fix: preserve compound expression equivalences through projection

datafusion/physical-expr/src/equivalence/class.rs

@@ -624,8 +624,8 @@ impl EquivalenceGroup {
             .filter(|req| !self.is_uniform_constant(&req.expr))
     }
 
-    /// Perform an indirect projection of `expr` by consulting the equivalence
-    /// classes.
+    /// Projects `expr` through its children or an equivalent expression in the
+    /// mapping, preferring to preserve its structure.
     fn project_expr_indirect(
         aug_mapping: &AugmentedMapping,
         expr: &Arc<dyn PhysicalExpr>,
@@ -635,8 +635,32 @@ impl EquivalenceGroup {
             return Some(Arc::clone(expr));
         }
 
-        // The given expression is not inside the mapping, so we try to project
-        // indirectly using equivalence classes.
+        // Preserve expression structure when its children can be projected.
+        // Otherwise, projecting a class such as a + 1 = b can collapse it to
+        // b = b, losing the equivalence even when a survives the projection.
+        let children = expr.children();
+        if !children.is_empty() {
+            let projected_children = children
+                .into_iter()
+                .map(|child| {
+                    // First, we try to project children with an exact match. If
+                    // we are unable to do this, we consult equivalence classes.
+                    if let Some((targets, _)) = aug_mapping.get(child) {
+                        // If we match the source, we can project directly:
+                        let (target, _) = targets.first();
+                        Some(Arc::clone(target))
+                    } else {
+                        Self::project_expr_indirect(aug_mapping, child)
+                    }
+                })
+                .collect::<Option<Vec<_>>>();
+            if let Some(children) = projected_children {
+                return Some(Arc::clone(expr).with_new_children(children).unwrap());
+            }
+        }
+
+        // Project expressions whose children could not be projected by consulting
+        // equivalence classes.
         for (targets, eq_class) in aug_mapping.values() {
             // If we match an equivalent expression to a source expression in
             // the mapping, then we can project. For example, if we have the
@@ -647,27 +671,7 @@ impl EquivalenceGroup {
                 return Some(Arc::clone(target));
             }
         }
-        // Project a non-leaf expression by projecting its children.
-        let children = expr.children();
-        if children.is_empty() {
-            // A leaf expression should be inside the mapping.
-            return None;
-        }
-        children
-            .into_iter()
-            .map(|child| {
-                // First, we try to project children with an exact match. If
-                // we are unable to do this, we consult equivalence classes.
-                if let Some((targets, _)) = aug_mapping.get(child) {
-                    // If we match the source, we can project directly:
-                    let (target, _) = targets.first();
-                    Some(Arc::clone(target))
-                } else {
-                    Self::project_expr_indirect(aug_mapping, child)
-                }
-            })
-            .collect::<Option<Vec<_>>>()
-            .map(|children| Arc::clone(expr).with_new_children(children).unwrap())
+        None
     }
 
     fn augment_projection_mapping<'a>(

datafusion/sqllogictest/test_files/joins.slt

@@ -6299,3 +6299,62 @@ reset datafusion.optimizer.repartition_file_scans;
 
 statement ok
 set datafusion.execution.target_partitions = 4;
+
+# Projecting away a temporary join-filter column must preserve the equivalence
+# between a struct field and its extracted scalar column. The parallel sort and
+# merge can use different members of that equivalence class.
+statement ok
+set datafusion.execution.target_partitions = 2;
+
+statement ok
+CREATE TABLE projection_events (id BIGINT) AS VALUES (1), (2), (3);
+
+# UNION ALL gives the right input two partitions, so ordering needs a merge.
+statement ok
+CREATE TABLE projection_networks AS
+SELECT 'b' AS label UNION ALL SELECT 'a' AS label;
+
+query TT
+EXPLAIN SELECT * FROM projection_networks;
+----
+logical_plan TableScan: projection_networks projection=[label]
+physical_plan DataSourceExec: partitions=2, partition_sizes=[1, 1]
+
+query I?
+SELECT id, n
+FROM projection_events
+LEFT JOIN (
+  SELECT named_struct('label', label) AS n FROM projection_networks
+) ON id > 1
+ORDER BY id, n['label'];
+----
+1 NULL
+2 {label: a}
+2 {label: b}
+3 {label: a}
+3 {label: b}
+
+# Also exercise the equivalence with an explicitly projected scalar alias.
+query I?
+SELECT id, n
+FROM projection_events
+LEFT JOIN (
+  SELECT named_struct('label', label) AS n, label AS extracted
+  FROM projection_networks
+) ON id > 1
+ORDER BY id, extracted;
+----
+1 NULL
+2 {label: a}
+2 {label: b}
+3 {label: a}
+3 {label: b}
+
+statement ok
+DROP TABLE projection_events;
+
+statement ok
+DROP TABLE projection_networks;
+
+statement ok
+set datafusion.execution.target_partitions = 4;