diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 084ae59a76c23..e305a5b0a6e9e 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2604,7 +2604,6 @@ async fn test_push_with_required_input_ordering_prohibited() -> Result<()> { Ok(()) } -// test when the required input ordering is satisfied so could push through #[tokio::test] async fn test_push_with_required_input_ordering_allowed() -> Result<()> { let schema = create_test_schema3()?; @@ -2618,14 +2617,6 @@ 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" @@ -2640,7 +2631,28 @@ async fn test_push_with_required_input_ordering_allowed() -> Result<()> { 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 6cf05ecfd8fc4..566779afdf1ec 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, collect}; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ @@ -1137,6 +1139,29 @@ fn test_sort_pushdown_exact_no_fetch_no_limit() { ); } +#[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] fn test_sort_pushdown_exact_preserves_fetch_with_global_limit() { // When a source returns Exact but does NOT support with_fetch(), 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..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 @@ -49,7 +49,10 @@ 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 @@ -548,49 +551,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: @@ -635,14 +607,11 @@ fn pushdown_requirement_to_children( } 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) + || !maintains_input_order.iter().any(|maintains| *maintains) || 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 Ok(None) } else if is_sort_preserving_merge(plan) { let new_ordering = LexOrdering::from(parent_required.first().clone()); @@ -1000,14 +969,22 @@ fn handle_custom_pushdown( ) -> Result>>> { let plan_children = plan.children(); - // If the plan has no children, return early: - if plan_children.is_empty() { + // 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| { @@ -1051,6 +1028,7 @@ fn handle_custom_pushdown( // 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(); @@ -1073,7 +1051,7 @@ fn handle_custom_pushdown( .collect::>>()?; // Prepare the result, populating with the updated requirements for children that maintain order - let result = maintains_input_order + let result: Vec<_> = maintains_input_order .iter() .map(|&maintains_order| { if maintains_order { @@ -1085,12 +1063,101 @@ fn handle_custom_pushdown( }) .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( @@ -1247,10 +1314,16 @@ 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::{ + ScalarSubqueryExec, ScalarSubqueryLink, + }; const DESC: SortOptions = SortOptions { descending: true, @@ -1366,6 +1439,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(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/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 0c8c2c5add363..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::{ @@ -73,6 +74,7 @@ use log::info; use sqlparser::ast; use tempfile::TempDir; +mod custom_sort_pushdown; mod range_partitioning; /// Context for running tests @@ -113,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); @@ -176,12 +179,30 @@ 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( + 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..05f87630a0568 --- /dev/null +++ b/datafusion/sqllogictest/src/test_context/custom_sort_pushdown.rs @@ -0,0 +1,231 @@ +// 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::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, + ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, +}; +use datafusion::prelude::SessionContext; + +/// 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(); + 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(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(()) +} + +/// 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 CustomExec(Arc); + +#[async_trait] +impl TableProvider for CustomExec { + 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 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 CustomExec { + fn name(&self) -> &str { + "CustomExec" + } + + fn properties(&self) -> &Arc { + self.0.properties() + } + + fn children(&self) -> Vec<&Arc> { + self.0.children() + } + + 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(Arc::clone(&self.0).replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?))) + } + + 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..7a6d81f328237 100644 --- a/datafusion/sqllogictest/test_files/sort_pushdown.slt +++ b/datafusion/sqllogictest/test_files/sort_pushdown.slt @@ -15,6 +15,92 @@ # 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] + +# 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; + +statement ok +SET datafusion.execution.target_partitions = 4; + #Sort Pushdown for ordered Parquet files statement ok SET datafusion.execution.parquet.pushdown_filters = true;