From 2bf86aacf258cce75a91f91700b477bf4c7e2759 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 | 222 +++++------------- .../physical-plan/src/execution_plan.rs | 14 +- datafusion/sqllogictest/src/test_context.rs | 7 + .../src/test_context/custom_sort_pushdown.rs | 139 +++++++++++ .../sqllogictest/test_files/sort_pushdown.slt | 31 +++ 8 files changed, 288 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..2c57dab1b37a7 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; @@ -479,7 +481,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 +549,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 +604,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 +623,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 +954,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 +1110,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 +1233,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 58e41d7763e9f..0965e72239ea5 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,14 @@ 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. /// /// 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; From e8a763592b9e9e9761a1bf4a3cb2272e43ace2f8 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Mon, 26 Oct 2026 10:31:59 +0100 Subject: [PATCH 2/3] fix: validate sort pushdown through custom execution plans Preserve the generic fallback for transparent custom nodes. Rebuild the node over candidate sorted children and check input ordering, distribution, and output ordering before accepting pushdown. Keep computed keys, fetched nodes, and new blocking sorts on unbounded inputs outside this fallback. Restore compatible-sort consolidation, cover conflicting mandatory input orderings, and remove the obsolete opt-in documentation. Existing custom nodes do not need to implement another hook. --- .../physical_optimizer/enforce_sorting.rs | 42 ++- .../tests/physical_optimizer/pushdown_sort.rs | 28 +- .../enforce_sorting/sort_pushdown.rs | 292 ++++++++++++++++-- .../physical-plan/src/execution_plan.rs | 14 +- 4 files changed, 328 insertions(+), 48 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 49989a50b5f84..fbe22a42a4a1f 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(()) } -// Compatible input requirements do not opt a custom node into sort pushdown. +// test when the required input ordering is satisfied so could push through #[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(); @@ -2618,14 +2618,50 @@ async fn test_custom_node_keeps_sort_with_compatible_input_ordering() -> 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 / 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] + "); + // Should be able to push down + 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 2c57dab1b37a7..b884386a38ccf 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 @@ -481,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); @@ -550,6 +552,9 @@ fn pushdown_requirement_to_children( Ok(None) } } else if plan.is::() && plan.fetch().is_some() { + // Push down through a fetched merge when: + // - requirement is aligned with output ordering + // - it preserves ordering during execution // Refining a fetched merge must preserve the ordering that selected its rows. let Some(ordering) = plan.output_ordering() else { return Ok(None); @@ -604,7 +609,16 @@ 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()) + { + // Leaf nodes and nodes that preserve no input ordering cannot push + // requirements down to their children. + // Keep sorts above RepartitionExec even when a single input partition + // lets it preserve order; pushing the sort below it provides no benefit. Ok(None) } else if is_sort_preserving_merge(plan) { let new_ordering = LexOrdering::from(parent_required.first().clone()); @@ -623,21 +637,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) } } @@ -954,6 +955,245 @@ 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(); + + // Return early for leaf nodes or fetched custom nodes. + // Pushing a sort below an unordered fetch changes which rows it keeps. + // Even an ordered fetch can select different rows if its input is reordered. + // 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(); + + // A custom node can preserve row positions while reversing value order + // (e.g. `-a AS a`). Keep the outer sort unless the proposed child sorts + // satisfy the node's input requirements and give its output the + // requested ordering. + if !is_custom_sort_pushdown_valid(plan, &result, requirement)? { + return Ok(None); + } + Ok(Some(result)) + } else { + Ok(None) + } +} + +/// Checks that proposed child orderings satisfy the node's input requirements +/// and produce the requested output ordering. +fn is_custom_sort_pushdown_valid( + plan: &Arc, + child_requirements: &[Option], + output_requirement: LexRequirement, +) -> Result { + let plan_children = plan.children(); + + // 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(child_requirements) { + 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(false); + } + // 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(false); + } + 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); + } + + // The proposed sorts must not break the node's own input requirements. + let children = probe_children + .iter() + .map(|child| child.as_ref()) + .collect::>(); + if !sort_pushdown_inputs_satisfied(plan.as_ref(), &children)? { + return Ok(false); + } + + // Recompute properties from the proposed inputs: the old properties cannot + // tell us how these new orderings pass through the node. + let probe = Arc::clone(plan).replace_children( + probe_children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + // Rebuilding can change the node's input requirements. Check them again, + // then verify that the output has the requested ordering. + 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(output_requirement)? + { + return Ok(false); + } + Ok(true) +} + +/// 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( @@ -1094,9 +1334,15 @@ fn handle_projection_pushdown( .iter() .any(|expr| !expr.expr.is::()) { + // Computed projection expressions are not supported by this + // remapping, so keep the sort above the projection. return Ok(None); } + // A `ProjectionExec` 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 `score@1` unchanged + // could refer to a different input column and sort on the wrong key. Ok( remap_requirement_through_projection(projection_exec, parent_required) .map(|requirements| vec![Some(requirements)]), @@ -1117,7 +1363,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, @@ -1237,7 +1485,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 0965e72239ea5..58e41d7763e9f 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,14 +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. + /// 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( From f7fc68d6f682b32d4db588614d059011d73ebb31 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Mon, 26 Oct 2026 10:32:35 +0100 Subject: [PATCH 3/3] test: cover custom sort pushdown boundaries in SQL Check returned rows and EXPLAIN plans for reordered columns, computed keys after filtering, and a custom row limit. Express the inner ORDER BY and LIMIT in SQL so the selected rows are deterministic. Reuse native execution plans behind one thin custom wrapper, with a planning rule restricted to the dedicated ordered test source. --- datafusion/sqllogictest/src/test_context.rs | 16 ++- .../src/test_context/custom_sort_pushdown.rs | 128 +++++++++++++++--- .../sqllogictest/test_files/sort_pushdown.slt | 55 ++++++++ 3 files changed, 180 insertions(+), 19 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 a89979c4c614c..05f87630a0568 100644 --- a/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs +++ b/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs @@ -23,20 +23,26 @@ 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, 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_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::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, 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 +53,101 @@ 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)?).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_ordered_input", Arc::new(CustomExec(input)))?; Ok(()) } -/// Negates values while preserving row positions. Keeping a distinct plan type -/// exercises the generic custom-plan sort pushdown path. +/// 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)] -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 +176,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 +196,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..7a6d81f328237 100644 --- a/datafusion/sqllogictest/test_files/sort_pushdown.slt +++ b/datafusion/sqllogictest/test_files/sort_pushdown.slt @@ -40,6 +40,61 @@ 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 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; +---- +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 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 (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 (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)------CustomDataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 ASC NULLS LAST + statement ok SET datafusion.explain.physical_plan_only = false;