feat: add projection support to SortMergeJoinExec (#24517)
## Which issue does this PR close?
- Closes https://github.com/apache/datafusion/issues/24518
## Rationale for this change
`SortMergeJoinExec` emits every column of both inputs.
We can add projection support to only produce the output columns (saving
some compute / copy).
```sql
select * from t1 right join t2 on t1.c3 = t2.c3 -- c3 needs a cast
```
```text
ProjectionExec: expr=[c1@0 as c1, ..., c4@8 as c4]
SortMergeJoinExec: join_type=Right, on=[(CAST(t1.c3 AS Decimal128(10, 2))@4, c3@2)]
```
The cast column is only there to be joined on, and every operator above
the join
carries it until the projection removes it.
## What changes are included in this PR?
`SortMergeJoinExec` takes an optional projection, set with
`with_projection`, the same
shape as `HashJoinExec`'s.
`try_swapping_with_projection` embeds the projection into the join when
it cannot push
it into the children, so the query above becomes:
```text
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]
```
The change doesn't bring a big speedup but helps aligning with other
join types and helping simplify join optimization in other areas (e.g.
join enumeration).
## Are these changes tested?
Yes:
- a projection pushdown test for a projection that interleaves the two
sides, which
the existing pushdown cannot handle
- a serialization round trip
- existing sqllogictests, whose plans lose a `ProjectionExec` in four
places
## Are there any user-facing changes?
Explain / proto changes.
---------
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs
index 113552c..4e8d479 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::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 @@
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<dyn ExecutionPlan> = 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 dc8540f..b2f643f 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
@@ -29,24 +29,27 @@
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};
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_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 sort_options: Vec<SortOptions>,
/// 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<Vec<usize>>,
/// Cache holding plan properties like equivalences, output partitioning etc.
cache: Arc<PlanProperties>,
}
@@ -187,7 +195,7 @@
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,31 @@
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<Vec<usize>>) -> Result<Self> {
+ 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)
+ })
+ }
+
/// 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 +310,32 @@
fn compute_properties(
left: &Arc<dyn ExecutionPlan>,
right: &Arc<dyn ExecutionPlan>,
- schema: SchemaRef,
+ schema: &SchemaRef,
join_type: JoinType,
join_on: JoinOnRef,
+ projection: Option<&[usize]>,
) -> Result<PlanProperties> {
// 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,
@@ -326,10 +363,16 @@
self.join_type().swap(),
self.sort_options.clone(),
self.null_equality,
- )?;
+ )?
+ .with_projection(swap_join_projection(
+ left.schema().fields().len(),
+ right.schema().fields().len(),
+ self.projection.as_deref(),
+ &self.join_type(),
+ ))?;
- // TODO: OR this condition with having a built-in projection (like
- // ordinary hash join) when we support it.
+ // 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
@@ -338,7 +381,8 @@
| 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())
@@ -346,6 +390,12 @@
}
}
+impl EmbeddedProjection for SortMergeJoinExec {
+ fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> {
+ self.with_projection(projection)
+ }
+}
+
impl DisplayAs for SortMergeJoinExec {
fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
match t {
@@ -362,9 +412,24 @@
} else {
""
};
+ let display_projections = match &self.projection {
+ Some(projection) => format!(
+ ", projection=[{}]",
+ projection
+ .iter()
+ .map(|index| format!(
+ "{}@{}",
+ self.schema.field(*index).name(),
+ index
+ ))
+ .collect::<Vec<_>>()
+ .join(", ")
+ ),
+ None => String::new(),
+ };
write!(
f,
- "{}: join_type={:?}, on=[{}]{}{}",
+ "{}: join_type={:?}, on=[{}]{}{}{}",
Self::static_name(),
self.join_type,
on,
@@ -373,6 +438,7 @@
|f| format!(", filter={}", f.expression())
),
display_null_equality,
+ display_projections,
)
}
DisplayFormatType::TreeRender => {
@@ -467,15 +533,18 @@
}))
}
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 +615,7 @@
)
.with_compression_type(context.session_config().spill_compression());
- if matches!(
+ let joined = if matches!(
self.join_type,
JoinType::LeftSemi
| JoinType::LeftAnti
@@ -589,7 +658,16 @@
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<MetricsSet> {
@@ -614,14 +692,18 @@
// - `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 +713,9 @@
&self,
projection: &ProjectionExec,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+ if self.projection.is_some() {
+ 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 +727,15 @@
&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 +744,7 @@
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 +783,7 @@
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 +834,20 @@
filter,
sort_options,
null_equality: null_equality.into(),
+ // 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()
+ }
+ },
},
)),
),
@@ -781,6 +883,7 @@
filter,
sort_options,
null_equality,
+ projection,
} = &**sort_join;
let left =
@@ -832,14 +935,24 @@
})
.collect();
- Ok(Arc::new(Self::try_new(
- left,
- right,
- on,
- filter,
- join_type,
- sort_options,
- null_equality,
- )?))
+ // 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(
+ left,
+ right,
+ on,
+ filter,
+ join_type,
+ sort_options,
+ null_equality,
+ )?
+ .with_projection(projection)?,
+ ))
}
}
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 91d1b89..0189a0e 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 @@
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-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto
index d98b67a..7b8f8e3 100644
--- a/datafusion/proto-models/proto/datafusion.proto
+++ b/datafusion/proto-models/proto/datafusion.proto
@@ -1691,6 +1691,7 @@
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 21309bb..bc88f5c 100644
--- a/datafusion/proto-models/src/generated/pbjson.rs
+++ b/datafusion/proto-models/src/generated/pbjson.rs
@@ -25603,6 +25603,9 @@
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 @@
.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 @@
"sortOptions",
"null_equality",
"nullEquality",
+ "projection",
];
#[allow(clippy::enum_variant_names)]
@@ -25660,6 +25667,7 @@
Filter,
SortOptions,
NullEquality,
+ Projection,
}
impl<'de> serde::Deserialize<'de> for GeneratedField {
fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error>
@@ -25688,6 +25696,7 @@
"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 @@
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 @@
}
null_equality__ = Some(map_.next_value::<super::datafusion_common::NullEquality>()? as i32);
}
+ GeneratedField::Projection => {
+ if projection__.is_some() {
+ return Err(serde::de::Error::duplicate_field("projection"));
+ }
+ projection__ =
+ Some(map_.next_value::<Vec<::pbjson::private::NumberDeserialize<_>>>()?
+ .into_iter().map(|x| x.0).collect())
+ ;
+ }
}
}
Ok(SortMergeJoinExecNode {
@@ -25768,6 +25787,7 @@
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 d830624..5bc46df 100644
--- a/datafusion/proto-models/src/generated/prost.rs
+++ b/datafusion/proto-models/src/generated/prost.rs
@@ -2559,6 +2559,8 @@
pub sort_options: ::prost::alloc::vec::Vec<SortExprNode>,
#[prost(enumeration = "super::datafusion_common::NullEquality", tag = "7")]
pub null_equality: i32,
+ #[prost(uint32, repeated, tag = "8")]
+ pub projection: ::prost::alloc::vec::Vec<u32>,
}
#[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 941e883..9bfe172 100644
--- a/datafusion/proto/tests/cases/plans/joins.rs
+++ b/datafusion/proto/tests/cases/plans/joins.rs
@@ -348,6 +348,53 @@
}
#[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 _,
+ )];
+
+ // 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::<SortMergeJoinExec>().unwrap();
+ assert_eq!(result.projection, projection);
+ }
+
+ Ok(())
+}
+
+#[test]
fn roundtrip_sort_merge_join() -> Result<()> {
let ctx = SessionContext::new();
let codec = DefaultPhysicalExtensionCodec {};
diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt
index 7a70683..e6a4af4 100644
--- a/datafusion/sqllogictest/test_files/joins.slt
+++ b/datafusion/sqllogictest/test_files/joins.slt
@@ -2838,16 +2838,15 @@
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 dd80fde..e4c6265 100644
--- a/datafusion/sqllogictest/test_files/range_partitioning.slt
+++ b/datafusion/sqllogictest/test_files/range_partitioning.slt
@@ -1041,12 +1041,11 @@
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=<slt:ignore>, 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=<slt:ignore>, 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 @@
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=<slt:ignore>, 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=<slt:ignore>, 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 69bb718..ab5fb15 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 @@
)
----
Plan with Metrics
-<slt:ignore>SortMergeJoinExec: join_type=LeftSemi, on=[(k@0, k@0)], filter=v@1 <= x@0 - 300, metrics=[output_rows=1,<slt:ignore>spill_count=1, spilled_bytes=<slt:ignore> KB, spilled_rows=<slt:ignore> K, peak_mem_used=<slt:ignore>
+<slt:ignore>SortMergeJoinExec: join_type=LeftSemi, on=[(k@0, k@0)], filter=v@1 <= x@0 - 300, projection=[k@0], metrics=[output_rows=1,<slt:ignore>spill_count=1, spilled_bytes=<slt:ignore> KB, spilled_rows=<slt:ignore> K, peak_mem_used=<slt:ignore>
# The same query must retain the matching first slice after later overflows.
query I
