jayzhan211 commented on code in PR #26118:
URL: https://github.com/apache/datafusion/pull/26118#discussion_r4377400751
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:
##########
@@ -1000,14 +976,24 @@ 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() {
+ // 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() {
Review Comment:
This early return also catches native `OutputRequirementExec`, which carries
the subquery's `fetch` and reaches this fallback, so `group_by.slt` now fails
at lines 2598, 2636 and 2709. The `(country, amount DESC)` requirement is no
longer merged into the TopK, and a second full `SortExec` stays above
`TopK(country)`. The required CI checks haven't run on this head yet, but it
reproduces locally, and the file passes with the base `sort_pushdown.rs`.
Applying the same rule as the native limit branch fixes it. Refining a
`SortExec` directly below the fetch only changes tie-breaking, so your
custom-limit slt and `test_custom_plan_with_fetch_keeps_outer_sort` still pass,
and the full sqllogictest suite is green with this change:
```diff
- // 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() {
+ if plan_children.is_empty() {
return Ok(None);
}
+ // Sorting below a fetch changes which rows it keeps. As with native
+ // limits, only refine a `SortExec` directly below the fetch.
+ if plan.fetch().is_some() {
+ let refines_sort_below = plan_children.len() == 1
+ && plan_children[0].is::<SortExec>()
+ && plan.output_ordering().is_some_and(|ordering| {
+ plan.equivalence_properties().requirements_compatible(
+ parent_required.first().clone(),
+ ordering.clone().into(),
+ )
+ });
+ if !refines_sort_below {
+ return Ok(None);
+ }
+ }
```
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:
##########
@@ -1085,12 +1072,128 @@ fn handle_custom_pushdown(
})
.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<dyn ExecutionPlan>,
+ child_requirements: &[Option<OrderingRequirements>],
+ output_requirement: LexRequirement,
+) -> Result<bool> {
+ 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()
Review Comment:
No test reaches this unbounded guard or the narrow-schema guard above it:
disabling both leaves every optimizer test and `sort_pushdown.slt` passing. A
snapshot in `enforce_sorting.rs` for each would keep them from silently going
away. For example, use a bounded custom node over an unsorted unbounded source,
and a node with two order-maintaining children of different widths. Fine to do
this in a follow-up.
##########
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<dyn ExecutionPlan>,
+ _: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ plan.transform_up(|plan| {
+ let Some(sort) = plan.downcast_ref::<SortExec>() 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::<ProjectionExec>()
+ .map_or(input, |projection| projection.input());
+ let Some(source) = source.downcast_ref::<CustomExec>() else {
+ return Ok(Transformed::no(plan));
+ };
+ // Only the dedicated ordered source is a wrapped memory scan.
+ if !source.0.is::<TestMemoryExec>()
+ || 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<dyn ExecutionPlan> =
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<dyn ExecutionPlan>);
+
+#[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<usize>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(self.clone());
+ if let Some(projection) = projection {
Review Comment:
Because `scan` always wraps the plan in a `ProjectionExec`, the `ORDER BY t,
1 / x` query is stopped by `handle_projection_pushdown` before it reaches
`handle_custom_pushdown`. It passes on base, and removing the new computed-key
guard changes no test. Skipping the identity projection makes the test reach
the guard. With this change, removing the guard fails the query with `Divide by
zero error`, and the PR as written passes:
```diff
- if let Some(projection) = projection {
+ if let Some(projection) = projection
+ && !projection
+ .iter()
+ .copied()
+ .eq(0..self.0.schema().fields().len())
+ {
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
