fix: Apply the OVER clause to multi-argument aggregates

crates/polars-sql/src/functions.rs

@@ -1860,12 +1860,14 @@ impl SQLFunctionVisitor<'_> {
             })
             .collect::<PolarsResult<Vec<_>>>()?;
 
-        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<Expr> {
@@ -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<Expr> {
@@ -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<Arg: FromSQLExpr>(
@@ -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<Expr> {
@@ -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<Expr> {

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]],
+        },
+    )