comphead commented on code in PR #26214:
URL: https://github.com/apache/datafusion/pull/26214#discussion_r4429707816
##########
datafusion/optimizer/src/decorrelate_predicate_subquery.rs:
##########
@@ -631,6 +632,34 @@ impl InValue {
scope_filters,
))
}
+
+ /// Can either side hold a NULL that its schema does not show? See
+ /// [`type_may_hide_null`].
+ fn may_hide_null(
Review Comment:
`may_hide_null` reads each side's type after coercion, so a dictionary under
a cast slips through. In `multi_column_not_in_sees_null_dictionary_values` both
sides are `Dictionary(Int32, Int32)`, and coercion leaves them alone. If `y` is
`Dictionary(Int32, Int64)` or a plain `Int64` instead, `comparison_coercion`
resolves the pair to `Int64` and wraps the dictionary side in `CAST(.. AS
Int64)`. The constant-element EXPLAIN in `null_aware_anti_join.slt` shows this
kind of cast on `y`. Both types are then `Int64`, so `b = y` becomes a scope
key. The cast turns the hidden NULL into a null bit, and `matchable_join_keys`
then puts the row out of scope.
From reading the code, I expect the `a = 1` row to return `true` again
instead of NULL. I have not run it. Checking the types of the columns each side
reads (at least through `strip_casts`) instead of the output type would close
this. A variant of the existing test with an `Int64` `y` would pin 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:
Now that the scope map is built lazily, `on_left` is only used for
`on_left[value_keys..].is_empty()`, which is the same as `keys.len() >
value_keys`. Dropping the parameter here and in `NotInBuild::try_new` also
makes the `PhysicalExprRef` import unused.
##########
datafusion/optimizer/src/decorrelate_predicate_subquery.rs:
##########
@@ -1037,6 +1079,51 @@ fn build_join(
}
let left = projected_left.as_ref().unwrap_or(left);
+ // The first `V` keys of a null-aware join are its `NOT IN` value keys.
+ //
+ // A tuple element that cannot be NULL in scope on either side never makes
+ // the tuple comparison UNKNOWN: compared with a definite value it is
either
+ // equal, which leaves the other elements to decide, or a definite
mismatch,
+ // which makes the comparison FALSE whatever the other elements are. It
+ // therefore only selects which subquery rows a NULL in another element is
+ // compared with, which is what a correlation scope key does. Moving such
+ // elements after the ones that can be NULL turns them into scope keys, so
+ // the hash join finds the rows a NULL-valued row is compared with through
a
+ // hash lookup instead of pairing it with every row of the other side
+ // (<https://github.com/apache/datafusion/issues/26088>).
+ //
+ // A null-aware tuple has at least one element that can be NULL, so `V` is
+ // at least one. With exactly one, the join is a scalar `NOT IN` on that
+ // element, correlated by the other elements.
+ //
+ // An element whose type can hide a NULL from its schema (see
+ // `type_may_hide_null`) stays a value key too. As a scope key, such a NULL
+ // would put its row outside every scope, which makes the comparison FALSE
+ // instead of UNKNOWN; a value key still sees it.
+ let null_aware_value_keys = if null_aware && in_values.len() > 1 {
+ let keeps_value_key = in_values
+ .iter()
+ .zip(elements_may_be_null)
+ .map(|(in_value, may_be_null)| {
+ Ok(may_be_null
+ || in_value.may_hide_null(left.schema(),
sub_query_alias.schema())?)
+ })
+ .collect::<Result<Vec<_>>>()?;
+ let (nullable, non_nullable): (Vec<_>, Vec<_>) = in_values
+ .into_iter()
+ .zip(keeps_value_key)
+ .partition(|(_, keeps_value_key)| *keeps_value_key);
+ let null_aware_value_keys = nullable.len();
+ in_values = nullable
+ .into_iter()
+ .chain(non_nullable)
+ .map(|(in_value, _)| in_value)
+ .collect();
Review Comment:
This reorder builds four `Vec`s. A stable sort on the zipped flags keeps the
order within each group in one `Vec`:
```rust
let mut keyed = in_values
.into_iter()
.zip(elements_may_be_null)
.map(|(in_value, may_be_null)| {
let keeps = may_be_null
|| in_value.may_hide_null(left.schema(),
sub_query_alias.schema())?;
Ok((keeps, in_value))
})
.collect::<Result<Vec<_>>>()?;
keyed.sort_by_key(|(keeps_value_key, _)| !keeps_value_key);
let null_aware_value_keys = keyed.iter().take_while(|(keeps, _)|
*keeps).count();
in_values = keyed.into_iter().map(|(_, in_value)| in_value).collect();
```
##########
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:
This reads as if a dictionary element made the join null-aware. `null_aware`
still comes only from the elements that can be NULL in scope. A dictionary,
run-end encoded or union element only stays a value key once another element
already makes the join null-aware. Maybe: "In a null-aware join, dictionary,
run-end encoded and union elements stay value keys, because they can hold a
NULL that their field's nullability does not show." The guide could also say
that a NULL-rejecting correlated filter turns a nullable element into a scope
key.
##########
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:
This repeats `update_hash` in `joins/utils.rs` step by step:
`create_hashes`, `extend_zero`, `matchable_join_keys(NullEqualsNothing)`, the
`is_none_or` filter, and `update_from_iter(.. .rev())`. The only step it leaves
out is evaluating `on` against the batch. Could the body of `update_hash` move
into a function that takes the evaluated `&[ArrayRef]`, called from both
places? That keeps a single copy of the NULL-skipping rule for join maps.
##########
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:
Two small repeats:
- `count_groups(&rest)` walks `rest` once, and the `chunk_by(..).collect()`
below walks it again. Collecting `groups` once and using `groups.len()` would
drop `count_groups`.
- When the sample says yes but the real group count says no, the sorted
`rest` is thrown away. A suffix of it is then sorted again after the second
`lookup_each_row`. Since `probe_rows` is increasing, `rest.retain(|(_, row)|
*row >= probe_rows[group_from])` keeps the hash order and saves the second
sort. This one is optional, since that path should be rare.
##########
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:
This test and `null_aware_scope_map_built_only_for_null_probe_values` run
the same no-NULL baseline and assert the same 9,899 rows. One test with the
three runs would cover both.
`fallback_context(limit)` further down already builds a task context with an
optional memory limit, so `run_scope_key_anti_join` could probably reuse it.
The `prepare_task_ctx(8192, false)` knobs don't apply here, because a two-key
join never takes the perfect hash path.
--
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]
