richacode007-byte opened a new pull request, #1783:
URL: https://github.com/apache/datafusion-python/pull/1783
# 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
