comphead commented on code in PR #26447:
URL: https://github.com/apache/datafusion/pull/26447#discussion_r4425166866
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -3398,6 +3398,18 @@ async fn collect_left_input(
offset += batch.num_rows();
}
+ // Store the rows of each repeated key next to each other, so that a
+ // probe reads its matches in order. Keep the chain if the memory for
+ // the groups is not available.
+ if let Some(size) = hashmap.grouped_rows_size()
Review Comment:
Grouping is optional, but it runs before mandatory reservations.
`concat_build_batches`, the key arrays and the bitmaps all call `try_grow(..)?`
afterwards. Grouping keeps one extra index per distinct hash value reserved, so
a build that fits on main can now fail there with `ResourcesExhausted` instead
of falling back to the chain. Running it after those reservations would let the
fallback cover them.
The caller also has to call `grouped_rows_size`, `try_grow`, `group_rows`
and `shrink` in that order, and `group_rows` repeats the size check. One method
that takes the reservation would keep that protocol in one place. The
`JoinHashMapU32` and `JoinHashMapU64` impls of `group_rows` are identical apart
from `size_of`, so their body could move into the free `group_rows`.
##########
datafusion/physical-plan/src/joins/join_hash_map.rs:
##########
@@ -396,6 +641,33 @@ where
let zero = T::try_from(0).unwrap();
let one = T::try_from(1).unwrap();
+ if !groups.is_empty() {
Review Comment:
Only the symmetric hash join calls `get_matched_indices`, and its
`PruningJoinHashMap` always passes `&[]`. So this branch, including its
`deleted_offset` handling, only runs in `test_grouped_rows_match_the_chain`.
Dispatching to `get_grouped_matched_indices` inside the `JoinHashMapU32` and
`JoinHashMapU64` impls would leave both public free functions unchanged. That
drops this branch, the `stream_join_utils.rs` hunk and those test lines. A
`debug_assert!(self.groups.is_empty())` in `get_matched_indices` would guard
the contract. `get_grouped_matched_indices` could then take `limit` by value,
since nothing reads `remaining_output` after it returns.
##########
datafusion/physical-plan/src/joins/join_hash_map.rs:
##########
@@ -165,29 +179,56 @@ pub trait JoinHashMapType: Send + Sync {
/// Returns the number of entries in the join hash map.
fn len(&self) -> usize;
+
+ /// Returns the number of bytes that [`Self::group_rows`] would allocate,
+ /// or `None` if it would not group the rows: when the hash values have
+ /// fewer than two build rows each on average, the chains are too short for
+ /// grouping to pay for itself.
+ fn grouped_rows_size(&self) -> Option<usize> {
+ None
+ }
+
+ /// Stores the build rows of each hash value next to each other, in chain
+ /// order, and drops the chain. A probe row then reads its matches from one
+ /// place, instead of following the chain with a dependent load for every
+ /// match, which misses the cache when the build side is large. Call this
+ /// once every build row has been inserted. Returns the number of bytes
+ /// freed.
+ fn group_rows(&mut self) -> usize {
+ 0
+ }
}
pub struct JoinHashMapU32 {
- // Stores hash value to last row index
+ // Stores hash value to last row index, or to its group in `groups`
map: HashTable<(u64, u32)>,
// Stores indices in chained list data structure
next: Vec<u32>,
// Whether any hash value maps to more than one row
chained: bool,
+ // The build rows of each hash value once `group_rows` has grouped them,
+ // after which `next` is empty
+ groups: Vec<u32>,
Review Comment:
After `group_rows`, a map value is a group position, but `update_from_iter`
doesn't know that. An insert on an existing hash would index the empty `next`
and panic, and a new hash would store a `row + 1` among the group positions. No
caller does this today. An enum like `KeyRows` in `ArrayMap`, for example
`Chain { next, chained }` and `Grouped { groups }`, would make the state
explicit and replace the `groups.is_empty()` checks. At least a `debug_assert!`
in `update_from_iter` would catch misuse.
##########
datafusion/physical-plan/src/joins/chain.rs:
##########
@@ -75,3 +75,43 @@ where
match_row_idx = next - one;
}
}
+
+/// Emits the build rows `rows[pos..end]` of probe row `prob_idx`, at most
+/// `remaining` of them, for maps that store the build rows of each key next to
+/// each other instead of in a chain. Returns `Some(offset)` if the limit was
+/// reached and there are more results to process, as [`traverse_chain`] does.
+/// A [`MidChain`] offset then holds the position in `rows` to resume from,
+/// which is never `0` because at least one row was emitted.
+#[expect(clippy::too_many_arguments)]
+#[inline(always)]
+pub(crate) fn emit_grouped_rows<T>(
Review Comment:
`end` is only used as the end of `rows`. Passing `&rows[..end]` drops a
parameter, which brings it to 7, so the `#[expect(clippy::too_many_arguments)]`
can go.
##########
datafusion/physical-plan/src/joins/join_hash_map.rs:
##########
@@ -337,11 +415,175 @@ impl JoinHashMapType for JoinHashMapU64 {
fn len(&self) -> usize {
self.map.len()
}
+
+ fn grouped_rows_size(&self) -> Option<usize> {
+ grouped_rows_size(&self.map, &self.next, self.chained)
+ }
+
+ fn group_rows(&mut self) -> usize {
+ if self.grouped_rows_size().is_none() {
+ return 0;
+ }
+ let freed = self.next.capacity() * size_of::<u64>();
+ self.groups = group_rows(&mut self.map, &self.next);
+ self.next = vec![];
+ freed
+ }
}
-use crate::joins::chain::traverse_chain;
+use crate::joins::chain::{emit_grouped_rows, traverse_chain};
use crate::joins::join_hash_map::ProbeOffset::{AtRow, MidChain};
+/// Returns the number of bytes that [`group_rows`] allocates for `map` and
+/// `next`, or `None` if the hash values have fewer than two rows each on
+/// average, or if a position in the groups would not fit in `T`.
+fn grouped_rows_size<T>(
+ map: &HashTable<(u64, T)>,
+ next: &[T],
+ chained: bool,
+) -> Option<usize>
+where
+ T: ArrowNativeType,
+{
+ // Following a chain costs a cache miss for each row after the first, which
+ // only adds up when the average chain is long enough. `next.len()` counts
+ // the rows left out of the map, such as NULL keys, too.
+ if !chained || next.len() < map.len().saturating_mul(2) {
Review Comment:
Two cases where the average misleads:
- Skew. One key with 500K rows plus 1M unique keys gives `next.len()` of
1.5M, below `2 * map.len()`, so nothing is grouped. Every probe row that hits
the hot key still walks a 500K-long chain, which is the q72 shape. Could the
benchmarks include a skewed case like this?
- NULL keys. `next.len()` counts every build row, but `update_hash` leaves
NULL keys out of the map. A build with 60% NULL keys and nearly unique non-NULL
keys passes this check and pays for grouping. `groups` also gets capacity for
the NULL rows, which stays reserved and unused for the whole join.
Counting the rows `update_from_iter` inserts would make both the threshold
and the capacity exact.
##########
datafusion/physical-plan/src/joins/join_hash_map.rs:
##########
@@ -692,6 +1057,177 @@ mod tests {
assert_eq!(match_indices, vec![1, 0]);
}
+ /// Collects every match of `hash_values`, `limit` at a time.
+ fn all_matches(
+ hash_map: &dyn JoinHashMapType,
+ hash_values: &[u64],
+ valid_keys: Option<&NullBuffer>,
+ limit: usize,
+ ) -> Vec<(u32, u64)> {
+ let mut matches = vec![];
+ let mut input_indices = vec![];
+ let mut match_indices = vec![];
+ let mut offset = AtRow { row: 0 };
+ loop {
+ let next_offset = hash_map.get_matched_indices_with_limit_offset(
+ hash_values,
+ valid_keys,
+ limit,
+ offset,
+ &mut input_indices,
+ &mut match_indices,
+ );
+ assert!(input_indices.len() <= limit);
+ assert_eq!(input_indices.len(), match_indices.len());
+ matches.extend(
+ input_indices
+ .iter()
+ .copied()
+ .zip(match_indices.iter().copied()),
+ );
+ match next_offset {
+ Some(next) => offset = next,
+ None => return matches,
+ }
+ }
+ }
+
+ /// Builds two maps from the same rows, inserted batch by batch and last
+ /// to first as `update_hash` does, with row 7 left out as a NULL key would
+ /// be, and groups the rows of the second.
+ fn chained_and_grouped<M: JoinHashMapType>(
+ new_map: impl Fn(usize) -> M,
+ build_hashes: &[u64],
+ ) -> (M, M) {
+ let insert = |hash_map: &mut M| {
+ for (batch_idx, hashes) in
build_hashes.chunks(64).enumerate().rev() {
+ let offset = batch_idx * 64;
+ hash_map.update_from_iter(
+ Box::new(
+ hashes
+ .iter()
+ .enumerate()
+ .map(move |(i, hash)| (offset + i, hash))
+ .filter(|(row, _)| *row != 7)
+ .rev(),
+ ),
+ 0,
+ );
+ }
+ };
+ let mut chained = new_map(build_hashes.len());
+ insert(&mut chained);
+ let mut grouped = new_map(build_hashes.len());
+ insert(&mut grouped);
+ assert!(grouped.grouped_rows_size().is_some());
+ assert!(grouped.group_rows() > 0);
+ (chained, grouped)
+ }
+
+ #[test]
+ fn test_grouped_rows_match_the_chain() {
+ // 13 hash values over 300 build rows, and probe rows that repeat,
+ // miss (13 to 16) or have a NULL key.
+ let build_hashes: Vec<u64> = (0..300u64).map(|i| (i * 7919) %
13).collect();
+ let probe_hashes: Vec<u64> = (0..50u64).map(|i| (i * 31) %
17).collect();
+ let valid_keys =
+ NullBuffer::from((0..50).map(|i| i % 9 != 4).collect::<Vec<_>>());
+
+ let (chained, grouped) =
+ chained_and_grouped(JoinHashMapU32::with_capacity, &build_hashes);
+ let (chained_u64, grouped_u64) =
+ chained_and_grouped(JoinHashMapU64::with_capacity, &build_hashes);
+ for limit in [1, 2, 3, 5, 23, 64, 8192] {
+ for valid in [None, Some(&valid_keys)] {
+ let expected = all_matches(&chained, &probe_hashes, valid,
limit);
+ assert!(!expected.is_empty());
+ assert_eq!(all_matches(&grouped, &probe_hashes, valid, limit),
expected);
+ assert_eq!(
+ all_matches(&chained_u64, &probe_hashes, valid, limit),
+ expected
+ );
+ assert_eq!(
+ all_matches(&grouped_u64, &probe_hashes, valid, limit),
+ expected
+ );
+ }
+ }
+
+ let iter = || Box::new(probe_hashes.iter().enumerate());
+ let expected = chained.get_matched_indices(iter(), None);
+ assert_eq!(grouped.get_matched_indices(iter(), None), expected);
+ assert_eq!(grouped_u64.get_matched_indices(iter(), None), expected);
+ assert_eq!(
+ grouped.contain_hashes(&probe_hashes),
+ chained.contain_hashes(&probe_hashes)
+ );
+ }
+
+ #[test]
+ fn test_grouped_rows_match_the_chain_for_long_chains() {
+ // Hash values 0 and 1 have chains too long for grouping to hold their
+ // rows until their end, among 2,000 short chains.
+ let build_hashes: Vec<u64> = (0..30_000u64)
+ .map(|i| if i % 3 == 0 { i % 2 } else { 2 + i % 2_000 })
+ .collect();
+ let probe_hashes: Vec<u64> = vec![0, 5, 1, 7, 2_001, 3_000, 0];
+
+ let (chained, grouped) =
+ chained_and_grouped(JoinHashMapU32::with_capacity, &build_hashes);
+ for limit in [1, 7, 4_096, 5_000, 8_192, 100_000] {
+ let expected = all_matches(&chained, &probe_hashes, None, limit);
+ assert_eq!(expected.len(), 2 * 5_000 + 5_000 + 3 * 10);
+ assert_eq!(all_matches(&grouped, &probe_hashes, None, limit),
expected);
+ }
+ }
+
+ #[test]
+ fn test_unique_hash_values_are_not_grouped() {
Review Comment:
This test and `test_short_chains_are_not_grouped` make the same assertions
on two inputs, and their `all_matches` checks re-test the chain and unique
paths that `test_get_matched_indices_skips_invalid_keys*` and
`test_unique_fast_path_with_uninserted_row` already cover. One table-driven
test over `(hashes, expected grouped_rows_size)` could cover both, together
with the size check in `test_group_rows_replaces_the_chain`. In
`test_grouped_rows_match_the_chain`, `chained_u64` re-tests the existing chain
path, so only `grouped_u64` adds coverage.
##########
datafusion/physical-plan/src/joins/array_map.rs:
##########
@@ -180,60 +197,88 @@ impl ArrayMap {
}
let size = (range + 1) as usize;
- let mut data: Vec<u32> = vec![0; size];
- let mut next: Vec<u32> = vec![];
- let mut num_of_distinct_key = 0;
-
downcast_supported_integer!(
array.data_type() => (
- fill_data,
+ build,
array,
min_val,
- &mut data,
- &mut next,
- &mut num_of_distinct_key
+ size
)
- )?;
-
- Ok(Self {
- data,
- offset: min_val,
- next,
- num_of_distinct_key,
- })
+ )
}
- fn fill_data<T: ArrowNumericType>(
+ fn build<T: ArrowNumericType>(
array: &ArrayRef,
- offset_val: u64,
- data: &mut [u32],
- next: &mut Vec<u32>,
- num_of_distinct_key: &mut usize,
- ) -> Result<()>
+ offset: u64,
+ size: usize,
+ ) -> Result<Self>
where
T::Native: AsPrimitive<u64>,
{
let arr = array.as_primitive::<T>();
- // Iterate in reverse to maintain FIFO order when there are duplicate
keys.
- for (i, val) in arr.iter().enumerate().rev() {
- if let Some(val) = val {
- let key: u64 = val.as_();
- let Some(idx) = Self::key_to_index(key, offset_val,
data.len()) else {
- return internal_err!("failed build Array idx >=
data.len()");
- };
+ let mut data: Vec<u32> = vec![0; size];
+ let mut num_of_distinct_key = 0;
+ let mut unique = true;
+ for (i, val) in arr.iter().enumerate() {
+ let Some(val) = val else { continue };
+ let Some(idx) = Self::key_to_index(val.as_(), offset, size) else {
+ return internal_err!("failed build Array idx >= data.len()");
+ };
+ if data[idx] != 0 {
+ unique = false;
+ break;
+ }
+ data[idx] = i as u32 + 1;
+ num_of_distinct_key += 1;
+ }
+ if unique {
+ return Ok(Self {
+ rows: KeyRows::Unique { data },
+ offset,
+ num_of_distinct_key,
+ });
+ }
- if data[idx] != 0 {
- if next.is_empty() {
- *next = vec![0; array.len()]
- }
- next[i] = data[idx]
- } else {
- *num_of_distinct_key += 1;
- }
- data[idx] = (i) as u32 + 1;
+ // Some key repeats, so group the build rows by key with a counting
+ // sort, reusing the allocation. First count each key's rows...
+ let mut starts = data;
+ starts.fill(0);
+ // One more slot, for the end of the last key's rows
+ starts.reserve_exact(1);
+ starts.push(0);
Review Comment:
`data` was allocated with exactly `size` capacity, so `reserve_exact(1)`
reallocates and may copy the whole slot array, right after `fill(0)` has swept
it. Dropping `data` and allocating `vec![0; size + 1]` here skips both the
sweep and the copy, and zeroed allocations can come straight from the OS. The
slot array scales with the key range, so this grows with it.
##########
datafusion/physical-plan/src/joins/array_map.rs:
##########
@@ -274,14 +324,27 @@ impl ArrayMap {
)
}
- /// Looks up `key` (a raw probe value cast to `u64`) in the build side,
- /// returning the 1-based build-side slot if the key maps to a non-empty
- /// bucket, or `None` otherwise.
+ /// Returns the slot of `key` (a raw probe value cast to `u64`), or `None`
+ /// if it is outside the build side's key range.
+ #[inline]
+ fn slot(&self, key: u64) -> Option<usize> {
Review Comment:
`slot` and `contains` match on `self.rows` for every row of `contain_keys`,
which backs the dynamic filter on every scanned probe row. Matching once
outside `collect_bool`, with one closure per variant, removes the per-row
branch and both helpers. The lookup paths already call `key_to_index` directly.
--
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]
