diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 7c0e89ee96418..e903b8af48a5e 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/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()); +} diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 1a631aac980ab..0c839d9dba3ea 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -775,9 +775,7 @@ impl CrossJoinStream { 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 { schema.fields().iter().map(|f| f.name().clone()).collect() diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 515dcc2931c05..0240540c68fc7 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -17,6 +17,7 @@ //! [`NestedLoopJoinExec`]: joins without equijoin (equality predicates). +use std::collections::HashMap; use std::fmt::Formatter; use std::ops::{BitOr, ControlFlow}; use std::sync::Arc; @@ -24,8 +25,8 @@ 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, projection: Option, + metadata: Option>, } 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>) -> Self { + self.metadata = metadata; + self + } + /// Build resulting execution plan. pub fn build(self) -> Result { 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> { 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 diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 654d873ae1b0e..547a8934fe269 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/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, ) -> (Schema, Vec) { 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) { 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