Skip to content

Commit d66f4b9

Browse files
jayzhan211jayzhan211
authored
fix: hash join dynamic filter with dictionary-encoded join keys (#24587)
## Rationale for this change Hash joins on a dictionary-encoded integer join key misbehave when dynamic filter pushdown is active. A dictionary array can carry NULL in its *values* while its key bitmap has no physical nulls (`null_count() == 0` but `logical_null_count() > 0`), and two places in the hash join only looked at the physical count: 1. **Wrong results:** for a `NullEqualsNull` join, `keys_have_null` was computed with `null_count()`, so the pushed dynamic filter was not widened with `key IS NULL` and pruned the probe-side NULL row that should have null-matched the build-side NULL. 2. **Query error:** `try_create_array_map` inferred "the key type is supported by `ArrayMap`" from the presence of min/max bounds. But bounds are also collected for dynamic filters on *any* key type, so with a dictionary key and a filter consumer the build side failed with `Unsupported type for ArrayMap: Dictionary(Int32, Int32)` instead of falling back to the regular hash map. ## What changes are included in this PR? - `stream.rs`: compute `keys_have_null` with `logical_null_count()`. - `exec.rs`: gate `ArrayMap` (perfect hash join) creation on the key type via `is_perfect_hash_join_candidate` rather than on the presence of bounds; also use `logical_null_count()` for the `NullEqualsNull` check in the same function. - Rename `should_collect_min_max_for_perfect_hash` → `is_perfect_hash_join_candidate` with a doc comment, since it describes what is checked rather than a side effect. ## Are these changes tested? Yes. A new test, `test_null_equal_dynamic_filter_keeps_probe_nulls_for_build_logical_null`, builds a `NullEqualsNull` inner join over dictionary keys with a dynamic filter consumer on the probe side and asserts both that the pushed filter contains the `IS NULL` disjunct and that the NULL–NULL row is produced. Before this PR it fails with the `ArrayMap` error; with only the `ArrayMap` gate it fails on the missing `IS NULL`. Existing hash join unit tests and the `joins`, `push_down_filter_parquet`, and `explain_analyze` sqllogictests pass. ## Are there any user-facing changes? No API changes. Joins on dictionary-encoded integer keys with dynamic filter pushdown now return correct results instead of erroring / dropping rows.
1 parent 2a7a1e3 commit d66f4b9

3 files changed

Lines changed: 317 additions & 8 deletions

File tree

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

Lines changed: 231 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,9 @@ fn try_create_array_map(
118118
perfect_hash_join_min_key_density: f64,
119119
null_equality: NullEquality,
120120
) -> Result<Option<(ArrayMap, RecordBatch, Vec<ArrayRef>)>> {
121-
if on_left.len() != 1 {
121+
// `bounds` are also collected for dynamic filters, on any key type, so the
122+
// key type cannot be inferred from their presence.
123+
if !is_perfect_hash_join_candidate(on_left, schema)? {
122124
return Ok(None);
123125
}
124126

@@ -2542,7 +2544,12 @@ impl BuildSideState {
25422544
}
25432545
}
25442546

2545-
fn should_collect_min_max_for_perfect_hash(
2547+
/// Returns whether the build side could be joined with an [`ArrayMap`]
2548+
/// (perfect hash join): a single join key of a supported integer type.
2549+
///
2550+
/// Only a candidate: the final decision also depends on the observed key
2551+
/// range and density, see [`try_create_array_map`].
2552+
fn is_perfect_hash_join_candidate(
25462553
on_left: &[PhysicalExprRef],
25472554
schema: &SchemaRef,
25482555
) -> Result<bool> {
@@ -2600,15 +2607,14 @@ async fn collect_left_input(
26002607
) -> Result<JoinLeftData> {
26012608
let schema = left_stream.schema();
26022609

2603-
let should_collect_min_max_for_phj =
2604-
should_collect_min_max_for_perfect_hash(&on_left, &schema)?;
2610+
let is_phj_candidate = is_perfect_hash_join_candidate(&on_left, &schema)?;
26052611

26062612
let initial = BuildSideState::try_new(
26072613
metrics,
26082614
reservation,
26092615
on_left.clone(),
26102616
&schema,
2611-
should_compute_dynamic_filters || should_collect_min_max_for_phj,
2617+
should_compute_dynamic_filters || is_phj_candidate,
26122618
)?;
26132619

26142620
let state = left_stream
@@ -2768,7 +2774,7 @@ async fn collect_left_input(
27682774
}
27692775
};
27702776

2771-
if should_collect_min_max_for_phj && !should_compute_dynamic_filters {
2777+
if is_phj_candidate && !should_compute_dynamic_filters {
27722778
bounds = None;
27732779
}
27742780

@@ -7535,6 +7541,225 @@ mod tests {
75357541
Ok(())
75367542
}
75377543

7544+
/// Regression test for a `NullEqualsNull` join whose build-side join key is
7545+
/// a dictionary with NULL stored in the dictionary *values*.
7546+
///
7547+
/// Such a key is only logically NULL: `null_count() == 0` because the key
7548+
/// bitmap has no physical nulls, but `logical_null_count() > 0`. The hash
7549+
/// join must still treat it as a NULL key in two places:
7550+
///
7551+
/// 1. The build side must not be turned into an `ArrayMap` (perfect hash
7552+
/// join), which only supports plain integer keys. Bounds are collected
7553+
/// for the dynamic filter regardless of key type, so the `ArrayMap`
7554+
/// path must check the key type itself instead of inferring it from
7555+
/// the presence of bounds; otherwise the build fails with
7556+
/// `Unsupported type for ArrayMap`.
7557+
/// 2. The dynamic filter pushed to the probe side must be widened with
7558+
/// `c2 IS NULL OR ...`, so the probe row with a NULL key still reaches
7559+
/// the join and null-matches the build NULL. Without it the filter
7560+
/// `c2 >= 1 AND c2 <= 2 AND c2 IN (...)` evaluates to NULL for that
7561+
/// row and silently drops it.
7562+
///
7563+
/// Setup: build keys `[1, NULL, 2]`, probe keys `[1, NULL, 4]`, inner join
7564+
/// with a `FilterExec` consuming the join's dynamic filter on the probe side.
7565+
/// Expected output: the `1 = 1` row and the `NULL = NULL` row.
7566+
#[tokio::test]
7567+
async fn test_null_equal_dynamic_filter_keeps_probe_nulls_for_build_logical_null()
7568+
-> Result<()> {
7569+
let task_ctx = Arc::new(TaskContext::default());
7570+
7571+
// Dictionary values: [1, NULL, 2]; keys: [0, 1, 2] => logical [1, NULL, 2]
7572+
let left = build_table_dict_key(
7573+
"c1",
7574+
vec![Some(1), None, Some(2)],
7575+
vec![0, 1, 2],
7576+
"d1",
7577+
vec![Some(100), Some(200), Some(300)],
7578+
);
7579+
let left_key = left
7580+
.execute(0, Arc::clone(&task_ctx))?
7581+
.next()
7582+
.await
7583+
.unwrap()?;
7584+
assert_eq!(left_key.column(0).null_count(), 0);
7585+
assert_eq!(left_key.column(0).logical_null_count(), 1);
7586+
7587+
// Dictionary values: [1, NULL, 4]; keys: [0, 1, 2] => logical [1, NULL, 4]
7588+
let right = build_table_dict_key(
7589+
"c2",
7590+
vec![Some(1), None, Some(4)],
7591+
vec![0, 1, 2],
7592+
"d2",
7593+
vec![Some(10), Some(20), Some(40)],
7594+
);
7595+
7596+
let on = vec![(
7597+
Arc::new(Column::new_with_schema("c1", &left.schema())?) as _,
7598+
Arc::new(Column::new_with_schema("c2", &right.schema())?) as _,
7599+
)];
7600+
7601+
// Wire the join's dynamic filter into a FilterExec over the probe side so
7602+
// that the filter is actually built (bounds are only collected when the
7603+
// probe subtree contains a consumer) and applied to the probe rows.
7604+
let dynamic_filter = HashJoinExec::create_dynamic_filter(&on);
7605+
let consumer: Arc<dyn PhysicalExpr> = Arc::clone(&dynamic_filter) as _;
7606+
let right = Arc::new(FilterExecBuilder::new(consumer, right).build()?);
7607+
let mut join = HashJoinExec::try_new(
7608+
left,
7609+
right,
7610+
on,
7611+
None,
7612+
&JoinType::Inner,
7613+
None,
7614+
PartitionMode::CollectLeft,
7615+
NullEquality::NullEqualsNull,
7616+
false,
7617+
)?;
7618+
join.dynamic_filter = Some(HashJoinExecDynamicFilter {
7619+
filter: Arc::clone(&dynamic_filter),
7620+
build_accumulator: OnceLock::new(),
7621+
});
7622+
7623+
// (1) Building the hash table must not fail on the dictionary key.
7624+
let stream = join.execute(0, task_ctx)?;
7625+
let batches = common::collect(stream).await?;
7626+
7627+
// (2) The build-side logical NULL must widen the pushed filter with
7628+
// `IS NULL`, because the join is NullEqualsNull.
7629+
dynamic_filter.wait_complete().await;
7630+
let filter = dynamic_filter.current()?.to_string();
7631+
assert!(
7632+
filter.contains("c2@0 IS NULL OR"),
7633+
"expected an IS NULL disjunct in the pushed filter, got: {filter}"
7634+
);
7635+
7636+
// The NULL-NULL row (d1 = 200, d2 = 20) is only produced if the probe
7637+
// NULL survived the pushed filter.
7638+
allow_duplicates! {
7639+
assert_snapshot!(batches_to_sort_string(&batches), @r"
7640+
+----+-----+----+----+
7641+
| c1 | d1 | c2 | d2 |
7642+
+----+-----+----+----+
7643+
| | 200 | | 20 |
7644+
| 1 | 100 | 1 | 10 |
7645+
+----+-----+----+----+
7646+
");
7647+
}
7648+
Ok(())
7649+
}
7650+
7651+
/// The [`PartitionMode::Partitioned`] counterpart of
7652+
/// [`test_null_equal_dynamic_filter_keeps_probe_nulls_for_build_logical_null`].
7653+
///
7654+
/// Here the build-side logical NULL lives in exactly one build partition, so
7655+
/// the `keys_have_null` flag has to survive the per-partition aggregation in
7656+
/// `build_partitioned_filter` and widen the *routed* filter as a whole: the
7657+
/// probe NULL is routed to a single `CASE` branch, and OR-ing `IS NULL` into
7658+
/// that one branch would still drop it for every other route.
7659+
#[tokio::test]
7660+
async fn test_partitioned_null_equal_dynamic_filter_keeps_probe_nulls_for_build_logical_null()
7661+
-> Result<()> {
7662+
let mut session_config = SessionConfig::default();
7663+
session_config
7664+
.options_mut()
7665+
.optimizer
7666+
.enable_dynamic_filter_pushdown = true;
7667+
let task_ctx =
7668+
Arc::new(TaskContext::default().with_session_config(session_config));
7669+
7670+
let partition_count = 4;
7671+
7672+
// Dictionary values: [1, NULL, 2]; keys: [0, 1, 2] => logical [1, NULL, 2]
7673+
let left = build_table_dict_key(
7674+
"c1",
7675+
vec![Some(1), None, Some(2)],
7676+
vec![0, 1, 2],
7677+
"d1",
7678+
vec![Some(100), Some(200), Some(300)],
7679+
);
7680+
// Dictionary values: [1, NULL, 4]; keys: [0, 1, 2] => logical [1, NULL, 4]
7681+
let right = build_table_dict_key(
7682+
"c2",
7683+
vec![Some(1), None, Some(4)],
7684+
vec![0, 1, 2],
7685+
"d2",
7686+
vec![Some(10), Some(20), Some(40)],
7687+
);
7688+
7689+
let left_key =
7690+
Arc::new(Column::new_with_schema("c1", &left.schema())?) as PhysicalExprRef;
7691+
let right_key =
7692+
Arc::new(Column::new_with_schema("c2", &right.schema())?) as PhysicalExprRef;
7693+
let on = vec![(Arc::clone(&left_key), Arc::clone(&right_key))];
7694+
7695+
// Hash-repartition both sides so the join runs with real partition routing:
7696+
// the build NULL and the probe NULL land in the same partition, and every
7697+
// other partition reports a NULL-free build.
7698+
let left_repartitioned: Arc<dyn ExecutionPlan> =
7699+
Arc::new(RepartitionExec::try_new(
7700+
left,
7701+
Partitioning::Hash(vec![left_key], partition_count),
7702+
)?);
7703+
let right_repartitioned: Arc<dyn ExecutionPlan> =
7704+
Arc::new(RepartitionExec::try_new(
7705+
right,
7706+
Partitioning::Hash(vec![Arc::clone(&right_key)], partition_count),
7707+
)?);
7708+
7709+
// The FilterExec sits above the repartition so each probe partition is only
7710+
// read after the join has waited for every partition's build report, which
7711+
// makes the finalized (routed and widened) filter the one actually applied.
7712+
let dynamic_filter = HashJoinExec::create_dynamic_filter(&on);
7713+
let consumer: Arc<dyn PhysicalExpr> = Arc::clone(&dynamic_filter) as _;
7714+
let right_filtered: Arc<dyn ExecutionPlan> =
7715+
Arc::new(FilterExecBuilder::new(consumer, right_repartitioned).build()?);
7716+
7717+
let mut join = HashJoinExec::try_new(
7718+
left_repartitioned,
7719+
right_filtered,
7720+
on,
7721+
None,
7722+
&JoinType::Inner,
7723+
None,
7724+
PartitionMode::Partitioned,
7725+
NullEquality::NullEqualsNull,
7726+
false,
7727+
)?;
7728+
join.dynamic_filter = Some(HashJoinExecDynamicFilter {
7729+
filter: Arc::clone(&dynamic_filter),
7730+
build_accumulator: OnceLock::new(),
7731+
});
7732+
7733+
let batches = crate::execution_plan::collect(Arc::new(join), task_ctx).await?;
7734+
7735+
dynamic_filter.wait_complete().await;
7736+
let filter = dynamic_filter.current()?.to_string();
7737+
// The aggregated build-side NULL must widen the routed filter once, at the
7738+
// top: `c2 IS NULL OR CASE hash(c2) % 4 WHEN ... END`.
7739+
assert!(
7740+
filter.contains("CASE"),
7741+
"expected a routed partitioned filter, got: {filter}"
7742+
);
7743+
assert!(
7744+
filter.starts_with("c2@0 IS NULL OR"),
7745+
"expected the IS NULL disjunct to wrap the whole routed filter, got: {filter}"
7746+
);
7747+
7748+
// The NULL-NULL row (d1 = 200, d2 = 20) is only produced if the probe
7749+
// NULL survived the pushed filter on its route.
7750+
allow_duplicates! {
7751+
assert_snapshot!(batches_to_sort_string(&batches), @r"
7752+
+----+-----+----+----+
7753+
| c1 | d1 | c2 | d2 |
7754+
+----+-----+----+----+
7755+
| | 200 | | 20 |
7756+
| 1 | 100 | 1 | 10 |
7757+
+----+-----+----+----+
7758+
");
7759+
}
7760+
Ok(())
7761+
}
7762+
75387763
/// Null-aware RightAnti must drop outer rows whose dictionary key is only
75397764
/// logically NULL (key points at a null dictionary value).
75407765
#[apply(hash_join_exec_configs)]

‎datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1446,6 +1446,88 @@ mod tests {
14461446
)
14471447
}
14481448

1449+
fn reported_with_null_keys(
1450+
pushdown: PushdownStrategy,
1451+
bounds: PartitionBounds,
1452+
) -> PartitionStatus {
1453+
PartitionStatus::Reported(PartitionData {
1454+
pushdown,
1455+
bounds,
1456+
keys_have_null: true,
1457+
})
1458+
}
1459+
1460+
/// A partitioned accumulator whose probe key is nullable, so that a build-side
1461+
/// NULL can actually widen the filter.
1462+
fn null_equal_partitioned_accumulator(
1463+
num_partitions: usize,
1464+
) -> SharedBuildAccumulator {
1465+
let mut acc = make_accumulator_for_test(
1466+
AccumulatedBuildData::Partitioned {
1467+
partitions: vec![PartitionStatus::Pending; num_partitions],
1468+
completed_partitions: 0,
1469+
},
1470+
test_on_right(),
1471+
);
1472+
acc.probe_schema = Arc::new(Schema::new(vec![Field::new(
1473+
"probe_key",
1474+
DataType::Int32,
1475+
true,
1476+
)]));
1477+
acc.null_equality = NullEquality::NullEqualsNull;
1478+
acc
1479+
}
1480+
1481+
/// The IS NULL disjunct must wrap the routed `CASE` as a whole. A probe NULL
1482+
/// takes exactly one route, so widening only the branch of the partition that
1483+
/// reported the NULL would still drop it on every other route.
1484+
#[test]
1485+
fn partitioned_null_keys_in_one_partition_widen_whole_routed_filter() {
1486+
let acc = null_equal_partitioned_accumulator(2);
1487+
1488+
acc.build_filter(FinalizeInput::Partitioned(vec![
1489+
reported(in_list(&[1]), no_bounds()),
1490+
reported_with_null_keys(in_list(&[2]), no_bounds()),
1491+
]))
1492+
.unwrap();
1493+
1494+
let expr = current_expr(&acc);
1495+
assert_top_binary_op(&expr, Operator::Or);
1496+
let widened = binary_expr(&expr);
1497+
assert!(
1498+
widened.left().downcast_ref::<IsNullExpr>().is_some(),
1499+
"expected the IS NULL disjunct first, got: {expr}"
1500+
);
1501+
case_expr(widened.right());
1502+
assert_eq!(
1503+
format!("{expr}").matches("IS NULL").count(),
1504+
1,
1505+
"the routed filter must be widened once, not per branch"
1506+
);
1507+
}
1508+
1509+
/// A canceled partition's build content is unknown, so it may hold a NULL key:
1510+
/// the aggregated flag must be permissive even when no partition reported one.
1511+
#[test]
1512+
fn partitioned_canceled_partition_widens_routed_filter() {
1513+
let acc = null_equal_partitioned_accumulator(2);
1514+
1515+
acc.build_filter(FinalizeInput::Partitioned(vec![
1516+
PartitionStatus::CanceledUnknown,
1517+
reported(in_list(&[2]), no_bounds()),
1518+
]))
1519+
.unwrap();
1520+
1521+
let expr = current_expr(&acc);
1522+
assert_top_binary_op(&expr, Operator::Or);
1523+
assert!(
1524+
binary_expr(&expr)
1525+
.left()
1526+
.downcast_ref::<IsNullExpr>()
1527+
.is_some()
1528+
);
1529+
}
1530+
14491531
#[test]
14501532
fn preserve_probe_nulls_only_widens_nullable_keys() {
14511533
let probe_schema = Arc::new(Schema::new(vec![

‎datafusion/physical-plan/src/joins/hash_join/stream.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -559,11 +559,13 @@ impl HashJoinStream {
559559
.bounds
560560
.clone()
561561
.unwrap_or_else(|| PartitionBounds::new(vec![]));
562-
// Arrow tracks null counts per array, so this costs no data scan.
562+
// Use the logical null count: a dictionary key whose entry points at a
563+
// NULL dictionary value is a NULL key even though the key bitmap has no
564+
// physical nulls (`null_count() == 0` but `logical_null_count() > 0`).
563565
let keys_have_null = left_data
564566
.values()
565567
.iter()
566-
.any(|array| array.null_count() > 0);
568+
.any(|array| array.logical_null_count() > 0);
567569

568570
let build_data = match self.mode {
569571
PartitionMode::Partitioned => PartitionBuildData::Partitioned {

0 commit comments

Comments
 (0)