feat(expr-ir): Support group_by, utilize pyarrow.acero
#3143
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Mapping things out a bit, no compliant yet
There's a few gaps, but overall surprised how much was reusable 🥳
lol didn't realise it was just describing python dict behavior
Everything here seems to be working already? 😱
May as well show it off
```py
>>> df.group_by("a", nwp.nth(2, 8)).agg(nwp.mean("d", "e", "g").name.suffix("_mean"))
NotImplementedError: TODO: `GroupBy.agg` needs a `CompliantGroupBy` to dispatch to:
keys:
(a=col('a'), c=col('c'), i=col('i'))
aggs:
(d_mean=col('d').mean(), e_mean=col('e').mean(), g_mean=col('g').mean())
result_schema:
FrozenSchema([
('a', String),
('c', Int64),
('i', Unknown),
('d_mean', Unknown),
('e_mean', Unknown),
('g_mean', Unknown),
])
```
Quite different to current version(s)
75 tasks
Gonna need space for the mini translator
woops Making it a separate node rather than having a flag https://github.com/pola-rs/polars/blob/cdd247aaba8db3332be0bd031e0f31bc3fc33f77/crates/polars-plan/src/dsl/mod.rs#L872-L889
`pyarrow` has the same behavior as `polars`
Wasn't expecting so much to be working already 🥳 🥳 🥳
Just pushing this as tests are working. Useful changes to follow: - Column renaming stuff will be avoidable - we just use `ArrowAggExpr.output_name` - Awkward stuff `first`, `last`, `_ensure_single_thread` can be avoided - `use_threads` was always available on `Declaration.to_table` - Whether we need to use can just be an `__ior__`
`ArrowDataFrame.drop_nulls` is shorter and waaaaaaay more efficient than `main`
In theory, we should be able to compose `over()` using combinations of: - aggregate - both scalar and hash - order_by - project - hashjoin
Will make for an easier diff in the PR that splits up this mess
dangotbanned
commented
Sep 29, 2025
dangotbanned
changed the title
group_bygroup_by, utilize pyarrow.acero
dangotbanned
commented
Sep 30, 2025
well, move it out of the docstring is more accurate i suppose
Still need to do `temp.column_names` as well That one is quite different to #3147, but was needed in this PR
This was referenced Oct 2, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What type of PR is this? (check all applicable)
Related issues
ExprIR #2572Description
This one spiralled into something pretty interesting IMO 😏
Acero is used under-the-hood for
pyarrowfeatures like:Table.group_by(most important to this PR)Table.filterTable.{join,join_asof}But it also has limited support for other operations like:
projectaggregateorder_byI've been experimenting with those a bit in (https://github.com/narwhals-dev/narwhals/blob/e0a3684ba4f77a59fe45d2914c74b4cff25cf344/narwhals/_plan/arrow/acero.py).
Anticipating that combining those nodes into one big soup will result in
over(...)😄Tasks
tests/frame/group_by_testtest_double_same_aggregationtest_group_by_len_1_columngroup_by(drop_null_keys=True)drop_null_keyswithExprExprkeystest_group_by_exprcase threadGroupBy.__iter__ArrowGroupBy.__iter____iter__pc.Expressioninstead of eager predicateremove_columnArrowDataFrame.row__iter__Expr.uniquegroup_by testsfirst,lastgroup_by testscol)df.group_by(**named_by)ArrowAggExprAggSpecArrowGroupByby_named_irsmay be reusable for everything but polarstemp(thread)temp.column_nameerror messagetemp.column_nametemp.column_namesteststest_temp_column_names_failed_uniquetemp.column_nametemp.column_namesitertools.cycle)