diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index b48bc19a5517c..7f927cee3538a 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -67,6 +67,7 @@ use datafusion_common::HashMap as DFHashMap; use datafusion_common::display::ToStringifiedPlan; use datafusion_common::format::ExplainAnalyzeCategories; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion_common::utils::has_float_leaf; use datafusion_common::{ DFSchema, DFSchemaRef, ScalarValue, exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_err, @@ -1744,6 +1745,22 @@ impl DefaultPhysicalPlanner { planning_ctx, )?; + if join_key_blocks_sort_based_join( + &on_left, + &physical_left.schema(), + )? || join_key_blocks_sort_based_join( + &on_right, + &physical_right.schema(), + )? { + return Ok(Arc::new(NestedLoopJoinExec::try_new( + physical_left, + physical_right, + join_filter, + join_type, + None, + )?)); + } + Arc::new(PiecewiseMergeJoinExec::try_new( physical_left, physical_right, @@ -1766,6 +1783,11 @@ impl DefaultPhysicalPlanner { && session_state.config().repartition_joins() && !prefer_hash_join && !*null_aware + && !has_nested_float_join_keys( + &join_on, + &physical_left.schema(), + &physical_right.schema(), + )? // Null-aware joins (e.g. `NOT IN` with a nullable subquery) must // use the CollectLeft HashJoin below: SortMergeJoinExec does not // implement null-aware anti-join semantics and would return wrong @@ -2212,6 +2234,29 @@ fn get_null_physical_expr_pair( Ok((Arc::new(null_value), physical_name)) } +fn join_key_blocks_sort_based_join( + expr: &Arc, + schema: &Schema, +) -> Result { + let key_type = expr.data_type(schema)?; + Ok(key_type.is_nested() && has_float_leaf(&key_type)) +} + +fn has_nested_float_join_keys( + join_on: &[(Arc, Arc)], + left_schema: &Schema, + right_schema: &Schema, +) -> Result { + for (left_key, right_key) in join_on { + if join_key_blocks_sort_based_join(left_key, left_schema)? + || join_key_blocks_sort_based_join(right_key, right_schema)? + { + return Ok(true); + } + } + Ok(false) +} + /// Qualifies the fields in a join schema with "left" and "right" qualifiers /// without mutating the original schema. This function should only be used when /// the join inputs have already been requalified earlier in `try_new_with_project_input`. diff --git a/datafusion/sqllogictest/test_files/negative_zero.slt b/datafusion/sqllogictest/test_files/negative_zero.slt index fe4df57e179cc..55fe07c66228c 100644 --- a/datafusion/sqllogictest/test_files/negative_zero.slt +++ b/datafusion/sqllogictest/test_files/negative_zero.slt @@ -254,3 +254,37 @@ SELECT l.id, r.id FROM nested_zeros l JOIN nested_zeros r ON l.a = r.a ORDER BY statement ok DROP TABLE nested_zeros; + +##### +## Sort-based joins with multi-element nested float keys +##### + +statement ok +CREATE TABLE fa AS SELECT * FROM (VALUES + (1, make_array(arrow_cast(-0.0, 'Float64'), 1.0)), + (2, make_array(0.0, 0.0)), + (3, make_array(0.0, 1.0))) t(id, k); + +statement ok +CREATE TABLE fb AS SELECT * FROM (VALUES + (10, make_array(0.0, 1.0)), + (11, make_array(0.0, 0.0))) t(id, k); + +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +query II +SELECT fa.id, fb.id FROM fa JOIN fb ON fa.k = fb.k ORDER BY fa.id, fb.id; +---- +1 10 +2 11 +3 10 + +statement ok +reset datafusion.optimizer.prefer_hash_join; + +statement ok +DROP TABLE fa; + +statement ok +DROP TABLE fb;