jayzhan211 commented on code in PR #27191:
URL: https://github.com/apache/datafusion/pull/27191#discussion_r4501885060


##########
datafusion/physical-expr/src/equivalence/properties/walk.rs:
##########
@@ -43,19 +43,39 @@ pub(super) struct SortKeyWalk<'a> {
     /// found by [`EquivalenceProperties::implied_ordering`] as the walk goes.
     implied: Vec<LexOrdering>,
     /// For each ordering (the normal orderings of `props`, then `implied`),
-    /// the index of its head. Empty until the first pin, so that a walk with
-    /// nothing pinned allocates nothing; until then every head is at index 0.
+    /// the index of its head. Empty until a head moves, so that a walk that
+    /// moves no head allocates nothing; until then every head is at index 0.
     head_idxs: Vec<usize>,
 }
 
 impl<'a> SortKeyWalk<'a> {
     pub(super) fn new(props: &'a EquivalenceProperties) -> Self {
-        Self {
+        let mut walk = Self {
             props,
             pinned: HashSet::new(),
             implied: vec![],
             head_idxs: vec![],
+        };
+        // A normal ordering keeps a constant whose value differs between
+        // partitions (see `EquivalenceGroup::normalize_sort_exprs`), and such
+        // a constant may lead it. Within a partition it has one value, so it
+        // is fixed and the head is the first expression after it.
+        let normal = &props.oeq_cache.normal_cls;
+        if normal
+            .iter()
+            .any(|ordering| walk.is_fixed(&ordering[0].expr))

Review Comment:
   Same answer as on the comment about the scan in `SortKeyWalk::new`, a few 
lines up: the planning benchmarks are flat, so I'd rather not add a cache.



##########
datafusion/physical-expr/src/equivalence/properties/walk.rs:
##########
@@ -43,19 +43,39 @@ pub(super) struct SortKeyWalk<'a> {
     /// found by [`EquivalenceProperties::implied_ordering`] as the walk goes.
     implied: Vec<LexOrdering>,
     /// For each ordering (the normal orderings of `props`, then `implied`),
-    /// the index of its head. Empty until the first pin, so that a walk with
-    /// nothing pinned allocates nothing; until then every head is at index 0.
+    /// the index of its head. Empty until a head moves, so that a walk that
+    /// moves no head allocates nothing; until then every head is at index 0.
     head_idxs: Vec<usize>,
 }
 
 impl<'a> SortKeyWalk<'a> {
     pub(super) fn new(props: &'a EquivalenceProperties) -> Self {
-        Self {
+        let mut walk = Self {
             props,
             pinned: HashSet::new(),
             implied: vec![],
             head_idxs: vec![],
+        };
+        // A normal ordering keeps a constant whose value differs between
+        // partitions (see `EquivalenceGroup::normalize_sort_exprs`), and such
+        // a constant may lead it. Within a partition it has one value, so it
+        // is fixed and the head is the first expression after it.
+        let normal = &props.oeq_cache.normal_cls;
+        if normal

Review Comment:
   When the heads are columns, this is two lookups per ordering, and the 
`sql_planner` bot run on `c777864f7c` is flat: `physical_plan_tpch_all` 
unchanged, `physical_plan_clickbench_all` +1%, `physical_plan_tpcds_all` -3%, 
worst single query +3%. A cached list of heads would have to be kept correct by 
every path that changes constants or orderings (`add_constants`, 
`add_orderings`, `update_oeq_cache`, `clear_orderings`, 
`clear_per_partition_constants`). That's more state for no visible gain, so I'd 
leave it out unless a benchmark shows a cost.



##########
datafusion/physical-expr/src/equivalence/properties/walk.rs:
##########
@@ -43,19 +43,39 @@ pub(super) struct SortKeyWalk<'a> {
     /// found by [`EquivalenceProperties::implied_ordering`] as the walk goes.
     implied: Vec<LexOrdering>,
     /// For each ordering (the normal orderings of `props`, then `implied`),
-    /// the index of its head. Empty until the first pin, so that a walk with
-    /// nothing pinned allocates nothing; until then every head is at index 0.
+    /// the index of its head. Empty until a head moves, so that a walk that
+    /// moves no head allocates nothing; until then every head is at index 0.
     head_idxs: Vec<usize>,
 }
 
 impl<'a> SortKeyWalk<'a> {
     pub(super) fn new(props: &'a EquivalenceProperties) -> Self {
-        Self {
+        let mut walk = Self {
             props,
             pinned: HashSet::new(),
             implied: vec![],
             head_idxs: vec![],
+        };
+        // A normal ordering keeps a constant whose value differs between
+        // partitions (see `EquivalenceGroup::normalize_sort_exprs`), and such
+        // a constant may lead it. Within a partition it has one value, so it
+        // is fixed and the head is the first expression after it.
+        let normal = &props.oeq_cache.normal_cls;
+        if normal
+            .iter()
+            .any(|ordering| walk.is_fixed(&ordering[0].expr))
+        {
+            walk.head_idxs = normal

Review Comment:
   I closed this in discovery rather than in the walk: `add_orderings` and 
`discover_new_orderings` now look past leading constants. `SortKeyWalk::new` 
runs on every ordering check and its result is dropped when the walk ends, so a 
`ProjectionExec` above couldn't use it. With the ordering stored, a window on 
`c` over a TopK on `(k, a, c)` with `a` not selected no longer needs a sort at 
one partition.



##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:
##########
@@ -1318,10 +1318,28 @@ fn handle_projection_pushdown(
     // (expressed in the projection's output schema) must be remapped into the
     // child schema before being pushed down. Forwarding `score@1` unchanged
     // could refer to a different input column and sort on the wrong key.
-    Ok(
+    let Some(requirements) =
         remap_requirement_through_projection(projection_exec, parent_required)
-            .map(|requirements| vec![Some(requirements)]),
-    )
+    else {
+        return Ok(None);
+    };
+    // If the input already has this ordering, no sort will be added below the
+    // projection. The parent still needs the projection to prove it, which it
+    // may not be able to do: the proof can depend on a column that it drops,
+    // such as a leading per-partition constant. Keep the sort above the
+    // projection in that case.
+    if projection_exec

Review Comment:
   This can't happen with the default optimizer rules: the physical planner 
builds `HashJoinExec` without a projection, and `ProjectionPushdown`, which 
embeds one, runs after `EnsureRequirements`. I ran the shape, a window on `(g, 
c)` over a join whose right side is a TopK on `(g, a, c)` with `a` not 
selected, and it plans at 1 and 4 partitions. A custom rule order that embeds 
the projection earlier could hit this. If anyone has one, I'm happy to open an 
issue for it.



##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:
##########
@@ -1318,10 +1318,28 @@ fn handle_projection_pushdown(
     // (expressed in the projection's output schema) must be remapped into the
     // child schema before being pushed down. Forwarding `score@1` unchanged
     // could refer to a different input column and sort on the wrong key.
-    Ok(
+    let Some(requirements) =
         remap_requirement_through_projection(projection_exec, parent_required)
-            .map(|requirements| vec![Some(requirements)]),
-    )
+    else {
+        return Ok(None);
+    };
+    // If the input already has this ordering, no sort will be added below the
+    // projection. The parent still needs the projection to prove it, which it
+    // may not be able to do: the proof can depend on a column that it drops,
+    // such as a leading per-partition constant. Keep the sort above the
+    // projection in that case.
+    if projection_exec
+        .input()
+        .equivalence_properties()
+        .ordering_satisfy_requirement(requirements.first().clone())?
+        && !projection_exec

Review Comment:
   After the change for your comment on alternatives, this checks the 
alternative that gets pushed down, which isn't always the first one that 
`satisfy_parent` checked. Reusing it would also add a parameter to 
`pushdown_requirement_to_children` that only this branch reads. The check only 
runs when the input already has the ordering, so I'd keep it.



##########
datafusion/physical-expr/src/equivalence/properties/mod.rs:
##########


Review Comment:
   These aren't the same computation: `discover_new_orderings` derives from an 
ordering, and `add_orderings` then derives from what that ordering implied. The 
recursion is what finds chains. With `d = CAST(a)`, `a = CAST(c)` and ordering 
`[d, a, c]`, `[c]` is found with it and missed with one level, and 
`implied_orderings_chain` covers that. It only runs when something was implied.



-- 
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]

Reply via email to