fix: Preserve metadata when nl-join is swapped

datafusion/core/tests/sql/joins.rs

@@ -16,11 +16,14 @@
 // under the License.
 
 use insta::assert_snapshot;
+use std::collections::HashMap;
 
 use datafusion::assert_batches_eq;
 use datafusion::catalog::MemTable;
 use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable};
 use datafusion::test_util::register_unbounded_file_with_ordering;
+use datafusion_physical_plan::joins::{CrossJoinExec, NestedLoopJoinExec};
+use datafusion_physical_plan::test::TestMemoryExec;
 use datafusion_sql::unparser::plan_to_sql;
 
 use super::*;
@@ -299,3 +302,35 @@ async fn unparse_cross_join() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_swap_joins_on_conflicting_metadata() {
+    let input = |field: &str, meta_value: &str| {
+        let schema = Arc::new(
+            Schema::new(vec![Field::new(field, DataType::Int32, false)]).with_metadata(
+                HashMap::from([(String::from("metadata_key"), String::from(meta_value))]),
+            ),
+        );
+        TestMemoryExec::try_new_exec(&[vec![]], schema, None).unwrap()
+    };
+
+    let join = CrossJoinExec::new(input("a", "left value"), input("b", "right value"));
+
+    let swapped_join = join.swap_inputs().unwrap();
+
+    // The metadata of the join and the swapped version must be the same
+    assert_eq!(join.schema().metadata(), swapped_join.schema().metadata());
+
+    let join = NestedLoopJoinExec::try_new(
+        input("a", "left value"),
+        input("b", "right value"),
+        None,
+        &JoinType::Inner,
+        None,
+    )
+    .unwrap();
+
+    let swapped_join = join.swap_inputs().unwrap();
+
+    assert_eq!(join.schema().metadata(), swapped_join.schema().metadata());
+}

datafusion/physical-plan/src/joins/cross_join.rs

@@ -775,9 +775,7 @@ impl<T: BatchTransformer> CrossJoinStream<T> {
 mod tests {
     use super::*;
     use crate::common;
-    use crate::test::{TestMemoryExec, assert_join_metrics, build_table_scan_i32};
-    use arrow_schema::{DataType, Field};
-    use std::collections::HashMap;
+    use crate::test::{assert_join_metrics, build_table_scan_i32};
 
     use datafusion_common::{assert_contains, test_util::batches_to_sort_string};
     use datafusion_execution::runtime_env::RuntimeEnvBuilder;
@@ -1070,28 +1068,6 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn test_swapped_cross_join_schema_on_conflicting_metadata() {
-        let input = |field: &str, meta_value: &str| {
-            let schema = Arc::new(
-                Schema::new(vec![Field::new(field, DataType::Int32, false)])
-                    .with_metadata(HashMap::from([(
-                        String::from("metadata_key"),
-                        String::from(meta_value),
-                    )])),
-            );
-            TestMemoryExec::try_new_exec(&[vec![]], schema, None).unwrap()
-        };
-        // Conflicting metadata on left and right input, right side wins "metadata_key" -> "right value"
-        let join =
-            CrossJoinExec::new(input("a", "left value"), input("b", "right value"));
-
-        let swapped_join = join.swap_inputs().unwrap();
-
-        // The metadata of the cross-join and the swapped cross-join (with projection on top) must be the same
-        assert_eq!(join.schema().metadata(), swapped_join.schema().metadata());
-    }
-
     /// Returns the column names on the schema
     fn columns(schema: &Schema) -> Vec<String> {
         schema.fields().iter().map(|f| f.name().clone()).collect()

datafusion/physical-plan/src/joins/nested_loop_join.rs

@@ -17,15 +17,16 @@
 
 //! [`NestedLoopJoinExec`]: joins without equijoin (equality predicates).
 
+use std::collections::HashMap;
 use std::fmt::Formatter;
 use std::ops::{BitOr, ControlFlow};
 use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::task::Poll;
 
 use super::utils::{
-    asymmetric_join_output_partitioning, need_produce_result_in_final,
-    reorder_output_after_swap, swap_join_projection,
+    asymmetric_join_output_partitioning, build_join_schema_with_metadata,
+    need_produce_result_in_final, reorder_output_after_swap, swap_join_projection,
 };
 use crate::common::can_project;
 use crate::execution_plan::{EmissionType, boundedness_from_children};
@@ -234,6 +235,7 @@ pub struct NestedLoopJoinExecBuilder {
     join_type: JoinType,
     filter: Option<JoinFilter>,
     projection: Option<ProjectionRef>,
+    metadata: Option<HashMap<String, String>>,
 }
 
 impl NestedLoopJoinExecBuilder {
@@ -249,6 +251,7 @@ impl NestedLoopJoinExecBuilder {
             join_type,
             filter: None,
             projection: None,
+            metadata: None,
         }
     }
 
@@ -269,6 +272,12 @@ impl NestedLoopJoinExecBuilder {
         self
     }
 
+    /// Set metadata for the schema.
+    pub fn with_metadata(mut self, metadata: Option<HashMap<String, String>>) -> Self {
+        self.metadata = metadata;
+        self
+    }
+
     /// Build resulting execution plan.
     pub fn build(self) -> Result<NestedLoopJoinExec> {
         let Self {
@@ -277,13 +286,23 @@ impl NestedLoopJoinExecBuilder {
             join_type,
             filter,
             projection,
+            metadata,
         } = self;
 
         let left_schema = left.schema();
         let right_schema = right.schema();
         check_join_is_valid(&left_schema, &right_schema, &[])?;
-        let (join_schema, column_indices) =
-            build_join_schema(&left_schema, &right_schema, &join_type);
+
+        let (join_schema, column_indices) = if let Some(metadata) = metadata {
+            build_join_schema_with_metadata(
+                &left_schema,
+                &right_schema,
+                &join_type,
+                &metadata,
+            )
+        } else {
+            build_join_schema(&left_schema, &right_schema, &join_type)
+        };
         let join_schema = Arc::new(join_schema);
         let cache = NestedLoopJoinExec::compute_properties(
             &left,
@@ -316,6 +335,7 @@ impl From<&NestedLoopJoinExec> for NestedLoopJoinExecBuilder {
             join_type: exec.join_type,
             filter: exec.filter.clone(),
             projection: exec.projection.clone(),
+            metadata: None,
         }
     }
 }
@@ -453,18 +473,24 @@ impl NestedLoopJoinExec {
     pub fn swap_inputs(&self) -> Result<Arc<dyn ExecutionPlan>> {
         let left = self.left();
         let right = self.right();
-        let new_join = NestedLoopJoinExec::try_new(
+
+        let swapped_projection = swap_join_projection(
+            left.schema().fields().len(),
+            right.schema().fields().len(),
+            self.projection.as_deref(),
+            self.join_type(),
+        );
+
+        let new_join = NestedLoopJoinExecBuilder::new(
             Arc::clone(right),
             Arc::clone(left),
-            self.filter().map(JoinFilter::swap),
-            &self.join_type().swap(),
-            swap_join_projection(
-                left.schema().fields().len(),
-                right.schema().fields().len(),
-                self.projection.as_deref(),
-                self.join_type(),
-            ),
-        )?;
+            self.join_type().swap(),
+        )
+        .with_projection(swapped_projection)
+        .with_filter(self.filter().map(JoinFilter::swap))
+        // Preserve existing metadata
+        .with_metadata(Some(self.join_schema.metadata.clone()))
+        .build()?;
 
         // For Semi/Anti joins, swap result will produce same output schema,
         // no need to wrap them into additional projection

datafusion/physical-plan/src/joins/utils.rs

@@ -18,7 +18,7 @@
 //! Join related functionality used both on logical and physical plans
 
 use std::cmp::{Ordering, min};
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{self, Debug};
 use std::future::Future;
 use std::iter::once;
@@ -265,12 +265,13 @@ fn output_join_field(old_field: &Field, join_type: &JoinType, is_left: bool) ->
     }
 }
 
-/// Creates a schema for a join operation.
+/// Creates a schema for a join operation and use existing metadata.
 /// The fields from the left side are first
-pub fn build_join_schema(
+pub fn build_join_schema_with_metadata(
     left: &Schema,
     right: &Schema,
     join_type: &JoinType,
+    metadata: &HashMap<String, String>,
 ) -> (Schema, Vec<ColumnIndex>) {
     let left_fields = || {
         left.fields()
@@ -334,6 +335,19 @@ pub fn build_join_schema(
         }
     };
 
+    (
+        fields.finish().with_metadata(metadata.clone()),
+        column_indices,
+    )
+}
+
+/// Creates a schema for a join operation.
+/// The fields from the left side are first
+pub fn build_join_schema(
+    left: &Schema,
+    right: &Schema,
+    join_type: &JoinType,
+) -> (Schema, Vec<ColumnIndex>) {
     let (schema1, schema2) = match join_type {
         JoinType::Right
         | JoinType::RightSemi
@@ -349,7 +363,7 @@ pub fn build_join_schema(
         .chain(schema2.metadata().clone())
         .collect();
 
-    (fields.finish().with_metadata(metadata), column_indices)
+    build_join_schema_with_metadata(left, right, join_type, &metadata)
 }
 
 /// A [`OnceAsync`] runs an `async` closure once, where multiple calls to