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

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

@@ -36,8 +36,9 @@ use crate::windows::{
 };
 use crate::{
     ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
-    ExecutionPlanProperties, InputOrderMode, PlanProperties, RecordBatchStream,
-    SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties,
+    ExecutionPlanProperties, InputDistributionRequirements, InputOrderMode,
+    PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr,
+    check_if_same_properties,
 };
 
 use arrow::compute::take_record_batch;
@@ -324,13 +325,16 @@ impl ExecutionPlan for BoundedWindowAggExec {
         self.input_distribution_requirements().into_per_child()
     }
 
-    fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements {
-        crate::InputDistributionRequirements::new(if self.partition_keys().is_empty() {
+    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
+        if self.partition_keys().is_empty() {
             debug!("No partition defined for BoundedWindowAggExec!!!");
-            vec![Distribution::SinglePartition]
+            InputDistributionRequirements::new(vec![Distribution::SinglePartition])
         } else {
-            vec![Distribution::KeyPartitioned(self.partition_keys().clone())]
-        })
+            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

@@ -32,8 +32,9 @@ use crate::windows::{
 };
 use crate::{
     ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
-    ExecutionPlanProperties, PhysicalExpr, PlanProperties, RecordBatchStream,
-    SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties,
+    ExecutionPlanProperties, InputDistributionRequirements, PhysicalExpr, PlanProperties,
+    RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr,
+    check_if_same_properties,
 };
 
 use arrow::array::ArrayRef;
@@ -233,12 +234,15 @@ impl ExecutionPlan for WindowAggExec {
         self.input_distribution_requirements().into_per_child()
     }
 
-    fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements {
-        crate::InputDistributionRequirements::new(if self.partition_keys().is_empty() {
-            vec![Distribution::SinglePartition]
+    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
+        if self.partition_keys().is_empty() {
+            InputDistributionRequirements::new(vec![Distribution::SinglePartition])
         } else {
-            vec![Distribution::KeyPartitioned(self.partition_keys())]
-        })
+            InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
+                self.partition_keys(),
+            )])
+            .allow_range_satisfaction_for_key_partitioning()
+        }
     }
 
     fn with_new_children(