Skip to content

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

Open
richacode007-byte wants to merge 3 commits into
apache:mainfrom
richacode007-byte:feat/1757-aggregate-column-names
Open

richacode007-byte wants to merge 3 commits into
apache:mainfrom
richacode007-byte:feat/1757-aggregate-column-names

Conversation

@richacode007-byte

Copy link
Copy Markdown

Which issue does this PR close?

Closes #1757.

Rationale for this change

Aggregate functions currently require an Expr for their column inputs, so users have to write F.sum(col("a")) even though F.sum("a") can only mean "sum column a". DataFrame methods such as select, sort and aggregate already accept column names as strings, and aggregate order_by= arguments do too. Accepting them for aggregate column inputs makes the API more consistent and less verbose.

The scope is limited on purpose to the column inputs of aggregate functions. The project's guidance (.ai/skills/make-pythonic/SKILL.md, "Category C") warns against treating strings as column names in functions.py, because for scalar functions a string is often a literal (concat, replace, split_part, ...). An aggregate's input, by contrast, is always a column or expression, so a string is unambiguous there. Scalar, string, array and window functions are unchanged.

What changes are included in this PR?

python/datafusion/functions/__init__.py: the column arguments are now typed Expr | str and converted with the existing _to_raw_expr helper from expr.py. A string becomes Expr.column(name), and any other non-Expr value still raises TypeError.

  • expression in approx_distinct, approx_median, array_agg, avg, max, median, min, sum, stddev, stddev_pop, var_pop, var_samp, first_value, last_value, nth_value, bit_and, bit_or, bit_xor, bool_and, bool_or, string_agg
  • value_y / value_x in corr, covar_pop, covar_samp
  • y / x in all nine regr_* functions
  • count: accepts Expr | str | list[Expr | str] | None
  • approx_percentile_cont, approx_percentile_cont_with_weight, percentile_cont: sort_expression accepts a column name (turned into a column before the default sort is applied), and weight accepts one too
  • Aliases updated so their type hints match their targets: mean, var, var_sample, var_population, stddev_samp, covar, quantile_cont
  • Each changed argument's docstring notes that it accepts a column name, and sum gains a doctest that passes a plain string

Unchanged: count_star, all filter= arguments, string_agg's delimiter, and nth_value's n.

.ai/skills/make-pythonic/SKILL.md: one paragraph in "Category C" recording aggregate column inputs as the documented exception, so future audits don't undo it.

python/tests/test_aggregation.py:

  • string-input versions of the existing test_aggregation_stats cases (avg, corr, count, covar, max, mean, sum, stddev_samp, var), checked against the same numpy results
  • test_aggregate_accepts_column_name: string and col(...) inputs give the same result for count([...]), regr_slope, approx_percentile_cont, approx_percentile_cont_with_weight, percentile_cont, first_value, array_agg, bool_and
  • test_string_agg_accepts_column_name
  • test_aggregate_rejects_non_column_input: inputs such as f.sum(1) still raise TypeError

Testing:

  • pytest python/tests: 1053 passed, 8 skipped
  • pytest python/datafusion/functions --doctest-modules: 323 passed
  • ruff check and ruff format --check: clean

Separate existing bug, not fixed here: on main, F.count([col("a"), col("b")]) already fails with TypeError: count() got multiple values for argument 'distinct', because the Rust count binding takes a single expression. The new tests use a single-item list. I can open a separate issue for this.

Are there any user-facing changes?

Yes, and they're additive. Aggregate functions now accept a column name wherever they accepted an Expr for a column input:

df.aggregate([], [F.sum("a"), F.corr("a", "b"), F.approx_percentile_cont("b", 0.5)])

Existing code that passes Expr behaves exactly as before. No public API is removed or renamed, so I don't think this needs the api change label. The docstrings for the affected functions are updated.

@richacode007-byte

Copy link
Copy Markdown
Author

@davisp @kou @viirya Pls review the code

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Allow "str" for function arguments that accept "col" expressions

1 participant