jayzhan211 commented on code in PR #26647:
URL: https://github.com/apache/datafusion/pull/26647#discussion_r4365598740
##########
datafusion/physical-plan/src/joins/array_map.rs:
##########
@@ -245,6 +245,9 @@ impl ArrayMap {
self.data.capacity() * size_of::<u32>() + self.next.capacity() *
size_of::<u32>()
}
+ // Kept out of line: inlined into `HashJoinStream::process_probe_batch`,
+ // this loop measured ~4-6% slower on high-fanout probes.
+ #[inline(never)]
Review Comment:
I only measured these on aarch64, so I've left both pins out. The
Performance section of the description now says what they recovered there and
that it isn't measured on x86_64. If one turns out to be worth keeping, I'll
send it as its own PR.
##########
datafusion/physical-plan/src/joins/hash_join/not_in.rs:
##########
@@ -0,0 +1,972 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! `NOT IN` (null-aware) semantics of the hash join.
+//!
+//! `x NOT IN (subquery)` is `NOT (x IN (subquery))` under SQL three-valued
+//! logic: `x IN (..)` is TRUE when some subquery value equals `x`, UNKNOWN
+//! (NULL) when none does but a NULL on either side might, and FALSE
+//! otherwise. `NOT IN` keeps a row only where `IN` is FALSE.
+//!
+//! The regular hash join finds the equal pairs; what NULLs change in the
+//! rows a `NOT IN` emits is here. The join reaches it through:
+//! - [`NotIn::try_new`] when the plan is built or executed: is this join
+//! legal, and which form is it;
+//! - [`NotInBuild::try_new`] once per build, from `collect_left_input`;
+//! - [`NotInProbe::before_lookup`] and [`NotInProbe::after_lookup`] around
+//! each probe batch's lookup;
+//! - [`NotInBuild::final_chunk`] on each chunk of build rows emitted after
+//! the probe side is exhausted.
+
+use std::sync::Arc;
+
+use parking_lot::Mutex;
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, RecordBatch,
+ UInt32Array, UInt64Array,
+};
+use arrow::buffer::{BooleanBuffer, NullBuffer};
+use arrow::compute::{filter, take};
+use arrow::util::bit_util;
+use datafusion_common::hash_utils::{RandomState, create_hashes};
+use datafusion_common::{JoinSide, JoinType, NullEquality, Result, plan_err};
+use datafusion_execution::memory_pool::MemoryReservation;
+use datafusion_expr::{ColumnarValue, Operator};
+use datafusion_physical_expr::PhysicalExprRef;
+use datafusion_physical_expr_common::datum::apply_cmp;
+
+use crate::joins::PartitionMode;
+use crate::joins::SharedBitmapBuilder;
+use crate::joins::hash_join::exec::{JoinLeftData, new_join_hashmap};
+use crate::joins::hash_join::probe_completion::ProbeSideSummary;
+use crate::joins::hash_join::stream::lookup_join_hashmap;
+use crate::joins::join_hash_map::JoinHashMapType;
+use crate::joins::join_hash_map::ProbeOffset::AtRow;
+use crate::joins::utils::{
+ BuildProbeJoinMetrics, JoinFilter, apply_join_filter_to_indices,
update_hash,
+};
+
+/// Which `NOT IN` a null-aware join computes, derived from
+/// `HashJoinExec::null_aware`, the join type, the keys and the filter.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(super) enum NotIn {
+ /// `build.key NOT IN (probe.key)`: the final stage gives every build row
+ /// its `IN` mark, from the visited bitmap and the NULLs seen.
+ BuildRows {
+ scope: NotInScope,
+ output: NotInOutput,
+ },
+ /// `probe.key NOT IN (build.key)`, single key, `CollectLeft`, no filter
+ /// (`RightAnti`): emits probe rows, none of them once any build key is
+ /// NULL, and never a probe row whose key is NULL.
+ ProbeRows,
+}
+
+/// Which probe rows a NULL can make a build row's mark UNKNOWN for.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(super) enum NotInScope {
+ /// One value key, no correlation keys, no filter: a NULL anywhere on the
+ /// probe side makes every unmatched mark UNKNOWN, so the probe partitions
+ /// share one summary of what they saw (`ProbeCompletion`).
+ Uncorrelated,
+ /// Correlation scope keys (`on[value_keys..]`), a multi-column value key
+ /// (`on[..value_keys]`), or a filter: only the probe rows in a build
+ /// row's scope count, so UNKNOWN is decided per build row.
+ Correlated { value_keys: usize },
+}
+
+/// What a [`NotIn::BuildRows`] join emits.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(super) enum NotInOutput {
+ /// `LeftAnti`: the build rows whose mark is FALSE.
+ Rows,
+ /// `LeftMark`: every build row with its mark, a nullable Boolean column.
+ Mark,
+}
+
+impl NotIn {
+ /// Validates that `null_aware` may be set for this join and returns its
+ /// form. The errors are those of the legality table this replaced.
+ pub(super) fn try_new(
+ join_type: JoinType,
+ partition_mode: PartitionMode,
+ num_keys: usize,
+ num_value_keys: usize,
+ has_filter: bool,
+ ) -> Result<Self> {
+ if num_value_keys == 0 || num_value_keys > num_keys {
+ return plan_err!(
+ "null_aware {join_type} join needs between 1 and {num_keys}
`NOT IN` value keys, got {num_value_keys}"
+ );
+ }
+ // `num_keys > 1` also covers a multi-column value key without any
+ // correlation: a NULL in one tuple element leaves the comparison FALSE
+ // whenever another element is a definite mismatch, so it too must be
+ // decided per build row.
+ let scope = if num_keys > 1 || has_filter {
+ NotInScope::Correlated {
+ value_keys: num_value_keys,
+ }
+ } else {
+ NotInScope::Uncorrelated
+ };
+ let output = match (join_type, partition_mode) {
+ (JoinType::LeftAnti, _) => NotInOutput::Rows,
+ (JoinType::LeftMark, _) => NotInOutput::Mark,
+ // `PartitionMode::CollectLeft` is safe because `RightAnti` is
probe-driven
+ (JoinType::RightAnti, PartitionMode::CollectLeft) => {
+ if num_keys != 1 {
+ return plan_err!(
+ "null_aware {join_type} joins only support single
column join key, got {num_keys} columns"
+ );
+ }
+ if has_filter {
+ return plan_err!(
+ "null_aware RightAnti join does not support a join
filter"
+ );
+ }
+ return Ok(Self::ProbeRows);
+ }
+ _ => {
+ return plan_err!(
+ "null_aware can only be true for LeftAnti joins and
RightAnti joins with `CollectLeft` `PartitionMode`, or LeftMark joins, got
{join_type} with {partition_mode}"
+ );
+ }
+ };
+ // Build-driven joins decide in the final stage from state shared
+ // across the probe partitions: the per-build-row null bitmap for
+ // the correlated forms, and the probe-side NULL flag for the
+ // uncorrelated ones. `Partitioned` gives each output partition its
+ // own build side, so neither is global there.
+ if partition_mode == PartitionMode::Partitioned {
+ return plan_err!(
+ "null_aware joins require PartitionMode::CollectLeft, got
PartitionMode::Partitioned"
+ );
+ }
+ Ok(Self::BuildRows { scope, output })
+ }
+
+ /// Number of leading keys that are `NOT IN` value keys; the keys after
+ /// them are correlation scope keys.
+ pub(super) fn value_keys(self) -> usize {
+ match self {
+ Self::BuildRows {
+ scope: NotInScope::Correlated { value_keys },
+ ..
+ } => value_keys,
+ _ => 1,
+ }
+ }
+}
+
+/// What a build side keeps for its `NOT IN`, one variant per form.
+pub(super) enum NotInBuild {
+ /// Correlated [`NotIn::BuildRows`]: UNKNOWN is recorded per build row.
+ Correlated {
+ value_keys: usize,
+ output: NotInOutput,
+ /// One bit per build row, set once its mark is UNKNOWN; shared by
+ /// every probe partition, and never cleared.
+ unknown: SharedBitmapBuilder,
+ /// Hash table over the scope keys (`on[value_keys..]`) of all build
+ /// rows, probed with NULL-valued probe rows. `None` without scope
keys.
+ scope_map: Option<Box<dyn JoinHashMapType>>,
+ /// The build rows with a NULL value key, `None` when there are none.
+ null_value_rows: Option<NullValueBuildRows>,
+ },
+ /// Uncorrelated [`NotIn::BuildRows`]: what every probe partition saw
+ /// decides every unmatched mark.
+ Uncorrelated {
+ output: NotInOutput,
+ /// NULLs of the value key, `None` when it has none. A NULL build
+ /// key's mark is UNKNOWN once the probe side has a row.
+ key_nulls: Option<NullBuffer>,
+ },
+ /// [`NotIn::ProbeRows`].
+ ProbeRows {
+ /// Whether a build key is NULL: then no probe row is emitted.
+ build_has_null: bool,
+ },
+}
+
+impl NotInBuild {
+ /// Builds the `NOT IN` state of a collected build side and charges it to
+ /// `reservation`. `keys` are the build keys as evaluated, `normalized`
+ /// with `-0.0` rewritten (see `JoinLeftData::values`).
+ #[expect(clippy::too_many_arguments)]
+ pub(super) fn try_new(
+ not_in: NotIn,
+ on_left: &[PhysicalExprRef],
+ batch: &RecordBatch,
+ keys: &[ArrayRef],
+ normalized: &[ArrayRef],
+ random_state: &RandomState,
+ reservation: &mut MemoryReservation,
+ metrics: &BuildProbeJoinMetrics,
+ ) -> Result<Self> {
+ let num_rows = batch.num_rows();
+ match not_in {
+ NotIn::ProbeRows => Ok(Self::ProbeRows {
+ build_has_null: keys[0].logical_null_count() > 0,
+ }),
+ NotIn::BuildRows {
+ scope: NotInScope::Uncorrelated,
+ output,
+ } => {
+ // `LeftMark` uses logical NULLs; `LeftAnti` still uses
+ // physical ones (#26568).
+ let key_nulls = match output {
+ NotInOutput::Rows => keys[0].nulls().cloned(),
+ NotInOutput::Mark => keys[0].logical_nulls(),
+ }
+ .filter(|n| n.null_count() > 0);
+ if let Some(nulls) = &key_nulls {
+ // `logical_nulls` reuses an array's physical null buffer
+ // when it can, which the batch reservation already holds.
+ // Charge only a newly materialized one (dictionaries).
+ let reuses_input_buffer = keys[0]
+ .nulls()
+ .is_some_and(|input|
input.inner().ptr_eq(nulls.inner()));
+ if !reuses_input_buffer {
+ let size = nulls.buffer().capacity();
+ reservation.try_grow(size)?;
+ metrics.build_mem_used.add(size);
+ }
+ }
+ Ok(Self::Uncorrelated { output, key_nulls })
+ }
+ NotIn::BuildRows {
+ scope: NotInScope::Correlated { value_keys },
+ output,
+ } => {
+ let bitmap_size = bit_util::ceil(num_rows, 8);
+ reservation.try_grow(bitmap_size)?;
+ metrics.build_mem_used.add(bitmap_size);
+ let mut unknown = BooleanBufferBuilder::new(num_rows);
+ unknown.append_n(num_rows, false);
+
+ let (scope_map, null_value_rows) = correlated_row_state(
+ on_left,
+ batch,
+ keys,
+ normalized,
+ value_keys,
+ random_state,
+ reservation,
+ metrics,
+ )?;
+ Ok(Self::Correlated {
+ value_keys,
+ output,
+ unknown: Mutex::new(unknown),
+ scope_map,
+ null_value_rows,
+ })
+ }
+ }
+ }
+
+ /// The rows to emit for one chunk of final build rows, once every probe
+ /// partition has finished (`summary`), with the mark column for
+ /// [`NotInOutput::Mark`].
+ ///
+ /// A matched build row's `IN` mark is TRUE. An unmatched one's is UNKNOWN
+ /// (NULL) or FALSE (the "null bucket" and `hadNull` of Neumann, Leis and
+ /// Kemper, "The Complete Story of Joins (in HyPer)", Section 5.6):
+ /// - correlated: UNKNOWN where the bitmap recorded it;
+ /// - uncorrelated, a probe key was NULL: UNKNOWN;
+ /// - uncorrelated, the probe side was empty: FALSE (`x IN (empty)`);
+ /// - uncorrelated otherwise: UNKNOWN exactly for a NULL build key.
+ pub(super) fn final_chunk(
+ &self,
+ summary: ProbeSideSummary,
+ left_side: UInt64Array,
+ right_side: UInt32Array,
+ ) -> (UInt64Array, UInt32Array, Option<ArrayRef>) {
+ match self {
+ Self::ProbeRows { .. } => (left_side, right_side, None),
+ Self::Correlated {
+ output, unknown, ..
+ } => {
+ let unknown = unknown.lock();
+ emit(*output, left_side, right_side, |row|
unknown.get_bit(row))
+ }
+ // A probe batch with a NULL key also has rows: `has_null` implies
+ // `non_empty`.
+ Self::Uncorrelated { output, .. } if summary.has_null => {
+ emit_uniform(*output, left_side, right_side, true)
+ }
+ Self::Uncorrelated { output, .. } if !summary.non_empty => {
+ emit_uniform(*output, left_side, right_side, false)
+ }
+ Self::Uncorrelated { output, key_nulls } => match key_nulls {
+ Some(nulls) => {
+ emit(*output, left_side, right_side, |row|
nulls.is_null(row))
+ }
+ None => emit_uniform(*output, left_side, right_side, false),
+ },
+ }
+ }
+}
+
+/// Emits a chunk of final build rows given which unmatched rows have an
+/// UNKNOWN mark (see [`NotInBuild::final_chunk`]). `LeftAnti` chunks hold
+/// only unmatched rows and keep those whose mark is FALSE; `LeftMark` chunks
+/// hold every row, a valid probe index marking a matched one.
+fn emit(
+ output: NotInOutput,
+ left_side: UInt64Array,
+ right_side: UInt32Array,
+ is_unknown: impl Fn(usize) -> bool,
+) -> (UInt64Array, UInt32Array, Option<ArrayRef>) {
+ match output {
+ NotInOutput::Rows => {
+ let kept = UInt64Array::from_iter_values(
+ left_side
+ .values()
+ .iter()
+ .copied()
+ .filter(|row| !is_unknown(*row as usize)),
+ );
+ let right_side = UInt32Array::new_null(kept.len());
+ (kept, right_side, None)
+ }
+ NotInOutput::Mark => {
+ let marks: BooleanArray = left_side
+ .values()
+ .iter()
+ .zip(right_side.iter())
+ .map(|(row, probe)| {
+ if probe.is_some() {
+ Some(true)
+ } else {
+ (!is_unknown(*row as usize)).then_some(false)
+ }
+ })
+ .collect();
+ (left_side, right_side, Some(Arc::new(marks)))
+ }
+ }
+}
+
+/// [`emit`] when every unmatched row's mark is the same: an anti join then
+/// keeps all of the chunk or none of it.
+fn emit_uniform(
+ output: NotInOutput,
+ left_side: UInt64Array,
+ right_side: UInt32Array,
+ unknown: bool,
+) -> (UInt64Array, UInt32Array, Option<ArrayRef>) {
+ match (output, unknown) {
+ (NotInOutput::Rows, true) => {
+ (UInt64Array::new_null(0), UInt32Array::new_null(0), None)
+ }
+ (NotInOutput::Rows, false) => (left_side, right_side, None),
+ (NotInOutput::Mark, _) => emit(output, left_side, right_side, |_|
unknown),
+ }
+}
+
+/// One probe stream's `NOT IN` work, with scratch space for the correlated
+/// marking pass, reused across probe batches.
+#[derive(Debug, Default)]
+pub(super) struct NotInProbe {
+ hashes: Vec<u64>,
+ probe_indices: Vec<u32>,
+ build_indices: Vec<u64>,
+}
+
+impl NotInProbe {
+ /// Before a lookup of a probe batch. `first_chunk` is whether this is the
+ /// batch's first lookup: a batch may resume lookup several times, but
+ /// what it contributes to `NOT IN` is recorded once.
+ ///
+ /// - [`NotIn::ProbeRows`] emits probe rows, so a NULL build key skips
+ /// every batch.
+ /// - Uncorrelated: record what the batch saw for the final stage. The
+ /// anti join skips once any partition saw a NULL probe key: every
+ /// unmatched mark is then UNKNOWN. This is a best-effort early exit,
+ /// checked on every resumed lookup; the final stage re-checks.
+ /// - Correlated: mark the build rows this batch makes UNKNOWN. Must also
+ /// run when the build map is empty: an all-NULL-key build side still
+ /// needs its marks.
+ ///
+ /// Returns whether the batch cannot contribute any output row, so its
+ /// lookup is skipped. A no-op returning `false` for other joins.
+ #[expect(clippy::too_many_arguments)]
+ pub(super) fn before_lookup(
+ &mut self,
+ left_data: &JoinLeftData,
+ probe_batch: &RecordBatch,
+ probe_values: &[ArrayRef],
+ first_chunk: bool,
+ filter: Option<&JoinFilter>,
+ random_state: &RandomState,
+ batch_size: usize,
+ ) -> Result<bool> {
+ let Some(not_in) = left_data.not_in() else {
+ return Ok(false);
+ };
+ match not_in {
+ NotInBuild::ProbeRows { build_has_null } => Ok(*build_has_null),
+ NotInBuild::Uncorrelated { output, .. } => {
+ if first_chunk {
+ // `on[0]` is the `NOT IN` value key. Only batches with
rows
+ // count: `NULL NOT IN (empty)` is TRUE.
+ let probe_has_null = probe_values[0].logical_null_count()
> 0;
+ left_data
+ .record_probe_batch(probe_batch.num_rows() > 0,
probe_has_null);
+ }
+ Ok(*output == NotInOutput::Rows &&
left_data.probe_side_has_null_hint())
+ }
+ NotInBuild::Correlated {
+ value_keys,
+ output,
+ unknown,
+ scope_map,
+ null_value_rows,
+ } => {
+ // The join filter needs the join type the output implies.
+ let join_type = match output {
+ NotInOutput::Rows => JoinType::LeftAnti,
+ NotInOutput::Mark => JoinType::LeftMark,
+ };
+ if first_chunk {
+ mark_null_candidates_for_probe_batch(
+ left_data,
+ unknown,
+ scope_map.as_deref(),
+ null_value_rows.as_ref(),
+ probe_batch,
+ probe_values,
+ *value_keys,
+ filter,
+ join_type,
+ random_state,
+ batch_size,
+ &mut self.hashes,
+ &mut self.probe_indices,
+ &mut self.build_indices,
+ )?;
+ }
+ Ok(false)
+ }
+ }
+ }
+
+ /// After a lookup: [`NotIn::ProbeRows`] drops the probe rows whose key is
+ /// NULL, since `NULL NOT IN (..)` is never TRUE. `valid_keys` is the probe
+ /// keys' validity (`matchable_join_keys`), when the batch has one.
+ pub(super) fn after_lookup(
Review Comment:
After the rebase onto #26213, `after_lookup` also drops the probe rows that
`before_lookup` marked UNKNOWN for a correlated `RightAnti`. That bitmap is
per-batch scratch, so it lives in `NotInProbe`, and `after_lookup` is now a
`&self` method there that uses it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
