Merge branch 'main' into metadata-join-swap

Author: mkleen

datafusion/physical-optimizer/src/window_topn.rs

@@ -61,6 +61,7 @@ use datafusion_physical_expr::window::StandardWindowExpr;
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::sorts::partitioned_topk::{
     PartitionedTopKExec, WindowFnKind,
 };
@@ -140,24 +141,25 @@ impl WindowTopN {
         // Step 2: Extract limit from predicate (rn <= K, rn < K, etc.)
         let (col_idx, limit_n) = extract_window_limit(filter.predicate())?;
 
-        // Step 3: Walk through optional ProjectionExec to find BoundedWindowAggExec
+        // Step 3: Walk through optional ProjectionExec and RepartitionExec to find BoundedWindowAggExec
         let child = filter.input();
-        let (window_exec, proj_between) = find_window_below(child)?;
+        let (window_exec, intermediates) = find_window_below(child)?;
 
         // Step 4: Verify col_idx references a supported window function output column
-        let input_field_count = window_exec.input().schema().fields().len();
+        let window_exec_typed = window_exec.downcast_ref::<BoundedWindowAggExec>()?;
+        let sort_exec = window_exec_typed.input().downcast_ref::<SortExec>()?;
+        let input_field_count = window_exec_typed.input().schema().fields().len();
         if col_idx < input_field_count {
             return None; // Filter is on an input column, not a window column
         }
         let window_expr_idx = col_idx - input_field_count;
-        let window_exprs = window_exec.window_expr();
+        let window_exprs = window_exec_typed.window_expr();
         if window_expr_idx >= window_exprs.len() {
             return None;
         }
         let fn_kind = supported_window_fn(&window_exprs[window_expr_idx])?;
 
-        // Step 5: Verify child of window is SortExec
-        let sort_exec = window_exec.input().downcast_ref::<SortExec>()?;
+        // Step 5: child of window is SortExec (verified above)
         let sort_child = sort_exec.input();
 
         // Step 6: Determine partition_prefix_len from the window expression
@@ -190,28 +192,19 @@ impl WindowTopN {
         .ok()?;
 
         // Step 8: Rebuild window with new child
-        let new_window = Arc::clone(&child_as_arc(window_exec))
+        let mut result = window_exec
             .with_new_children(vec![Arc::new(partitioned_topk)])
             .ok()?;
 
-        // Step 9: If ProjectionExec was between Filter and Window, rebuild it
-        let result = match proj_between {
-            Some(proj) => Arc::clone(&child_as_arc(proj))
-                .with_new_children(vec![new_window])
-                .ok()?,
-            None => new_window,
-        };
+        // Step 9: Rebuild intermediate nodes (ProjectionExec/RepartitionExec)
+        for node in intermediates.into_iter().rev() {
+            result = node.with_new_children(vec![result]).ok()?;
+        }
 
         Some(result)
     }
 }
 
-/// Helper to get an `Arc<dyn ExecutionPlan>` from a reference.
-/// We need this because `with_new_children` takes `Arc<Self>`.
-fn child_as_arc<T: ExecutionPlan + Clone + 'static>(plan: &T) -> Arc<dyn ExecutionPlan> {
-    Arc::new(plan.clone())
-}
-
 impl PhysicalOptimizerRule for WindowTopN {
     fn optimize(
         &self,
@@ -336,29 +329,32 @@ fn supported_window_fn(
     }
 }
 
+type PlanAndIntermediates = (Arc<dyn ExecutionPlan>, Vec<Arc<dyn ExecutionPlan>>);
+
 /// Walk below a plan node looking for a [`BoundedWindowAggExec`].
 ///
-/// Handles two cases:
-/// - Direct child: `FilterExec → BoundedWindowAggExec`
-/// - With projection: `FilterExec → ProjectionExec → BoundedWindowAggExec`
+/// Handles sequences of `ProjectionExec` and `RepartitionExec`.
+/// This is safe because `PartitionedTopKExec` can be pushed below them:
+/// projections only provide aliases, and pushing the limit below repartitions
+/// is safe because the limit is computed per-partition.
 ///
-/// Returns the window exec and an optional `ProjectionExec` in between,
-/// or `None` if no `BoundedWindowAggExec` is found within one or two levels.
-fn find_window_below(
-    plan: &Arc<dyn ExecutionPlan>,
-) -> Option<(&BoundedWindowAggExec, Option<&ProjectionExec>)> {
-    // Direct child is BoundedWindowAggExec
-    if let Some(window) = plan.downcast_ref::<BoundedWindowAggExec>() {
-        return Some((window, None));
-    }
-
-    // Child is ProjectionExec with BoundedWindowAggExec below
-    if let Some(proj) = plan.downcast_ref::<ProjectionExec>() {
-        let proj_child = proj.input();
-        if let Some(window) = proj_child.downcast_ref::<BoundedWindowAggExec>() {
-            return Some((window, Some(proj)));
+/// Returns the window exec and a list of intermediate nodes to rebuild,
+/// or `None` if no `BoundedWindowAggExec` is found.
+fn find_window_below(plan: &Arc<dyn ExecutionPlan>) -> Option<PlanAndIntermediates> {
+    let mut current = Arc::clone(plan);
+    let mut intermediates = Vec::new();
+
+    loop {
+        if current.downcast_ref::<BoundedWindowAggExec>().is_some() {
+            return Some((current, intermediates));
+        } else if current.downcast_ref::<ProjectionExec>().is_some()
+            || current.downcast_ref::<RepartitionExec>().is_some()
+        {
+            let next = Arc::clone(current.children().first()?);
+            intermediates.push(current);
+            current = next;
+        } else {
+            return None;
         }
     }
-
-    None
 }

datafusion/physical-plan/src/joins/utils.rs

@@ -33,7 +33,8 @@ use crate::metrics::{
 };
 use crate::projection::{ProjectionExec, ProjectionExpr};
 use crate::{
-    ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, Statistics,
+    ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning,
+    RangePartitioning, Statistics,
 };
 // compatibility
 pub use super::join_filter::JoinFilter;
@@ -68,7 +69,7 @@ use datafusion_common::stats::Precision;
 use datafusion_common::utils::normalize_float_zero;
 use datafusion_common::{
     DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult,
-    not_impl_err, plan_err,
+    internal_datafusion_err, not_impl_err, plan_err,
 };
 use datafusion_expr::Operator;
 use datafusion_expr::interval_arithmetic::Interval;
@@ -146,9 +147,19 @@ pub fn adjust_right_output_partitioning(
             Partitioning::Hash(new_exprs, *size)
         }
         Partitioning::Range(range) => {
-            // Range partitioning optimizer propagation is tracked in
-            // https://github.com/apache/datafusion/issues/22395
-            Partitioning::UnknownPartitioning(range.partition_count())
+            let ordering = add_offset_to_physical_sort_exprs(
+                range.ordering().iter().cloned(),
+                left_columns_len as _,
+            )?;
+            let ordering = LexOrdering::new(ordering).ok_or_else(|| {
+                internal_datafusion_err!(
+                    "Offsetting range partitioning produced an empty ordering"
+                )
+            })?;
+            Partitioning::Range(RangePartitioning::new(
+                ordering,
+                range.split_points().to_vec(),
+            ))
         }
         result => result.clone(),
     };
@@ -2468,7 +2479,7 @@ mod tests {
     use arrow::datatypes::{DataType, Fields};
     use arrow::error::{ArrowError, Result as ArrowResult};
     use datafusion_common::stats::Precision::{Absent, Exact, Inexact};
-    use datafusion_common::{ScalarValue, arrow_datafusion_err, arrow_err};
+    use datafusion_common::{ScalarValue, SplitPoint, arrow_datafusion_err, arrow_err};
     use datafusion_physical_expr::PhysicalSortExpr;
 
     use rstest::rstest;
@@ -4131,6 +4142,53 @@ mod tests {
         assert_eq!(result.column_statistics[1].byte_size, Inexact(256));
     }
 
+    #[test]
+    fn test_adjust_right_output_partitioning_preserves_range() -> Result<()> {
+        let split_points = vec![
+            SplitPoint::new(vec![
+                ScalarValue::Int32(Some(10)),
+                ScalarValue::Int32(Some(100)),
+            ]),
+            SplitPoint::new(vec![
+                ScalarValue::Int32(Some(20)),
+                ScalarValue::Int32(Some(50)),
+            ]),
+        ];
+        let range = RangePartitioning::try_new(
+            LexOrdering::new([
+                PhysicalSortExpr::new(
+                    Arc::new(Column::new("a", 0)),
+                    SortOptions::new(false, true),
+                ),
+                PhysicalSortExpr::new(
+                    Arc::new(Column::new("b", 2)),
+                    SortOptions::new(true, false),
+                ),
+            ])
+            .unwrap(),
+            split_points.clone(),
+        )?;
+
+        let adjusted = adjust_right_output_partitioning(&Partitioning::Range(range), 3)?;
+        let expected = Partitioning::Range(RangePartitioning::new(
+            LexOrdering::new([
+                PhysicalSortExpr::new(
+                    Arc::new(Column::new("a", 3)),
+                    SortOptions::new(false, true),
+                ),
+                PhysicalSortExpr::new(
+                    Arc::new(Column::new("b", 5)),
+                    SortOptions::new(true, false),
+                ),
+            ])
+            .unwrap(),
+            split_points,
+        ));
+
+        assert_eq!(adjusted, expected);
+        Ok(())
+    }
+
     #[test]
     fn test_calculate_join_output_ordering() -> Result<()> {
         let left_ordering = LexOrdering::new(vec![

datafusion/physical-plan/src/sorts/cursor.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use std::cmp::Ordering;
+use std::fmt::Debug;
 use std::sync::Arc;
 
 use arrow::array::{
@@ -32,7 +33,7 @@ use datafusion_execution::memory_pool::MemoryReservation;
 ///
 /// This is a trait as there are several specialized implementations, such as for
 /// single columns or for normalized multi column keys ([`Rows`])
-pub trait CursorValues {
+pub trait CursorValues: Debug + Sync + Send {
     fn len(&self) -> usize;
 
     /// Returns true if `l[l_idx] == r[r_idx]`
@@ -298,6 +299,7 @@ impl<T: ArrowNativeTypeOp> CursorValues for PrimitiveValues<T> {
     }
 }
 
+#[derive(Debug)]
 pub struct ByteArrayValues<T: OffsetSizeTrait> {
     offsets: OffsetBuffer<T>,
     values: Buffer,