From 88e7f149a2a24d1e8de39c915b68fcbb7b7dbef4 Mon Sep 17 00:00:00 2001
From: Mingfa Lim <346975989+XiaoMingFa@users.noreply.github.com>
Date: Fri, 23 Oct 2026 20:27:34 +0200
Subject: [PATCH] fix: preserve can_repartition in
 WindowAggExec::replace_children (#26626)

## Which issue does this PR close?

- Closes #26626.

## Rationale for this change

When `datafusion.execution.target_partitions = 1`, the physical planner constructs `WindowAggExec` with `can_repartition = false` so that the window stage requires `SinglePartition` distribution and inserts a `CoalescePartitionsExec`.

During physical optimization, replacing children via `replace_children` under `ChildrenPropertiesMode::Recompute` rebuilt `WindowAggExec` with hardcoded `can_repartition = true` instead of propagating `self.can_repartition`. This flipped the distribution requirement from `SinglePartition` to `KeyPartitioned`, stripping the required coalesce and triggering a `SanityCheckPlan` distribution error.

`BoundedWindowAggExec` already correctly propagates `self.can_repartition` during `replace_children`. This change aligns `WindowAggExec` with `BoundedWindowAggExec`.

## What changes are included in this PR?

- Pass `self.can_repartition` instead of `true` in `WindowAggExec::replace_children` when `ChildrenPropertiesMode::Recompute` is used.
- Added unit test `test_window_agg_replace_children_preserves_can_repartition` in `window_agg_exec.rs` verifying preservation for both `false` and `true`.
- Added sqllogictest regression test in `window.slt` verifying query execution with `target_partitions = 1`.

## Are these changes tested?

Yes:
- Local unit test `test_window_agg_replace_children_preserves_can_repartition` passed.
- Integration test added to `window.slt`.
---
 .../src/windows/window_agg_exec.rs            | 58 ++++++++++++++++++-
 datafusion/sqllogictest/test_files/window.slt | 15 +++++
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs
index 425af3a3c057d..c28cb2f1a674a 100644
--- a/datafusion/physical-plan/src/windows/window_agg_exec.rs
+++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs
@@ -277,7 +277,7 @@ impl ExecutionPlan for WindowAggExec {
             ChildrenPropertiesMode::Recompute => Ok(Arc::new(WindowAggExec::try_new(
                 self.window_expr.clone(),
                 children.swap_remove(0),
-                true,
+                self.can_repartition,
             )?)),
         }
     }
@@ -679,4 +679,60 @@ mod tests {
         ));
         Ok(())
     }
+
+    #[test]
+    fn test_window_agg_replace_children_preserves_can_repartition() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)]));
+        let input: Arc<dyn ExecutionPlan> =
+            Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?);
+        let args = vec![crate::expressions::col("a", &schema)?];
+        let window_expr = create_window_expr(
+            &WindowFunctionDefinition::AggregateUDF(count_udaf()),
+            "count(a)".to_string(),
+            &args,
+            &[],
+            &[],
+            Arc::new(WindowFrame::new_bounds(
+                WindowFrameUnits::Rows,
+                WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
+                WindowFrameBound::CurrentRow,
+            )),
+            Arc::clone(&schema),
+            false,
+            false,
+            None,
+        )?;
+
+        // can_repartition = false should be preserved after replace_children with Recompute
+        let window = Arc::new(WindowAggExec::try_new(
+            vec![window_expr.clone()],
+            Arc::clone(&input),
+            false,
+        )?);
+        assert!(!window.can_repartition);
+
+        let new_child: Arc<dyn ExecutionPlan> =
+            Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?);
+        let replaced = window.clone().replace_children(
+            vec![new_child],
+            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+        )?;
+        let replaced_window = replaced.downcast_ref::<WindowAggExec>().unwrap();
+        assert_eq!(replaced_window.can_repartition, false);
+
+        // can_repartition = true should also be preserved
+        let window_true =
+            Arc::new(WindowAggExec::try_new(vec![window_expr], input, true)?);
+        assert!(window_true.can_repartition);
+        let new_child2: Arc<dyn ExecutionPlan> =
+            Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?);
+        let replaced_true = window_true.clone().replace_children(
+            vec![new_child2],
+            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+        )?;
+        let replaced_window_true = replaced_true.downcast_ref::<WindowAggExec>().unwrap();
+        assert_eq!(replaced_window_true.can_repartition, true);
+
+        Ok(())
+    }
 }
diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt
index 864a6e56bf574..7471d767006db 100644
--- a/datafusion/sqllogictest/test_files/window.slt
+++ b/datafusion/sqllogictest/test_files/window.slt
@@ -7747,3 +7747,18 @@ DROP TABLE range_null_peers;
 
 statement ok
 DROP TABLE range_nan_peers;
+
+# Test issue #26626: WindowAggExec replaces children while preserving can_repartition under target_partitions = 1
+statement ok
+SET datafusion.execution.target_partitions = 1;
+
+query II
+SELECT count(*) OVER (PARTITION BY a + NULL) AS x,
+       count(*) OVER (PARTITION BY a) AS y
+FROM (SELECT 1 AS a UNION ALL SELECT 2);
+----
+2 1
+2 1
+
+statement ok
+RESET datafusion.execution.target_partitions;