Skip to content

Commit 019e9c8

Browse files
committed
fix: account for volatility and nullability in join ordering proofs
1 parent 4e2eabf commit 019e9c8

2 files changed

Lines changed: 280 additions & 22 deletions

File tree

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

Lines changed: 253 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3497,22 +3497,55 @@ async fn test_passthrough_wrapper_projection_keeps_ordering() -> Result<()> {
34973497

34983498
#[tokio::test]
34993499
async fn test_unique_build_hash_join_unbounded_limit() -> Result<()> {
3500-
check_unique_build_join_unbounded_limit(JoinType::Inner, false).await
3500+
check_unique_build_join_unbounded_limit(
3501+
JoinType::Inner,
3502+
false,
3503+
NullEquality::NullEqualsNothing,
3504+
)
3505+
.await
35013506
}
35023507

35033508
#[tokio::test]
35043509
async fn test_unique_build_right_join_unbounded_limit() -> Result<()> {
3505-
check_unique_build_join_unbounded_limit(JoinType::Right, false).await
3510+
check_unique_build_join_unbounded_limit(
3511+
JoinType::Right,
3512+
false,
3513+
NullEquality::NullEqualsNothing,
3514+
)
3515+
.await
35063516
}
35073517

35083518
#[tokio::test]
35093519
async fn test_unique_build_left_sort_merge_join_unbounded_limit() -> Result<()> {
3510-
check_unique_build_join_unbounded_limit(JoinType::Left, false).await
3520+
check_unique_build_join_unbounded_limit(
3521+
JoinType::Left,
3522+
false,
3523+
NullEquality::NullEqualsNothing,
3524+
)
3525+
.await
35113526
}
35123527

35133528
#[tokio::test]
35143529
async fn test_unique_build_computed_key_join_unbounded_limit() -> Result<()> {
3515-
check_unique_build_join_unbounded_limit(JoinType::Inner, true).await
3530+
check_unique_build_join_unbounded_limit(
3531+
JoinType::Inner,
3532+
true,
3533+
NullEquality::NullEqualsNothing,
3534+
)
3535+
.await
3536+
}
3537+
3538+
#[tokio::test]
3539+
async fn test_nullable_unique_build_nonnullable_probe_unbounded_limit() -> Result<()> {
3540+
for join_type in [JoinType::Inner, JoinType::Right, JoinType::Left] {
3541+
check_unique_build_join_unbounded_limit(
3542+
join_type,
3543+
false,
3544+
NullEquality::NullEqualsNull,
3545+
)
3546+
.await?;
3547+
}
3548+
Ok(())
35163549
}
35173550

35183551
fn unique_join_build() -> Result<Arc<DataSourceExec>> {
@@ -3541,6 +3574,7 @@ fn unique_join_build() -> Result<Arc<DataSourceExec>> {
35413574
async fn check_unique_build_join_unbounded_limit(
35423575
join_type: JoinType,
35433576
computed_key: bool,
3577+
null_equality: NullEquality,
35443578
) -> Result<()> {
35453579
#[derive(Debug)]
35463580
struct PendingPartition(RecordBatch);
@@ -3558,9 +3592,25 @@ async fn check_unique_build_join_unbounded_limit(
35583592
}
35593593

35603594
let build = unique_join_build()?;
3595+
let build =
3596+
if null_equality == NullEquality::NullEqualsNull {
3597+
assert!(build.schema().field(0).is_nullable());
3598+
Arc::new(build.as_ref().clone().with_constraints(
3599+
Constraints::new_unverified(vec![Constraint::Unique(vec![0])]),
3600+
))
3601+
} else {
3602+
build
3603+
};
35613604

3562-
let probe_batch = record_batch!(("probe_k", Int32, [1, 1]))?;
3563-
let probe_schema = probe_batch.schema();
3605+
let probe_schema = Arc::new(Schema::new(vec![Field::new(
3606+
"probe_k",
3607+
DataType::Int32,
3608+
false,
3609+
)]));
3610+
let probe_batch = RecordBatch::try_new(
3611+
Arc::clone(&probe_schema),
3612+
vec![Arc::new(Int32Array::from(vec![1, 1]))],
3613+
)?;
35643614
let probe_ordering = LexOrdering::new([sort_expr("probe_k", &probe_schema)]).unwrap();
35653615
let probe = Arc::new(StreamingTableExec::try_new(
35663616
probe_schema,
@@ -3589,7 +3639,7 @@ async fn check_unique_build_join_unbounded_limit(
35893639
None,
35903640
join_type,
35913641
vec![SortOptions::default()],
3592-
NullEquality::NullEqualsNothing,
3642+
null_equality,
35933643
)?)
35943644
} else {
35953645
Arc::new(
@@ -3599,6 +3649,7 @@ async fn check_unique_build_join_unbounded_limit(
35993649
vec![(Arc::new(Column::new("build_k", 0)), probe_key)],
36003650
join_type,
36013651
)
3652+
.with_null_equality(null_equality)
36023653
.with_partition_mode(PartitionMode::CollectLeft)
36033654
.build()?,
36043655
)
@@ -3897,3 +3948,198 @@ async fn test_nullable_unique_build_null_equals_null_requires_sort() -> Result<(
38973948
);
38983949
Ok(())
38993950
}
3951+
3952+
/// A filter can equate a volatile expression with a column for its own output,
3953+
/// but evaluating the expression again as a join key can produce different values.
3954+
#[tokio::test]
3955+
async fn test_unique_build_volatile_key_equivalence_limit() -> Result<()> {
3956+
use datafusion_expr::{Volatility, create_udf};
3957+
use datafusion_physical_expr::ScalarFunctionExpr;
3958+
use datafusion_physical_plan::filter::FilterExec;
3959+
use std::sync::atomic::{AtomicUsize, Ordering};
3960+
3961+
for original_key_is_volatile in [true, false] {
3962+
let calls = AtomicUsize::new(0);
3963+
let flip = create_udf(
3964+
"volatile_flip",
3965+
vec![DataType::Int32],
3966+
DataType::Int32,
3967+
Volatility::Volatile,
3968+
Arc::new(move |args| {
3969+
// The filter keeps every row. On reevaluation in the join, equal
3970+
// probe keys alternate between two different unique build rows.
3971+
if original_key_is_volatile && calls.fetch_add(1, Ordering::SeqCst) == 0 {
3972+
return Ok(args[0].clone());
3973+
}
3974+
let values = if original_key_is_volatile {
3975+
[0, 1, 0, 0]
3976+
} else {
3977+
[1, 2, 1, 2]
3978+
};
3979+
let len = match &args[0] {
3980+
ColumnarValue::Array(array) => array.len(),
3981+
ColumnarValue::Scalar(_) => 1,
3982+
};
3983+
Ok(ColumnarValue::Array(Arc::new(
3984+
Int32Array::from_iter_values(values.into_iter().cycle().take(len)),
3985+
)))
3986+
}),
3987+
);
3988+
let probe_batch = record_batch!(
3989+
("probe_k", Int32, [1, 1, 1, 2]),
3990+
("probe_q", Int32, [1, 2, 1, 2])
3991+
)?;
3992+
let probe_schema = probe_batch.schema();
3993+
let probe_key = col("probe_k", &probe_schema)?;
3994+
let flip: Arc<dyn PhysicalExpr> = Arc::new(ScalarFunctionExpr::try_new(
3995+
Arc::new(flip),
3996+
vec![Arc::clone(&probe_key)],
3997+
&probe_schema,
3998+
Arc::new(ConfigOptions::default()),
3999+
)?);
4000+
let probe = DataSourceExec::from_data_source(
4001+
MemorySourceConfig::try_new(
4002+
&[vec![probe_batch]],
4003+
Arc::clone(&probe_schema),
4004+
None,
4005+
)?
4006+
.try_with_sort_information(vec![
4007+
LexOrdering::new([sort_expr("probe_k", &probe_schema)]).unwrap(),
4008+
])?,
4009+
);
4010+
let (predicate, join_key): (_, Arc<dyn PhysicalExpr>) =
4011+
if original_key_is_volatile {
4012+
(
4013+
BinaryExpr::new(
4014+
Arc::clone(&probe_key),
4015+
Operator::Eq,
4016+
Arc::clone(&flip),
4017+
),
4018+
Arc::new(BinaryExpr::new(probe_key, Operator::Plus, flip)),
4019+
)
4020+
} else {
4021+
let q = col("probe_q", &probe_schema)?;
4022+
(BinaryExpr::new(flip, Operator::Eq, Arc::clone(&q)), q)
4023+
};
4024+
let probe = Arc::new(FilterExec::try_new(Arc::new(predicate), probe)?);
4025+
// Equivalence can either hide a volatile child or introduce one through an alias.
4026+
let normalized = probe
4027+
.properties()
4028+
.equivalence_properties()
4029+
.eq_group()
4030+
.normalize_expr(Arc::clone(&join_key));
4031+
assert_eq!(
4032+
datafusion_physical_expr_common::physical_expr::is_volatile(&normalized),
4033+
!original_key_is_volatile,
4034+
);
4035+
let join = HashJoinExecBuilder::new(
4036+
unique_join_build()?,
4037+
probe,
4038+
vec![(Arc::new(Column::new("build_k", 0)), join_key)],
4039+
JoinType::Inner,
4040+
)
4041+
.with_partition_mode(PartitionMode::CollectLeft)
4042+
.build()?;
4043+
let required = LexOrdering::new([
4044+
sort_expr("probe_k", &join.schema()),
4045+
sort_expr("build_k", &join.schema()),
4046+
sort_expr("build_v", &join.schema()),
4047+
])
4048+
.unwrap();
4049+
let plan = Arc::new(SortExec::new(required, Arc::new(join)).with_fetch(Some(2)));
4050+
let ctx = SessionContext::new_with_config(
4051+
SessionConfig::new().with_target_partitions(1),
4052+
);
4053+
let batches = datafusion_physical_plan::collect(plan, ctx.task_ctx()).await?;
4054+
assert_batches_eq!(
4055+
[
4056+
"+---------+---------+---------+---------+",
4057+
"| build_k | build_v | probe_k | probe_q |",
4058+
"+---------+---------+---------+---------+",
4059+
"| 1 | 100 | 1 | 1 |",
4060+
"| 1 | 100 | 1 | 1 |",
4061+
"+---------+---------+---------+---------+",
4062+
],
4063+
&batches
4064+
);
4065+
}
4066+
4067+
Ok(())
4068+
}
4069+
4070+
#[tokio::test]
4071+
async fn test_unique_build_volatile_build_key_equivalence_limit() -> Result<()> {
4072+
use datafusion_expr::{Volatility, create_udf};
4073+
use datafusion_physical_expr::ScalarFunctionExpr;
4074+
use datafusion_physical_plan::filter::FilterExec;
4075+
use std::sync::atomic::{AtomicUsize, Ordering};
4076+
4077+
let calls = AtomicUsize::new(0);
4078+
let flip = create_udf(
4079+
"volatile_build_key",
4080+
vec![DataType::Int32],
4081+
DataType::Int32,
4082+
Volatility::Volatile,
4083+
Arc::new(move |args| {
4084+
// The filter sees the unique keys [1, 2], but the join sees [1, 1].
4085+
if calls.fetch_add(1, Ordering::SeqCst) == 0 {
4086+
Ok(args[0].clone())
4087+
} else {
4088+
Ok(ColumnarValue::Scalar(1_i32.into()))
4089+
}
4090+
}),
4091+
);
4092+
let build = unique_join_build()?;
4093+
let build_key = col("build_k", &build.schema())?;
4094+
let flip: Arc<dyn PhysicalExpr> = Arc::new(ScalarFunctionExpr::try_new(
4095+
Arc::new(flip),
4096+
vec![Arc::clone(&build_key)],
4097+
&build.schema(),
4098+
Arc::new(ConfigOptions::default()),
4099+
)?);
4100+
let build = Arc::new(FilterExec::try_new(
4101+
Arc::new(BinaryExpr::new(build_key, Operator::Eq, Arc::clone(&flip))),
4102+
build,
4103+
)?);
4104+
let probe_batch = record_batch!(("probe_k", Int32, [1, 1]))?;
4105+
let probe_schema = probe_batch.schema();
4106+
let probe = DataSourceExec::from_data_source(
4107+
MemorySourceConfig::try_new(
4108+
&[vec![probe_batch]],
4109+
Arc::clone(&probe_schema),
4110+
None,
4111+
)?
4112+
.try_with_sort_information(vec![
4113+
LexOrdering::new([sort_expr("probe_k", &probe_schema)]).unwrap(),
4114+
])?,
4115+
);
4116+
let join = HashJoinExecBuilder::new(
4117+
build,
4118+
probe,
4119+
vec![(flip, col("probe_k", &probe_schema)?)],
4120+
JoinType::Inner,
4121+
)
4122+
.with_partition_mode(PartitionMode::CollectLeft)
4123+
.build()?;
4124+
let required = LexOrdering::new([
4125+
sort_expr("probe_k", &join.schema()),
4126+
sort_expr("build_k", &join.schema()),
4127+
sort_expr("build_v", &join.schema()),
4128+
])
4129+
.unwrap();
4130+
let plan = Arc::new(SortExec::new(required, Arc::new(join)).with_fetch(Some(2)));
4131+
let batches =
4132+
datafusion_physical_plan::collect(plan, SessionContext::new().task_ctx()).await?;
4133+
assert_batches_eq!(
4134+
[
4135+
"+---------+---------+---------+",
4136+
"| build_k | build_v | probe_k |",
4137+
"+---------+---------+---------+",
4138+
"| 1 | 100 | 1 |",
4139+
"| 1 | 100 | 1 |",
4140+
"+---------+---------+---------+",
4141+
],
4142+
&batches
4143+
);
4144+
Ok(())
4145+
}

0 commit comments

Comments
 (0)