feat(expr-ir): Support group_by
#3143
Draft
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)
60 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`
Didn't show up as an issue until trying to drop multiple keys Surprised pyarrow doesnt support this?
dangotbanned
commented
Sep 23, 2025
dangotbanned
commented
Sep 23, 2025
Comment on lines
+130
to
+144
| class temp: # noqa: N801 | ||
| """Temporary mini namespace for temporary utils.""" | ||
|
|
||
| _MAX_ITERATIONS: ClassVar[int] = 100 | ||
|
|
||
| @classmethod | ||
| def column_name( | ||
| cls, | ||
| source: _StoresColumns | Iterable[str], | ||
| /, | ||
| *, | ||
| prefix: str = "nw", | ||
| n_chars: int = 16, | ||
| ) -> str: | ||
| """Generate a single, unique column name that is not present in `source`.""" |
There was a problem hiding this comment.
TODO: Document + test all the temp stuff
Rushed this a bit since I needed temp.column_names.
Would be nice to have something to point to about why this:
nw.temp.column_name(df, prefix="count_")is preferable to this:
subset = df.columns
nw.generate_temporary_column_name(n_bytes=8, columns=subset, prefix="count_")RE: #3147
dangotbanned
commented
Sep 23, 2025
narwhals/_plan/dataframe.py
Outdated
Comment on lines
83
to
88
| # NOTE: Want to be able to call `with_columns` at compliant level - and still get the right schema | ||
| # - Currently it acts like select in `group_by` | ||
| # - Doing some gymnastics to workaround for now | ||
| def with_columns(self, *exprs: OneOrIterable[IntoExpr], **named_exprs: Any) -> Self: | ||
| named_irs, _ = self._project(exprs, named_exprs, ExprContext.WITH_COLUMNS) | ||
| return self._from_compliant(self._compliant.with_columns(named_irs)) |
There was a problem hiding this comment.
Important
Don't merge back into (#2572) without addressing this
There was a problem hiding this comment.
Mostly addressed in:
As in it'll do for now, but I really wish we could have something more like
Added these rules before I realised it described `dict`
dangotbanned
commented
Sep 23, 2025
Still need to decide what to do with named tuple
…ojection` Can accomodate - any narwhals/compliant-level frame - any reasonable transformation on a schema - the result of `freeze_schema` can be passed *back* into itself for free
- Currently just replaces the narwhals-level bits - Next part is making use of them for compliant-level
Tricky to do this in a way that doesn't end up duplicating lots of logic/layers
- Enough to pass the remaining tests - Hoping to reduce the complexity in `__iter__` next
Bug isnt present on `main`, probably need to redo the resolve stuff to fix
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 #2572Tasks
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__Expr.uniquegroup_by testsfirst,lastgroup_by testscol)df.group_by(**named_by)ArrowAggExprArrowGroupByby_named_irsmay be reusable for everything but polars