viirya commented on code in PR #26214:
URL: https://github.com/apache/datafusion/pull/26214#discussion_r4431021755
##########
datafusion/physical-plan/src/joins/hash_join/not_in.rs:
##########
@@ -649,29 +631,98 @@ fn correlated_row_state(
reservation.try_grow(retained_size)?;
metrics.build_mem_used.add(retained_size);
- let null_scope_map = if scope_keys.is_empty() {
- None
+ let scope = if scope_keys.is_empty() {
+ NullValueScope::Unscoped
} else {
- let scope_values = normalized[value_keys..]
- .iter()
- .map(|values| Ok(filter(values.as_ref(), &null_mask)?))
- .collect::<Result<Vec<_>>>()?;
let null_rows = build_indices.len();
- let mut map = new_join_hashmap(null_rows, reservation, metrics)?;
- let mut hashes_buffer = vec![0; null_rows];
- create_hashes(&scope_values, random_state, &mut hashes_buffer)?;
- map.update_from_iter(Box::new(hashes_buffer.iter().enumerate().rev()),
0);
- Some(map)
+ match new_join_hashmap(null_rows, reservation, metrics) {
+ Ok(mut map) => {
+ let scope_values = normalized[value_keys..]
+ .iter()
+ .map(|values| Ok(filter(values.as_ref(), &null_mask)?))
+ .collect::<Result<Vec<_>>>()?;
+ let mut hashes_buffer = vec![0; null_rows];
+ create_hashes(&scope_values, random_state, &mut
hashes_buffer)?;
+
map.update_from_iter(Box::new(hashes_buffer.iter().enumerate().rev()), 0);
+ NullValueScope::Map(map)
+ }
+ // The map only speeds up the pairing, so the join goes on without
+ // it rather than fail.
+ Err(DataFusionError::ResourcesExhausted(_)) =>
NullValueScope::Unindexed,
+ Err(e) => return Err(e),
+ }
};
Ok((
scope_map,
Some(NullValueBuildRows {
- scope_map: null_scope_map,
+ scope,
build_indices,
}),
))
}
+/// Hash table over the correlation scope keys (`on[value_keys..]`) of all
+/// build rows, used by correlated `NOT IN` (see [`NotInScope::Correlated`]).
+///
+/// Only the probe rows with a NULL value key look it up, so it is built on
+/// the first such probe row rather than with the build side: a probe side
+/// without NULLs never pays for a second hash table over the build side.
+///
+/// The map only speeds up the pairing of those probe rows with the build
+/// rows in their scope. So when the memory pool cannot hold it, the join
+/// pairs them with every build row instead, and compares their scope keys.
+/// That is the pairing of a join whose tuple elements are all value keys,
+/// which needs no map either.
+pub(super) struct BuildScopeMap {
+ /// `None` when the memory pool could not hold the map.
+ map: OnceLock<Option<Box<dyn JoinHashMapType>>>,
+ /// Charged for the map once it is built. Its lock also makes the probe
+ /// partitions build the map only once.
+ reservation: Mutex<MemoryReservation>,
+ metrics: BuildProbeJoinMetrics,
+}
+
+impl BuildScopeMap {
+ /// Returns the map, building it on first use from `scope_values`, the
+ /// build side's scope keys (`JoinLeftData::values()[value_keys..]`).
+ /// Returns `None` when the memory pool cannot hold it.
+ fn get_or_try_init(
+ &self,
+ scope_values: &[ArrayRef],
+ random_state: &RandomState,
+ ) -> Result<Option<&dyn JoinHashMapType>> {
+ if let Some(map) = self.map.get() {
+ return Ok(map.as_deref());
+ }
+ let mut reservation = self.reservation.lock();
+ if let Some(map) = self.map.get() {
+ return Ok(map.as_deref());
+ }
+ // Scope-only NULL marking uses a HashMap (the primary join map may
+ // use ArrayMap for full-key matches, but scope keys have arbitrary
+ // shape).
+ let num_rows = scope_values[0].len();
+ let mut map = match new_join_hashmap(num_rows, &mut reservation,
&self.metrics) {
+ Ok(map) => map,
+ Err(DataFusionError::ResourcesExhausted(_)) => {
+ return Ok(self.map.get_or_init(|| None).as_deref());
+ }
+ Err(e) => return Err(e),
+ };
+ let mut hashes_buffer = vec![0; num_rows];
+ create_hashes(scope_values, random_state, &mut hashes_buffer)?;
+ map.extend_zero(num_rows);
+ // A build row whose scope keys hold a NULL is in no scope.
+ let valid = matchable_join_keys(scope_values,
NullEquality::NullEqualsNothing);
+ let hashes = hashes_buffer
+ .iter()
+ .enumerate()
+ .filter(|(i, _)| valid.as_ref().is_none_or(|valid|
valid.is_valid(*i)));
+ map.update_from_iter(Box::new(hashes.rev()), 0);
Review Comment:
Done in 84d39903c. The body of `update_hash` moved into
`update_hash_from_values`, which takes the evaluated keys. `update_hash` and
the lazy scope map both call it.
##########
datafusion/physical-plan/src/joins/hash_join/not_in.rs:
##########
@@ -601,42 +602,23 @@ fn retain_probe_rows(
}
/// The scope map and NULL-valued build rows of a correlated `NOT IN`.
-#[expect(clippy::too_many_arguments, clippy::type_complexity)]
fn correlated_row_state(
on_left: &[PhysicalExprRef],
- batch: &RecordBatch,
keys: &[ArrayRef],
normalized: &[ArrayRef],
value_keys: usize,
random_state: &RandomState,
reservation: &mut MemoryReservation,
metrics: &BuildProbeJoinMetrics,
-) -> Result<(Option<Box<dyn JoinHashMapType>>, Option<NullValueBuildRows>)> {
- let num_rows = batch.num_rows();
+) -> Result<(Option<BuildScopeMap>, Option<NullValueBuildRows>)> {
// `on_left[..V]` are the value keys and `on_left[V..]` the (possibly
// empty) correlation scope keys.
let scope_keys = &on_left[value_keys..];
Review Comment:
Done in 84d39903c. `correlated_row_state` and `NotInBuild::try_new` no
longer take `on_left`. The check is `keys.len() > value_keys` now.
##########
datafusion/physical-plan/src/joins/hash_join/not_in.rs:
##########
@@ -1245,34 +1375,36 @@ where
self.batch_size,
)
};
- let rest_hashes = &hashes[group_from..];
// Estimate the groups from a sample first, because a sort of all the
// rest that finds no group worth it costs as much as the wasted
// matches did.
- let worth_sorting = saves_work(rest_hashes.len(),
estimate_groups(rest_hashes));
- let mut rest: Vec<(u64, u32)> = vec![];
- let mut groups = vec![];
- if worth_sorting {
- rest = rest_hashes
- .iter()
- .copied()
- .zip(probe_rows[group_from..].iter().copied())
- .collect();
- rest.sort_unstable_by_key(|(hash, _)| *hash);
- groups = rest.chunk_by(|a, b| a.0 == b.0).collect();
+ let mut rest = vec![];
+ if saves_work(
+ probe_rows.len() - group_from,
+ estimate_groups(&hashes[group_from..]),
+ ) {
+ rest = sort_by_hash(&hashes[group_from..],
&probe_rows[group_from..]);
}
- if !worth_sorting || !saves_work(rest.len(), groups.len()) {
- self.lookup_each_row(
+ if rest.is_empty() || !saves_work(rest.len(), count_groups(&rest)) {
Review Comment:
Both done in 84d39903c. The groups are collected once. When the real group
count rejects the sort, the second lookup keeps the suffix with `retain`, so
the rows are not sorted again. `hot_scope_among_rare_ones_is_grouped` takes
this path.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -10577,6 +10579,125 @@ mod tests {
Ok(())
}
+ /// Runs a null-aware `LeftAnti` join with the value key `id` and the
+ /// scope key `grp`, over 10,000 build rows `(2 * i, i % 100)`, with a
+ /// NULL `id` for each multiple of `null_build_every`, and the probe rows
+ /// `(probe_ids[i], i)`, in a memory pool of `memory_limit` bytes or an
+ /// unbounded one. Returns the output rows and the `build_mem_used` metric.
+ async fn run_scope_key_anti_join(
+ null_build_every: Option<i32>,
+ probe_ids: Vec<Option<i32>>,
+ memory_limit: Option<usize>,
+ ) -> Result<(usize, usize)> {
+ let build_rows = 10_000;
+ let build_id = |i: i32| match null_build_every {
+ Some(every) if i % every == 0 => None,
+ _ => Some(i * 2),
+ };
+ let left = build_table_two_cols(
+ ("id", &(0..build_rows).map(build_id).collect()),
+ ("grp", &(0..build_rows).map(|i| Some(i % 100)).collect()),
+ );
+ let probe_grps = (0..probe_ids.len() as i32).map(Some).collect();
+ let right = build_table_two_cols(("id", &probe_ids), ("grp",
&probe_grps));
+ let on = vec![
+ (
+ Arc::new(Column::new_with_schema("id", &left.schema())?) as _,
+ Arc::new(Column::new_with_schema("id", &right.schema())?) as _,
+ ),
+ (
+ Arc::new(Column::new_with_schema("grp", &left.schema())?) as _,
+ Arc::new(Column::new_with_schema("grp", &right.schema())?) as
_,
+ ),
+ ];
+ let join = HashJoinExec::try_new(
+ left,
+ right,
+ on,
+ None,
+ &JoinType::LeftAnti,
+ None,
+ PartitionMode::CollectLeft,
+ NullEquality::NullEqualsNothing,
+ true,
+ )?;
+ let mut task_ctx = prepare_task_ctx(8192, false);
+ if let Some(limit) = memory_limit {
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_limit(limit, 1.0)
+ .build_arc()?;
+ task_ctx = Arc::new(
+ TaskContext::default()
+ .with_session_config(task_ctx.session_config().clone())
+ .with_runtime(runtime),
+ );
+ }
+ let batches = common::collect(join.execute(0, task_ctx)?).await?;
+ let rows = batches.iter().map(|b| b.num_rows()).sum();
+ let build_mem_used = join
+ .metrics()
+ .unwrap()
+ .sum_by_name("build_mem_used")
+ .unwrap()
+ .as_usize();
+ Ok((rows, build_mem_used))
+ }
+
+ /// The hash table over the build rows' correlation scope keys is only
+ /// looked up for probe rows with a NULL value key, so a probe side
+ /// without one never builds it, and the result is the same either way.
+ #[tokio::test]
+ async fn null_aware_scope_map_built_only_for_null_probe_values() ->
Result<()> {
+ // `(0, 0)` matches the build row `(0, 0)`; `(1, 1)` matches nothing.
+ let (rows, without_null) =
+ run_scope_key_anti_join(None, vec![Some(0), Some(1)], None).await?;
+ assert_eq!(rows, 9_999);
+ // `(NULL, 1)` makes the 100 build rows with `grp = 1` UNKNOWN.
+ let (rows, with_null) =
+ run_scope_key_anti_join(None, vec![Some(0), None], None).await?;
+ assert_eq!(rows, 9_899);
+ // The scope map holds a bucket and a chain entry for every build row.
+ assert!(
+ with_null >= without_null + 10_000 * size_of::<u32>(),
+ "{without_null} -> {with_null}"
+ );
+ Ok(())
+ }
+
+ /// A memory pool that holds the build side but not the scope map: the
+ /// join pairs the NULL-valued probe rows with every build row instead,
+ /// and gives the same result.
+ #[tokio::test]
+ async fn null_aware_scope_map_falls_back_when_memory_is_short() ->
Result<()> {
Review Comment:
`run_scope_key_anti_join` uses `fallback_context` now (8f2cb1257). Its batch
size of 16 also runs the fallback pairing over many chunks.
I kept the two tests apart. One checks that the map is built lazily and the
other that the join falls back, and the separate names show which one broke.
The baseline run they share takes a few milliseconds.
##########
docs/source/library-user-guide/upgrading/56.0.0.md:
##########
@@ -934,9 +934,13 @@ Wire compatibility is directional:
### Null-aware joins record their `NOT IN` value key count
Multi-column `(a, b) NOT IN (SELECT x, y ...)` is now planned as a null-aware
-join. Its first equi-join keys are the tuple elements, followed by any
-correlation keys, so a null-aware join now records how many of its leading keys
-are `NOT IN` value keys:
+join. Its first equi-join keys are the tuple elements that can be NULL. A
+dictionary, run-end encoded or union element counts as one that can be NULL,
+because it can hold a NULL that its field's nullability does not show. The
Review Comment:
Reworded in 8f2cb1257. The guide now says that these types, `Null` included,
stay among the first keys, and that a correlated filter rejecting NULL makes a
nullable element one that cannot be NULL.
--
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]
