From 2672c49d465495debaa15e57d2dc9460b382faf2 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Wed, 7 Oct 2026 18:03:09 +0200 Subject: [PATCH 1/3] fix: keep required sorts above custom execution plans Keep sorts above operators without explicit sort-pushdown handling. Preserve known built-in input mappings and make custom-node support opt-in through the existing try_pushdown_sort API. Cover correct query results and sort placement for a value-transforming custom projection, built-in input mappings, and explicit custom support. Closes #26656. --- .../physical_optimizer/enforce_sorting.rs | 20 +- .../tests/physical_optimizer/pushdown_sort.rs | 53 ++-- .../tests/physical_optimizer/test_utils.rs | 4 + .../enforce_sorting/sort_pushdown.rs | 226 +++++------------- .../physical-plan/src/execution_plan.rs | 16 +- datafusion/sqllogictest/src/test_context.rs | 7 + .../src/test_context/custom_sort_pushdown.rs | 139 +++++++++++ .../sqllogictest/test_files/sort_pushdown.slt | 31 +++ .../library-user-guide/upgrading/56.0.0.md | 22 ++ 9 files changed, 316 insertions(+), 202 deletions(-) create mode 100644 datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 084ae59a76c23..49989a50b5f84 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2604,9 +2604,9 @@ async fn test_push_with_required_input_ordering_prohibited() -> Result<()> { Ok(()) } -// test when the required input ordering is satisfied so could push through +// Compatible input requirements do not opt a custom node into sort pushdown. #[tokio::test] -async fn test_push_with_required_input_ordering_allowed() -> Result<()> { +async fn test_custom_node_keeps_sort_with_compatible_input_ordering() -> Result<()> { let schema = create_test_schema3()?; let ordering_a: LexOrdering = [sort_expr("a", &schema)].into(); let ordering_ab = [sort_expr("a", &schema), sort_expr("b", &schema)].into(); @@ -2618,29 +2618,15 @@ async fn test_push_with_required_input_ordering_allowed() -> Result<()> { .into_arc(); let plan = sort_exec(ordering_ab, plan); - /* - let expected_input = [ - "SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[false]", // <-- can push this down (as it is compatible with the required input ordering) - " RequiredInputOrderingExec", // <-- this requires input sorted by a, and preserves the input order - " SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - " DataSourceExec: partitions=1, partition_sizes=[0]", - ]; - */ let test = EnforceSortingTest::new(plan.clone()).with_repartition_sorts(true); assert_snapshot!(test.run(), @r" - Input Plan: + Input / Optimized Plan: SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[false] RequiredInputOrderingExec SortExec: expr=[a@0 ASC], preserve_partitioning=[false] DataSourceExec: partitions=1, partition_sizes=[0] - - Optimized Plan: - RequiredInputOrderingExec - SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[false] - DataSourceExec: partitions=1, partition_sizes=[0] "); - // Should be able to push down Ok(()) } diff --git a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs index 6cf05ecfd8fc4..68d3418345d64 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs @@ -34,8 +34,10 @@ use datafusion_physical_expr::expressions; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_optimizer::PhysicalOptimizerRule; +use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; use datafusion_physical_optimizer::pushdown_sort::PushdownSort; -use datafusion_physical_plan::collect; +use datafusion_physical_plan::sorts::sort::SortExec; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties, collect}; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ @@ -1111,9 +1113,7 @@ fn test_sort_pushdown_with_test_scan_arbitrary_ordering() { // ============================================================================ #[test] -fn test_sort_pushdown_exact_no_fetch_no_limit() { - // When a source returns Exact (without fetch), the SortExec should be - // removed entirely with no GlobalLimitExec wrapper. +fn test_custom_plan_can_opt_into_sort_pushdown() -> Result<()> { let schema = schema(); let a = sort_expr("a", &schema); let b = sort_expr("b", &schema); @@ -1121,20 +1121,39 @@ fn test_sort_pushdown_exact_no_fetch_no_limit() { Arc::new(TestScan::new(schema.clone(), vec![]).with_exact_pushdown(true)); let ordering = LexOrdering::new(vec![a, b.reverse()]).unwrap(); - let plan = sort_exec(ordering, source); + let plan = sort_exec(ordering.clone(), source); + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 1; + let plan = EnsureRequirements::new().optimize(plan, &config)?; + assert!(plan.is::()); - insta::assert_snapshot!( - OptimizationTest::new(plan, PushdownSort::new(), true), - @r" - OptimizationTest: - input: - - SortExec: expr=[a@0 ASC, b@1 DESC NULLS LAST], preserve_partitioning=[false] - - TestScan - output: - Ok: - - TestScan: requested_ordering=[a@0 ASC, b@1 DESC NULLS LAST] - " - ); + let optimized = PushdownSort::new().optimize(plan, &config)?; + assert!(optimized.is::()); + assert_eq!(optimized.output_ordering(), Some(&ordering)); + Ok(()) +} + +#[test] +fn test_custom_plan_with_fetch_keeps_outer_sort() -> Result<()> { + let schema = schema(); + let a = sort_expr("a", &schema); + let source = TestScan::new(schema.clone(), vec![[a.clone()].into()]) + .with_supports_fetch(true) + .with_fetch(Some(10)) + .unwrap(); + let ordering = LexOrdering::new(vec![a, sort_expr("b", &schema)]).unwrap(); + let plan = sort_exec(ordering.clone(), source); + + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 1; + let optimized = EnsureRequirements::new().optimize(plan, &config)?; + let sort = optimized + .downcast_ref::() + .expect("a fetched custom plan must retain the requested outer sort"); + assert_eq!(sort.expr(), &ordering); + assert!(sort.input().is::()); + assert_eq!(sort.input().fetch(), Some(10)); + Ok(()) } #[test] diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 3733dba39c97f..9be1c001a0710 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -1074,6 +1074,10 @@ impl ExecutionPlan for TestScan { Ok(Arc::new(Statistics::new_unknown(&self.schema))) } + fn supports_limit_pushdown(&self) -> bool { + self.supports_fetch + } + fn with_fetch(&self, fetch: Option) -> Option> { if self.supports_fetch { let mut new_scan = self.clone(); diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 2aeddd23aaffe..510e8df83a154 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -18,6 +18,7 @@ use std::fmt::Debug; use std::sync::Arc; +use crate::output_requirements::OutputRequirementExec; use crate::utils::{ add_sort_above_with_distribution, is_sort, is_sort_preserving_merge, is_union, is_window, @@ -37,15 +38,16 @@ use datafusion_physical_expr_common::sort_expr::{ PhysicalSortRequirement, }; use datafusion_physical_plan::aggregates::AggregateExec; +use datafusion_physical_plan::buffer::BufferExec; +use datafusion_physical_plan::coop::CooperativeExec; use datafusion_physical_plan::execution_plan::CardinalityEffect; -use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::{ ColumnIndex, calculate_join_output_ordering, }; use datafusion_physical_plan::joins::{HashJoinExec, SortMergeJoinExec}; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; -use datafusion_physical_plan::repartition::RepartitionExec; +use datafusion_physical_plan::scalar_subquery::ScalarSubqueryExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::tree_node::PlanContext; @@ -447,6 +449,10 @@ fn pushdown_sorts_helper( /// Calculate the pushdown ordering requirements for children. /// If sort cannot be pushed down, return None. +/// +/// Pushdown requires a known mapping from output sort keys to child expressions, +/// provided by a specialized handler or a known passthrough case. Preserving +/// row sequence alone does not establish that mapping. fn pushdown_requirement_to_children( plan: &Arc, parent_required: OrderingRequirements, @@ -479,7 +485,6 @@ fn pushdown_requirement_to_children( return Ok(None); } - let maintains_input_order = plan.maintains_input_order(); if is_window(plan) { let mut required_input_ordering = plan.required_input_ordering(); let maybe_child_requirement = required_input_ordering.swap_remove(0); @@ -548,49 +553,18 @@ fn pushdown_requirement_to_children( } else { Ok(None) } - } else if plan.fetch().is_some() - && plan.supports_limit_pushdown() - && plan - .maintains_input_order() - .into_iter() - .all(|maintain| maintain) - { - // Push down through operator with fetch when: - // - requirement is aligned with output ordering - // - it preserves ordering during execution - // - // A `ProjectionExec` reports a `fetch()` forwarded from its input and - // can renumber/reorder columns, so the requirement (expressed in the - // projection's output schema) must be remapped into the child schema - // before being pushed down — forwarding it unchanged would let a key - // such as `score@1` (valid in the output schema) refer to a different - // column in the child schema, producing a `SortExec` whose key points - // at the wrong column ("does not satisfy order requirements ... - // Child-0 order: []"). If a required column maps to a computed - // (non-`Column`) projection expression it cannot be expressed below the - // projection, so the sort is kept above it. - let child_required = - if let Some(projection) = plan.downcast_ref::() { - match remap_requirement_through_projection(projection, &parent_required) { - Some(remapped) => remapped, - None => return Ok(None), - } - } else { - parent_required.clone() - }; - - // Pushing a sort below an unordered fetch changes which rows it keeps. - let Some(ordering) = plan.properties().output_ordering() else { + } else if plan.is::() && plan.fetch().is_some() { + // Refining a fetched merge must preserve the ordering that selected its rows. + let Some(ordering) = plan.output_ordering() else { return Ok(None); }; - if plan.properties().eq_properties.requirements_compatible( - parent_required.first().clone(), - ordering.clone().into(), - ) { - Ok(Some(vec![Some(child_required)])) - } else { - Ok(None) - } + Ok(plan + .equivalence_properties() + .requirements_compatible( + parent_required.first().clone(), + ordering.clone().into(), + ) + .then(|| vec![Some(parent_required)])) } else if is_union(plan) { // `UnionExec` does not have real sort requirements for its input, we // just propagate the sort requirements down: @@ -634,15 +608,7 @@ fn pushdown_requirement_to_children( handle_aggregate_pushdown(aggregate_exec, parent_required) } else if let Some(projection_exec) = plan.downcast_ref::() { handle_projection_pushdown(projection_exec, &parent_required) - } else if maintains_input_order.is_empty() - || !maintains_input_order.iter().any(|o| *o) - || plan.is::() - || plan.is::() - || pushdown_would_violate_requirements(&parent_required, plan.as_ref()) - { - // If the current plan is a leaf node or can not maintain any of the input ordering, can not pushed down requirements. - // For RepartitionExec, we always choose to not push down the sort requirements even the RepartitionExec(input_partition=1) could maintain input ordering. - // Pushing down is not beneficial + } else if pushdown_would_violate_requirements(&parent_required, plan.as_ref()) { Ok(None) } else if is_sort_preserving_merge(plan) { let new_ordering = LexOrdering::from(parent_required.first().clone()); @@ -661,8 +627,21 @@ fn pushdown_requirement_to_children( } } else if let Some(hash_join) = plan.downcast_ref::() { handle_hash_join(hash_join, parent_required) + } else if plan.is::() + || plan.is::() + || plan.is::() + || plan.is::() + { + // These operators expose the first child's columns without transformation. + // Scalar subquery children compute side inputs and do not contribute rows. + let mut child_requirements = vec![None; plan.children().len()]; + child_requirements[0] = Some(parent_required); + Ok(Some(child_requirements)) } else { - handle_custom_pushdown(plan, parent_required, &maintains_input_order) + // Preserving row positions does not identify the input expressions needed + // to order transformed output values. Custom nodes opt in through + // ExecutionPlan::try_pushdown_sort, handled by the PushdownSort rule. + Ok(None) } } @@ -979,118 +958,6 @@ fn expr_source_side( } } -/// Handles the custom pushdown of parent-required sorting requirements down to -/// the child execution plans, considering whether the input order is maintained. -/// -/// # Arguments -/// -/// * `plan` - A reference to an `ExecutionPlan` for which the pushdown will be applied. -/// * `parent_required` - The sorting requirements expected by the parent node. -/// * `maintains_input_order` - A vector of booleans indicating whether each child -/// maintains the input order. -/// -/// # Returns -/// -/// Returns `Ok(Some(Vec>))` if the sorting requirements can be -/// pushed down, `Ok(None)` if not. On error, returns a `Result::Err`. -fn handle_custom_pushdown( - plan: &Arc, - parent_required: OrderingRequirements, - maintains_input_order: &[bool], -) -> Result>>> { - let plan_children = plan.children(); - - // If the plan has no children, return early: - if plan_children.is_empty() { - return Ok(None); - } - - // Collect all unique column indices used in the parent-required sorting - // expression: - let requirement = parent_required.into_single(); - let all_indices: HashSet = requirement - .iter() - .flat_map(|order| { - collect_columns(&order.expr) - .iter() - .map(|col| col.index()) - .collect::>() - }) - .collect(); - - // Get the number of fields in each child's schema: - let children_schema_lengths: Vec = plan_children - .iter() - .map(|c| c.schema().fields().len()) - .collect(); - - // Find the index of the order-maintaining child: - let Some(maintained_child_idx) = maintains_input_order - .iter() - .enumerate() - .find(|(_, m)| **m) - .map(|pair| pair.0) - else { - return Ok(None); - }; - - // Check if all required columns come from the order-maintaining child: - let start_idx = children_schema_lengths[..maintained_child_idx] - .iter() - .sum::(); - let end_idx = start_idx + children_schema_lengths[maintained_child_idx]; - let all_from_maintained_child = - all_indices.iter().all(|i| i >= &start_idx && i < &end_idx); - - // If all columns are from the maintained child, update the parent requirements: - if all_from_maintained_child { - let sub_offset = children_schema_lengths - .iter() - .take(maintained_child_idx) - .sum::(); - // Transform the parent-required expression for the child schema by - // adjusting columns: - let updated_parent_req = requirement - .into_iter() - .map(|req| { - let child_schema = plan_children[maintained_child_idx].schema(); - let updated_columns = req - .expr - .transform_up(|expr| { - if let Some(col) = expr.downcast_ref::() { - let new_index = col.index() - sub_offset; - Ok(Transformed::yes(Arc::new(Column::new( - child_schema.field(new_index).name(), - new_index, - )))) - } else { - Ok(Transformed::no(expr)) - } - })? - .data; - Ok(PhysicalSortRequirement::new(updated_columns, req.options)) - }) - .collect::>>()?; - - // Prepare the result, populating with the updated requirements for children that maintain order - let result = maintains_input_order - .iter() - .map(|&maintains_order| { - if maintains_order { - LexRequirement::new(updated_parent_req.clone()) - .map(OrderingRequirements::new) - } else { - None - } - }) - .collect(); - - Ok(Some(result)) - } else { - Ok(None) - } -} - // For hash join we only maintain the input order for the right child // for join type: Inner, Right, RightSemi, RightAnti, RightMark fn handle_hash_join( @@ -1247,10 +1114,14 @@ mod tests { use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_expr::Operator; + use datafusion_expr::physical_planning_context::{ + ScalarSubqueryResults, SubqueryIndex, + }; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::{BinaryExpr, col}; use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::limit::GlobalLimitExec; + use datafusion_physical_plan::scalar_subquery::ScalarSubqueryLink; const DESC: SortOptions = SortOptions { descending: true, @@ -1366,6 +1237,29 @@ mod tests { LexRequirement::new(reqs).unwrap() } + #[test] + fn scalar_subquery_sort_applies_only_to_the_main_input() -> Result<()> { + let schema = child_schema(); + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let requirement = OrderingRequirements::new_soft(lex([ + req("c", &schema, DESC), + req("a", &schema, ASC), + ])); + let plan: Arc = Arc::new(ScalarSubqueryExec::new( + Arc::clone(&input), + vec![ScalarSubqueryLink { + plan: Arc::clone(&input), + index: SubqueryIndex::new(0), + }], + ScalarSubqueryResults::new(1), + )); + assert_eq!( + pushdown_requirement_to_children(&plan, requirement.clone(), None)?, + Some(vec![Some(requirement), None]) + ); + Ok(()) + } + #[test] fn remap_single_hard_requirement_through_reordering_projection() { let projection = reordering_projection(); diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index ee15653f7df59..e9cc3244f2ad6 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -233,6 +233,10 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// DataFusion uses this metadata to apply certain optimizations /// such as automatically repartitioning correctly. /// + /// This describes row sequence, not how input expressions map to outputs. + /// Custom nodes that support sort pushdown should implement + /// [`Self::try_pushdown_sort`]; returning `true` here alone does not opt in. + /// /// The default implementation returns `false` /// /// WARNING: if you override this default, you *MUST* ensure that @@ -983,8 +987,16 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// - Return [`SortOrderPushdownResult::Unsupported`] if the node cannot optimize /// for the ordering /// - /// For transparent nodes (that preserve ordering), implement this to delegate to - /// children and wrap the result with a new instance of this node. + /// Custom nodes opt into sort pushdown by implementing this method. Requested + /// expressions refer to this node's output schema; map them to input + /// expressions before delegating to a child. Decline pushdown when a correct + /// mapping is unavailable. + /// + /// Nodes that preserve the full input row sequence and column values can + /// delegate to their child and wrap the returned plan in a new instance of + /// the node. Delegation relies on the child supporting this method. If the + /// child returns [`SortOrderPushdownResult::Unsupported`], delegation alone + /// does not insert or move sorts to satisfy the requested ordering. /// /// Default implementation returns `Unsupported`. fn try_pushdown_sort( diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 0c8c2c5add363..dd7b9fb1e2222 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -73,6 +73,7 @@ use log::info; use sqlparser::ast; use tempfile::TempDir; +mod custom_sort_pushdown; mod range_partitioning; /// Context for running tests @@ -182,6 +183,12 @@ impl TestContext { let file_name = relative_path.file_name().unwrap().to_str().unwrap(); match file_name { + "sort_pushdown.slt" => { + custom_sort_pushdown::register_custom_sort_pushdown( + test_ctx.session_ctx(), + ) + .expect("register custom projection table"); + } "parquet_missing_bounds.slt" => { register_parquet_missing_bounds(&mut test_ctx).await; } diff --git a/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs b/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs new file mode 100644 index 0000000000000..a89979c4c614c --- /dev/null +++ b/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs @@ -0,0 +1,139 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{self, Formatter}; +use std::sync::Arc; + +use arrow::array::record_batch; +use arrow::datatypes::{self as arrow_schema, SchemaRef}; +use async_trait::async_trait; +use datafusion::catalog::{Session, TableProvider}; +use datafusion::common::Result; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::{Expr, TableType}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::{col, negative}; +use datafusion::physical_expr::projection::ProjectionExprs; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::test::TestMemoryExec; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, +}; +use datafusion::prelude::SessionContext; + +/// Register a fixed custom operator that negates its input column. +pub(super) fn register_custom_sort_pushdown(ctx: &SessionContext) -> Result<()> { + let batch = record_batch!(("a", Int32, [2, 1, 3]))?; + let schema = batch.schema(); + let input = TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + let projection = ProjectionExec::try_new( + [(negative(col("a", &schema)?, &schema)?, "a".to_string())], + input, + )?; + ctx.register_table( + "custom_negated_projection", + Arc::new(CustomProjection(projection)), + )?; + Ok(()) +} + +/// Negates values while preserving row positions. Keeping a distinct plan type +/// exercises the generic custom-plan sort pushdown path. +#[derive(Debug, Clone)] +struct CustomProjection(ProjectionExec); + +#[async_trait] +impl TableProvider for CustomProjection { + fn schema(&self) -> SchemaRef { + self.0.schema() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _: &dyn Session, + projection: Option<&[usize]>, + _: &[Expr], + _: Option, + ) -> Result> { + let plan: Arc = Arc::new(self.clone()); + if let Some(projection) = projection { + let expressions = ProjectionExprs::from_indices(projection, &self.0.schema()); + Ok(Arc::new(ProjectionExec::try_new( + expressions.iter().cloned(), + plan, + )?)) + } else { + Ok(plan) + } + } +} + +impl DisplayAs for CustomProjection { + fn fmt_as(&self, display_type: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "Custom")?; + self.0.fmt_as(display_type, f) + } +} + +impl ExecutionPlan for CustomProjection { + fn name(&self) -> &str { + "CustomProjectionExec" + } + + fn properties(&self) -> &Arc { + self.0.properties() + } + + fn children(&self) -> Vec<&Arc> { + self.0.children() + } + + // Row positions are preserved; projected values have the opposite ordering. + fn maintains_input_order(&self) -> Vec { + self.0.maintains_input_order() + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self(ProjectionExec::try_new( + self.0.expr().to_vec(), + Arc::clone(&children[0]), + )?))) + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.0.apply_expressions(f) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.0.execute(partition, context) + } +} diff --git a/datafusion/sqllogictest/test_files/sort_pushdown.slt b/datafusion/sqllogictest/test_files/sort_pushdown.slt index da583f127c2b7..30deeaf9aee82 100644 --- a/datafusion/sqllogictest/test_files/sort_pushdown.slt +++ b/datafusion/sqllogictest/test_files/sort_pushdown.slt @@ -15,6 +15,37 @@ # specific language governing permissions and limitations # under the License. +# The custom operator evaluates -a AS a, preserving row positions while +# reversing the ordering of values. Its input is [2, 1, 3]. +statement ok +SET datafusion.execution.target_partitions = 1; + +query I +SELECT a FROM custom_negated_projection ORDER BY a; +---- +-3 +-2 +-1 + +# The sort must stay above the custom projection. +statement ok +SET datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT a FROM custom_negated_projection ORDER BY a; +---- +physical_plan +01)SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--CustomProjectionExec: expr=[(- a@0) as a] +03)----CooperativeExec +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +SET datafusion.explain.physical_plan_only = false; + +statement ok +SET datafusion.execution.target_partitions = 4; + #Sort Pushdown for ordered Parquet files statement ok SET datafusion.execution.parquet.pushdown_filters = true; diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 2ec01c99479c4..047e64c03f89f 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -25,6 +25,28 @@ in this section pertains to features and changes that have already been merged to the main branch and are awaiting release in this version. +### Sort pushdown through custom execution plans + +Returning `true` from `ExecutionPlan::maintains_input_order()` alone no longer +enables sort pushdown through a custom node. Preserving row sequence does not +establish how output sort keys map to input expressions; moving a sort through +a value transformation could return incorrectly sorted results. Custom nodes +that previously relied on implicit pushdown may therefore retain extra sorts. + +**Migration guide:** + +Implement `ExecutionPlan::try_pushdown_sort` to optimize for the requested output +ordering. Map sort keys to input expressions before delegating, and return +`Exact` only when the replacement plan guarantees the requested ordering. + +Transparent nodes can delegate to a child and wrap the returned plan when that +child supports the request. Delegation does not automatically move or +consolidate sorts through unsupported children, such as joins, unions, windows, +or `SortExec`. + +A node's reported output ordering can still eliminate a redundant sort without +implementing this hook. + ### Session time zone now applies to mixed timestamp operations When `datafusion.execution.time_zone` is set, DataFusion now uses it to From a36fdc57a2443961fe7d47071047065a94b795b2 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Sat, 24 Oct 2026 22:45:29 +0200 Subject: [PATCH 2/3] fix: validate sort pushdown through custom execution plans --- .../physical_optimizer/enforce_sorting.rs | 32 ++- .../tests/physical_optimizer/pushdown_sort.rs | 28 +- .../enforce_sorting/sort_pushdown.rs | 254 ++++++++++++++++-- .../physical-plan/src/execution_plan.rs | 16 +- .../src/test_context/custom_sort_pushdown.rs | 76 ++++-- .../sqllogictest/test_files/sort_pushdown.slt | 49 ++++ .../library-user-guide/upgrading/56.0.0.md | 22 -- 7 files changed, 384 insertions(+), 93 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 49989a50b5f84..e305a5b0a6e9e 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2604,9 +2604,8 @@ async fn test_push_with_required_input_ordering_prohibited() -> Result<()> { Ok(()) } -// Compatible input requirements do not opt a custom node into sort pushdown. #[tokio::test] -async fn test_custom_node_keeps_sort_with_compatible_input_ordering() -> Result<()> { +async fn test_push_with_required_input_ordering_allowed() -> Result<()> { let schema = create_test_schema3()?; let ordering_a: LexOrdering = [sort_expr("a", &schema)].into(); let ordering_ab = [sort_expr("a", &schema), sort_expr("b", &schema)].into(); @@ -2621,11 +2620,38 @@ async fn test_custom_node_keeps_sort_with_compatible_input_ordering() -> Result< let test = EnforceSortingTest::new(plan.clone()).with_repartition_sorts(true); assert_snapshot!(test.run(), @r" - Input / Optimized Plan: + Input Plan: SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[false] RequiredInputOrderingExec SortExec: expr=[a@0 ASC], preserve_partitioning=[false] DataSourceExec: partitions=1, partition_sizes=[0] + + Optimized Plan: + RequiredInputOrderingExec + SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[0] + "); + Ok(()) +} + +#[tokio::test] +async fn test_push_with_required_input_ordering_conflicting_suffix() -> Result<()> { + let schema = create_test_schema3()?; + let required: LexOrdering = [sort_expr("a", &schema), sort_expr("b", &schema)].into(); + let requested = [sort_expr("a", &schema), sort_expr("c", &schema)].into(); + let input = sort_exec(required.clone(), memory_exec(&schema)); + let plan = RequirementsTestExec::new(input) + .with_required_input_ordering(Some(required)) + .into_arc(); + let plan = sort_exec(requested, plan); + let test = EnforceSortingTest::new(plan).with_repartition_sorts(true); + + assert_snapshot!(test.run(), @r" + Input / Optimized Plan: + SortExec: expr=[a@0 ASC, c@2 ASC], preserve_partitioning=[false] + RequiredInputOrderingExec + SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[0] "); Ok(()) } diff --git a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs index 68d3418345d64..566779afdf1ec 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs @@ -37,7 +37,7 @@ use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; use datafusion_physical_optimizer::pushdown_sort::PushdownSort; use datafusion_physical_plan::sorts::sort::SortExec; -use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties, collect}; +use datafusion_physical_plan::{ExecutionPlan, collect}; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ @@ -1113,7 +1113,9 @@ fn test_sort_pushdown_with_test_scan_arbitrary_ordering() { // ============================================================================ #[test] -fn test_custom_plan_can_opt_into_sort_pushdown() -> Result<()> { +fn test_sort_pushdown_exact_no_fetch_no_limit() { + // When a source returns Exact (without fetch), the SortExec should be + // removed entirely with no GlobalLimitExec wrapper. let schema = schema(); let a = sort_expr("a", &schema); let b = sort_expr("b", &schema); @@ -1121,16 +1123,20 @@ fn test_custom_plan_can_opt_into_sort_pushdown() -> Result<()> { Arc::new(TestScan::new(schema.clone(), vec![]).with_exact_pushdown(true)); let ordering = LexOrdering::new(vec![a, b.reverse()]).unwrap(); - let plan = sort_exec(ordering.clone(), source); - let mut config = ConfigOptions::new(); - config.execution.target_partitions = 1; - let plan = EnsureRequirements::new().optimize(plan, &config)?; - assert!(plan.is::()); + let plan = sort_exec(ordering, source); - let optimized = PushdownSort::new().optimize(plan, &config)?; - assert!(optimized.is::()); - assert_eq!(optimized.output_ordering(), Some(&ordering)); - Ok(()) + insta::assert_snapshot!( + OptimizationTest::new(plan, PushdownSort::new(), true), + @r" + OptimizationTest: + input: + - SortExec: expr=[a@0 ASC, b@1 DESC NULLS LAST], preserve_partitioning=[false] + - TestScan + output: + Ok: + - TestScan: requested_ordering=[a@0 ASC, b@1 DESC NULLS LAST] + " + ); } #[test] diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 510e8df83a154..acb1d84f7fa4b 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -18,7 +18,6 @@ use std::fmt::Debug; use std::sync::Arc; -use crate::output_requirements::OutputRequirementExec; use crate::utils::{ add_sort_above_with_distribution, is_sort, is_sort_preserving_merge, is_union, is_window, @@ -38,20 +37,22 @@ use datafusion_physical_expr_common::sort_expr::{ PhysicalSortRequirement, }; use datafusion_physical_plan::aggregates::AggregateExec; -use datafusion_physical_plan::buffer::BufferExec; -use datafusion_physical_plan::coop::CooperativeExec; use datafusion_physical_plan::execution_plan::CardinalityEffect; +use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::{ ColumnIndex, calculate_join_output_ordering, }; use datafusion_physical_plan::joins::{HashJoinExec, SortMergeJoinExec}; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; -use datafusion_physical_plan::scalar_subquery::ScalarSubqueryExec; +use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::tree_node::PlanContext; -use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion_physical_plan::{ + ChildSatisfactionOptions, ChildrenPropertiesMode, ExecutionPlan, + ExecutionPlanProperties, ReplaceChildrenOptions, +}; /// "Data class" used by sort pushdown (now driven from `EnsureRequirements`) /// to push down [`SortExec`] in the plan. In some cases the total @@ -449,10 +450,6 @@ fn pushdown_sorts_helper( /// Calculate the pushdown ordering requirements for children. /// If sort cannot be pushed down, return None. -/// -/// Pushdown requires a known mapping from output sort keys to child expressions, -/// provided by a specialized handler or a known passthrough case. Preserving -/// row sequence alone does not establish that mapping. fn pushdown_requirement_to_children( plan: &Arc, parent_required: OrderingRequirements, @@ -485,6 +482,7 @@ fn pushdown_requirement_to_children( return Ok(None); } + let maintains_input_order = plan.maintains_input_order(); if is_window(plan) { let mut required_input_ordering = plan.required_input_ordering(); let maybe_child_requirement = required_input_ordering.swap_remove(0); @@ -608,7 +606,12 @@ fn pushdown_requirement_to_children( handle_aggregate_pushdown(aggregate_exec, parent_required) } else if let Some(projection_exec) = plan.downcast_ref::() { handle_projection_pushdown(projection_exec, &parent_required) - } else if pushdown_would_violate_requirements(&parent_required, plan.as_ref()) { + } else if maintains_input_order.is_empty() + || !maintains_input_order.iter().any(|maintains| *maintains) + || plan.is::() + || plan.is::() + || pushdown_would_violate_requirements(&parent_required, plan.as_ref()) + { Ok(None) } else if is_sort_preserving_merge(plan) { let new_ordering = LexOrdering::from(parent_required.first().clone()); @@ -627,21 +630,8 @@ fn pushdown_requirement_to_children( } } else if let Some(hash_join) = plan.downcast_ref::() { handle_hash_join(hash_join, parent_required) - } else if plan.is::() - || plan.is::() - || plan.is::() - || plan.is::() - { - // These operators expose the first child's columns without transformation. - // Scalar subquery children compute side inputs and do not contribute rows. - let mut child_requirements = vec![None; plan.children().len()]; - child_requirements[0] = Some(parent_required); - Ok(Some(child_requirements)) } else { - // Preserving row positions does not identify the input expressions needed - // to order transformed output values. Custom nodes opt in through - // ExecutionPlan::try_pushdown_sort, handled by the PushdownSort rule. - Ok(None) + handle_custom_pushdown(plan, parent_required, &maintains_input_order) } } @@ -958,6 +948,216 @@ fn expr_source_side( } } +/// Handles the custom pushdown of parent-required sorting requirements down to +/// the child execution plans, considering whether the input order is maintained. +/// +/// # Arguments +/// +/// * `plan` - A reference to an `ExecutionPlan` for which the pushdown will be applied. +/// * `parent_required` - The sorting requirements expected by the parent node. +/// * `maintains_input_order` - A vector of booleans indicating whether each child +/// maintains the input order. +/// +/// # Returns +/// +/// Returns `Ok(Some(Vec>))` if the sorting requirements can be +/// pushed down, `Ok(None)` if not. On error, returns a `Result::Err`. +fn handle_custom_pushdown( + plan: &Arc, + parent_required: OrderingRequirements, + maintains_input_order: &[bool], +) -> Result>>> { + let plan_children = plan.children(); + + // A fetched custom node already selects rows; reordering its input can + // change that selection. Native limit and merge cases are handled earlier. + if plan_children.is_empty() || plan.fetch().is_some() { + return Ok(None); + } + + // Collect all unique column indices used in the parent-required sorting + // expression: + let requirement = parent_required.into_single(); + // Evaluating a computed key earlier can fail on rows this node would discard. + if requirement + .iter() + .any(|sort| sort.expr.downcast_ref::().is_none()) + { + return Ok(None); + } + let all_indices: HashSet = requirement + .iter() + .flat_map(|order| { + collect_columns(&order.expr) + .iter() + .map(|col| col.index()) + .collect::>() + }) + .collect(); + + // Get the number of fields in each child's schema: + let children_schema_lengths: Vec = plan_children + .iter() + .map(|c| c.schema().fields().len()) + .collect(); + + // Find the index of the order-maintaining child: + let Some(maintained_child_idx) = maintains_input_order + .iter() + .enumerate() + .find(|(_, m)| **m) + .map(|pair| pair.0) + else { + return Ok(None); + }; + + // Check if all required columns come from the order-maintaining child: + let start_idx = children_schema_lengths[..maintained_child_idx] + .iter() + .sum::(); + let end_idx = start_idx + children_schema_lengths[maintained_child_idx]; + let all_from_maintained_child = + all_indices.iter().all(|i| i >= &start_idx && i < &end_idx); + + // If all columns are from the maintained child, update the parent requirements: + if all_from_maintained_child { + let sub_offset = children_schema_lengths + .iter() + .take(maintained_child_idx) + .sum::(); + // Transform the parent-required expression for the child schema by + // adjusting columns: + let updated_parent_req = requirement + .clone() + .into_iter() + .map(|req| { + let child_schema = plan_children[maintained_child_idx].schema(); + let updated_columns = req + .expr + .transform_up(|expr| { + if let Some(col) = expr.downcast_ref::() { + let new_index = col.index() - sub_offset; + Ok(Transformed::yes(Arc::new(Column::new( + child_schema.field(new_index).name(), + new_index, + )))) + } else { + Ok(Transformed::no(expr)) + } + })? + .data; + Ok(PhysicalSortRequirement::new(updated_columns, req.options)) + }) + .collect::>>()?; + + // Prepare the result, populating with the updated requirements for children that maintain order + let result: Vec<_> = maintains_input_order + .iter() + .map(|&maintains_order| { + if maintains_order { + LexRequirement::new(updated_parent_req.clone()) + .map(OrderingRequirements::new) + } else { + None + } + }) + .collect(); + + // Test the proposed child orderings without moving the real sorts yet. + let mut probe_children = Vec::with_capacity(plan_children.len()); + for (child, required) in plan_children.iter().zip(&result) { + let child = if let Some(required) = required { + // Other order-maintaining inputs may have narrower schemas. + if required.first().iter().any(|sort| { + collect_columns(&sort.expr) + .iter() + .any(|column| column.index() >= child.schema().fields().len()) + }) { + return Ok(None); + } + // A bounded operator can stop an unbounded input; a new blocking + // sort below it cannot. An already ordered input needs no new sort. + if child.boundedness().is_unbounded() + && !child + .equivalence_properties() + .ordering_satisfy_requirement(required.first().clone())? + { + return Ok(None); + } + Arc::new( + SortExec::new(required.first().clone().into(), Arc::clone(child)) + .with_preserve_partitioning(true), + ) as Arc + } else { + Arc::clone(child) + }; + probe_children.push(child); + } + let children = probe_children + .iter() + .map(|child| child.as_ref()) + .collect::>(); + if !sort_pushdown_inputs_satisfied(plan.as_ref(), &children)? { + return Ok(None); + } + let probe = Arc::clone(plan).replace_children( + probe_children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let children = probe + .children() + .into_iter() + .map(|child| child.as_ref()) + .collect::>(); + if !sort_pushdown_inputs_satisfied(probe.as_ref(), &children)? + || !probe + .equivalence_properties() + .ordering_satisfy_requirement(requirement)? + { + return Ok(None); + } + Ok(Some(result)) + } else { + Ok(None) + } +} + +/// A sort probe must satisfy the node's input requirements as well as its output order. +fn sort_pushdown_inputs_satisfied( + plan: &dyn ExecutionPlan, + children: &[&dyn ExecutionPlan], +) -> Result { + let ordering = plan.required_input_ordering(); + let distribution = plan.input_distribution_requirements(); + if ordering.len() != children.len() + || distribution.per_child_distributions().len() != children.len() + { + return Ok(false); + } + for (index, (child, required)) in children.iter().zip(ordering).enumerate() { + if let Some(required) = required + && !child + .equivalence_properties() + .ordering_satisfy_requirement(required.first().clone())? + { + return Ok(false); + } + if !distribution + .child_satisfaction( + index, + *child, + ChildSatisfactionOptions::new().with_allow_subset(true), + )? + .is_satisfied() + { + return Ok(false); + } + } + Ok(distribution + .unsatisfied_co_partitioned_children(plan.name(), children)? + .is_empty()) +} + // For hash join we only maintain the input order for the right child // for join type: Inner, Right, RightSemi, RightAnti, RightMark fn handle_hash_join( @@ -1121,7 +1321,9 @@ mod tests { use datafusion_physical_expr::expressions::{BinaryExpr, col}; use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::limit::GlobalLimitExec; - use datafusion_physical_plan::scalar_subquery::ScalarSubqueryLink; + use datafusion_physical_plan::scalar_subquery::{ + ScalarSubqueryExec, ScalarSubqueryLink, + }; const DESC: SortOptions = SortOptions { descending: true, @@ -1241,7 +1443,7 @@ mod tests { fn scalar_subquery_sort_applies_only_to_the_main_input() -> Result<()> { let schema = child_schema(); let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let requirement = OrderingRequirements::new_soft(lex([ + let requirement = OrderingRequirements::new(lex([ req("c", &schema, DESC), req("a", &schema, ASC), ])); diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index e9cc3244f2ad6..ee15653f7df59 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -233,10 +233,6 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// DataFusion uses this metadata to apply certain optimizations /// such as automatically repartitioning correctly. /// - /// This describes row sequence, not how input expressions map to outputs. - /// Custom nodes that support sort pushdown should implement - /// [`Self::try_pushdown_sort`]; returning `true` here alone does not opt in. - /// /// The default implementation returns `false` /// /// WARNING: if you override this default, you *MUST* ensure that @@ -987,16 +983,8 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// - Return [`SortOrderPushdownResult::Unsupported`] if the node cannot optimize /// for the ordering /// - /// Custom nodes opt into sort pushdown by implementing this method. Requested - /// expressions refer to this node's output schema; map them to input - /// expressions before delegating to a child. Decline pushdown when a correct - /// mapping is unavailable. - /// - /// Nodes that preserve the full input row sequence and column values can - /// delegate to their child and wrap the returned plan in a new instance of - /// the node. Delegation relies on the child supporting this method. If the - /// child returns [`SortOrderPushdownResult::Unsupported`], delegation alone - /// does not insert or move sorts to satisfy the requested ordering. + /// For transparent nodes (that preserve ordering), implement this to delegate to + /// children and wrap the result with a new instance of this node. /// /// Default implementation returns `Unsupported`. fn try_pushdown_sort( diff --git a/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs b/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs index a89979c4c614c..709579facfb7d 100644 --- a/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs +++ b/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs @@ -25,18 +25,21 @@ use datafusion::catalog::{Session, TableProvider}; use datafusion::common::Result; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; -use datafusion::logical_expr::{Expr, TableType}; -use datafusion::physical_expr::PhysicalExpr; -use datafusion::physical_expr::expressions::{col, negative}; +use datafusion::logical_expr::{Expr, Operator, TableType}; +use datafusion::physical_expr::expressions::{binary, col, lit, negative}; use datafusion::physical_expr::projection::ProjectionExprs; +use datafusion::physical_expr::{PhysicalExpr, PhysicalSortExpr}; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::limit::LocalLimitExec; use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::test::TestMemoryExec; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, }; use datafusion::prelude::SessionContext; -/// Register a fixed custom operator that negates its input column. +/// Register custom plans whose native implementations are hidden from the optimizer. pub(super) fn register_custom_sort_pushdown(ctx: &SessionContext) -> Result<()> { let batch = record_batch!(("a", Int32, [2, 1, 3]))?; let schema = batch.schema(); @@ -47,18 +50,54 @@ pub(super) fn register_custom_sort_pushdown(ctx: &SessionContext) -> Result<()> )?; ctx.register_table( "custom_negated_projection", - Arc::new(CustomProjection(projection)), + Arc::new(CustomExec(Arc::new(projection))), + )?; + + let batch = record_batch!(("a", Int32, [2, 1, 3]), ("b", Int32, [10, 30, 20]))?; + let schema = batch.schema(); + let input = TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + let projection = ProjectionExec::try_new( + [ + (col("b", &schema)?, "a".to_string()), + (col("a", &schema)?, "b".to_string()), + ], + input, + )?; + ctx.register_table( + "custom_reordered_projection", + Arc::new(CustomExec(Arc::new(projection))), + )?; + + let batch = record_batch!(("t", Int32, [1, 2, 3]), ("x", Int32, [0, 2, 4]))?; + let schema = batch.schema(); + let input = TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + let predicate = binary(col("x", &schema)?, Operator::NotEq, lit(0i32), &schema)?; + ctx.register_table( + "custom_filtered_input", + Arc::new(CustomExec(Arc::new(FilterExec::try_new(predicate, input)?))), + )?; + + let batch = record_batch!(("a", Int32, [1, 2, 3]), ("b", Int32, [30, 20, 10]))?; + let schema = batch.schema(); + let ordering = [PhysicalSortExpr::new_default(col("a", &schema)?)].into(); + let input = Arc::new( + TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?, + ); + ctx.register_table( + "custom_limited_input", + Arc::new(CustomExec(Arc::new(LocalLimitExec::new(input, 2)))), )?; Ok(()) } -/// Negates values while preserving row positions. Keeping a distinct plan type -/// exercises the generic custom-plan sort pushdown path. +/// Exposes a native plan's children and properties under a distinct plan type, +/// exercising the generic custom-plan sort pushdown path. #[derive(Debug, Clone)] -struct CustomProjection(ProjectionExec); +struct CustomExec(Arc); #[async_trait] -impl TableProvider for CustomProjection { +impl TableProvider for CustomExec { fn schema(&self) -> SchemaRef { self.0.schema() } @@ -87,16 +126,16 @@ impl TableProvider for CustomProjection { } } -impl DisplayAs for CustomProjection { +impl DisplayAs for CustomExec { fn fmt_as(&self, display_type: DisplayFormatType, f: &mut Formatter) -> fmt::Result { write!(f, "Custom")?; self.0.fmt_as(display_type, f) } } -impl ExecutionPlan for CustomProjection { +impl ExecutionPlan for CustomExec { fn name(&self) -> &str { - "CustomProjectionExec" + "CustomExec" } fn properties(&self) -> &Arc { @@ -107,18 +146,21 @@ impl ExecutionPlan for CustomProjection { self.0.children() } - // Row positions are preserved; projected values have the opposite ordering. fn maintains_input_order(&self) -> Vec { self.0.maintains_input_order() } + fn fetch(&self) -> Option { + self.0.fetch() + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(Self(ProjectionExec::try_new( - self.0.expr().to_vec(), - Arc::clone(&children[0]), + Ok(Arc::new(Self(Arc::clone(&self.0).replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), )?))) } diff --git a/datafusion/sqllogictest/test_files/sort_pushdown.slt b/datafusion/sqllogictest/test_files/sort_pushdown.slt index 30deeaf9aee82..e1035ff38abe6 100644 --- a/datafusion/sqllogictest/test_files/sort_pushdown.slt +++ b/datafusion/sqllogictest/test_files/sort_pushdown.slt @@ -40,6 +40,55 @@ physical_plan 03)----CooperativeExec 04)------DataSourceExec: partitions=1, partition_sizes=[1] +# Reordering columns must not make the generic path sort the wrong input column. +query II +SELECT a, b FROM custom_reordered_projection ORDER BY a; +---- +10 2 +20 3 +30 1 + +query TT +EXPLAIN SELECT a, b FROM custom_reordered_projection ORDER BY a; +---- +physical_plan +01)SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--CustomProjectionExec: expr=[b@1 as a, a@0 as b] +03)----CooperativeExec +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# Sorting before the custom filter would evaluate 1 / x on the rejected zero row. +query II +SELECT t, x FROM custom_filtered_input ORDER BY t, 1 / x; +---- +2 2 +3 4 + +query TT +EXPLAIN SELECT t, x FROM custom_filtered_input ORDER BY t, 1 / x; +---- +physical_plan +01)SortExec: expr=[t@0 ASC NULLS LAST, 1 / CAST(x@1 AS Int64) ASC NULLS LAST], preserve_partitioning=[false] +02)--CustomFilterExec: x@1 != 0 +03)----CooperativeExec +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# The input is ordered by a. Sorting by b must keep the first two selected rows. +query II +SELECT a, b FROM custom_limited_input ORDER BY b; +---- +2 20 +1 30 + +query TT +EXPLAIN SELECT a, b FROM custom_limited_input ORDER BY b; +---- +physical_plan +01)SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false] +02)--CustomLocalLimitExec: fetch=2 +03)----CooperativeExec +04)------DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 ASC + statement ok SET datafusion.explain.physical_plan_only = false; diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 047e64c03f89f..2ec01c99479c4 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -25,28 +25,6 @@ in this section pertains to features and changes that have already been merged to the main branch and are awaiting release in this version. -### Sort pushdown through custom execution plans - -Returning `true` from `ExecutionPlan::maintains_input_order()` alone no longer -enables sort pushdown through a custom node. Preserving row sequence does not -establish how output sort keys map to input expressions; moving a sort through -a value transformation could return incorrectly sorted results. Custom nodes -that previously relied on implicit pushdown may therefore retain extra sorts. - -**Migration guide:** - -Implement `ExecutionPlan::try_pushdown_sort` to optimize for the requested output -ordering. Map sort keys to input expressions before delegating, and return -`Exact` only when the replacement plan guarantees the requested ordering. - -Transparent nodes can delegate to a child and wrap the returned plan when that -child supports the request. Delegation does not automatically move or -consolidate sorts through unsupported children, such as joins, unions, windows, -or `SortExec`. - -A node's reported output ordering can still eliminate a redundant sort without -implementing this hook. - ### Session time zone now applies to mixed timestamp operations When `datafusion.execution.time_zone` is set, DataFusion now uses it to From ec3badb56a99f9b3175443425e11b3be2b0c5861 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Sun, 25 Oct 2026 07:17:41 +0100 Subject: [PATCH 3/3] test: cover ordered selection in custom sort pushdown --- datafusion/sqllogictest/src/test_context.rs | 16 ++++- .../src/test_context/custom_sort_pushdown.rs | 66 ++++++++++++++++--- .../sqllogictest/test_files/sort_pushdown.slt | 16 +++-- 3 files changed, 84 insertions(+), 14 deletions(-) diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index dd7b9fb1e2222..0d0e424266e48 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -46,6 +46,7 @@ use datafusion::logical_expr::{ ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, create_udf, }; +use datafusion::physical_optimizer::optimizer::PhysicalOptimizer; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::*; use datafusion::{ @@ -114,6 +115,7 @@ impl TestContext { /// If `None` is returned (e.g. because some needed feature is not /// enabled), the file should be skipped pub async fn try_new_for_test_file(relative_path: &Path) -> Option { + let file_name = relative_path.file_name().unwrap().to_str().unwrap(); let config = SessionConfig::new() // hardcode target partitions so plans are deterministic .with_target_partitions(4); @@ -177,11 +179,23 @@ impl TestContext { state_builder = state_builder.with_statistics_registry(registry); } + if file_name == "sort_pushdown.slt" { + let mut rules = PhysicalOptimizer::new().rules; + let ensure_requirements = rules + .iter() + .position(|rule| rule.name() == "EnsureRequirements") + .expect("default optimizer includes EnsureRequirements"); + rules.insert( + ensure_requirements, + Arc::new(custom_sort_pushdown::CustomLimitRule), + ); + state_builder = state_builder.with_physical_optimizer_rules(rules); + } + let state = state_builder.build(); let mut test_ctx = TestContext::new(SessionContext::new_with_state(state)); - let file_name = relative_path.file_name().unwrap().to_str().unwrap(); match file_name { "sort_pushdown.slt" => { custom_sort_pushdown::register_custom_sort_pushdown( diff --git a/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs b/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs index 709579facfb7d..05f87630a0568 100644 --- a/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs +++ b/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs @@ -23,19 +23,22 @@ use arrow::datatypes::{self as arrow_schema, SchemaRef}; use async_trait::async_trait; use datafusion::catalog::{Session, TableProvider}; use datafusion::common::Result; -use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::logical_expr::{Expr, Operator, TableType}; use datafusion::physical_expr::expressions::{binary, col, lit, negative}; use datafusion::physical_expr::projection::ProjectionExprs; use datafusion::physical_expr::{PhysicalExpr, PhysicalSortExpr}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::limit::LocalLimitExec; use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::test::TestMemoryExec; use datafusion::physical_plan::{ - ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, - ReplaceChildrenOptions, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, }; use datafusion::prelude::SessionContext; @@ -79,18 +82,65 @@ pub(super) fn register_custom_sort_pushdown(ctx: &SessionContext) -> Result<()> let batch = record_batch!(("a", Int32, [1, 2, 3]), ("b", Int32, [30, 20, 10]))?; let schema = batch.schema(); - let ordering = [PhysicalSortExpr::new_default(col("a", &schema)?)].into(); + let ordering = + [PhysicalSortExpr::new_default(col("a", &schema)?).nulls_last()].into(); let input = Arc::new( TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? .try_with_sort_information(vec![ordering])?, ); - ctx.register_table( - "custom_limited_input", - Arc::new(CustomExec(Arc::new(LocalLimitExec::new(input, 2)))), - )?; + ctx.register_table("custom_ordered_input", Arc::new(CustomExec(input)))?; Ok(()) } +/// Plans `ORDER BY ... LIMIT` as a custom node to test generic sort pushdown. +#[derive(Debug)] +pub(super) struct CustomLimitRule; + +impl PhysicalOptimizerRule for CustomLimitRule { + fn optimize( + &self, + plan: Arc, + _: &ConfigOptions, + ) -> Result> { + plan.transform_up(|plan| { + let Some(sort) = plan.downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + let Some(fetch) = sort.fetch() else { + return Ok(Transformed::no(plan)); + }; + let input = sort.input(); + let source = input + .downcast_ref::() + .map_or(input, |projection| projection.input()); + let Some(source) = source.downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + // Only the dedicated ordered source is a wrapped memory scan. + if !source.0.is::() + || input.output_partitioning().partition_count() != 1 + || !input + .equivalence_properties() + .ordering_satisfy(sort.expr().clone())? + { + return Ok(Transformed::no(plan)); + } + let limit = LocalLimitExec::new(Arc::clone(input), fetch); + let plan: Arc = Arc::new(CustomExec(Arc::new(limit))); + Ok(Transformed::yes(plan)) + }) + .map(|transformed| transformed.data) + } + + fn name(&self) -> &str { + "CustomLimitRule" + } + + fn schema_check(&self) -> bool { + true + } +} + /// Exposes a native plan's children and properties under a distinct plan type, /// exercising the generic custom-plan sort pushdown path. #[derive(Debug, Clone)] diff --git a/datafusion/sqllogictest/test_files/sort_pushdown.slt b/datafusion/sqllogictest/test_files/sort_pushdown.slt index e1035ff38abe6..7a6d81f328237 100644 --- a/datafusion/sqllogictest/test_files/sort_pushdown.slt +++ b/datafusion/sqllogictest/test_files/sort_pushdown.slt @@ -57,7 +57,8 @@ physical_plan 03)----CooperativeExec 04)------DataSourceExec: partitions=1, partition_sizes=[1] -# Sorting before the custom filter would evaluate 1 / x on the rejected zero row. +# Sorting before the custom filter would evaluate 1 / x on the rejected row +# where x = 0, failing the query with an integer division-by-zero error. query II SELECT t, x FROM custom_filtered_input ORDER BY t, 1 / x; ---- @@ -73,21 +74,26 @@ physical_plan 03)----CooperativeExec 04)------DataSourceExec: partitions=1, partition_sizes=[1] -# The input is ordered by a. Sorting by b must keep the first two selected rows. +# The inner ORDER BY determines which rows LIMIT selects. The outer sort +# must keep those rows, even when LIMIT is implemented by a custom node. query II -SELECT a, b FROM custom_limited_input ORDER BY b; +SELECT a, b +FROM (SELECT a, b FROM custom_ordered_input ORDER BY a LIMIT 2) s +ORDER BY b; ---- 2 20 1 30 query TT -EXPLAIN SELECT a, b FROM custom_limited_input ORDER BY b; +EXPLAIN SELECT a, b +FROM (SELECT a, b FROM custom_ordered_input ORDER BY a LIMIT 2) s +ORDER BY b; ---- physical_plan 01)SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false] 02)--CustomLocalLimitExec: fetch=2 03)----CooperativeExec -04)------DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 ASC +04)------CustomDataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 ASC NULLS LAST statement ok SET datafusion.explain.physical_plan_only = false;