fix: keep existing options when chaining aggregate and window builders

Upstream ExprFunctionExt methods on an Expr start from an empty builder,
so build() resets every option not set again. The Python function
wrappers already apply their keyword options, so calls such as
string_agg(..., order_by=...).distinct().build() silently dropped the
ordering, and .filter() dropped order_by, and so on.

Seed the builder from the expression's existing params instead. A window
frame equal to the default for its order_by is left unset so build()
derives it again from the final order_by.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: timsaucer

crates/core/src/expr.rs

@@ -625,34 +625,38 @@ impl PyExpr {
     // Expression Function Builder functions
 
     pub fn order_by(&self, order_by: Vec<PySortExpr>) -> PyExprFuncBuilder {
-        self.expr
-            .clone()
+        builder_from_expr(&self.expr)
             .order_by(to_sort_expressions(order_by))
             .into()
     }
 
     pub fn filter(&self, filter: PyExpr) -> PyExprFuncBuilder {
-        self.expr.clone().filter(filter.expr.clone()).into()
+        builder_from_expr(&self.expr)
+            .filter(filter.expr.clone())
+            .into()
     }
 
     pub fn distinct(&self) -> PyExprFuncBuilder {
-        self.expr.clone().distinct().into()
+        builder_from_expr(&self.expr).distinct().into()
     }
 
     pub fn null_treatment(&self, null_treatment: NullTreatment) -> PyExprFuncBuilder {
-        self.expr
-            .clone()
+        builder_from_expr(&self.expr)
             .null_treatment(Some(null_treatment.into()))
             .into()
     }
 
     pub fn partition_by(&self, partition_by: Vec<PyExpr>) -> PyExprFuncBuilder {
         let partition_by = partition_by.iter().map(|e| e.expr.clone()).collect();
-        self.expr.clone().partition_by(partition_by).into()
+        builder_from_expr(&self.expr)
+            .partition_by(partition_by)
+            .into()
     }
 
     pub fn window_frame(&self, window_frame: PyWindowFrame) -> PyExprFuncBuilder {
-        self.expr.clone().window_frame(window_frame.into()).into()
+        builder_from_expr(&self.expr)
+            .window_frame(window_frame.into())
+            .into()
     }
 
     #[pyo3(signature = (partition_by=None, window_frame=None, order_by=None, null_treatment=None))]
@@ -743,6 +747,57 @@ impl PyExpr {
     }
 }
 
+/// Start an [`ExprFuncBuilder`] that keeps the options already set on `expr`.
+///
+/// Upstream's `ExprFunctionExt` methods on an `Expr` start from an empty
+/// builder, so `build()` would reset every option not set again. The Python
+/// function wrappers already apply their keyword options, so chaining another
+/// builder method onto their result must not discard them.
+fn builder_from_expr(expr: &Expr) -> ExprFuncBuilder {
+    match expr {
+        Expr::AggregateFunction(agg) => {
+            let params = &agg.params;
+            let mut builder = expr.clone().null_treatment(params.null_treatment);
+            if !params.order_by.is_empty() {
+                builder = builder.order_by(params.order_by.clone());
+            }
+            if let Some(filter) = &params.filter {
+                builder = builder.filter(filter.as_ref().clone());
+            }
+            if params.distinct {
+                builder = builder.distinct();
+            }
+            builder
+        }
+        Expr::WindowFunction(window) => {
+            let params = &window.params;
+            let mut builder = expr.clone().null_treatment(params.null_treatment);
+            if !params.partition_by.is_empty() {
+                builder = builder.partition_by(params.partition_by.clone());
+            }
+            let has_order_by = !params.order_by.is_empty();
+            if has_order_by {
+                builder = builder.order_by(params.order_by.clone());
+            }
+            // A frame equal to the default `build()` derived from the order-by is
+            // left unset, so it is derived again from the final order-by.
+            if params.window_frame
+                != datafusion::logical_expr::WindowFrame::new(has_order_by.then_some(true))
+            {
+                builder = builder.window_frame(params.window_frame.clone());
+            }
+            if let Some(filter) = &params.filter {
+                builder = builder.filter(filter.as_ref().clone());
+            }
+            if params.distinct {
+                builder = builder.distinct();
+            }
+            builder
+        }
+        _ => expr.clone().null_treatment(None),
+    }
+}
+
 #[pyclass(
     from_py_object,
     frozen,

python/tests/test_expr.py

@@ -53,6 +53,8 @@
     TransactionEnd,
     TransactionStart,
     Values,
+    Window,
+    WindowFrame,
     coerce_to_expr,
     coerce_to_expr_list,
     coerce_to_expr_or_none,
@@ -1251,3 +1253,76 @@ def test_expr_to_bytes_no_ctx_default_codec() -> None:
     restored = Expr.from_bytes(blob, ctx=fresh)
 
     assert restored.canonical_name() == original.canonical_name()
+
+
+@pytest.fixture
+def builder_df():
+    ctx = SessionContext()
+    return ctx.from_pydict(
+        {"g": [1, 1, 1, 2], "s": ["y", "x", "z", "w"], "v": [3, 1, 2, 4]}
+    )
+
+
+@pytest.mark.parametrize(
+    ("build_expr", "expected"),
+    [
+        pytest.param(
+            lambda: (
+                functions.array_agg(col("s"), order_by="s")
+                .filter(col("v") > lit(1))
+                .build()
+            ),
+            ["w", "y", "z"],
+            id="order_by_kept_after_filter",
+        ),
+        pytest.param(
+            lambda: (
+                functions.array_agg(col("s"), filter=col("v") > lit(1))
+                .order_by(col("s").sort(ascending=False))
+                .build()
+            ),
+            ["z", "y", "w"],
+            id="filter_kept_after_order_by",
+        ),
+        pytest.param(
+            lambda: (
+                functions.string_agg(col("s"), ",", order_by="s").distinct().build()
+            ),
+            "w,x,y,z",
+            id="order_by_kept_after_distinct",
+        ),
+        pytest.param(
+            lambda: (
+                functions.first_value(col("s"), order_by="v")
+                .filter(col("v") > lit(1))
+                .build()
+            ),
+            "z",
+            id="first_value_order_by_kept_after_filter",
+        ),
+    ],
+)
+def test_aggregate_builder_keeps_existing_options(builder_df, build_expr, expected):
+    result = builder_df.aggregate([], [build_expr().alias("r")])
+    assert result.collect_column("r")[0].as_py() == expected
+
+
+def test_window_builder_keeps_existing_options(builder_df):
+    expr = functions.lead(col("v"), order_by="v").partition_by(col("g")).build()
+    result = builder_df.select(col("v"), expr.alias("r")).sort(col("v"))
+    assert result.collect_column("r").to_pylist() == [2, 3, None, None]
+
+
+def test_window_builder_keeps_explicit_frame(builder_df):
+    window = Window(order_by=col("v"), window_frame=WindowFrame("rows", 1, 0))
+    expr = functions.sum(col("v")).over(window).partition_by(col("g")).build()
+    result = builder_df.select(col("v"), expr.alias("r")).sort(col("v"))
+    assert result.collect_column("r").to_pylist() == [1, 3, 5, 4]
+
+
+def test_window_builder_rederives_default_frame(builder_df):
+    # No order_by means a whole-partition frame; adding one later must switch
+    # to the running frame rather than keep the whole-partition default.
+    expr = functions.sum(col("v")).over(Window()).order_by(col("v")).build()
+    result = builder_df.select(col("v"), expr.alias("r")).sort(col("v"))
+    assert result.collect_column("r").to_pylist() == [1, 3, 6, 10]