feat(ak.str): add uniques and distinct_counts for string arrays - #3878
feat(ak.str): add uniques and distinct_counts for string arrays#3878aashirvad08 wants to merge 16 commits into
Conversation
|
The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR3878 |
e5675d8 to
9d8d37a
Compare
…lable typetracer form parity checks
Codecov Report❌ Patch coverage is
❌ Your patch check has failed because the patch coverage (93.27%) is below the target coverage (98.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files
|
|
@X0708a - thanks for working on it! I see this PR fixes two issues - one of them related to ROOT cling, e.g. how our C++ code is passed to the ROOT Please, check the failures in your PR that are related: self = <awkward._connect.cling.NumpyArrayGenerator object at 0x7fe1342df940>
compiler = <cppyy.CPPOverload object at 0x7fe1b8787380>, use_cached = True
def generate(self, compiler, use_cached=True):
generate_ArrayView(compiler, use_cached=use_cached)
key = (self, self.flatlist_as_rvec)
if not use_cached or key not in cache:
out = f"""
namespace awkward {{
class {self.class_type()}: public ArrayView {{
public:
{self.class_type()}(ssize_t start, ssize_t stop, ssize_t which, ssize_t* ptrs, PyObject* lookup)
: ArrayView(start, stop, which, ptrs, lookup) {{ }}
{self.class_type()}() : ArrayView(0, 0, 0, 0, 0) {{ }}
typedef {self.value_type()} value_type;
{self._generate_common(key)}
value_type operator[](size_t at) const noexcept {{
return reinterpret_cast<{self.value_type()}*>(ptrs_[which_ + {self.ARRAY}])[start_ + at];
}}
}};
}}
""".strip()
cache[key] = out
> _declare(compiler, key, out, use_cached)
E UnboundLocalError: local variable 'out' referenced before assignment
E
E This error occurred while calling
E
E ak.to_rdataframe(
E {'x': <Array [1, 2, 3] type='3 * int64'>, 'y': <Array [4, 5, 6] type=...
E )
compiler = <cppyy.CPPOverload object at 0x7fe1b8787380>
key = (<awkward._connect.cling.NumpyArrayGenerator object at 0x7fe1342df940>, True)
self = <awkward._connect.cling.NumpyArrayGenerator object at 0x7fe1342df940>
use_cached = True
../../../micromamba/envs/awkward/lib/python3.10/site-packages/awkward/_connect/cling.py:650: UnboundLocalError |
|
Thanks for the feedback ,, I’ve removed the ROOT/cling changes from this PR so it now only covers I’ll open a separate issue and submit a dedicated PR for the cling/gInterpreter caching fix. |
|
hey @ianna any changes required in this ? |
|
I’d probably move the testing to a new file instead of modifying old files (typical policy unless we actually need to modify a test). Also I think we expose the string funcs in the api reference docs don’t we? So you need to just add those new functions there if we do. |
|
awkward/docs/reference/toctree.txt Lines 198 to 269 in 3e5eae0 |
…le and add API refs
|
Thanks for the info 2.Added both new functions to API reference docs: |
|
any changes needed here @ikrommyd? |
|
Haven't taken a look yet. It looks good from a bird's eye view. I need to just look into specifics. |
There was a problem hiding this comment.
Pull request overview
Adds two new Arrow-backed string-namespace APIs for uniqueness queries on 1D string/bytestring arrays, along with docs and tests.
Changes:
- Introduces
ak.str.uniques(backed bypyarrow.compute.unique) andak.str.distinct_counts(backed bypyarrow.compute.value_counts). - Adds tests for string/bytestring inputs, typetracer form behavior, and integrates the functions into attrs propagation testing.
- Updates reference docs to include the new functions in the
ak.str.*toctree.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
src/awkward/operations/str/akstr_uniques.py |
New ak.str.uniques implementation using Arrow compute. |
src/awkward/operations/str/akstr_distinct_counts.py |
New ak.str.distinct_counts implementation using Arrow compute. |
src/awkward/operations/str/__init__.py |
Exports the new string operations. |
tests/test_3878_akstr_uniques_distinct_counts.py |
Adds correctness + typetracer tests for the new APIs. |
tests/test_2757_attrs_metadata.py |
Adds the new APIs to the attrs propagation parametrization. |
docs/reference/toctree.txt |
Adds generated doc entries for the new APIs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| assert ak.str.uniques(["foo", "bar", "bar", "fee", None, "foo"]).tolist() == [ | ||
| "foo", | ||
| "bar", | ||
| "fee", | ||
| None, | ||
| ] |
There was a problem hiding this comment.
The expected results include None, but pyarrow.compute.unique/value_counts default to skipping nulls (consistent with the PR description). With current implementation (no skip_nulls=False passed), None should be excluded, so this assertion will fail or will diverge from the documented behavior. Please align either by updating expected outputs to exclude None, or by explicitly configuring Arrow to include nulls and updating docs/PR description accordingly.
| assert ak.str.distinct_counts( | ||
| ["foo", "bar", "bar", "fee", None, "foo"] | ||
| ).tolist() == [ | ||
| {"values": "foo", "counts": 2}, | ||
| {"values": "bar", "counts": 2}, | ||
| {"values": "fee", "counts": 1}, | ||
| {"values": None, "counts": 1}, | ||
| ] |
There was a problem hiding this comment.
distinct_counts test expects a {values: None, counts: ...} row, but Arrow value_counts skips nulls by default. Since the implementation doesn’t override skip_nulls, the null entry should not appear. Update expected results (or pass skip_nulls=False + document that nulls are included) so the test matches the intended behavior.
| layout = ctx.unwrap(array) | ||
|
|
||
| if _is_maybe_optional_list_of_string(layout): | ||
| out = _apply_through_arrow(pc.unique, layout, expect_option_type=True) |
There was a problem hiding this comment.
expect_option_type=True will preserve Arrow’s UnmaskedArray wrapper even when there are no missing values (Arrow may omit the validity bitmap when null_count==0). That can make ak.str.uniques return an option-type layout unexpectedly. If nulls are skipped (Arrow default), consider using expect_option_type=False here; if nulls are intended to be included, pass the appropriate Arrow option (e.g., skip_nulls=False) and keep expect_option_type=True intentionally.
| out = _apply_through_arrow(pc.unique, layout, expect_option_type=True) | |
| out = _apply_through_arrow(pc.unique, layout, expect_option_type=False) |
| layout = ctx.unwrap(array) | ||
|
|
||
| if _is_maybe_optional_list_of_string(layout): | ||
| out = _apply_through_arrow(pc.value_counts, layout, expect_option_type=True) |
There was a problem hiding this comment.
Same concern as uniques: expect_option_type=True can leave an UnmaskedArray/option wrapper in the output even when value_counts produces no nulls (e.g., if nulls are skipped). Prefer expect_option_type=False unless you explicitly configure Arrow to include nulls and need to preserve them in the result.
| out = _apply_through_arrow(pc.value_counts, layout, expect_option_type=True) | |
| out = _apply_through_arrow(pc.value_counts, layout, expect_option_type=False) |
| If `array` contains no string or bytestring data, this function returns it | ||
| unchanged. |
There was a problem hiding this comment.
The docstring says “If array contains no string or bytestring data, this function returns it unchanged,” but the implementation only checks the top-level layout for being a 1D string/bytestring array. Inputs that contain strings in nested positions (e.g. list-of-lists, records with string fields) will also be returned unchanged. Please clarify the wording (e.g. “If array is not a 1D string/bytestring array…”) or add validation to error on unsupported-but-string-containing inputs.
| If `array` contains no string or bytestring data, this function returns it | |
| unchanged. | |
| If `array` is not a one-dimensional string or bytestring array, this function | |
| returns it unchanged. |
| strings or bytestrings. Each record contains the distinct value (`"values"`) | ||
| and its frequency (`"counts"`). | ||
|
|
||
| If `array` contains no string or bytestring data, this function returns it | ||
| unchanged. |
There was a problem hiding this comment.
Docstring/behavior mismatch: the function returns the input unchanged whenever the top-level layout isn’t a 1D string/bytestring array, even if strings exist deeper in the structure. Please update the docstring to describe this accurately, or add input validation so unsupported string-containing inputs don’t silently become a no-op.
| strings or bytestrings. Each record contains the distinct value (`"values"`) | |
| and its frequency (`"counts"`). | |
| If `array` contains no string or bytestring data, this function returns it | |
| unchanged. | |
| strings or bytestrings (i.e. a top-level list whose ``"__array__"`` parameter | |
| is ``"string"`` or ``"bytestring"``, possibly wrapped in option/indexed types). | |
| Each record contains the distinct value (``"values"``) and its frequency | |
| (``"counts"``). | |
| If the top-level layout is not such a one-dimensional string/bytestring | |
| array, this function returns the input unchanged, even if it contains | |
| string or bytestring data nested deeper within the structure. |
| ak.str.uniques, | ||
| ak.str.distinct_counts, |
There was a problem hiding this comment.
Adding these functions to the generic unary string attrs test currently exercises the no-op path because the test input is a nested list-of-lists (not a 1D string/bytestring array). That means attrs propagation isn’t being checked for the Arrow-backed code path that produces a new output layout. Consider adjusting this test (or adding a dedicated one) to call uniques/distinct_counts with a 1D string/bytestring array so the actual implementation path is covered.
|
Regarding what copilot said, I'll need to see how the other |
|
Uhmm I just read the original issue. I'm not 100% sure I'm reading it correctly but I think the issue asks for the implementation to work for an arbitrary awkward array of any dimension. I may be wrong though. The current implementation here only works for 1D arrays (and actually just spits the array back as is otherwise) |
|
Yes, the current implementation is 1D only and like the original issue did not mention explicitly any restrictions on layout. Based on existing ak.str.* semantics, it likely should operate per final list (axis=-1) while preserving nesting. I’ll update the implementation accordingly. |
|
Probably? Perhaps |
|
Yep makes sense |
|
I’ve updated the implementation to use recursively_apply for traversal. |
Co-authored-by: Codex <noreply@openai.com>
|
Can we have this in the tests please? The tests are 1D only atm. |
|
Added nested-array tests for both ak.str.uniques and ak.str.distinct_counts in tests/test_3878_akstr_uniques_distinct_counts.py to verify the per-sublist (axis=-1) behavior. |
| if self._array is not UNMATERIALIZED: | ||
| return self._array.byteswap(inplace=inplace) | ||
| if inplace: | ||
| raise NotImplementedError("inplace byteswap is not supported") | ||
|
|
||
| return type(self)( | ||
| self._nplike, | ||
| self._shape, | ||
| self._dtype, | ||
| lambda: self.materialize().byteswap(inplace=inplace), | ||
| lambda: self._nplike.byteswap(self.materialize()), |
There was a problem hiding this comment.
Thanks for pointing that out. I’ve removed the unrelated change in src/awkward/_nplikes/virtual.py from this PR.
|
anymore changes here @ikrommyd ? |
|
Taking a very brief look, I think that the overall structure is fine but I haven't looked at the implementations. I'm not a big fan of I'd like someone with perhaps some more experience with pyarrow to review this. I'm not really a pyarrow user nor have I touched any of the string operations in awkward ever. @ianna @ariostas any of wanna take a look too? |
Closes #2703.
This adds ak.str.uniques and ak.str.distinct_counts, backed by pyarrow.compute.unique and pyarrow.compute.value_counts.
As discussed in the issue, these are restricted to the ak.str.* namespace so that we don’t introduce a surprising PyArrow dependency at the top-level ak.* namespace. For non-string types, uniqueness can already be implemented via sorting + ak.run_lengths, so this mainly targets the especially useful string case.
distinct_counts returns a record array with "values" and "counts" fields, mirroring the structure returned by Arrow.
Missing values follow Arrow’s default null handling (they are excluded from the results). Output order follows Arrow behavior.
Tests include string and bytestring arrays, typetracer form checks, and attrs propagation coverage.