fix: preserve equivalence join API compatibility

Author: Prajwal-k-tech

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

@@ -814,12 +814,40 @@ impl EquivalenceGroup {
 
     /// Combine equivalence groups of the given join children.
     pub fn join(
+        &self,
+        right_equivalences: &Self,
+        join_type: &JoinType,
+        left_size: usize,
+        on: &[(PhysicalExprRef, PhysicalExprRef)],
+    ) -> Result<Self> {
+        self.join_with_optional_schema(right_equivalences, join_type, left_size, None, on)
+    }
+
+    /// Combine equivalence groups for a join with its output schema available.
+    pub(crate) fn join_with_schema(
         &self,
         right_equivalences: &Self,
         join_type: &JoinType,
         left_size: usize,
         join_schema: &SchemaRef,
         on: &[(PhysicalExprRef, PhysicalExprRef)],
+    ) -> Result<Self> {
+        self.join_with_optional_schema(
+            right_equivalences,
+            join_type,
+            left_size,
+            Some(join_schema),
+            on,
+        )
+    }
+
+    fn join_with_optional_schema(
+        &self,
+        right_equivalences: &Self,
+        join_type: &JoinType,
+        left_size: usize,
+        join_schema: Option<&SchemaRef>,
+        on: &[(PhysicalExprRef, PhysicalExprRef)],
     ) -> Result<Self> {
         let group = match join_type {
             JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right => {
@@ -836,21 +864,25 @@ impl EquivalenceGroup {
                 // null-extended row; expressions such as coalesce(a, 0) do
                 // not, even when they were equivalent to a column before the
                 // join.
-                match join_type {
-                    JoinType::Left => {
-                        right = right
-                            .with_null_preserving_expressions(&null_batch(join_schema)?);
-                    }
-                    JoinType::Right => {
-                        left = left
-                            .with_null_preserving_expressions(&null_batch(join_schema)?);
+                if let Some(join_schema) = join_schema {
+                    match join_type {
+                        JoinType::Left => {
+                            right = right.with_null_preserving_expressions(&null_batch(
+                                join_schema,
+                            )?);
+                        }
+                        JoinType::Right => {
+                            left = left.with_null_preserving_expressions(&null_batch(
+                                join_schema,
+                            )?);
+                        }
+                        JoinType::Full => {
+                            let batch = null_batch(join_schema)?;
+                            left = left.with_null_preserving_expressions(&batch);
+                            right = right.with_null_preserving_expressions(&batch);
+                        }
+                        _ => {}
                     }
-                    JoinType::Full => {
-                        let batch = null_batch(join_schema)?;
-                        left = left.with_null_preserving_expressions(&batch);
-                        right = right.with_null_preserving_expressions(&batch);
-                    }
-                    _ => {}
                 }
 
                 let mut result =
@@ -1033,7 +1065,7 @@ mod tests {
         let left = EquivalenceGroup::new(vec![]);
         let right = EquivalenceGroup::new(vec![]);
 
-        let result = left.join(&right, &JoinType::Full, 0, &schema, &[])?;
+        let result = left.join_with_schema(&right, &JoinType::Full, 0, &schema, &[])?;
 
         assert_eq!(result.iter().count(), 0);
         Ok(())

datafusion/physical-expr/src/equivalence/properties/joins.rs

@@ -32,9 +32,13 @@ pub fn join_equivalence_properties(
     on: &[(PhysicalExprRef, PhysicalExprRef)],
 ) -> Result<EquivalenceProperties> {
     let left_size = left.schema.fields.len();
-    let equivalences =
-        left.eq_group()
-            .join(right.eq_group(), join_type, left_size, &join_schema, on)?;
+    let equivalences = left.eq_group().join_with_schema(
+        right.eq_group(),
+        join_type,
+        left_size,
+        &join_schema,
+        on,
+    )?;
     let mut result = EquivalenceProperties::new(join_schema);
     result.add_equivalence_group(equivalences)?;