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


##########
datafusion/physical-expr/src/equivalence/class.rs:
##########
@@ -204,6 +204,34 @@ impl EquivalenceClass {
         self.exprs.insert(expr);
     }
 
+    /// Makes sure the canonical expression does not contain another expression
+    /// of this class. Otherwise, with `a + b` as the canonical expression of a
+    /// class that also holds `a`, normalization would replace the `a` inside
+    /// `a + b` and turn `a + b` into `(a + b) + b`.
+    fn promote_sub_expressions(&mut self) {

Review Comment:
   @abp002 thanks for the fix! Classes built through `EquivalenceGroup::new` or 
merged by `EquivalenceGroup::extend` (via `bridge_classes`) skip this 
promotion, so `EquivalenceGroup::new([EquivalenceClass::new([a + b, a])])` 
still normalizes `a` to `(a + b) + b`. Both are public. Could you run the 
promotion inside `EquivalenceClass::push` and `extend` instead? Then every way 
of building a class gets it, and the four calls in `add_equal_conditions` go 
away:
   
   ```diff
            self.exprs.insert(expr);
   +        self.promote_sub_expressions();
        }
   @@
        pub fn extend(&mut self, other: Self) {
            self.exprs.extend(other.exprs);
   +        self.promote_sub_expressions();
   @@
                    self.classes[first_idx].extend(other_class);
   -                self.classes[first_idx].promote_sub_expressions();
   @@
                    self.classes[group_idx].push(right);
   -                self.classes[group_idx].promote_sub_expressions();
   @@
                    self.classes[group_idx].push(left);
   -                self.classes[group_idx].promote_sub_expressions();
   @@
   -                let mut class = EquivalenceClass::new([left, right]);
   -                class.promote_sub_expressions();
   +                let class = EquivalenceClass::new([left, right]);
   ```
   
   With this change on current main, the physical-expr unit tests and the full 
sqllogictest suite pass.
   
   <details>
   <summary>Unit test for both paths (fails on the current PR, passes with the 
change)</summary>
   
   ```rs
       #[test]
       fn test_sub_expression_is_canonical_on_every_path() -> Result<()> {
           let schema = abc_schema();
           let (a, b, c) = (col("a", &schema)?, col("b", &schema)?, col("c", 
&schema)?);
           let a_plus_b = binary(Arc::clone(&a), Operator::Plus, b, &schema)?;
   
           // A class built directly, with the larger expression first
           let group = EquivalenceGroup::new([EquivalenceClass::new([
               Arc::clone(&a_plus_b),
               Arc::clone(&a),
           ])]);
           assert!(group.normalize_expr(Arc::clone(&a_plus_b)).eq(&a));
   
           // Two classes bridged by `extend` through the shared `c`
           let mut group = EquivalenceGroup::default();
           group.add_equal_conditions(Arc::clone(&a_plus_b), Arc::clone(&c));
           let mut other = EquivalenceGroup::default();
           other.add_equal_conditions(c, Arc::clone(&a));
           group.extend(other);
           assert!(group.normalize_expr(a_plus_b).eq(&a));
           Ok(())
       }
   ```
   
   </details>



##########
datafusion/sqllogictest/test_files/equivalence_properties.slt:
##########
@@ -204,6 +204,49 @@ SELECT count(*) FROM t_double WHERE a = sqrt(b);
 statement ok
 DROP TABLE t_double;
 
+###
+### Equalities
+###
+
+statement ok
+CREATE TABLE eq_t (a BIGINT, b BIGINT, c BIGINT) AS VALUES (1, 1, 2), (3, 3, 
4), (5, 5, 0), (2, 2, 3);
+
+# The filter puts a and b in one class, and c and a + 1 in another. The inner
+# ORDER BY a + 1 LIMIT 100 sorts by a member of the second class, so the outer
+# ORDER BY a + 1 is already satisfied and the plan has no second sort.
+query TT
+EXPLAIN SELECT * FROM (SELECT * FROM eq_t ORDER BY a + 1 LIMIT 100) WHERE a = 
b AND c = a + 1 ORDER BY a + 1;
+----
+physical_plan
+01)FilterExec: b@1 = a@0 AND c@2 = a@0 + 1
+02)--SortExec: TopK(fetch=100), expr=[a@0 + 1 ASC NULLS LAST], 
preserve_partitioning=[false]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+
+query III
+SELECT * FROM (SELECT * FROM eq_t ORDER BY a + 1 LIMIT 100) WHERE a = b AND c 
= a + 1 ORDER BY a + 1;
+----
+1 1 2
+2 2 3
+3 3 4
+
+statement ok
+DROP TABLE eq_t;
+
+# The two IN lists hold the same values in a different order, so they compare
+# equal but are kept as two members of the class of b. Planning must still
+# finish.
+statement ok
+CREATE TABLE eq_in (a INT, b BOOLEAN) AS VALUES (1, true), (9, false);
+
+query IB rowsort
+SELECT * FROM eq_in WHERE (a IN (1, 2, 3, 4)) = b AND (a IN (4, 3, 2, 1)) = b;
+----
+1 true
+9 false
+
+statement ok
+DROP TABLE eq_in;

Review Comment:
   @abp002 the same bug also breaks a window on main, without a sort-merge 
join: `PARTITION BY a + b ORDER BY a` under `WHERE a + b = a` fails at runtime 
with `Expects PARTITION BY expression to be ordered`. This PR fixes it, but 
nothing tests that path. Could you add it?
   
   ```suggestion
   statement ok
   DROP TABLE eq_in;
   
   # With a + b = a, the window's PARTITION BY a + b is the same key as the
   # ORDER BY a below it
   statement ok
   CREATE TABLE eq_sub (a INT, b INT) AS VALUES (1, 0), (2, 3), (1, 0), (5, 0);
   
   query II
   SELECT a, sum(b) OVER (PARTITION BY a + b ORDER BY a) FROM eq_sub WHERE a + 
b = a ORDER BY a;
   ----
   1 0
   1 0
   5 0
   
   statement ok
   DROP TABLE eq_sub;
   ```



##########
datafusion/sqllogictest/test_files/sort_merge_join.slt:
##########
@@ -1129,6 +1129,20 @@ DROP TABLE smj_key_r;
 statement ok
 RESET datafusion.execution.batch_size;
 
+# Regression test for #26623: a filter that equates an expression with one of
+# its own sub-expressions must not make the sort below the join fail the
+# sanity check
+statement ok
+CREATE TABLE smj_self_eq (a INT, b INT) AS VALUES (1, 0), (2, 3);
+
+query IIII
+SELECT * FROM smj_self_eq x JOIN smj_self_eq y ON x.a = y.a AND x.b = y.b 
WHERE x.a + x.b = x.a;
+----
+1 0 1 0
+
+statement ok
+DROP TABLE smj_self_eq;

Review Comment:
   Removing the promotion after the merge (`extend`) or after `push(right)` in 
`add_equal_conditions` still passes the tests here, because the issue query 
only reaches the new-class arm. These two queries reach the other arms. Both 
fail with `SanityCheckPlan` on main, and each catches one of those removals:
   
   ```suggestion
   DROP TABLE smj_self_eq;
   
   statement ok
   CREATE TABLE smj_self_eq_cd (a INT, b INT, c INT, d INT) AS VALUES (1, 0, 1, 
1), (2, 3, 5, 2);
   
   # a joins the class of a + b after the class exists
   query IIII
   SELECT x.a, x.b, y.a, y.b FROM smj_self_eq_cd x JOIN smj_self_eq_cd y ON x.a 
= y.a AND x.b = y.b WHERE x.a + x.b = x.c AND x.c = x.a;
   ----
   1 0 1 0
   
   # a reaches the class of a + b when two classes merge
   query IIII
   SELECT x.a, x.b, y.a, y.b FROM smj_self_eq_cd x JOIN smj_self_eq_cd y ON x.a 
= y.a AND x.b = y.b WHERE x.a + x.b = x.c AND x.a = x.d AND x.d = x.c;
   ----
   1 0 1 0
   
   statement ok
   DROP TABLE smj_self_eq_cd;
   ```



##########
datafusion/sqllogictest/test_files/equivalence_properties.slt:
##########
@@ -204,6 +204,49 @@ SELECT count(*) FROM t_double WHERE a = sqrt(b);
 statement ok
 DROP TABLE t_double;
 
+###
+### Equalities
+###
+
+statement ok
+CREATE TABLE eq_t (a BIGINT, b BIGINT, c BIGINT) AS VALUES (1, 1, 2), (3, 3, 
4), (5, 5, 0), (2, 2, 3);
+
+# The filter puts a and b in one class, and c and a + 1 in another. The inner
+# ORDER BY a + 1 LIMIT 100 sorts by a member of the second class, so the outer
+# ORDER BY a + 1 is already satisfied and the plan has no second sort.
+query TT
+EXPLAIN SELECT * FROM (SELECT * FROM eq_t ORDER BY a + 1 LIMIT 100) WHERE a = 
b AND c = a + 1 ORDER BY a + 1;
+----
+physical_plan
+01)FilterExec: b@1 = a@0 AND c@2 = a@0 + 1
+02)--SortExec: TopK(fetch=100), expr=[a@0 + 1 ASC NULLS LAST], 
preserve_partitioning=[false]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+
+query III
+SELECT * FROM (SELECT * FROM eq_t ORDER BY a + 1 LIMIT 100) WHERE a = b AND c 
= a + 1 ORDER BY a + 1;
+----
+1 1 2
+2 2 3
+3 3 4
+
+statement ok
+DROP TABLE eq_t;
+
+# The two IN lists hold the same values in a different order, so they compare

Review Comment:
   Optional cleanup: this test doesn't reach the loop bound in 
`promote_sub_expressions`. With the bound replaced by `loop`, the query still 
finishes, because neither `IN` list is inside the other. A promoted expression 
is always strictly inside the previous canonical expression, so it can't cycle. 
I'd drop this test, or change its comment so it doesn't claim to cover 
termination.



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