jayzhan211 commented on code in PR #26459:
URL: https://github.com/apache/datafusion/pull/26459#discussion_r4324943045
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs:
##########
@@ -232,6 +235,71 @@ pub fn replace_with_partial_sort(
Ok(plan)
}
+/// Adds a [`SortExec`] below every node whose ordering requirement is not
+/// satisfied by its child's equivalence properties.
+///
+/// [`pushdown_sorts`] removes a sort when it can map the requirement onto a
+/// descendant that satisfies it. That mapping is computed separately from the
+/// equivalence properties of the operators in between, and the two can
+/// disagree. The plan may then be correct, but its properties cannot show it,
+/// and `SanityCheckPlan` rejects it. Adding the sort back keeps every plan
+/// produced by `EnsureRequirements` valid by the same check.
+///
+/// A requirement that a new [`SortExec`] still cannot satisfy is left as is,
+/// for `SanityCheckPlan` to report.
+///
+/// [`pushdown_sorts`]: sort_pushdown::pushdown_sorts
+pub fn add_sort_for_unproven_ordering(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ let requirements = plan.required_input_ordering();
+ if requirements.iter().all(Option::is_none) {
+ return Ok(Transformed::no(plan));
+ }
+ let mut children: Vec<_> = plan.children().into_iter().cloned().collect();
+ let mut changed = false;
+ for (child, requirement) in children.iter_mut().zip(requirements) {
+ let Some(requirement) = requirement else {
+ continue;
+ };
+ let requirement = requirement.into_single();
+ let eq_properties = child.equivalence_properties();
+ if eq_properties.ordering_satisfy_requirement(requirement.clone())? {
+ continue;
+ }
+ let sort_reqs: Vec<_> = requirement.clone().into();
+ let sort_exprs = sort_reqs
+ .into_iter()
+ .filter(|req| eq_properties.is_expr_constant(&req.expr).is_none())
+ .map(Into::into);
+ let Some(ordering) = LexOrdering::new(sort_exprs) else {
+ continue;
+ };
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(
+ SortExec::new(ordering,
Arc::clone(child)).with_preserve_partitioning(
Review Comment:
This no longer applies, since the step is removed. For the record,
`LimitPushdown` ran after it and turned the added sort into a TopK under a
`LIMIT`.
##########
datafusion/sqllogictest/test_files/order.slt:
##########
@@ -2276,3 +2276,79 @@ SELECT i, -i AS n FROM (
-128 -128
-128 -128
NULL NULL
+
+# EnsureRequirements removes the sort on `k` because `k` is constant on the
+# probe side of the join, but the join output does not report `k` as constant.
+# The sort must then be kept so the plan can be checked and run.
+# https://github.com/apache/datafusion/issues/26283
+statement ok
+CREATE TABLE probe_constant_t(a INT, b INT) AS VALUES (1, 1), (2, 2), (3, 3),
(4, 4);
+
+statement ok
+CREATE TABLE probe_constant_u(a INT) AS VALUES (1);
+
+statement ok
+CREATE TABLE probe_constant_v(x INT) AS VALUES (1), (2);
+
+query II
+SELECT s.k, s.a
+FROM (
+ SELECT probe_constant_t.a, -2 AS k
+ FROM probe_constant_t JOIN probe_constant_v ON probe_constant_t.b <>
probe_constant_v.x
+) s
+JOIN probe_constant_u ON s.a = probe_constant_u.a
+ORDER BY s.k DESC;
+----
+-2 1
+
+query TT
+EXPLAIN SELECT s.k, s.a
+FROM (
+ SELECT probe_constant_t.a, -2 AS k
+ FROM probe_constant_t JOIN probe_constant_v ON probe_constant_t.b <>
probe_constant_v.x
+) s
+JOIN probe_constant_u ON s.a = probe_constant_u.a
+ORDER BY s.k DESC;
+----
+logical_plan
+01)Sort: s.k DESC NULLS FIRST
+02)--Projection: s.k, s.a
+03)----Inner Join: s.a = probe_constant_u.a
+04)------SubqueryAlias: s
+05)--------Projection: probe_constant_t.a, Int64(-2) AS k
+06)----------Inner Join: Filter: probe_constant_t.b != probe_constant_v.x
+07)------------TableScan: probe_constant_t projection=[a, b]
+08)------------TableScan: probe_constant_v projection=[x]
+09)------TableScan: probe_constant_u projection=[a]
+physical_plan
+01)SortPreservingMergeExec: [k@0 DESC]
+02)--SortExec: expr=[k@0 DESC], preserve_partitioning=[true]
+03)----HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, a@0)],
projection=[k@2, a@1]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+05)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+06)--------ProjectionExec: expr=[a@0 as a, -2 as k]
+07)----------NestedLoopJoinExec: join_type=Inner, filter=b@0 != x@1,
projection=[a@1]
+08)------------DataSourceExec: partitions=1, partition_sizes=[1]
+09)------------DataSourceExec: partitions=1, partition_sizes=[1]
+
+statement ok
+SET datafusion.execution.target_partitions = 1;
+
+query ITI
+SELECT s.a, s.k, row_number() OVER (PARTITION BY s.a ORDER BY s.k) AS r
Review Comment:
The window query has an `EXPLAIN` now. The Rust test went away with the step
it tested. A unit test in `class.rs` covers the join change, and the slt
queries run the whole optimizer, `SanityCheckPlan` included.
--
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]
