diff --git a/crates/polars-sql/src/functions.rs b/crates/polars-sql/src/functions.rs index db0b3f197cff..92c9b1fba4c2 100644 --- a/crates/polars-sql/src/functions.rs +++ b/crates/polars-sql/src/functions.rs @@ -1860,12 +1860,14 @@ impl SQLFunctionVisitor<'_> { }) .collect::>>()?; - Ok(self + let expr = self .ctx .function_registry .get_udf(func_name)? .ok_or_else(|| polars_err!(SQLInterface: "UDF {} not found", func_name))? - .call(args)) + .call(args); + + self.apply_window_spec(expr, &self.func.over) } /// Validate window frame specifications. @@ -2031,6 +2033,7 @@ impl SQLFunctionVisitor<'_> { .dot(self.parse_array_inner_product_arg(rhs)?)), _ => self.not_supported_error(), } + .and_then(|e| self.apply_window_spec(e, &self.func.over)) } fn visit_unary(&mut self, f: impl Fn(Expr) -> Expr) -> PolarsResult { @@ -2107,6 +2110,7 @@ impl SQLFunctionVisitor<'_> { }, _ => self.not_supported_error(), } + .and_then(|e| self.apply_window_spec(e, &self.func.over)) } fn visit_variadic(&mut self, f: impl Fn(&[Expr]) -> Expr) -> PolarsResult { @@ -2126,7 +2130,7 @@ impl SQLFunctionVisitor<'_> { return self.not_supported_error(); }; } - f(&expr_args) + f(&expr_args).and_then(|e| self.apply_window_spec(e, &self.func.over)) } fn try_visit_ternary( @@ -2147,6 +2151,7 @@ impl SQLFunctionVisitor<'_> { }, _ => self.not_supported_error(), } + .and_then(|e| self.apply_window_spec(e, &self.func.over)) } fn visit_nullary(&self, f: impl Fn() -> Expr) -> PolarsResult { @@ -2216,7 +2221,7 @@ impl SQLFunctionVisitor<'_> { sql_expr, "ARRAY_AGG", )?; - Ok(base.implode(true)) + self.apply_window_spec(base.implode(true), &self.func.over) }, _ => { polars_bail!(SQLSyntax: "ARRAY_AGG must have exactly one argument; found {}", args.len()) @@ -2264,9 +2269,12 @@ impl SQLFunctionVisitor<'_> { .list() .join(separator, true); - Ok(when(base.clone().null_count().lt(base.len())) - .then(joined) - .otherwise(lit(LiteralValue::untyped_null()))) + self.apply_window_spec( + when(base.clone().null_count().lt(base.len())) + .then(joined) + .otherwise(lit(LiteralValue::untyped_null())), + &self.func.over, + ) } fn visit_arr_to_string(&mut self) -> PolarsResult { diff --git a/py-polars/tests/unit/sql/test_window_functions.py b/py-polars/tests/unit/sql/test_window_functions.py index 69eae4dd7f88..0c511dd20fa8 100644 --- a/py-polars/tests/unit/sql/test_window_functions.py +++ b/py-polars/tests/unit/sql/test_window_functions.py @@ -758,3 +758,53 @@ def test_window_order_by_mixed_nulls_placement_error() -> None: ctx.execute( "SELECT ROW_NUMBER() OVER (ORDER BY a NULLS LAST, b NULLS FIRST) FROM df" ) + + +@pytest.mark.parametrize( + "agg", + [ + "CORR(a, b)", + "COVAR_POP(a, b)", + "COVAR_SAMP(a, b)", + "QUANTILE_CONT(a, 0.5)", + "QUANTILE_DISC(a, 0.5)", + "STRING_AGG(g, '-')", + ], +) +def test_window_multi_arg_aggregate_partition_by(agg: str) -> None: + lf = pl.LazyFrame( + { + "i": [0, 1, 2, 3, 4], + "g": ["a", "a", "a", "b", "b"], + "a": [1, 2, 3, 4, 5], + "b": [1, 3, 2, 10, 20], + } + ) + assert_sql_matches( + {"df": lf}, + query=f""" + SELECT i, g, {agg} OVER (PARTITION BY g) AS res + FROM df + ORDER BY i + """, + compare_with="duckdb", + ) + + +def test_window_array_agg_partition_by() -> None: + lf = pl.LazyFrame( + { + "i": [0, 1, 2, 3, 4], + "g": ["a", "a", "a", "b", "b"], + "a": [1, 2, 3, 4, 5], + } + ) + assert_sql_matches( + lf, + query="SELECT i, ARRAY_AGG(a) OVER (PARTITION BY g) AS res FROM self ORDER BY i", + compare_with=None, + expected={ + "i": [0, 1, 2, 3, 4], + "res": [[1, 2, 3], [1, 2, 3], [1, 2, 3], [4, 5], [4, 5]], + }, + )