feat: accept column names for aggregate function inputs by richacode007-byte · Pull Request #1783 · apache/datafusion-python

richacode007-byte · 2026-10-03T21:37:12Z

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 like select and aggregate already accept plain column names. This PR lets aggregate functions take a column name too, so F.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 like concat or replace a 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?

Aggregate functions (sum, avg, min, max, count, corr, regr_*, first_value, bit_and, string_agg, etc.) now accept a column name wherever they took an Expr for a column. Strings go through the existing _to_raw_expr helper, and anything that isn't an Expr or a string still raises TypeError.
count also accepts a list of names, and the percentile functions accept a name for sort_expression and weight.
Things that aren't columns are left alone: filter=, string_agg's delimiter, nth_value's n, and count_star.
Added a short note to the make-pythonic skill so this exception doesn't get undone later.
Added tests showing string inputs give the same results as col(...), and that bad inputs like F.sum(1) still raise.
Merged latest main. The new any_value also accepts a column name, and the new distinct= 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") or F.corr("a", "b"), and existing code that passes Expr works exactly as before.

richacode007-byte · 2026-10-03T21:44:56Z

@davisp @kou @viirya Pls review the code

timsaucer · 2026-10-05T12:53:05Z

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.

richacode007-byte · 2026-10-05T16:00:57Z

@timsaucer Issue are resolved , Pls review them

timsaucer

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 · 2026-10-05T16:18:45Z

 **IMPORTANT:** In `functions.py`, string arguments almost never mean column names. Functions operate on expressions, and column references should use `col()`. Category C applies mainly to DataFrame methods and context APIs, not to scalar/aggregate/window functions. Do NOT convert string arguments to column expressions in `functions.py` unless there is a very clear reason to do so.

+The documented exception is the column inputs of aggregate functions (`sum`, `avg`, `count`, `corr`, the `regr_*` family, and so on). There a string can only mean a column, so they accept `Expr | str` via `_to_raw_expr()`. Their literal arguments (for example `string_agg`'s `delimiter` or `nth_value`'s `n`) are unaffected.
+


It sounds like we've got a fairly large list of "exceptions" to the rule here, so maybe the rule needs revisiting instead. Maybe we need to rewrite this to explain when it is appropriate to turn bare strings into literals and when to treat it as column names.

timsaucer · 2026-10-05T16:19:20Z

+def test_aggregate_rejects_non_column_input(build) -> None:
+    with pytest.raises(TypeError, match="Expected Expr or column name"):
+        build()


What do you think is the value added of this kind of test?

timsaucer · 2026-10-05T16:20:58Z

+def test_string_agg_accepts_column_name() -> None:
+    ctx = SessionContext()
+    df = ctx.from_pydict({"a": ["one", "two", "three"], "b": [2, 0, 1]})
+
+    result = df.aggregate([], [f.string_agg("a", ",", order_by="b").alias("v")])
+
+    assert result.collect()[0].to_pydict() == {"v": ["two,three,one"]}


Could this not have been one of the variants in the previous test?

timsaucer · 2026-10-05T16:23:26Z

+        args = [_to_raw_expr(arg) for arg in expressions]
    else:
-        args = [expressions.expr]
+        args = [_to_raw_expr(expressions)]


My agent picked up on something I didn't. It's review:

The one to decide: count("*"). SQL users will write this. With the PR, "*" becomes a column named * and fails at execution with Schema error: No field named "*". Before, it failed right away with AttributeError. It's not a regression, but the PR invites the call. Two options: - Treat "*" as count() (same as count(lit(1))). - Reject it at build time with a TypeError pointing to count(). I'd pick the first: it matches SQL and costs one line. Either way, add a test.

richacode007-byte added 3 commits October 4, 2026 03:04

Update __init__.py

9840d67

Update test_aggregation.py

32b1bbf

Update SKILL.md

786efe2

Merge branch 'main' into feat/1757-aggregate-column-names

75d6f27

timsaucer reviewed Oct 5, 2026

View reviewed changes

Navigation Menu

feat: accept column names for aggregate function inputs - #1783

feat: accept column names for aggregate function inputs#1783
richacode007-byte wants to merge 4 commits into
apache:mainfrom
richacode007-byte:feat/1757-aggregate-column-names

richacode007-byte commented Oct 3, 2026 •

edited

Loading

Uh oh!

richacode007-byte commented Oct 3, 2026

Uh oh!

timsaucer commented Oct 5, 2026

Uh oh!

richacode007-byte commented Oct 5, 2026

Uh oh!

timsaucer left a comment

Uh oh!

timsaucer Oct 5, 2026

Uh oh!

timsaucer Oct 5, 2026

Uh oh!

timsaucer Oct 5, 2026

Uh oh!

timsaucer Oct 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		IMPORTANT: In `functions.py`, string arguments almost never mean column names. Functions operate on expressions, and column references should use `col()`. Category C applies mainly to DataFrame methods and context APIs, not to scalar/aggregate/window functions. Do NOT convert string arguments to column expressions in `functions.py` unless there is a very clear reason to do so.

		The documented exception is the column inputs of aggregate functions (`sum`, `avg`, `count`, `corr`, the `regr_*` family, and so on). There a string can only mean a column, so they accept `Expr \| str` via `_to_raw_expr()`. Their literal arguments (for example `string_agg`'s `delimiter` or `nth_value`'s `n`) are unaffected.

Conversation

richacode007-byte commented Oct 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Uh oh!

richacode007-byte commented Oct 3, 2026

Uh oh!

timsaucer commented Oct 5, 2026

Uh oh!

richacode007-byte commented Oct 5, 2026

Uh oh!

timsaucer left a comment

Choose a reason for hiding this comment

Uh oh!

timsaucer Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

timsaucer Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

timsaucer Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

timsaucer Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

richacode007-byte commented Oct 3, 2026 •

edited

Loading