jayzhan211 commented on code in PR #26118:
URL: https://github.com/apache/datafusion/pull/26118#discussion_r4388577078
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:
##########
@@ -1034,61 +1029,178 @@ fn handle_custom_pushdown(
return Ok(None);
};
- // Check if all required columns come from the order-maintaining child:
+ // Matching column positions is only a candidate; prove output ordering
below.
let start_idx = children_schema_lengths[..maintained_child_idx]
.iter()
.sum::<usize>();
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);
+ let fits_maintained_input = columns
+ .iter()
+ .all(|col| (start_idx..end_idx).contains(&col.index()));
- // If all columns are from the maintained child, update the parent
requirements:
- if all_from_maintained_child {
- let sub_offset = children_schema_lengths
+ let child_requirements = if fits_maintained_input {
Review Comment:
Optional cleanup: an early return here would flatten the remap and proposal
by one level:
```diff
- let child_requirements = if fits_maintained_input {
- // Remap the columns into the maintained child's schema.
- let child_schema = plan_children[maintained_child_idx].schema();
- let mapped_requirement = columns
- .iter()
- .zip(output_requirement.iter())
- .map(|(col, req)| {
- let index = col.index() - start_idx;
- let col = Column::new(child_schema.field(index).name(),
index);
- PhysicalSortRequirement::new(Arc::new(col), req.options)
- })
- .collect::<Vec<_>>();
-
- // Propose the same ordering for every order-maintaining input.
- maintains_input_order
- .iter()
- .map(|&maintains_order| {
- if maintains_order {
- LexRequirement::new(mapped_requirement.clone())
- .map(OrderingRequirements::new)
- } else {
- None
- }
- })
- .collect::<Vec<_>>()
- } else {
+ if !fits_maintained_input {
return Ok(None);
- };
+ }
+
+ // Remap the columns into the maintained child's schema.
+ let child_schema = plan_children[maintained_child_idx].schema();
+ let mapped_requirement = columns
+ .iter()
+ .zip(output_requirement.iter())
+ .map(|(col, req)| {
+ let index = col.index() - start_idx;
+ let col = Column::new(child_schema.field(index).name(), index);
+ PhysicalSortRequirement::new(Arc::new(col), req.options)
+ })
+ .collect::<Vec<_>>();
+
+ // Propose the same ordering for every order-maintaining input.
+ let child_requirements = maintains_input_order
+ .iter()
+ .map(|&maintains_order| {
+ if maintains_order {
+ LexRequirement::new(mapped_requirement.clone())
+ .map(OrderingRequirements::new)
+ } else {
+ None
+ }
+ })
+ .collect::<Vec<_>>();
```
--
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]
