timsaucer commented on code in PR #1783:
URL:
https://github.com/apache/datafusion-python/pull/1783#discussion_r4186232969
##########
python/tests/test_aggregation.py:
##########
@@ -555,6 +565,57 @@ def test_string_agg(name, expr, result) -> None:
assert df.collect()[0].to_pydict() == expected
[email protected](
+ ("by_name", "by_expr"),
+ [
+ (f.count(["e"]), f.count([column("e")])),
+ (f.regr_slope("c", "a"), f.regr_slope(column("c"), column("a"))),
+ (
+ f.approx_percentile_cont("b", 0.5),
+ f.approx_percentile_cont(column("b"), 0.5),
+ ),
+ (
+ f.approx_percentile_cont_with_weight("b", "a", 0.5),
+ f.approx_percentile_cont_with_weight(column("b"), column("a"),
0.5),
+ ),
+ (f.percentile_cont("c", 0.5), f.percentile_cont(column("c"), 0.5)),
+ (
+ f.first_value("a", order_by="c"),
+ f.first_value(column("a"), order_by="c"),
+ ),
+ (f.array_agg("a", order_by="a"), f.array_agg(column("a"),
order_by="a")),
+ (f.bool_and("d"), f.bool_and(column("d"))),
+ ],
+)
+def test_aggregate_accepts_column_name(df, by_name, by_expr) -> None:
+ result = df.aggregate([], [by_name.alias("v")]).collect()[0]
+ expected = df.aggregate([], [by_expr.alias("v")]).collect()[0]
+ assert result == expected
+
+
+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"]}
+
+
[email protected](
+ "build",
+ [
+ lambda: f.sum(1),
+ lambda: f.corr("a", 1),
+ lambda: f.count(["a", 1]),
+ lambda: f.approx_percentile_cont_with_weight("a", 1, 0.5),
+ ],
+)
+def test_aggregate_rejects_non_column_input(build) -> None:
+ with pytest.raises(TypeError, match="Expected Expr or column name"):
+ build()
Review Comment:
What do you think is the value added of this kind of test?
##########
.ai/skills/make-pythonic/SKILL.md:
##########
@@ -338,6 +338,8 @@ def some_function(column: Expr | str) -> Expr:
**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.
+
Review Comment:
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.
##########
python/datafusion/functions/__init__.py:
##########
@@ -5920,14 +5935,16 @@ def count(
if expressions is None:
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)]
Review Comment:
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.
```
##########
python/tests/test_aggregation.py:
##########
@@ -555,6 +565,57 @@ def test_string_agg(name, expr, result) -> None:
assert df.collect()[0].to_pydict() == expected
[email protected](
+ ("by_name", "by_expr"),
+ [
+ (f.count(["e"]), f.count([column("e")])),
+ (f.regr_slope("c", "a"), f.regr_slope(column("c"), column("a"))),
+ (
+ f.approx_percentile_cont("b", 0.5),
+ f.approx_percentile_cont(column("b"), 0.5),
+ ),
+ (
+ f.approx_percentile_cont_with_weight("b", "a", 0.5),
+ f.approx_percentile_cont_with_weight(column("b"), column("a"),
0.5),
+ ),
+ (f.percentile_cont("c", 0.5), f.percentile_cont(column("c"), 0.5)),
+ (
+ f.first_value("a", order_by="c"),
+ f.first_value(column("a"), order_by="c"),
+ ),
+ (f.array_agg("a", order_by="a"), f.array_agg(column("a"),
order_by="a")),
+ (f.bool_and("d"), f.bool_and(column("d"))),
+ ],
+)
+def test_aggregate_accepts_column_name(df, by_name, by_expr) -> None:
+ result = df.aggregate([], [by_name.alias("v")]).collect()[0]
+ expected = df.aggregate([], [by_expr.alias("v")]).collect()[0]
+ assert result == expected
+
+
+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"]}
Review Comment:
Could this not have been one of the variants in the previous test?
--
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]
