Skip to content
Merged
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
197 changes: 197 additions & 0 deletions datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs
Original file line number Diff line number Diff line change
Expand Up @@ -742,6 +742,119 @@ pub(super) fn completed_partitions_for_test(acc: &SharedBuildAccumulator) -> usi
mod tests {
use super::*;

use arrow::array::{ArrayRef, Int32Array};
use datafusion_physical_expr::expressions::{Column, Literal};

fn test_on_right() -> Vec<PhysicalExprRef> {
vec![Arc::new(Column::new("probe_key", 0))]
}

fn test_probe_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![Field::new(
"probe_key",
DataType::Int32,
false,
)]))
}

fn test_dynamic_filter(
on_right: &[PhysicalExprRef],
) -> Arc<DynamicFilterPhysicalExpr> {
Arc::new(DynamicFilterPhysicalExpr::new(on_right.to_vec(), lit(true)))
}

fn make_accumulator_for_test(
data: AccumulatedBuildData,
on_right: Vec<PhysicalExprRef>,
) -> SharedBuildAccumulator {
let dynamic_filter = test_dynamic_filter(&on_right);
SharedBuildAccumulator {
inner: Mutex::new(AccumulatorState {
data,
completion: CompletionState::Pending,
}),
completion_notify: Notify::new(),
dynamic_filter,
on_right,
repartition_random_state: SeededRandomState::with_seed(1),
probe_schema: test_probe_schema(),
}
}

fn make_collect_left_accumulator_for_test() -> SharedBuildAccumulator {
make_accumulator_for_test(
AccumulatedBuildData::CollectLeft {
data: PartitionStatus::Pending,
reported_count: 0,
expected_reports: 1,
},
test_on_right(),
)
}

fn make_partitioned_expr_accumulator_for_test(
num_partitions: usize,
) -> SharedBuildAccumulator {
make_accumulator_for_test(
AccumulatedBuildData::Partitioned {
partitions: vec![PartitionStatus::Pending; num_partitions],
completed_partitions: 0,
},
test_on_right(),
)
}

fn in_list(values: &[i32]) -> PushdownStrategy {
PushdownStrategy::InList(Arc::new(Int32Array::from(values.to_vec())) as ArrayRef)
}

fn bounds(min: i32, max: i32) -> PartitionBounds {
PartitionBounds::new(vec![ColumnBounds::new(
ScalarValue::Int32(Some(min)),
ScalarValue::Int32(Some(max)),
)])
}

fn no_bounds() -> PartitionBounds {
PartitionBounds::new(vec![])
}

fn reported(pushdown: PushdownStrategy, bounds: PartitionBounds) -> PartitionStatus {
PartitionStatus::Reported(PartitionData { pushdown, bounds })
}

fn current_expr(acc: &SharedBuildAccumulator) -> PhysicalExprRef {
acc.dynamic_filter
.current()
.expect("dynamic filter current expression should be available")
}

fn in_list_expr(expr: &PhysicalExprRef) -> &InListExpr {
expr.downcast_ref::<InListExpr>()
.expect("expected InListExpr dynamic filter")
}

fn binary_expr(expr: &PhysicalExprRef) -> &BinaryExpr {
expr.downcast_ref::<BinaryExpr>()
.expect("expected BinaryExpr dynamic filter")
}

fn case_expr(expr: &PhysicalExprRef) -> &CaseExpr {
expr.downcast_ref::<CaseExpr>()
.expect("expected CaseExpr dynamic filter")
}

fn assert_literal_bool(expr: &PhysicalExprRef, expected: bool) {
let literal = expr
.downcast_ref::<Literal>()
.expect("expected literal bool dynamic filter");
assert_eq!(literal.value(), &ScalarValue::Boolean(Some(expected)));
}

fn assert_top_binary_op(expr: &PhysicalExprRef, expected: Operator) {
assert_eq!(binary_expr(expr).op(), &expected);
}

fn partitioned_state(acc: &SharedBuildAccumulator) -> (Vec<PartitionStatus>, usize) {
let guard = acc.inner.lock();
let AccumulatedBuildData::Partitioned {
Expand All @@ -754,6 +867,90 @@ mod tests {
(partitions.clone(), *completed_partitions)
}

#[test]
fn collect_left_updates_with_membership_only() {
let acc = make_collect_left_accumulator_for_test();

acc.build_filter(FinalizeInput::CollectLeft(reported(
in_list(&[1, 2, 3]),
no_bounds(),
)))
.unwrap();

let expr = current_expr(&acc);
in_list_expr(&expr);

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

after in_list_expr(&expr), also asserting the child column and the in-list values are [1,2,3] would be better.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

  • Updated collect_left_updates_with_membership_only to assert:
  • child column = probe_key@0
  • in-list values = [1, 2, 3]

}

#[test]
fn collect_left_updates_with_bounds_only() {
let acc = make_collect_left_accumulator_for_test();

acc.build_filter(FinalizeInput::CollectLeft(reported(
PushdownStrategy::Empty,
bounds(10, 20),
)))
.unwrap();

let expr = current_expr(&acc);
assert_top_binary_op(&expr, Operator::And);
}

#[test]
fn collect_left_empty_build_data_does_not_update_filter() {
let acc = make_collect_left_accumulator_for_test();
let initial_generation = acc.dynamic_filter.snapshot_generation();

acc.build_filter(FinalizeInput::CollectLeft(reported(
PushdownStrategy::Empty,
no_bounds(),
)))
.unwrap();

assert_eq!(
acc.dynamic_filter.snapshot_generation(),
initial_generation,
"empty CollectLeft input must not update with a no-op filter"
);
let expr = current_expr(&acc);
assert_literal_bool(&expr, true);
}

#[test]
fn partitioned_one_real_partition_with_rest_empty_skips_case() {
let acc = make_partitioned_expr_accumulator_for_test(3);

acc.build_filter(FinalizeInput::Partitioned(vec![
reported(PushdownStrategy::Empty, no_bounds()),
reported(in_list(&[2]), no_bounds()),
reported(PushdownStrategy::Empty, no_bounds()),
]))
.unwrap();

let expr = current_expr(&acc);
in_list_expr(&expr);
assert!(expr.downcast_ref::<CaseExpr>().is_none());
}

#[test]
fn partitioned_canceled_unknown_partitions_keep_unknown_routes_permissive() {
let acc = make_partitioned_expr_accumulator_for_test(2);

acc.build_filter(FinalizeInput::Partitioned(vec![
PartitionStatus::CanceledUnknown,
reported(PushdownStrategy::Empty, no_bounds()),
]))
.unwrap();

let expr = current_expr(&acc);
let case = case_expr(&expr);
assert_eq!(case.when_then_expr().len(), 1);
assert_literal_bool(&case.when_then_expr()[0].1, false);
assert_literal_bool(
case.else_expr().expect("expected permissive fallback"),
true,
);
}

// Regression guard for the build-report lifecycle fix: on `Drop`, a stream
// in `BuildReportState::ReportScheduled` still calls `report_canceled_partition`
// because it cannot tell whether the coordinator has already observed the
Expand Down
Loading