diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index f6ff292556c24..ce669509a7316 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -355,6 +355,11 @@ struct JoinBuildData { /// `build_side_has_null` were computed from the original keys in /// `collect_left_input`. values: Vec, + /// Logical nulls of the scalar key for uncorrelated null-aware `LeftMark`. + /// + /// Computed once while collecting the build side because dictionary + /// logical-null discovery can scan nested dictionary values. + left_mark_build_key_nulls: Option, /// Bounds computed from the build side; absent for an empty partition. bounds: Option, /// IN-list values or a hash-table reference used for filter pushdown. @@ -431,6 +436,10 @@ impl JoinLeftData { &self.build.values } + pub(super) fn left_mark_build_key_nulls(&self) -> Option<&NullBuffer> { + self.build.left_mark_build_key_nulls.as_ref() + } + /// returns a reference to the visited indices bitmap pub(super) fn visited_indices_bitmap(&self) -> &SharedBitmapBuilder { &self.visited_indices_bitmap @@ -3543,6 +3552,33 @@ async fn collect_left_input( && !left_values.is_empty() && left_values[0].logical_null_count() > 0; + let left_mark_build_key_nulls = if matches!( + null_aware, + Some(NullAwareMode::LeftMark { + correlated: false, + .. + }) + ) { + left_values[0] + .logical_nulls() + .filter(|nulls| nulls.null_count() > 0) + } else { + None + }; + if let Some(nulls) = &left_mark_build_key_nulls { + // `logical_nulls` reuses an array's physical null buffer when it can. + // Only charge newly materialized buffers, such as dictionary logical + // nulls, because the build batch's buffers were reserved above. + let reuses_input_buffer = left_values[0] + .nulls() + .is_some_and(|input| input.inner().ptr_eq(nulls.inner())); + if !reuses_input_buffer { + let retained_size = nulls.buffer().capacity(); + reservation.try_grow(retained_size)?; + metrics.build_mem_used.add(retained_size); + } + } + if prepared { drop(batches); // Prepared keys are direct columns. IN-list arrays share these batch @@ -3561,6 +3597,7 @@ async fn collect_left_input( map, batch, values: normalized_values, + left_mark_build_key_nulls, bounds, membership, reservation, @@ -9396,6 +9433,59 @@ mod tests { Ok(()) } + /// The dictionary counterpart of + /// [`test_null_aware_left_mark_probe_null_in_other_partition`], where the + /// probe NULL exists in a dictionary value instead of the key bitmap. + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_null_aware_left_mark_probe_logical_null_in_other_partition( + batch_size: usize, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, false); + + for order in [[0, 1], [1, 0]] { + let left = build_table_dict_key( + "c1", + vec![Some(1), Some(4)], + vec![0, 1], + "dummy", + vec![Some(10), Some(40)], + ); + let right = build_two_partition_dict_probe_with_logical_null_in_partition_0(); + + let on = vec![( + Arc::new(Column::new_with_schema("c1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("c2", &right.schema())?) as _, + )]; + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftMark, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + + let batches = collect_partitions_in_order(&join, &order, &task_ctx).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+-------+------+ + | c1 | dummy | mark | + +----+-------+------+ + | 1 | 10 | true | + | 4 | 40 | | + +----+-------+------+ + "); + } + } + Ok(()) + } + /// Test null-aware anti join when build side (left) contains NULL keys /// Expected: rows with NULL keys should not be output #[apply(hash_join_exec_configs)] @@ -10316,6 +10406,145 @@ mod tests { Ok(()) } + /// A dictionary key that points to a NULL dictionary value is a NULL build + /// key even when the dictionary key bitmap itself has no NULLs. + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_null_aware_left_mark_build_logical_null( + batch_size: usize, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, false); + + let left = build_table_dict_key( + "c1", + vec![Some(1), None, Some(4)], + vec![0, 1, 2], + "dummy", + vec![Some(10), Some(0), Some(40)], + ); + let left_key = left + .execute(0, Arc::clone(&task_ctx))? + .next() + .await + .unwrap()?; + assert_eq!(left_key.column(0).null_count(), 0); + assert_eq!(left_key.column(0).logical_null_count(), 1); + + let right = build_table_dict_key( + "c2", + vec![Some(1), Some(2), Some(3)], + vec![0, 1, 2], + "dummy", + vec![Some(100), Some(200), Some(300)], + ); + + let on = vec![( + Arc::new(Column::new_with_schema("c1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("c2", &right.schema())?) as _, + )]; + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftMark, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+-------+-------+ + | c1 | dummy | mark | + +----+-------+-------+ + | | 0 | | + | 1 | 10 | true | + | 4 | 40 | false | + +----+-------+-------+ + "); + } + Ok(()) + } + + fn build_single_column_table(name: &str, array: ArrayRef) -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new( + name, + array.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array]).unwrap(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + fn nested_dictionary_array( + leaf_values: Vec>, + inner_keys: Vec, + outer_keys: Vec, + ) -> ArrayRef { + let leaf_values: ArrayRef = Arc::new(Int32Array::from(leaf_values)); + let inner_keys = Int32Array::from(inner_keys); + let inner_values: ArrayRef = + Arc::new(DictionaryArray::::new(inner_keys, leaf_values)); + Arc::new(DictionaryArray::::new( + Int32Array::from(outer_keys), + inner_values, + )) + } + + /// A nested dictionary's logical null bitmap must be computed once for the + /// complete build key and remain indexed by the original build row. + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_null_aware_left_mark_nested_dictionary_logical_null( + batch_size: usize, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, false); + let left_key = nested_dictionary_array( + vec![Some(1), None, Some(4)], + vec![0, 1, 2], + vec![0, 1, 2], + ); + assert_eq!(left_key.logical_null_count(), 1); + let left = build_single_column_table("c1", left_key); + let right = build_single_column_table( + "c2", + nested_dictionary_array( + vec![Some(1), Some(2), Some(3)], + vec![0, 1, 2], + vec![0, 1, 2], + ), + ); + let on = vec![( + Arc::new(Column::new_with_schema("c1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("c2", &right.schema())?) as _, + )]; + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftMark, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + + let batches = common::collect(join.execute(0, task_ctx)?).await?; + let marks = batches + .iter() + .flat_map(|batch| batch.column(batch.num_columns() - 1).as_boolean().iter()) + .collect::>(); + assert_eq!(marks, vec![Some(true), None, Some(false)]); + Ok(()) + } + /// Scalar correlated null-aware left mark join where only the build side's /// value key contains NULL. /// diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index fb2263ebcca71..c2d3e79cbd7c0 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -1289,18 +1289,19 @@ fn null_aware_skip_probe_batch( | NullAwareMode::LeftMark { correlated: false, .. } => { - // `on[0]` is the `NOT IN` value key for both modes. - let probe_key_column = &state.values[0]; let is_anti = matches!(mode, NullAwareMode::LeftAnti { .. }); - let probe_has_null = if is_anti { - probe_key_column.logical_null_count() > 0 - } else { - probe_key_column.null_count() > 0 - }; - // Only batches with rows count: `NULL NOT IN (empty)` is TRUE. - left_data.record_probe_batch(state.batch.num_rows() > 0, probe_has_null); + if state.offset == (0, None) { + // `on[0]` is the `NOT IN` value key for both modes. A probe + // batch may resume lookup many times, but its summary is + // invariant and only needs to be computed once. + let probe_has_null = state.values[0].logical_null_count() > 0; + // Only batches with rows count: `NULL NOT IN (empty)` is TRUE. + left_data.record_probe_batch(state.batch.num_rows() > 0, probe_has_null); + } // Best-effort early exit; the final stage re-checks the flag - // through `report_probe_completed`. + // through `report_probe_completed`. Keep this outside the first- + // chunk guard so a sibling partition can stop an anti join between + // resumed lookups. is_anti && left_data.probe_side_has_null_hint() } } @@ -1388,13 +1389,19 @@ fn null_aware_left_mark_column( left_side: &UInt64Array, right_side: &UInt32Array, ) -> ArrayRef { - let build_key_column = &left_data.values()[0]; // Correlated joins precomputed the UNKNOWN decision per build row. let null_indices_bitmap = correlated.then(|| left_data.null_indices_bitmap().lock()); + // Uncorrelated joins use the logical-null bitmap computed once with the + // build side. Correlated joins carry the answer in `null_indices_bitmap`. + let build_key_nulls = if correlated { + None + } else { + left_data.left_mark_build_key_nulls() + }; build_null_aware_left_mark_column( left_side, right_side, - build_key_column.as_ref(), + build_key_nulls, null_indices_bitmap.as_deref(), probe_summary.has_null, probe_summary.non_empty, diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index a56aab6d250f9..c79c691fc9903 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -1382,12 +1382,14 @@ pub(crate) fn build_batch_from_indices( /// For correlated scalar `NOT IN`, `null_indices_bitmap` carries the same UNKNOWN /// decision per build row, scoped by the correlated equality keys. /// +/// When present, `build_key_nulls` is indexed by the original build row number. +/// /// This is the helper equivalent of the paper's "null bucket" and `hadNull` handling. /// It is intentionally scoped to scalar null-aware mark joins. pub(crate) fn build_null_aware_left_mark_column( build_indices: &UInt64Array, probe_indices: &UInt32Array, - build_key_column: &dyn Array, + build_key_nulls: Option<&NullBuffer>, null_indices_bitmap: Option<&BooleanBufferBuilder>, probe_side_has_null: bool, probe_side_non_empty: bool, @@ -1397,7 +1399,9 @@ pub(crate) fn build_null_aware_left_mark_column( // uncorrelated rules are cases 1-4 in the doc above. let unmatched_mark_is_null = |build_idx: usize| match null_indices_bitmap { Some(bitmap) => bitmap.get_bit(build_idx), - None if build_key_column.is_null(build_idx) => probe_side_non_empty, + None if build_key_nulls.is_some_and(|nulls| nulls.is_null(build_idx)) => { + probe_side_non_empty + } None => probe_side_has_null, };