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 = + 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 = + 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::().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 = + 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::().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;