From 2d55b49bb09c0623c33feb033268eeb81057fb6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 20 Aug 2026 12:42:13 +0200 Subject: [PATCH 1/3] feat: add projection support to SortMergeJoinExec A sort merge join emitted every column it joined, so a query selecting a subset needed a ProjectionExec above it, and one joining on an expression kept that expression in the output. Projection pushdown could only push into the children when each side's columns stayed together, and gave up otherwise. SortMergeJoinExec now carries an optional projection, as HashJoinExec does, and projection pushdown embeds into the join whatever it cannot push into the children. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_012YiAABcW4WSqij31zz2P6c --- .../physical_optimizer/projection_pushdown.rs | 47 ++++- .../src/joins/sort_merge_join/exec.rs | 169 ++++++++++++++---- .../proto-models/proto/datafusion.proto | 1 + .../proto-models/src/generated/pbjson.rs | 20 +++ .../proto-models/src/generated/prost.rs | 2 + datafusion/proto/tests/cases/plans/joins.rs | 44 +++++ datafusion/sqllogictest/test_files/joins.slt | 19 +- .../test_files/range_partitioning.slt | 26 ++- .../test_files/sort_merge_join_spill.slt | 2 +- 9 files changed, 269 insertions(+), 61 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 113552c462f76..4e8d47974057b 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -49,8 +49,8 @@ use datafusion_physical_plan::coop::CooperativeExec; use datafusion_physical_plan::filter::{FilterExec, FilterExecBuilder}; use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion_physical_plan::joins::{ - HashJoinExec, NestedLoopJoinExec, PartitionMode, StreamJoinPartitionMode, - SymmetricHashJoinExec, + HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, + StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr, update_expr}; use datafusion_physical_plan::repartition::RepartitionExec; @@ -1875,3 +1875,46 @@ fn test_filter_with_embedded_projection_after_renaming_projection() -> Result<() Ok(()) } + +#[test] +fn test_sort_merge_join_interleaved_projection_embeds() -> Result<()> { + // SELECT t1.c, t2.c, t1.a FROM t1 JOIN t2 ON t1.b = t2.c + // Taking a column from each side in turn leaves the sides interleaved, which + // cannot be pushed into the children, so the join applies it itself. + let join = Arc::new(SortMergeJoinExec::try_new( + create_simple_csv_exec(), + create_simple_csv_exec(), + vec![(Arc::new(Column::new("b", 1)), Arc::new(Column::new("c", 2)))], + None, + JoinType::Inner, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + )?); + let projection: Arc = Arc::new(ProjectionExec::try_new( + vec![ + ProjectionExpr::new(Arc::new(Column::new("c", 2)), "c_from_left"), + ProjectionExpr::new(Arc::new(Column::new("c", 7)), "c_from_right"), + ProjectionExpr::new(Arc::new(Column::new("a", 0)), "a_from_left"), + ], + join, + )?); + + let after_optimize = + ProjectionPushdown::new().optimize(projection, &ConfigOptions::new())?; + let actual = displayable(after_optimize.as_ref()) + .indent(true) + .to_string() + .trim() + .to_string(); + assert_snapshot!( + actual, + @r" + ProjectionExec: expr=[c@0 as c_from_left, c@1 as c_from_right, a@2 as a_from_left] + SortMergeJoinExec: join_type=Inner, on=[(b@1, c@2)], projection=[c@2, c@7, a@0] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false + " + ); + + Ok(()) +} diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index dc8540fe6df4d..eeb3979b6f1d5 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -34,19 +34,22 @@ use crate::joins::utils::{ }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet, SpillMetrics}; use crate::projection::{ - ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children, - physical_to_column_exprs, update_join_on, + EmbeddedProjection, ProjectionExec, join_allows_pushdown, join_table_borders, + new_join_children, physical_to_column_exprs, try_embed_projection, update_join_on, }; use crate::spill::spill_manager::SpillManager; use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::RecordBatchStreamAdapter; use crate::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, - ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, validate_child_count, + ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, common::can_project, + validate_child_count, }; use arrow::compute::SortOptions; use arrow::datatypes::SchemaRef; +use datafusion_common::project_schema; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err, internal_err, @@ -54,9 +57,12 @@ use datafusion_common::{ }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryConsumer; -use datafusion_physical_expr::equivalence::join_equivalence_properties; +use datafusion_physical_expr::equivalence::{ + ProjectionMapping, join_equivalence_properties, +}; use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql}; use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; +use futures::StreamExt; /// Join execution plan that executes equi-join predicates on multiple partitions using Sort-Merge /// join algorithm and applies an optional filter post join. Can be used to join arbitrarily large @@ -128,6 +134,8 @@ pub struct SortMergeJoinExec { pub sort_options: Vec, /// Defines the null equality for the join. pub null_equality: NullEquality, + /// The columns of `schema` to emit, in order. `None` emits all of them. + pub projection: Option>, /// Cache holding plan properties like equivalences, output partitioning etc. cache: Arc, } @@ -187,7 +195,7 @@ impl SortMergeJoinExec { let schema = Arc::new(build_join_schema(&left_schema, &right_schema, &join_type).0); let cache = - Self::compute_properties(&left, &right, Arc::clone(&schema), join_type, &on)?; + Self::compute_properties(&left, &right, &schema, join_type, &on, None)?; Ok(Self { left, right, @@ -200,10 +208,36 @@ impl SortMergeJoinExec { right_sort_exprs, sort_options, null_equality, + projection: None, + cache: Arc::new(cache), + }) + } + + /// Returns this join emitting only the columns in `projection`, in that order. + /// The indices address the join's own schema, before any projection. + pub fn with_projection(&self, projection: Option>) -> Result { + can_project(&self.schema, projection.as_deref())?; + let cache = Self::compute_properties( + &self.left, + &self.right, + &self.schema, + self.join_type, + &self.on, + projection.as_deref(), + )?; + Ok(Self { + projection, + metrics: ExecutionPlanMetricsSet::new(), cache: Arc::new(cache), + ..Self::clone(self) }) } + /// Whether the join emits fewer or reordered columns than it joins. + pub fn contains_projection(&self) -> bool { + self.projection.is_some() + } + /// Get probe side (e.g streaming side) information for this sort merge join. /// In current implementation, probe side is determined according to join type. pub fn probe_side(join_type: &JoinType) -> JoinSide { @@ -281,24 +315,32 @@ impl SortMergeJoinExec { fn compute_properties( left: &Arc, right: &Arc, - schema: SchemaRef, + schema: &SchemaRef, join_type: JoinType, join_on: JoinOnRef, + projection: Option<&[usize]>, ) -> Result { // Calculate equivalence properties: - let eq_properties = join_equivalence_properties( + let mut eq_properties = join_equivalence_properties( left.equivalence_properties().clone(), right.equivalence_properties().clone(), &join_type, - schema, + Arc::clone(schema), &Self::maintains_input_order(join_type), Some(Self::probe_side(&join_type)), join_on, )?; - let output_partitioning = + let mut output_partitioning = symmetric_join_output_partitioning(left, right, &join_type)?; + if let Some(projection) = projection { + let mapping = ProjectionMapping::from_indices(projection, schema)?; + let projected = project_schema(schema, Some(projection))?; + output_partitioning = output_partitioning.project(&mapping, &eq_properties); + eq_properties = eq_properties.project(&mapping, projected); + } + Ok(PlanProperties::new( eq_properties, output_partitioning, @@ -346,6 +388,12 @@ impl SortMergeJoinExec { } } +impl EmbeddedProjection for SortMergeJoinExec { + fn with_projection(&self, projection: Option>) -> Result { + self.with_projection(projection) + } +} + impl DisplayAs for SortMergeJoinExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { match t { @@ -362,9 +410,27 @@ impl DisplayAs for SortMergeJoinExec { } else { "" }; + let display_projections = if self.contains_projection() { + format!( + ", projection=[{}]", + self.projection + .as_ref() + .unwrap() + .iter() + .map(|index| format!( + "{}@{}", + self.schema.fields().get(*index).unwrap().name(), + index + )) + .collect::>() + .join(", ") + ) + } else { + "".to_string() + }; write!( f, - "{}: join_type={:?}, on=[{}]{}{}", + "{}: join_type={:?}, on=[{}]{}{}{}", Self::static_name(), self.join_type, on, @@ -373,6 +439,7 @@ impl DisplayAs for SortMergeJoinExec { |f| format!(", filter={}", f.expression()) ), display_null_equality, + display_projections, ) } DisplayFormatType::TreeRender => { @@ -467,15 +534,18 @@ impl ExecutionPlan for SortMergeJoinExec { })) } ChildrenPropertiesMode::Recompute => match &children[..] { - [left, right] => Ok(Arc::new(SortMergeJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - self.on.clone(), - self.filter.clone(), - self.join_type, - self.sort_options.clone(), - self.null_equality, - )?)), + [left, right] => Ok(Arc::new( + SortMergeJoinExec::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.filter.clone(), + self.join_type, + self.sort_options.clone(), + self.null_equality, + )? + .with_projection(self.projection.clone())?, + )), _ => internal_err!("SortMergeJoin wrong number of children"), }, } @@ -546,7 +616,7 @@ impl ExecutionPlan for SortMergeJoinExec { ) .with_compression_type(context.session_config().spill_compression()); - if matches!( + let joined = if matches!( self.join_type, JoinType::LeftSemi | JoinType::LeftAnti @@ -589,7 +659,16 @@ impl ExecutionPlan for SortMergeJoinExec { spill_manager, context.runtime_env(), ) - } + }?; + + let Some(projection) = self.projection.clone() else { + return Ok(joined); + }; + let schema = self.schema(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + joined.map(move |batch| Ok(batch?.project(&projection)?)), + ))) } fn metrics(&self) -> Option { @@ -614,14 +693,18 @@ impl ExecutionPlan for SortMergeJoinExec { // - `A LEFT JOIN B ON A.col=B.col` with `COUNT_DISTINCT(B.col)=COUNT(B.col)` let left_stats = input_stats[0].as_ref().clone(); let right_stats = input_stats[1].as_ref().clone(); - Ok(Arc::new(estimate_join_statistics( + let stats = estimate_join_statistics( left_stats, right_stats, &self.on, self.null_equality, &self.join_type, &self.schema, - )?)) + )?; + Ok(Arc::new(match &self.projection { + Some(projection) => stats.project(Some(projection)), + None => stats, + })) } /// Tries to swap the projection with its input [`SortMergeJoinExec`]. If it can be done, @@ -631,6 +714,9 @@ impl ExecutionPlan for SortMergeJoinExec { &self, projection: &ProjectionExec, ) -> Result>> { + if self.contains_projection() { + return Ok(None); + } // Convert projected PhysicalExpr's to columns. If not possible, we cannot proceed. let Some(projection_as_columns) = physical_to_column_exprs(projection.expr()) else { @@ -642,13 +728,15 @@ impl ExecutionPlan for SortMergeJoinExec { &projection_as_columns, ); + // Pushing into the children needs each side's columns to stay together, which + // an arbitrary projection does not. The join can apply that one itself. if !join_allows_pushdown( &projection_as_columns, &self.schema(), far_right_left_col_ind, far_left_right_col_ind, ) { - return Ok(None); + return try_embed_projection(projection, self); } let Some(new_on) = update_join_on( @@ -657,7 +745,7 @@ impl ExecutionPlan for SortMergeJoinExec { self.on(), self.left().schema().fields().len(), ) else { - return Ok(None); + return try_embed_projection(projection, self); }; let (new_left, new_right) = new_join_children( @@ -696,6 +784,7 @@ impl ExecutionPlan for SortMergeJoinExec { join_type, sort_options, null_equality, + projection, // derived from the children's schemas by `try_new` on decode schema: _, // runtime metrics, not part of the plan @@ -746,6 +835,11 @@ impl ExecutionPlan for SortMergeJoinExec { filter, sort_options, null_equality: null_equality.into(), + projection: projection + .iter() + .flatten() + .map(|index| *index as u32) + .collect(), }, )), ), @@ -781,6 +875,7 @@ impl SortMergeJoinExec { filter, sort_options, null_equality, + projection, } = &**sort_join; let left = @@ -832,14 +927,20 @@ impl SortMergeJoinExec { }) .collect(); - Ok(Arc::new(Self::try_new( - left, - right, - on, - filter, - join_type, - sort_options, - null_equality, - )?)) + let projection = (!projection.is_empty()) + .then(|| projection.iter().map(|index| *index as usize).collect()); + + Ok(Arc::new( + Self::try_new( + left, + right, + on, + filter, + join_type, + sort_options, + null_equality, + )? + .with_projection(projection)?, + )) } } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index d98b67a66e0a9..7b8f8e398539c 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1691,6 +1691,7 @@ message SortMergeJoinExecNode { JoinFilter filter = 5; repeated SortExprNode sort_options = 6; datafusion_common.NullEquality null_equality = 7; + repeated uint32 projection = 8; } message AsyncFuncExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 21309bb2d0941..bc88f5c9c2f45 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -25603,6 +25603,9 @@ impl serde::Serialize for SortMergeJoinExecNode { if self.null_equality != 0 { len += 1; } + if !self.projection.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.SortMergeJoinExecNode", len)?; if let Some(v) = self.left.as_ref() { struct_ser.serialize_field("left", v)?; @@ -25629,6 +25632,9 @@ impl serde::Serialize for SortMergeJoinExecNode { .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.null_equality)))?; struct_ser.serialize_field("nullEquality", &v)?; } + if !self.projection.is_empty() { + struct_ser.serialize_field("projection", &self.projection)?; + } struct_ser.end() } } @@ -25649,6 +25655,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { "sortOptions", "null_equality", "nullEquality", + "projection", ]; #[allow(clippy::enum_variant_names)] @@ -25660,6 +25667,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { Filter, SortOptions, NullEquality, + Projection, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -25688,6 +25696,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { "filter" => Ok(GeneratedField::Filter), "sortOptions" | "sort_options" => Ok(GeneratedField::SortOptions), "nullEquality" | "null_equality" => Ok(GeneratedField::NullEquality), + "projection" => Ok(GeneratedField::Projection), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -25714,6 +25723,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { let mut filter__ = None; let mut sort_options__ = None; let mut null_equality__ = None; + let mut projection__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Left => { @@ -25758,6 +25768,15 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { } null_equality__ = Some(map_.next_value::()? as i32); } + GeneratedField::Projection => { + if projection__.is_some() { + return Err(serde::de::Error::duplicate_field("projection")); + } + projection__ = + Some(map_.next_value::>>()? + .into_iter().map(|x| x.0).collect()) + ; + } } } Ok(SortMergeJoinExecNode { @@ -25768,6 +25787,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { filter: filter__, sort_options: sort_options__.unwrap_or_default(), null_equality: null_equality__.unwrap_or_default(), + projection: projection__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index d830624322e14..5bc46dfd7c2c8 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2559,6 +2559,8 @@ pub struct SortMergeJoinExecNode { pub sort_options: ::prost::alloc::vec::Vec, #[prost(enumeration = "super::datafusion_common::NullEquality", tag = "7")] pub null_equality: i32, + #[prost(uint32, repeated, tag = "8")] + pub projection: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct AsyncFuncExecNode { diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs index 941e8832952c6..f3d7b3f0188ea 100644 --- a/datafusion/proto/tests/cases/plans/joins.rs +++ b/datafusion/proto/tests/cases/plans/joins.rs @@ -347,6 +347,50 @@ fn roundtrip_sym_hash_join() -> Result<()> { Ok(()) } +#[test] +fn roundtrip_sort_merge_join_with_projection() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let schema_left = Arc::new(Schema::new(vec![Field::new( + "col_a", + DataType::Int64, + false, + )])); + let schema_right = Arc::new(Schema::new(vec![Field::new( + "col_b", + DataType::Int64, + false, + )])); + let on = vec![( + Arc::new(Column::new("col_a", 0)) as _, + Arc::new(Column::new("col_b", 0)) as _, + )]; + + let projection = Some(vec![1, 0]); + let result = roundtrip_test_and_return( + Arc::new( + SortMergeJoinExec::try_new( + Arc::new(EmptyExec::new(schema_left)), + Arc::new(EmptyExec::new(schema_right)), + on, + None, + JoinType::Inner, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + )? + .with_projection(projection.clone())?, + ), + &ctx, + &codec, + &proto_converter, + )?; + let result = result.downcast_ref::().unwrap(); + assert_eq!(result.projection, projection); + + Ok(()) +} + #[test] fn roundtrip_sort_merge_join() -> Result<()> { let ctx = SessionContext::new(); diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index 7a706836f44d6..e6a4af40e541e 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -2838,16 +2838,15 @@ logical_plan 04)--SubqueryAlias: t2 05)----TableScan: hashjoin_datatype_table_t2 projection=[c1, c2, c3, c4] physical_plan -01)ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, c3@2 as c3, c4@3 as c4, c1@5 as c1, c2@6 as c2, c3@7 as c3, c4@8 as c4] -02)--SortMergeJoinExec: join_type=Right, on=[(CAST(t1.c3 AS Decimal128(10, 2))@4, c3@2)] -03)----SortExec: expr=[CAST(t1.c3 AS Decimal128(10, 2))@4 ASC], preserve_partitioning=[true] -04)------RepartitionExec: partitioning=Hash([CAST(t1.c3 AS Decimal128(10, 2))@4], 2), input_partitions=2 -05)--------ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, c3@2 as c3, c4@3 as c4, CAST(c3@2 AS Decimal128(10, 2)) as CAST(t1.c3 AS Decimal128(10, 2))] -06)----------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1 -07)------------DataSourceExec: partitions=1, partition_sizes=[1] -08)----SortExec: expr=[c3@2 ASC], preserve_partitioning=[true] -09)------RepartitionExec: partitioning=Hash([c3@2], 2), input_partitions=1 -10)--------DataSourceExec: partitions=1, partition_sizes=[1] +01)SortMergeJoinExec: join_type=Right, on=[(CAST(t1.c3 AS Decimal128(10, 2))@4, c3@2)], projection=[c1@0, c2@1, c3@2, c4@3, c1@5, c2@6, c3@7, c4@8] +02)--SortExec: expr=[CAST(t1.c3 AS Decimal128(10, 2))@4 ASC], preserve_partitioning=[true] +03)----RepartitionExec: partitioning=Hash([CAST(t1.c3 AS Decimal128(10, 2))@4], 2), input_partitions=2 +04)------ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, c3@2 as c3, c4@3 as c4, CAST(c3@2 AS Decimal128(10, 2)) as CAST(t1.c3 AS Decimal128(10, 2))] +05)--------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1 +06)----------DataSourceExec: partitions=1, partition_sizes=[1] +07)--SortExec: expr=[c3@2 ASC], preserve_partitioning=[true] +08)----RepartitionExec: partitioning=Hash([c3@2], 2), input_partitions=1 +09)------DataSourceExec: partitions=1, partition_sizes=[1] # sort_merge_join_on_decimal right join on data type (Decimal) query DDRTDDRT rowsort diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index dd80fded93eb3..e4c6265bce511 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -1041,12 +1041,11 @@ FROM range_partitioned l JOIN range_partitioned r ON l.range_key = r.range_key; ---- physical_plan -01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] -02)--SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)] -03)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] -04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, sort_order_for_reorder=[range_key@0 ASC] -05)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] -06)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, sort_order_for_reorder=[range_key@0 ASC] +01)SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +03)----DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.parquet]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, sort_order_for_reorder=[range_key@0 ASC] +04)--SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +05)----DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.parquet]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, sort_order_for_reorder=[range_key@0 ASC] query III SELECT l.range_key, l.value, r.value @@ -1075,14 +1074,13 @@ FROM range_partitioned l JOIN range_partitioned_shifted r ON l.range_key = r.range_key; ---- physical_plan -01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] -02)--SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)] -03)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] -04)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 -05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet -06)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] -07)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 -08)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=parquet +01)SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +03)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.parquet]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +05)--SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +06)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +07)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned_shifted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned_shifted/part-1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned_shifted/part-2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned_shifted/part-3.parquet]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=parquet query III SELECT l.range_key, l.value, r.value diff --git a/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt b/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt index 69bb718bd8c1f..ab5fb15bbee4e 100644 --- a/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt +++ b/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt @@ -274,7 +274,7 @@ WHERE EXISTS ( ) ---- Plan with Metrics -SortMergeJoinExec: join_type=LeftSemi, on=[(k@0, k@0)], filter=v@1 <= x@0 - 300, metrics=[output_rows=1,spill_count=1, spilled_bytes= KB, spilled_rows= K, peak_mem_used= +SortMergeJoinExec: join_type=LeftSemi, on=[(k@0, k@0)], filter=v@1 <= x@0 - 300, projection=[k@0], metrics=[output_rows=1,spill_count=1, spilled_bytes= KB, spilled_rows= K, peak_mem_used= # The same query must retain the matching first slice after later overflows. query I From dc0e106a71642ade33fbaaedd3e293a25bfabb6d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 20 Aug 2026 15:23:41 +0200 Subject: [PATCH 2/3] Swap the projection with the inputs, and keep an empty one over the wire Swapping a projected join left the indices pointing at the columns the other side now occupies, so it emitted the wrong columns. An empty projection also decoded as an absent one, which is a different output schema. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_012YiAABcW4WSqij31zz2P6c --- .../src/joins/sort_merge_join/exec.rs | 46 +++++++--- .../src/joins/sort_merge_join/tests.rs | 83 +++++++++++++++++++ datafusion/proto/tests/cases/plans/joins.rs | 43 +++++----- 3 files changed, 139 insertions(+), 33 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index eeb3979b6f1d5..d29421cc3cfa9 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -29,7 +29,7 @@ use crate::execution_plan::{EmissionType, boundedness_from_children}; use crate::expressions::PhysicalSortExpr; use crate::joins::utils::{ JoinFilter, JoinOn, JoinOnRef, build_join_schema, check_join_is_valid, - estimate_join_statistics, reorder_output_after_swap, + estimate_join_statistics, reorder_output_after_swap, swap_join_projection, symmetric_join_output_partitioning, }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet, SpillMetrics}; @@ -368,10 +368,16 @@ impl SortMergeJoinExec { self.join_type().swap(), self.sort_options.clone(), self.null_equality, - )?; - - // TODO: OR this condition with having a built-in projection (like - // ordinary hash join) when we support it. + )? + .with_projection(swap_join_projection( + left.schema().fields().len(), + right.schema().fields().len(), + self.projection.as_deref(), + &self.join_type(), + ))?; + + // A semi, anti or mark join emits one side, and a projection already names the + // columns to emit, so in both cases swapping leaves the output order alone. if matches!( self.join_type(), JoinType::LeftSemi @@ -380,7 +386,8 @@ impl SortMergeJoinExec { | JoinType::RightAnti | JoinType::LeftMark | JoinType::RightMark - ) { + ) || self.projection.is_some() + { Ok(Arc::new(new_join)) } else { reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema()) @@ -835,11 +842,20 @@ impl ExecutionPlan for SortMergeJoinExec { filter, sort_options, null_equality: null_equality.into(), - projection: projection - .iter() - .flatten() - .map(|index| *index as u32) - .collect(), + // Proto3 `repeated` cannot distinguish `None` from + // `Some(vec![])`. `Some(vec![])` (reachable via + // `try_embed_projection` for e.g. `SELECT count(1) … JOIN …`) + // changes the output schema, so it is encoded with the + // single-element sentinel `[u32::MAX]` (never a valid column + // index); every other state is sent as-is. See + // `try_from_proto` for the matching decoder. + projection: match projection.as_ref() { + None => Vec::new(), + Some(indices) if indices.is_empty() => vec![u32::MAX], + Some(indices) => { + indices.iter().map(|index| *index as u32).collect() + } + }, }, )), ), @@ -927,8 +943,12 @@ impl SortMergeJoinExec { }) .collect(); - let projection = (!projection.is_empty()) - .then(|| projection.iter().map(|index| *index as usize).collect()); + // Preserve the empty-projection sentinel written by `try_to_proto`. + let projection = match projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|index| *index as usize).collect()), + }; Ok(Arc::new( Self::try_new( diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index 91d1b893f1b29..0189a0e390c9d 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -5985,3 +5985,86 @@ async fn bitwise_spill_pending_stream() -> Result<()> { Ok(()) } + +/// A projection names the columns to emit, so swapping the inputs must renumber it +/// rather than leave it pointing at the columns the other side now occupies. +#[tokio::test] +async fn swap_inputs_swaps_the_projection() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![10, 20, 30]), + ("c1", &vec![100, 200, 300]), + ); + let right = build_table( + ("a2", &vec![1, 2, 3]), + ("b2", &vec![11, 22, 33]), + ("c2", &vec![111, 222, 333]), + ); + let on: JoinOn = vec![( + Arc::new(Column::new("a1", 0)) as _, + Arc::new(Column::new("a2", 0)) as _, + )]; + // One column from each side, in an order that tells the two sides apart. + let join = SortMergeJoinExec::try_new( + left, + right, + on, + None, + Inner, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + )? + .with_projection(Some(vec![4, 2]))?; + + let swapped = join.swap_inputs()?; + assert_eq!( + swapped.schema().fields(), + join.schema().fields(), + "swapping must not change what the join emits" + ); + + let task_ctx = Arc::new(TaskContext::default()); + let expected = common::collect(join.execute(0, Arc::clone(&task_ctx))?).await?; + let actual = common::collect(swapped.execute(0, task_ctx)?).await?; + assert_eq!(expected, actual); + + Ok(()) +} + +/// An empty projection still changes the output schema, and the row count has to +/// survive it: `SELECT count(1)` over a join needs the rows but none of the columns. +#[tokio::test] +async fn an_empty_projection_keeps_the_rows() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![10, 20, 30]), + ("c1", &vec![100, 200, 300]), + ); + let right = build_table( + ("a2", &vec![1, 2, 3]), + ("b2", &vec![11, 22, 33]), + ("c2", &vec![111, 222, 333]), + ); + let on: JoinOn = vec![( + Arc::new(Column::new("a1", 0)) as _, + Arc::new(Column::new("a2", 0)) as _, + )]; + let join = SortMergeJoinExec::try_new( + left, + right, + on, + None, + Inner, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + )? + .with_projection(Some(vec![]))?; + + assert_eq!(join.schema().fields().len(), 0); + let batches = + common::collect(join.execute(0, Arc::new(TaskContext::default()))?).await?; + let rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(rows, 3); + + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs index f3d7b3f0188ea..9bfe172a7100f 100644 --- a/datafusion/proto/tests/cases/plans/joins.rs +++ b/datafusion/proto/tests/cases/plans/joins.rs @@ -367,26 +367,29 @@ fn roundtrip_sort_merge_join_with_projection() -> Result<()> { Arc::new(Column::new("col_b", 0)) as _, )]; - let projection = Some(vec![1, 0]); - let result = roundtrip_test_and_return( - Arc::new( - SortMergeJoinExec::try_new( - Arc::new(EmptyExec::new(schema_left)), - Arc::new(EmptyExec::new(schema_right)), - on, - None, - JoinType::Inner, - vec![SortOptions::default()], - NullEquality::NullEqualsNothing, - )? - .with_projection(projection.clone())?, - ), - &ctx, - &codec, - &proto_converter, - )?; - let result = result.downcast_ref::().unwrap(); - assert_eq!(result.projection, projection); + // An empty projection is not an absent one: it changes the output schema, and + // proto3 cannot tell the two apart without a sentinel. + for projection in [None, Some(vec![]), Some(vec![1, 0])] { + let result = roundtrip_test_and_return( + Arc::new( + SortMergeJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema_left))), + Arc::new(EmptyExec::new(Arc::clone(&schema_right))), + on.clone(), + None, + JoinType::Inner, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + )? + .with_projection(projection.clone())?, + ), + &ctx, + &codec, + &proto_converter, + )?; + let result = result.downcast_ref::().unwrap(); + assert_eq!(result.projection, projection); + } Ok(()) } From dd3a41e3a6e21c3f4c42c9477bfbd29dc651912e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 20 Aug 2026 15:43:56 +0200 Subject: [PATCH 3/3] Read the projection field directly instead of through a public helper Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_012YiAABcW4WSqij31zz2P6c --- .../src/joins/sort_merge_join/exec.rs | 22 ++++++------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index d29421cc3cfa9..b2f643ffa83ca 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -233,11 +233,6 @@ impl SortMergeJoinExec { }) } - /// Whether the join emits fewer or reordered columns than it joins. - pub fn contains_projection(&self) -> bool { - self.projection.is_some() - } - /// Get probe side (e.g streaming side) information for this sort merge join. /// In current implementation, probe side is determined according to join type. pub fn probe_side(join_type: &JoinType) -> JoinSide { @@ -417,23 +412,20 @@ impl DisplayAs for SortMergeJoinExec { } else { "" }; - let display_projections = if self.contains_projection() { - format!( + let display_projections = match &self.projection { + Some(projection) => format!( ", projection=[{}]", - self.projection - .as_ref() - .unwrap() + projection .iter() .map(|index| format!( "{}@{}", - self.schema.fields().get(*index).unwrap().name(), + self.schema.field(*index).name(), index )) .collect::>() .join(", ") - ) - } else { - "".to_string() + ), + None => String::new(), }; write!( f, @@ -721,7 +713,7 @@ impl ExecutionPlan for SortMergeJoinExec { &self, projection: &ProjectionExec, ) -> Result>> { - if self.contains_projection() { + if self.projection.is_some() { return Ok(None); } // Convert projected PhysicalExpr's to columns. If not possible, we cannot proceed.