fix: validate sort pushdown through custom execution plans

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(())
 }
 

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::<SortExec>()
+        .expect("a fetched custom plan must retain the requested outer sort");
+    assert_eq!(sort.expr(), &ordering);
+    assert!(sort.input().is::<TestScan>());
+    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(),

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<usize>) -> Option<Arc<dyn ExecutionPlan>> {
         if self.supports_fetch {
             let mut new_scan = self.clone();

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::<ProjectionExec>() {
-                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::<SortPreservingMergeExec>() && 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::<ProjectionExec>() {
         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::<RepartitionExec>()
         || plan.is::<FilterExec>()
         || 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<Option<Vec<Option<OrderingRequirements>>>> {
     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::<Column>().is_none())
+    {
+        return Ok(None);
+    }
     let all_indices: HashSet<usize> = 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::<Result<Vec<_>>>()?;
 
         // 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<dyn ExecutionPlan>
+            } else {
+                Arc::clone(child)
+            };
+            probe_children.push(child);
+        }
+        let children = probe_children
+            .iter()
+            .map(|child| child.as_ref())
+            .collect::<Vec<_>>();
+        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::<Vec<_>>();
+        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<bool> {
+    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<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::clone(&schema)));
+        let requirement = OrderingRequirements::new(lex([
+            req("c", &schema, DESC),
+            req("a", &schema, ASC),
+        ]));
+        let plan: Arc<dyn ExecutionPlan> = 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();