Fix groupby size dtype for masked series and apply group order for sort=False - #23260
Conversation
…rt=False Fixes the 6 failing tests in pandas' tests/groupby/methods/test_size.py under cudf.pandas (22/22 pass) and removes their plugin entries. * SeriesGroupBy.size on masked (Int*/UInt*/Float*/boolean) dtypes now returns Int64 like pandas (GH#54132), joining the existing string[pyarrow] -> Int64 and ArrowDtype -> int64[pyarrow] special cases. DataFrameGroupBy.size stays int64, also like pandas. * GroupBy.apply with sort=False now processes groups in order of first appearance: libcudf returns groups sorted by key, so the grouped layout (group names, offsets, keys and values) is permuted to appearance order before either engine runs, mirroring pandas' iteration order for every result shape.
github-actions
Bot
added
Python
galipremsagar
added
bug
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesGroupBy behavior alignment
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
python/cudf/cudf/core/groupby/groupby.py (1)
2494-2521: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate first-appearance-order computation; consider extracting a shared helper.
This block re-implements the exact same "gather earliest row position per group, then argsort" logic already present, unchanged, in
__iter__(lines 617-624). Extracting a small helper (e.g._first_appearance_group_order()) would avoid two independent copies of this logic drifting apart over time.Separately,
row_orderis built via a Python-level list comprehension overgroup_orderplusnp.concatenate, whereas the very similar permutation in_head_tail(lines 1635-1641) is fully vectorized withnp.repeat/np.cumsumand avoids a per-group Python loop. Since this path is only exercised whensort=Falseand there is more than one group, andapply()already warns above_MAX_GROUPS_BEFORE_WARNgroups for the iterative engine, the practical impact is limited, but the vectorized form would be more consistent with the rest of the file.♻️ Sketch of vectorized row_order construction (mirrors `_head_tail`)
first_pos = positions.take(as_column(pos_offsets[:-1])) group_order = first_pos.argsort().to_numpy() sizes = np.diff(np.asarray(offsets, dtype=SIZE_TYPE_DTYPE)) - row_order = as_column( - np.concatenate( - [ - np.arange( - offsets[i], offsets[i + 1], dtype=SIZE_TYPE_DTYPE - ) - for i in group_order - ] - ) - ) + starts = np.asarray(offsets, dtype=SIZE_TYPE_DTYPE)[group_order] + ordered_sizes = sizes[group_order] + row_order_arr = np.arange(ordered_sizes.sum(), dtype=SIZE_TYPE_DTYPE) + fixup = np.empty_like(ordered_sizes) + fixup[0] = 0 + np.cumsum(ordered_sizes[:-1], out=fixup[1:]) + row_order_arr += np.repeat(starts - fixup, ordered_sizes) + row_order = as_column(row_order_arr)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf/cudf/core/groupby/groupby.py` around lines 2494 - 2521, Extract the shared first-appearance ordering logic from __iter__ and this sort=False branch into a helper such as _first_appearance_group_order(), then reuse it in both call sites. In the same branch, replace the Python list-comprehension/np.concatenate construction of row_order with the vectorized np.repeat/np.cumsum approach used by _head_tail, preserving the existing group and row permutation results.python/cudf/cudf/tests/groupby/test_size.py (1)
24-35: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCorrect, but consider adding a null-containing case.
The test only uses non-null values (
[1, 1, 1]), so it doesn't exercise the masked dtype's actual nulls. Since the fix targets nullable/masked dtypes, a variant withpd.NA/Nonein the data (and possibly an empty or single-element group) would strengthen coverage of the edge cases this change is meant to support.As per path instructions,
python/**/test_*.py: "Ensure test files provide comprehensive edge case coverage (empty, all-null, single-element, mixed types) and do not depend on external datasets."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf/cudf/tests/groupby/test_size.py` around lines 24 - 35, Add null-containing coverage to test_size_series_masked_dtype using pd.NA or None in the nullable input, while retaining validation of the Int64 result dtype and groupby equality. Include relevant edge cases such as an all-null or single-element group without relying on external data.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@python/cudf/cudf/core/groupby/groupby.py`:
- Around line 2494-2521: Extract the shared first-appearance ordering logic from
__iter__ and this sort=False branch into a helper such as
_first_appearance_group_order(), then reuse it in both call sites. In the same
branch, replace the Python list-comprehension/np.concatenate construction of
row_order with the vectorized np.repeat/np.cumsum approach used by _head_tail,
preserving the existing group and row permutation results.
In `@python/cudf/cudf/tests/groupby/test_size.py`:
- Around line 24-35: Add null-containing coverage to
test_size_series_masked_dtype using pd.NA or None in the nullable input, while
retaining validation of the Int64 result dtype and groupby equality. Include
relevant edge cases such as an all-null or single-element group without relying
on external data.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4f26d38f-8bac-4414-849c-1cb83d825e18
📒 Files selected for processing (3)
python/cudf/cudf/core/groupby/groupby.pypython/cudf/cudf/pandas/scripts/pandas-testing-plugin.pypython/cudf/cudf/tests/groupby/test_size.py
💤 Files with no reviewable changes (1)
- python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py
|
/okay to test defe55a |
|
/okay to test 9eaba7c |
|
/merge |
galipremsagar
added
5 - Ready to Merge
4 participants
Description
Running pandas' own test suite under
cudf.pandas,tests/groupby/methods/test_size.pyhad 6 failing tests. This PR fixes them (22/22 pass) and removes the corresponding xfail entries from the pandas-testing plugin. Two independent root causes:SeriesGroupBy.size()dtype for masked seriespandas returns
Int64forsize()on masked (Int*/UInt*/Float*/boolean) series (pandas GH#54132). Verified matrix on pandas 3.0.3:.size()onInt64,Float64,boolean, ...)Int64string[pyarrow]seriesInt64ArrowDtypeseriesint64[pyarrow]object/str/string[python]seriesint64DataFrameGroupByint64cuDF already special-cased the Arrow and
string[pyarrow]rows; this adds the missing masked branch (Series-only,StringDtypeexcluded sostring[python]staysint64).GroupBy.applygroup order withsort=Falsepandas processes groups in order of first appearance when
sort=False. libcudf returns groups sorted by key, andapplyconsumed that layout directly, so results came out key-sorted. (size()was already correct — the failing tests comparesize()againstapply(lambda a: a.shape[0]), and theapplyside was the broken one.)The grouped layout (group names, offsets, grouped keys and values) is now permuted into first-appearance order before either engine runs — the same device-side first-position/argsort technique
GroupBy.__iter__already uses — so every result shape (scalar-per-group, frame-per-group, transforms) sees pandas' iteration order.Validation
tests/groupby/methods/test_size.py: 6 failed → 22/22 pass.tests/groupby/test_apply.py+ the fulltests/groupby/methods/directory: remaining failures are exactly the pre-existing known-failure entries (plus one pre-existing local-onlytest_nthfailure that reproduces identically onmainand is fixed by Rewrite GroupBy.nth as a pandas-compatible positional row filter #23257).Checklist