Skip to content

Commit dc0e106

Browse files
Dandandanclaude
Dandandan
and
claude
committed
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 <[email protected]> Claude-Session: https://claude.ai/code/session_012YiAABcW4WSqij31zz2P6c
1 parent 2d55b49 commit dc0e106

3 files changed

Lines changed: 139 additions & 33 deletions

File tree

‎datafusion/physical-plan/src/joins/sort_merge_join/exec.rs‎

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ use crate::execution_plan::{EmissionType, boundedness_from_children};
2929
use crate::expressions::PhysicalSortExpr;
3030
use crate::joins::utils::{
3131
JoinFilter, JoinOn, JoinOnRef, build_join_schema, check_join_is_valid,
32-
estimate_join_statistics, reorder_output_after_swap,
32+
estimate_join_statistics, reorder_output_after_swap, swap_join_projection,
3333
symmetric_join_output_partitioning,
3434
};
3535
use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet, SpillMetrics};
@@ -368,10 +368,16 @@ impl SortMergeJoinExec {
368368
self.join_type().swap(),
369369
self.sort_options.clone(),
370370
self.null_equality,
371-
)?;
372-
373-
// TODO: OR this condition with having a built-in projection (like
374-
// ordinary hash join) when we support it.
371+
)?
372+
.with_projection(swap_join_projection(
373+
left.schema().fields().len(),
374+
right.schema().fields().len(),
375+
self.projection.as_deref(),
376+
&self.join_type(),
377+
))?;
378+
379+
// A semi, anti or mark join emits one side, and a projection already names the
380+
// columns to emit, so in both cases swapping leaves the output order alone.
375381
if matches!(
376382
self.join_type(),
377383
JoinType::LeftSemi
@@ -380,7 +386,8 @@ impl SortMergeJoinExec {
380386
| JoinType::RightAnti
381387
| JoinType::LeftMark
382388
| JoinType::RightMark
383-
) {
389+
) || self.projection.is_some()
390+
{
384391
Ok(Arc::new(new_join))
385392
} else {
386393
reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema())
@@ -835,11 +842,20 @@ impl ExecutionPlan for SortMergeJoinExec {
835842
filter,
836843
sort_options,
837844
null_equality: null_equality.into(),
838-
projection: projection
839-
.iter()
840-
.flatten()
841-
.map(|index| *index as u32)
842-
.collect(),
845+
// Proto3 `repeated` cannot distinguish `None` from
846+
// `Some(vec![])`. `Some(vec![])` (reachable via
847+
// `try_embed_projection` for e.g. `SELECT count(1) … JOIN …`)
848+
// changes the output schema, so it is encoded with the
849+
// single-element sentinel `[u32::MAX]` (never a valid column
850+
// index); every other state is sent as-is. See
851+
// `try_from_proto` for the matching decoder.
852+
projection: match projection.as_ref() {
853+
None => Vec::new(),
854+
Some(indices) if indices.is_empty() => vec![u32::MAX],
855+
Some(indices) => {
856+
indices.iter().map(|index| *index as u32).collect()
857+
}
858+
},
843859
},
844860
)),
845861
),
@@ -927,8 +943,12 @@ impl SortMergeJoinExec {
927943
})
928944
.collect();
929945

930-
let projection = (!projection.is_empty())
931-
.then(|| projection.iter().map(|index| *index as usize).collect());
946+
// Preserve the empty-projection sentinel written by `try_to_proto`.
947+
let projection = match projection.as_slice() {
948+
[] => None,
949+
[u32::MAX] => Some(Vec::new()),
950+
indices => Some(indices.iter().map(|index| *index as usize).collect()),
951+
};
932952

933953
Ok(Arc::new(
934954
Self::try_new(

‎datafusion/physical-plan/src/joins/sort_merge_join/tests.rs‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5985,3 +5985,86 @@ async fn bitwise_spill_pending_stream() -> Result<()> {
59855985

59865986
Ok(())
59875987
}
5988+
5989+
/// A projection names the columns to emit, so swapping the inputs must renumber it
5990+
/// rather than leave it pointing at the columns the other side now occupies.
5991+
#[tokio::test]
5992+
async fn swap_inputs_swaps_the_projection() -> Result<()> {
5993+
let left = build_table(
5994+
("a1", &vec![1, 2, 3]),
5995+
("b1", &vec![10, 20, 30]),
5996+
("c1", &vec![100, 200, 300]),
5997+
);
5998+
let right = build_table(
5999+
("a2", &vec![1, 2, 3]),
6000+
("b2", &vec![11, 22, 33]),
6001+
("c2", &vec![111, 222, 333]),
6002+
);
6003+
let on: JoinOn = vec![(
6004+
Arc::new(Column::new("a1", 0)) as _,
6005+
Arc::new(Column::new("a2", 0)) as _,
6006+
)];
6007+
// One column from each side, in an order that tells the two sides apart.
6008+
let join = SortMergeJoinExec::try_new(
6009+
left,
6010+
right,
6011+
on,
6012+
None,
6013+
Inner,
6014+
vec![SortOptions::default()],
6015+
NullEquality::NullEqualsNothing,
6016+
)?
6017+
.with_projection(Some(vec![4, 2]))?;
6018+
6019+
let swapped = join.swap_inputs()?;
6020+
assert_eq!(
6021+
swapped.schema().fields(),
6022+
join.schema().fields(),
6023+
"swapping must not change what the join emits"
6024+
);
6025+
6026+
let task_ctx = Arc::new(TaskContext::default());
6027+
let expected = common::collect(join.execute(0, Arc::clone(&task_ctx))?).await?;
6028+
let actual = common::collect(swapped.execute(0, task_ctx)?).await?;
6029+
assert_eq!(expected, actual);
6030+
6031+
Ok(())
6032+
}
6033+
6034+
/// An empty projection still changes the output schema, and the row count has to
6035+
/// survive it: `SELECT count(1)` over a join needs the rows but none of the columns.
6036+
#[tokio::test]
6037+
async fn an_empty_projection_keeps_the_rows() -> Result<()> {
6038+
let left = build_table(
6039+
("a1", &vec![1, 2, 3]),
6040+
("b1", &vec![10, 20, 30]),
6041+
("c1", &vec![100, 200, 300]),
6042+
);
6043+
let right = build_table(
6044+
("a2", &vec![1, 2, 3]),
6045+
("b2", &vec![11, 22, 33]),
6046+
("c2", &vec![111, 222, 333]),
6047+
);
6048+
let on: JoinOn = vec![(
6049+
Arc::new(Column::new("a1", 0)) as _,
6050+
Arc::new(Column::new("a2", 0)) as _,
6051+
)];
6052+
let join = SortMergeJoinExec::try_new(
6053+
left,
6054+
right,
6055+
on,
6056+
None,
6057+
Inner,
6058+
vec![SortOptions::default()],
6059+
NullEquality::NullEqualsNothing,
6060+
)?
6061+
.with_projection(Some(vec![]))?;
6062+
6063+
assert_eq!(join.schema().fields().len(), 0);
6064+
let batches =
6065+
common::collect(join.execute(0, Arc::new(TaskContext::default()))?).await?;
6066+
let rows: usize = batches.iter().map(|batch| batch.num_rows()).sum();
6067+
assert_eq!(rows, 3);
6068+
6069+
Ok(())
6070+
}

‎datafusion/proto/tests/cases/plans/joins.rs‎

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -367,26 +367,29 @@ fn roundtrip_sort_merge_join_with_projection() -> Result<()> {
367367
Arc::new(Column::new("col_b", 0)) as _,
368368
)];
369369

370-
let projection = Some(vec![1, 0]);
371-
let result = roundtrip_test_and_return(
372-
Arc::new(
373-
SortMergeJoinExec::try_new(
374-
Arc::new(EmptyExec::new(schema_left)),
375-
Arc::new(EmptyExec::new(schema_right)),
376-
on,
377-
None,
378-
JoinType::Inner,
379-
vec![SortOptions::default()],
380-
NullEquality::NullEqualsNothing,
381-
)?
382-
.with_projection(projection.clone())?,
383-
),
384-
&ctx,
385-
&codec,
386-
&proto_converter,
387-
)?;
388-
let result = result.downcast_ref::<SortMergeJoinExec>().unwrap();
389-
assert_eq!(result.projection, projection);
370+
// An empty projection is not an absent one: it changes the output schema, and
371+
// proto3 cannot tell the two apart without a sentinel.
372+
for projection in [None, Some(vec![]), Some(vec![1, 0])] {
373+
let result = roundtrip_test_and_return(
374+
Arc::new(
375+
SortMergeJoinExec::try_new(
376+
Arc::new(EmptyExec::new(Arc::clone(&schema_left))),
377+
Arc::new(EmptyExec::new(Arc::clone(&schema_right))),
378+
on.clone(),
379+
None,
380+
JoinType::Inner,
381+
vec![SortOptions::default()],
382+
NullEquality::NullEqualsNothing,
383+
)?
384+
.with_projection(projection.clone())?,
385+
),
386+
&ctx,
387+
&codec,
388+
&proto_converter,
389+
)?;
390+
let result = result.downcast_ref::<SortMergeJoinExec>().unwrap();
391+
assert_eq!(result.projection, projection);
392+
}
390393

391394
Ok(())
392395
}

0 commit comments

Comments
 (0)