Skip to content

Commit e547fe5

Browse files
neilconwayneilconway
committed
.
1 parent dae03ee commit e547fe5

7 files changed

Lines changed: 363 additions & 12 deletions

File tree

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

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2015,6 +2015,7 @@ async fn collect_left_input(
20152015
&mut hashes_buffer,
20162016
0,
20172017
true,
2018+
null_equality,
20182019
)?;
20192020
offset += batch.num_rows();
20202021
}
@@ -3329,6 +3330,153 @@ mod tests {
33293330
Ok(())
33303331
}
33313332

3333+
/// Under NullEqualsNothing, NULL join keys are not inserted into the hash
3334+
/// map, so a build side whose keys are all NULL produces an empty map even
3335+
/// though it contains rows. Join types that emit unmatched build rows must
3336+
/// still produce them from the visited bitmap.
3337+
#[rstest]
3338+
#[tokio::test]
3339+
async fn join_all_null_build_keys(
3340+
#[values(PartitionMode::CollectLeft, PartitionMode::Partitioned)]
3341+
partition_mode: PartitionMode,
3342+
) -> Result<()> {
3343+
let left = build_table_two_cols(
3344+
("a1", &vec![Some(1), Some(2)]),
3345+
("b1", &vec![None, None]), // all build-side join keys are NULL
3346+
);
3347+
let right = build_table_two_cols(
3348+
("a2", &vec![Some(10), Some(20), Some(30)]),
3349+
("b1", &vec![Some(4), None, Some(6)]),
3350+
);
3351+
let on = vec![(
3352+
Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
3353+
Arc::new(Column::new_with_schema("b1", &right.schema())?) as _,
3354+
)];
3355+
3356+
for join_type in [
3357+
JoinType::Inner,
3358+
JoinType::Left,
3359+
JoinType::Right,
3360+
JoinType::Full,
3361+
JoinType::LeftSemi,
3362+
JoinType::LeftAnti,
3363+
JoinType::RightSemi,
3364+
JoinType::RightAnti,
3365+
JoinType::LeftMark,
3366+
JoinType::RightMark,
3367+
] {
3368+
let (_, batches, _) = join_collect_with_partition_mode(
3369+
Arc::clone(&left),
3370+
Arc::clone(&right),
3371+
on.clone(),
3372+
&join_type,
3373+
partition_mode,
3374+
NullEquality::NullEqualsNothing,
3375+
Arc::new(TaskContext::default()),
3376+
)
3377+
.await?;
3378+
3379+
match join_type {
3380+
JoinType::Inner | JoinType::LeftSemi | JoinType::RightSemi => {
3381+
let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3382+
assert_eq!(num_rows, 0, "unexpected rows for {join_type}");
3383+
}
3384+
JoinType::Left => {
3385+
allow_duplicates! {
3386+
assert_snapshot!(batches_to_sort_string(&batches), @r"
3387+
+----+----+----+----+
3388+
| a1 | b1 | a2 | b1 |
3389+
+----+----+----+----+
3390+
| 1 | | | |
3391+
| 2 | | | |
3392+
+----+----+----+----+
3393+
");
3394+
}
3395+
}
3396+
JoinType::Right => {
3397+
allow_duplicates! {
3398+
assert_snapshot!(batches_to_sort_string(&batches), @r"
3399+
+----+----+----+----+
3400+
| a1 | b1 | a2 | b1 |
3401+
+----+----+----+----+
3402+
| | | 10 | 4 |
3403+
| | | 20 | |
3404+
| | | 30 | 6 |
3405+
+----+----+----+----+
3406+
");
3407+
}
3408+
}
3409+
JoinType::Full => {
3410+
allow_duplicates! {
3411+
assert_snapshot!(batches_to_sort_string(&batches), @r"
3412+
+----+----+----+----+
3413+
| a1 | b1 | a2 | b1 |
3414+
+----+----+----+----+
3415+
| | | 10 | 4 |
3416+
| | | 20 | |
3417+
| | | 30 | 6 |
3418+
| 1 | | | |
3419+
| 2 | | | |
3420+
+----+----+----+----+
3421+
");
3422+
}
3423+
}
3424+
JoinType::LeftAnti => {
3425+
allow_duplicates! {
3426+
assert_snapshot!(batches_to_sort_string(&batches), @r"
3427+
+----+----+
3428+
| a1 | b1 |
3429+
+----+----+
3430+
| 1 | |
3431+
| 2 | |
3432+
+----+----+
3433+
");
3434+
}
3435+
}
3436+
JoinType::RightAnti => {
3437+
allow_duplicates! {
3438+
assert_snapshot!(batches_to_sort_string(&batches), @r"
3439+
+----+----+
3440+
| a2 | b1 |
3441+
+----+----+
3442+
| 10 | 4 |
3443+
| 20 | |
3444+
| 30 | 6 |
3445+
+----+----+
3446+
");
3447+
}
3448+
}
3449+
JoinType::LeftMark => {
3450+
allow_duplicates! {
3451+
assert_snapshot!(batches_to_sort_string(&batches), @r"
3452+
+----+----+-------+
3453+
| a1 | b1 | mark |
3454+
+----+----+-------+
3455+
| 1 | | false |
3456+
| 2 | | false |
3457+
+----+----+-------+
3458+
");
3459+
}
3460+
}
3461+
JoinType::RightMark => {
3462+
allow_duplicates! {
3463+
assert_snapshot!(batches_to_sort_string(&batches), @r"
3464+
+----+----+-------+
3465+
| a2 | b1 | mark |
3466+
+----+----+-------+
3467+
| 10 | 4 | false |
3468+
| 20 | | false |
3469+
| 30 | 6 | false |
3470+
+----+----+-------+
3471+
");
3472+
}
3473+
}
3474+
}
3475+
}
3476+
3477+
Ok(())
3478+
}
3479+
33323480
#[apply(hash_join_exec_configs)]
33333481
#[tokio::test]
33343482
async fn partitioned_join_left_one(
@@ -4459,6 +4607,7 @@ mod tests {
44594607
&[right_keys_values],
44604608
NullEquality::NullEqualsNothing,
44614609
&hashes_buffer,
4610+
None,
44624611
8192,
44634612
(0, None),
44644613
&mut probe_indices_buffer,
@@ -4520,6 +4669,7 @@ mod tests {
45204669
&[right_keys_values],
45214670
NullEquality::NullEqualsNothing,
45224671
&hashes_buffer,
4672+
None,
45234673
8192,
45244674
(0, None),
45254675
&mut probe_indices_buffer,

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

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ use crate::joins::hash_join::shared_bounds::{
3333
PartitionBounds, PartitionBuildData, SharedBuildAccumulator,
3434
};
3535
use crate::joins::utils::{
36-
OnceFut, equal_rows_arr, get_final_indices_from_shared_bitmap,
36+
OnceFut, equal_rows_arr, get_final_indices_from_shared_bitmap, matchable_join_keys,
3737
};
3838
use crate::stream::EmptyRecordBatchStream;
3939
use crate::{
@@ -48,6 +48,7 @@ use crate::{
4848
};
4949

5050
use arrow::array::{Array, ArrayRef, UInt32Array, UInt64Array};
51+
use arrow::buffer::NullBuffer;
5152
use arrow::datatypes::{Schema, SchemaRef};
5253
use arrow::record_batch::RecordBatch;
5354
use datafusion_common::{
@@ -156,6 +157,10 @@ pub(super) struct ProcessProbeBatchState {
156157
batch: RecordBatch,
157158
/// Probe-side on expressions values
158159
values: Vec<ArrayRef>,
160+
/// Combined validity of the probe-side key columns, set when NULL keys
161+
/// exist and cannot match (`NullEquality::NullEqualsNothing`); NULL rows
162+
/// are skipped during JoinHashMap lookups
163+
valid_keys: Option<NullBuffer>,
159164
/// Starting offset for JoinHashMap lookups
160165
offset: MapOffset,
161166
/// Max joined probe-side index from current batch
@@ -394,13 +399,15 @@ pub(super) fn lookup_join_hashmap(
394399
probe_side_values: &[ArrayRef],
395400
null_equality: NullEquality,
396401
hashes_buffer: &[u64],
402+
valid_keys: Option<&NullBuffer>,
397403
limit: usize,
398404
offset: MapOffset,
399405
probe_indices_buffer: &mut Vec<u32>,
400406
build_indices_buffer: &mut Vec<u64>,
401407
) -> Result<(UInt64Array, UInt32Array, Option<MapOffset>)> {
402408
let next_offset = build_hashmap.get_matched_indices_with_limit_offset(
403409
hashes_buffer,
410+
valid_keys,
404411
limit,
405412
offset,
406413
probe_indices_buffer,
@@ -516,7 +523,7 @@ impl HashJoinStream {
516523
join_type: JoinType,
517524
left_data: &JoinLeftData,
518525
) -> HashJoinStreamState {
519-
if left_data.map().is_empty()
526+
if left_data.batch().num_rows() == 0
520527
&& join_type.empty_build_side_produces_empty_result()
521528
{
522529
HashJoinStreamState::Completed
@@ -679,6 +686,7 @@ impl HashJoinStream {
679686
// Precalculate hash values for fetched batch
680687
let keys_values = evaluate_expressions_to_arrays(&self.on_right, &batch)?;
681688

689+
let mut valid_keys = None;
682690
if let Map::HashMap(_) = self.build_side.try_as_ready()?.left_data.map() {
683691
self.hashes_buffer.clear();
684692
self.hashes_buffer.resize(batch.num_rows(), 0);
@@ -687,6 +695,7 @@ impl HashJoinStream {
687695
&self.random_state,
688696
&mut self.hashes_buffer,
689697
)?;
698+
valid_keys = matchable_join_keys(&keys_values, self.null_equality);
690699
}
691700

692701
self.join_metrics.input_batches.add(1);
@@ -696,6 +705,7 @@ impl HashJoinStream {
696705
HashJoinStreamState::ProcessProbeBatch(ProcessProbeBatchState {
697706
batch,
698707
values: keys_values,
708+
valid_keys,
699709
offset: (0, None),
700710
joined_probe_idx: None,
701711
});
@@ -759,14 +769,9 @@ impl HashJoinStream {
759769
}
760770
}
761771

762-
// If the build side is empty, this stream only reaches ProcessProbeBatch for
763-
// join types whose output still depends on probe rows.
764772
let is_empty = build_side.left_data.map().is_empty();
765773

766774
if is_empty {
767-
// Invariant: state_after_build_ready should have already completed
768-
// join types whose result is fixed to empty when the build side is empty.
769-
debug_assert!(!self.join_type.empty_build_side_produces_empty_result());
770775
let result = build_batch_empty_build_side(
771776
&self.schema,
772777
build_side.left_data.batch(),
@@ -790,6 +795,7 @@ impl HashJoinStream {
790795
&state.values,
791796
self.null_equality,
792797
&self.hashes_buffer,
798+
state.valid_keys.as_ref(),
793799
self.batch_size,
794800
state.offset,
795801
&mut self.probe_indices_buffer,

0 commit comments

Comments
 (0)