Skip to content

Commit 0b6303e

Browse files
committed
Add type checks for aggregation expressions in DataFrame.aggregate method
1 parent bc6fb7c commit 0b6303e

File tree

1 file changed

+9
-1
lines changed

1 file changed

+9
-1
lines changed

python/datafusion/dataframe.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -549,7 +549,15 @@ def aggregate(
549549
group_by_exprs = [
550550
Expr.column(e).expr if isinstance(e, str) else e.expr for e in group_by_list
551551
]
552-
aggs_exprs = [e.expr for e in aggs_list]
552+
aggs_exprs = []
553+
for agg in aggs_list:
554+
if not isinstance(agg, Expr):
555+
msg = (
556+
f"Expected Expr, got {type(agg).__name__}. "
557+
"Use col() or lit() to construct expressions."
558+
)
559+
raise TypeError(msg)
560+
aggs_exprs.append(agg.expr)
553561
return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs))
554562

555563
def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:

0 commit comments

Comments
 (0)