geoffreyclaude commented on code in PR #26118:
URL: https://github.com/apache/datafusion/pull/26118#discussion_r4379002475
##########
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() {
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
Review Comment:
Done, thanks! This is much simpler.
##########
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()
+ && !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<dyn ExecutionPlan>
+ } else {
+ Arc::clone(child)
+ };
+ probe_children.push(child);
+ }
+
+ // The proposed sorts must not break the node's own input requirements.
Review Comment:
I removed the distribution checks before rebuilding. I kept the ordering
check because a custom node may reject incompatible input ordering during
rebuilding. The full checks still run afterward, since rebuilding may change
the input requirements.
--
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]
