Skip to content

Commit 31d1284

Browse files
lyne7-sclyne7-sc
committed
fix: consider join key subexpressions for probe ordering
1 parent 976250b commit 31d1284

2 files changed

Lines changed: 252 additions & 169 deletions

File tree

‎datafusion/core/tests/physical_optimizer/enforce_sorting.rs‎

Lines changed: 135 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -3634,139 +3634,145 @@ impl ExecutionPlan for JoinOrderingSource {
36343634
#[tokio::test]
36353635
async fn test_join_suffix_unbounded_limit() -> Result<()> {
36363636
// Either a unique probe row or a unique build match permits the build suffix.
3637-
for unique_probe in [true, false] {
3638-
for use_smj in [false, true] {
3639-
let build_values = if unique_probe { vec![1, 1] } else { vec![1, 2] };
3640-
let build_batch = record_batch!(
3641-
("build_k", Int32, [1, 1]),
3642-
("build_v", Int32, build_values),
3643-
("payload", Int32, [10, 20])
3644-
)?;
3645-
let build_schema = build_batch.schema();
3646-
let build = DataSourceExec::new(Arc::new(
3647-
MemorySourceConfig::try_new(
3648-
&[vec![build_batch]],
3649-
Arc::clone(&build_schema),
3650-
None,
3651-
)?
3652-
.try_with_sort_information(vec![
3653-
LexOrdering::new([
3654-
sort_expr("build_k", &build_schema),
3655-
sort_expr("build_v", &build_schema),
3656-
])
3657-
.unwrap(),
3658-
LexOrdering::new([sort_expr("payload", &build_schema)]).unwrap(),
3659-
])?,
3660-
))
3661-
.with_constraints(Constraints::new_unverified(
3662-
if unique_probe {
3663-
vec![]
3664-
} else {
3665-
vec![Constraint::PrimaryKey(vec![0, 1])]
3666-
},
3667-
));
3668-
let probe_batch = if unique_probe {
3669-
// UNIQUE permits repeated NULLs; ordinary equality excludes them.
3670-
record_batch!(
3671-
("probe_k", Int32, [None, None, Some(1)]),
3672-
("probe_v", Int32, [None, None, Some(1)]),
3673-
("extra", Int32, [1, 2, 3])
3674-
)?
3637+
for (unique_probe, computed_keys, use_smj) in [
3638+
(true, false, false),
3639+
(true, false, true),
3640+
(false, false, false),
3641+
(false, false, true),
3642+
// Modulo does not preserve the key ordering required by a streaming SMJ.
3643+
(false, true, false),
3644+
] {
3645+
let build_values = if unique_probe { vec![1, 1] } else { vec![1, 2] };
3646+
let build_batch = record_batch!(
3647+
("build_k", Int32, [1, 1]),
3648+
("build_v", Int32, build_values),
3649+
("payload", Int32, [10, 20])
3650+
)?;
3651+
let build_schema = build_batch.schema();
3652+
let build = DataSourceExec::new(Arc::new(
3653+
MemorySourceConfig::try_new(
3654+
&[vec![build_batch]],
3655+
Arc::clone(&build_schema),
3656+
None,
3657+
)?
3658+
.try_with_sort_information(vec![
3659+
LexOrdering::new([
3660+
sort_expr("build_k", &build_schema),
3661+
sort_expr("build_v", &build_schema),
3662+
])
3663+
.unwrap(),
3664+
LexOrdering::new([sort_expr("payload", &build_schema)]).unwrap(),
3665+
])?,
3666+
))
3667+
.with_constraints(Constraints::new_unverified(if unique_probe {
3668+
vec![]
3669+
} else {
3670+
vec![Constraint::PrimaryKey(vec![0, 1])]
3671+
}));
3672+
let probe_batch = if unique_probe {
3673+
// UNIQUE permits repeated NULLs; ordinary equality excludes them.
3674+
record_batch!(
3675+
("probe_k", Int32, [None, None, Some(1)]),
3676+
("probe_v", Int32, [None, None, Some(1)]),
3677+
("extra", Int32, [1, 2, 3])
3678+
)?
3679+
} else {
3680+
record_batch!(
3681+
("probe_k", Int32, [1, 1]),
3682+
("probe_v", Int32, [1, 1]),
3683+
("extra", Int32, [1, 2])
3684+
)?
3685+
};
3686+
let probe_schema = probe_batch.schema();
3687+
assert!(probe_schema.field(0).is_nullable());
3688+
// The unrelated ordering must not obscure the sufficient prefix.
3689+
let orderings = ["extra", "probe_k", "probe_v"]
3690+
.map(|name| LexOrdering::new([sort_expr(name, &probe_schema)]).unwrap())
3691+
.to_vec();
3692+
let probe = Arc::new(JoinOrderingSource::new(
3693+
probe_batch,
3694+
orderings,
3695+
Constraints::new_unverified(if unique_probe {
3696+
vec![Constraint::Unique(vec![0])]
36753697
} else {
3676-
record_batch!(
3677-
("probe_k", Int32, [1, 1]),
3678-
("probe_v", Int32, [1, 1]),
3679-
("extra", Int32, [1, 2])
3680-
)?
3681-
};
3682-
let probe_schema = probe_batch.schema();
3683-
assert!(probe_schema.field(0).is_nullable());
3684-
// The unrelated ordering must not obscure the sufficient prefix.
3685-
let orderings = ["extra", "probe_k", "probe_v"]
3686-
.map(|name| LexOrdering::new([sort_expr(name, &probe_schema)]).unwrap())
3687-
.to_vec();
3688-
let probe = Arc::new(JoinOrderingSource::new(
3689-
probe_batch,
3690-
orderings,
3691-
Constraints::new_unverified(if unique_probe {
3692-
vec![Constraint::Unique(vec![0])]
3693-
} else {
3694-
vec![]
3695-
}),
3696-
));
3697-
let on = vec![
3698-
(
3699-
col("build_k", &build_schema)?,
3700-
col("probe_k", &probe_schema)?,
3701-
),
3702-
(
3703-
col("build_v", &build_schema)?,
3704-
col("probe_v", &probe_schema)?,
3705-
),
3706-
];
3707-
let join: Arc<dyn ExecutionPlan> = if use_smj {
3708-
Arc::new(SortMergeJoinExec::try_new(
3709-
probe,
3710-
Arc::new(build),
3711-
on.into_iter().map(|(b, p)| (p, b)).collect(),
3712-
None,
3713-
JoinType::Inner,
3714-
vec![SortOptions::default(); 2],
3715-
NullEquality::NullEqualsNothing,
3716-
)?)
3698+
vec![]
3699+
}),
3700+
));
3701+
let probe_key = |name| -> Result<Arc<dyn PhysicalExpr>> {
3702+
let column = col(name, &probe_schema)?;
3703+
Ok(if computed_keys {
3704+
Arc::new(BinaryExpr::new(
3705+
column,
3706+
Operator::Modulo,
3707+
datafusion_physical_expr::expressions::lit(2_i32),
3708+
))
37173709
} else {
3718-
Arc::new(
3719-
HashJoinExecBuilder::new(Arc::new(build), probe, on, JoinType::Inner)
3720-
.with_partition_mode(PartitionMode::CollectLeft)
3721-
.build()?,
3722-
)
3723-
};
3724-
let mut required = vec![sort_expr("probe_k", &join.schema())];
3725-
if !unique_probe {
3726-
required.push(sort_expr("probe_v", &join.schema()));
3727-
}
3728-
required.push(sort_expr("payload", &join.schema()));
3729-
let mut plan: Arc<dyn ExecutionPlan> = Arc::new(
3730-
SortExec::new(LexOrdering::new(required).unwrap(), join)
3731-
.with_fetch(Some(1)),
3732-
);
3733-
let ctx = SessionContext::new_with_config(
3734-
SessionConfig::new()
3735-
.with_target_partitions(1)
3736-
.with_batch_size(2),
3737-
);
3738-
let state = ctx.state();
3739-
for optimizer in state.physical_optimizers() {
3740-
plan = optimizer.optimize(plan, state.config_options())?;
3741-
}
3742-
let batches = tokio::time::timeout(
3743-
Duration::from_secs(2),
3744-
datafusion_physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()),
3710+
column
3711+
})
3712+
};
3713+
let on = vec![
3714+
(col("build_k", &build_schema)?, probe_key("probe_k")?),
3715+
(col("build_v", &build_schema)?, probe_key("probe_v")?),
3716+
];
3717+
let join: Arc<dyn ExecutionPlan> = if use_smj {
3718+
Arc::new(SortMergeJoinExec::try_new(
3719+
probe,
3720+
Arc::new(build),
3721+
on.into_iter().map(|(b, p)| (p, b)).collect(),
3722+
None,
3723+
JoinType::Inner,
3724+
vec![SortOptions::default(); 2],
3725+
NullEquality::NullEqualsNothing,
3726+
)?)
3727+
} else {
3728+
Arc::new(
3729+
HashJoinExecBuilder::new(Arc::new(build), probe, on, JoinType::Inner)
3730+
.with_partition_mode(PartitionMode::CollectLeft)
3731+
.build()?,
37453732
)
3746-
.await
3747-
.unwrap_or_else(|_| {
3748-
panic!(
3749-
"LIMIT waited for more probe rows (unique_probe={unique_probe}, smj={use_smj}):\n{}",
3750-
displayable(plan.as_ref()).indent(true),
3751-
)
3752-
})?;
3753-
let payloads = batches
3754-
.iter()
3755-
.map(|batch| {
3756-
batch.project(&[batch.schema().index_of("payload").unwrap()])
3757-
})
3758-
.collect::<std::result::Result<Vec<_>, _>>()?;
3759-
assert_batches_eq!(
3760-
[
3761-
"+---------+",
3762-
"| payload |",
3763-
"+---------+",
3764-
"| 10 |",
3765-
"+---------+",
3766-
],
3767-
&payloads
3768-
);
3733+
};
3734+
let mut required = vec![sort_expr("probe_k", &join.schema())];
3735+
if !unique_probe {
3736+
required.push(sort_expr("probe_v", &join.schema()));
37693737
}
3738+
required.push(sort_expr("payload", &join.schema()));
3739+
let mut plan: Arc<dyn ExecutionPlan> = Arc::new(
3740+
SortExec::new(LexOrdering::new(required).unwrap(), join).with_fetch(Some(1)),
3741+
);
3742+
let ctx = SessionContext::new_with_config(
3743+
SessionConfig::new()
3744+
.with_target_partitions(1)
3745+
.with_batch_size(2),
3746+
);
3747+
let state = ctx.state();
3748+
for optimizer in state.physical_optimizers() {
3749+
plan = optimizer.optimize(plan, state.config_options())?;
3750+
}
3751+
let batches = tokio::time::timeout(
3752+
Duration::from_secs(2),
3753+
datafusion_physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()),
3754+
)
3755+
.await
3756+
.unwrap_or_else(|_| {
3757+
panic!(
3758+
"LIMIT waited for more probe rows (unique_probe={unique_probe}, computed_keys={computed_keys}, smj={use_smj}):\n{}",
3759+
displayable(plan.as_ref()).indent(true),
3760+
)
3761+
})?;
3762+
let payloads = batches
3763+
.iter()
3764+
.map(|batch| batch.project(&[batch.schema().index_of("payload").unwrap()]))
3765+
.collect::<std::result::Result<Vec<_>, _>>()?;
3766+
assert_batches_eq!(
3767+
[
3768+
"+---------+",
3769+
"| payload |",
3770+
"+---------+",
3771+
"| 10 |",
3772+
"+---------+",
3773+
],
3774+
&payloads
3775+
);
37703776
}
37713777
Ok(())
37723778
}

0 commit comments

Comments
 (0)