feat: accept column names for aggregate function inputs - #1783
Open
richacode007-byte wants to merge 3 commits into
Open
richacode007-byte wants to merge 3 commits into
richacode007-byte wants to merge 3 commits into
Conversation
Author
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
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.
Which issue does this PR close?
Closes #1757.
Rationale for this change
Aggregate functions currently require an
Exprfor their column inputs, so users have to writeF.sum(col("a"))even thoughF.sum("a")can only mean "sum columna". DataFrame methods such asselect,sortandaggregatealready accept column names as strings, and aggregateorder_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 infunctions.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 typedExpr | strand converted with the existing_to_raw_exprhelper fromexpr.py. A string becomesExpr.column(name), and any other non-Exprvalue still raisesTypeError.expressioninapprox_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_aggvalue_y/value_xincorr,covar_pop,covar_sampy/xin all nineregr_*functionscount: acceptsExpr | str | list[Expr | str] | Noneapprox_percentile_cont,approx_percentile_cont_with_weight,percentile_cont:sort_expressionaccepts a column name (turned into a column before the default sort is applied), andweightaccepts one toomean,var,var_sample,var_population,stddev_samp,covar,quantile_contsumgains a doctest that passes a plain stringUnchanged:
count_star, allfilter=arguments,string_agg'sdelimiter, andnth_value'sn..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:test_aggregation_statscases (avg,corr,count,covar,max,mean,sum,stddev_samp,var), checked against the same numpy resultstest_aggregate_accepts_column_name: string andcol(...)inputs give the same result forcount([...]),regr_slope,approx_percentile_cont,approx_percentile_cont_with_weight,percentile_cont,first_value,array_agg,bool_andtest_string_agg_accepts_column_nametest_aggregate_rejects_non_column_input: inputs such asf.sum(1)still raiseTypeErrorTesting:
pytest python/tests: 1053 passed, 8 skippedpytest python/datafusion/functions --doctest-modules: 323 passedruff checkandruff format --check: cleanSeparate existing bug, not fixed here: on
main,F.count([col("a"), col("b")])already fails withTypeError: count() got multiple values for argument 'distinct', because the Rustcountbinding 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
Exprfor a column input:Existing code that passes
Exprbehaves exactly as before. No public API is removed or renamed, so I don't think this needs theapi changelabel. The docstrings for the affected functions are updated.