feat: accept column names for aggregate function inputs - #1783
richacode007-byte wants to merge 7 commits into
Conversation
|
Thank you for the PR. The description reads like it's got quite a bit of AI slop in it. Could you make this more human facing? We've also got some conflicts that need to be resolved. |
|
@timsaucer Issue are resolved , Pls review them |
timsaucer
left a comment
timsaucer
left a comment
There was a problem hiding this comment.
Looks pretty good. A thing my agent picked up on:
- python/datafusion/functions/__init__.py:5755 (low): PR missed grouping().
Signature still expression: Expr, body still calls expression.expr
(line 5811). So F.grouping("a") fails with AttributeError: 'str' object has
no attribute 'expr', not the clear TypeError PR promises. Also contradicts
new note in .ai/skills/make-pythonic/SKILL.md saying aggregate column inputs
take Expr | str. Fix: type it Expr | str, call _to_raw_expr(expression).
|
@timsaucer : Thanks for the review, @timsaucer! I've pushed an update that addresses all of the comments |
There was a problem hiding this comment.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Copilot review overview
Review effort: Lite
Findings: 1
Open (2)
What changed in this PR
Adds support for passing plain column names (str) to aggregate function inputs (e.g., F.sum("a")) to match DataFrame method ergonomics and reduce boilerplate, while keeping non-column string parameters (like delimiters) as literals.
Changes:
- Expanded many aggregate function signatures to accept
Expr | strand routed inputs through_to_raw_expr. - Updated percentile/count variants to accept column names in additional positions and added test coverage to ensure parity with
col(...). - Documented the “column name vs literal string” rule in the make-pythonic skill guide to prevent future regressions.
| File | Description |
|---|---|
| python/tests/test_aggregation.py | Adds regression tests verifying string column-name inputs behave like col(...) for aggregates. |
| python/datafusion/functions/__init__.py | Updates aggregate APIs to accept str inputs and converts them via _to_raw_expr / column coercion. |
| .ai/skills/make-pythonic/SKILL.md | Clarifies guidance on when `Expr |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| if expressions is None or (isinstance(expressions, str) and expressions == "*"): | ||
| args = [Expr.literal(1).expr] | ||
| elif isinstance(expressions, list): | ||
| args = [arg.expr for arg in expressions] | ||
| args = [_to_raw_expr(arg) for arg in expressions] | ||
| else: | ||
| args = [expressions.expr] | ||
| args = [_to_raw_expr(expressions)] | ||
|
|
||
| return Expr(f.count(*args, distinct=distinct, filter=filter_raw)) |
|
|
||
| Args: | ||
| expression: Values to combine into an array | ||
| expression: Values to combine into an array (expression or column name) |
Which issue does this PR close?
Closes #1757.
Rationale for this change
Right now you have to write
F.sum(col("a"))to sum a column, even though DataFrame methods likeselectandaggregatealready accept plain column names. This PR lets aggregate functions take a column name too, soF.sum("a")just works.I kept this to aggregate inputs only. The make-pythonic skill warns against treating strings as columns in
functions.py, since for things likeconcatorreplacea string is usually a literal value. For an aggregate's input, a string can only mean a column, so there's no ambiguity.What changes are included in this PR?
sum,avg,min,max,count,corr,regr_*,first_value,bit_and,string_agg, etc.) now accept a column name wherever they took anExprfor a column. Strings go through the existing_to_raw_exprhelper, and anything that isn't anExpror a string still raisesTypeError.countalso accepts a list of names, and the percentile functions accept a name forsort_expressionandweight.filter=,string_agg's delimiter,nth_value'sn, andcount_star.col(...), and that bad inputs likeF.sum(1)still raise.main. The newany_valuealso accepts a column name, and the newdistinct=options are kept.All tests, doctests and ruff checks pass locally.
Side note: on
main,F.count([col("a"), col("b")])already fails because the Rust binding only takes one expression. I didn't touch that here, but I'm happy to open a separate issue.Are there any user-facing changes?
Yes, but nothing breaks. You can now write
F.sum("a")orF.corr("a", "b"), and existing code that passesExprworks exactly as before.