feat: allow Partitioning::Range to satisfy window Distribution::KeyPartitioned requirements

Windows got a hash repartition even when the input was already range
partitioned on the PARTITION BY keys. Opt WindowAggExec and
BoundedWindowAggExec into range satisfaction for KeyPartitioned, the
same way AggregateExec does. Incompatible keys still fall back to hash
repartitioning, and windows without PARTITION BY keep requiring a
single partition.

Covered by enforce_distribution and sanity_checker tests plus a window
section in range_partitioning.slt.

Closes #23289

Author: mithuncy

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -20,8 +20,8 @@ use std::ops::Deref;
 use std::sync::Arc;
 
 use crate::physical_optimizer::test_utils::{
-    check_integrity, coalesce_partitions_exec, parquet_exec_with_sort,
-    parquet_exec_with_stats, repartition_exec, schema, sort_exec,
+    bounded_window_exec_with_can_repartition, check_integrity, coalesce_partitions_exec,
+    parquet_exec_with_sort, parquet_exec_with_stats, repartition_exec, schema, sort_exec,
     sort_exec_with_preserve_partitioning, sort_merge_join_exec,
     sort_preserving_merge_exec, union_exec,
 };
@@ -870,6 +870,69 @@ fn range_inner_hash_join_rehashes_incompatible_range_partitioning() -> Result<()
     Ok(())
 }
 
+#[test]
+fn range_window_reuses_range_partitioning() -> Result<()> {
+    let input = parquet_exec_with_output_partitioning(range_partitioning(
+        "a",
+        [10, 20, 30],
+        SortOptions::default(),
+    )?);
+    let window = bounded_window_exec_with_can_repartition(
+        "a",
+        vec![],
+        &[col("a", &schema())?],
+        input,
+        true,
+    );
+
+    let plan = TestConfig::default()
+        .with_query_execution_partitions(4)
+        .to_plan(window, &DISTRIB_DISTRIB_SORT);
+
+    assert_plan!(
+        plan,
+        @r#"
+    BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+      SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true]
+        DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet
+    "#
+    );
+
+    Ok(())
+}
+
+#[test]
+fn range_window_rehashes_incompatible_range_partitioning() -> Result<()> {
+    let input = parquet_exec_with_output_partitioning(range_partitioning(
+        "a",
+        [10, 20, 30],
+        SortOptions::default(),
+    )?);
+    let window = bounded_window_exec_with_can_repartition(
+        "b",
+        vec![],
+        &[col("b", &schema())?],
+        input,
+        true,
+    );
+
+    let plan = TestConfig::default()
+        .with_query_execution_partitions(4)
+        .to_plan(window, &DISTRIB_DISTRIB_SORT);
+
+    assert_plan!(
+        plan,
+        @r#"
+    BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+      SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[true]
+        RepartitionExec: partitioning=Hash([b@1], 4), input_partitions=4
+          DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet
+    "#
+    );
+
+    Ok(())
+}
+
 #[test]
 fn multi_hash_joins() -> Result<()> {
     let left = parquet_exec();

datafusion/core/tests/physical_optimizer/sanity_checker.rs

@@ -19,9 +19,10 @@ use insta::assert_snapshot;
 use std::sync::Arc;
 
 use crate::physical_optimizer::test_utils::{
-    bounded_window_exec, global_limit_exec, hash_join_exec, local_limit_exec,
-    memory_exec, projection_exec, repartition_exec, sort_exec, sort_expr,
-    sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, union_exec,
+    bounded_window_exec, bounded_window_exec_with_can_repartition, global_limit_exec,
+    hash_join_exec, local_limit_exec, memory_exec, projection_exec, repartition_exec,
+    sort_exec, sort_exec_with_preserve_partitioning, sort_expr, sort_expr_options,
+    sort_merge_join_exec, sort_preserving_merge_exec, union_exec,
 };
 
 use arrow::compute::SortOptions;
@@ -501,6 +502,76 @@ async fn test_bounded_window_agg_no_sort_requirement() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+/// Tests that a window over a compatible range-partitioned input satisfies
+/// the window's key distribution requirement without a hash repartition.
+async fn test_bounded_window_agg_range_partitioning() -> Result<()> {
+    let schema = create_test_schema2();
+    let source = range_partitioned_exec(&schema, "a", [10, 20, 30])?;
+    let ordering: LexOrdering = [sort_expr_options(
+        "a",
+        &schema,
+        SortOptions {
+            descending: false,
+            nulls_first: false,
+        },
+    )]
+    .into();
+    let partition_by = vec![col("a", &schema)?];
+    let sort = sort_exec_with_preserve_partitioning(ordering, source);
+    let bw =
+        bounded_window_exec_with_can_repartition("a", vec![], &partition_by, sort, true);
+    let plan_str = displayable(bw.as_ref()).indent(true).to_string();
+    let actual = plan_str.trim();
+    assert_snapshot!(
+        actual,
+        @r#"
+    BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+      SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true]
+        RepartitionExec: partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), input_partitions=1
+          DataSourceExec: partitions=1, partition_sizes=[0]
+    "#
+    );
+    assert_sanity_check(&bw, true);
+    Ok(())
+}
+
+#[tokio::test]
+/// Tests that a window over an incompatible range-partitioned input fails
+/// the window's key distribution requirement.
+async fn test_bounded_window_agg_incompatible_range_partitioning() -> Result<()> {
+    let schema = create_test_schema2();
+    let source = range_partitioned_exec(&schema, "a", [10, 20, 30])?;
+    let ordering: LexOrdering = [sort_expr_options(
+        "b",
+        &schema,
+        SortOptions {
+            descending: false,
+            nulls_first: false,
+        },
+    )]
+    .into();
+    let partition_by = vec![col("b", &schema)?];
+    let sort = sort_exec_with_preserve_partitioning(ordering, source);
+    let bw =
+        bounded_window_exec_with_can_repartition("b", vec![], &partition_by, sort, true);
+    let plan_str = displayable(bw.as_ref()).indent(true).to_string();
+    let actual = plan_str.trim();
+    assert_snapshot!(
+        actual,
+        @r#"
+    BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+      SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[true]
+        RepartitionExec: partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), input_partitions=1
+          DataSourceExec: partitions=1, partition_sizes=[0]
+    "#
+    );
+    // Range([a]) does not colocate `b` values, so the window's key
+    // distribution requirement is not satisfied.
+    assert_sanity_check(&bw, false);
+    Ok(())
+}
+
 #[tokio::test]
 /// A valid when a single partition requirement
 /// is satisfied.

datafusion/core/tests/physical_optimizer/test_utils.rs

@@ -274,6 +274,22 @@ pub fn bounded_window_exec_with_partition(
     sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
     partition_by: &[Arc<dyn PhysicalExpr>],
     input: Arc<dyn ExecutionPlan>,
+) -> Arc<dyn ExecutionPlan> {
+    bounded_window_exec_with_can_repartition(
+        col_name,
+        sort_exprs,
+        partition_by,
+        input,
+        false,
+    )
+}
+
+pub fn bounded_window_exec_with_can_repartition(
+    col_name: &str,
+    sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
+    partition_by: &[Arc<dyn PhysicalExpr>],
+    input: Arc<dyn ExecutionPlan>,
+    can_repartition: bool,
 ) -> Arc<dyn ExecutionPlan> {
     let sort_exprs = sort_exprs.into_iter().collect::<Vec<_>>();
     let schema = input.schema();
@@ -296,7 +312,7 @@ pub fn bounded_window_exec_with_partition(
             vec![window_expr],
             Arc::clone(&input),
             InputOrderMode::Sorted,
-            false,
+            can_repartition,
         )
         .unwrap(),
     )

datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs

@@ -336,12 +336,15 @@ impl ExecutionPlan for BoundedWindowAggExec {
     }
 
     fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements {
-        crate::InputDistributionRequirements::new(if self.partition_keys().is_empty() {
+        if self.partition_keys().is_empty() {
             debug!("No partition defined for BoundedWindowAggExec!!!");
-            vec![Distribution::SinglePartition]
+            crate::InputDistributionRequirements::new(vec![Distribution::SinglePartition])
         } else {
-            vec![Distribution::KeyPartitioned(self.partition_keys().clone())]
-        })
+            crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
+                self.partition_keys(),
+            )])
+            .allow_range_satisfaction_for_key_partitioning()
+        }
     }
 
     fn maintains_input_order(&self) -> Vec<bool> {

datafusion/physical-plan/src/windows/window_agg_exec.rs

@@ -245,11 +245,14 @@ impl ExecutionPlan for WindowAggExec {
     }
 
     fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements {
-        crate::InputDistributionRequirements::new(if self.partition_keys().is_empty() {
-            vec![Distribution::SinglePartition]
+        if self.partition_keys().is_empty() {
+            crate::InputDistributionRequirements::new(vec![Distribution::SinglePartition])
         } else {
-            vec![Distribution::KeyPartitioned(self.partition_keys())]
-        })
+            crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
+                self.partition_keys(),
+            )])
+            .allow_range_satisfaction_for_key_partitioning()
+        }
     }
 
     fn with_new_children(