jayzhan211 commented on code in PR #27347:
URL: https://github.com/apache/datafusion/pull/27347#discussion_r4551779479
##########
datafusion/physical-expr/src/equivalence/properties/union.rs:
##########
@@ -104,6 +102,22 @@ pub fn calculate_union(
groupings.add_groupings(std::iter::once(grouping.iter().cloned()));
}
+ // An equality holds across the union only if it holds in every input
partition.
Review Comment:
@galaxytemple thanks for picking this up! Because the equalities are added
after the fold, each `calculate_union_binary` step still sees an accumulator
without `x = y`. So with 3 arms where the arm lacking the equality comes last,
the same SanityCheckPlan error remains:
```sql
CREATE TABLE t(a INT, b INT) AS VALUES (2, 20), (1, 30);
SELECT a AS x, a AS y FROM t
UNION ALL SELECT a AS x, a AS y FROM t
UNION ALL SELECT a AS x, b AS y FROM t
ORDER BY x, y;
-- SanityCheckPlan ... Child-0 order: [[x@0 ASC NULLS LAST]]
```
Computing the meet inside `calculate_union_binary` fixes it and lets you
drop the n-ary block, the emptiness guard and the `unwrap`. With this change
the physical-expr tests and the full sqllogictest suite pass, and the new
`union.slt` case below fails without it:
<details>
<summary>Diff for union.rs and union.slt</summary>
```diff
@@ fn calculate_union_binary(
let orderings = orderings.build();
+ // An equality holds at the output only if it holds in both inputs. Keep
+ // it in the result so the next fold step can still use it.
+ let mut eq_group = EquivalenceGroup::default();
+ for class in lhs.eq_group().iter() {
+ for (i, left) in class.iter().enumerate() {
+ for right in class.iter().skip(i + 1) {
+ if rhs.eq_group().exprs_equal(left, right) {
+ eq_group.add_equal_conditions(Arc::clone(left),
Arc::clone(right));
+ }
+ }
+ }
+ }
+
let mut eq_properties = EquivalenceProperties::new(lhs.schema);
+ eq_properties.add_equivalence_group(eq_group)?;
eq_properties.add_constants(constants)?;
@@ pub fn calculate_union(
- // An equality holds across the union only if it holds in every input
partition.
- let mut eq_group = EquivalenceGroup::default();
- if !eqps.is_empty() && eqps.iter().all(|p| !p.eq_group().is_empty()) {
- ...
- }
-
@@
- acc.add_equivalence_group(eq_group)?;
acc.add_groupings(groupings);
```
```diff
--- a/datafusion/sqllogictest/test_files/union.slt
+++ b/datafusion/sqllogictest/test_files/union.slt
@@ -1170,6 +1170,22 @@ ORDER BY x, y;
2 2
2 2
+# The arm without x = y is folded last, after the arms that have it.
+query II
+SELECT a AS x, a AS y FROM t_26625
+UNION ALL
+SELECT a AS x, a AS y FROM t_26625
+UNION ALL
+SELECT a AS x, (a % 2)::INT AS y FROM t_26625
+ORDER BY x, y;
+----
+1 1
+1 1
+1 1
+2 0
+2 2
+2 2
+
statement ok
DROP TABLE t_26625;
```
The `::INT` matters: without it, `a % 2` is coerced to Int64 and the first
two arms project `CAST(a AS Int64)`, so `x = y` disappears from them and the
test passes without the fix.
</details>
--
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]
