andygrove commented on code in PR #26447:
URL: https://github.com/apache/datafusion/pull/26447#discussion_r4432072877
##########
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:
Done: sorted rows get a new zeroed `starts`. Since @adriangb's review, an
`ArrayMap` only sorts its rows that way when it has at least two rows for each
slot, without building a chain first. Any other keeps main's chain and has its
rows grouped later, without a pass over the slots.
##########
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:
Done: `contain_keys` matches on the layout once, with a loop for each, and
`slot` and `contains` are gone.
##########
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:
Done: `test_which_build_rows_are_grouped` is a table of build hashes and the
length of their groups, which covers both tests and the size check, and it also
checks the reservation and the matches. `test_grouped_rows_match_the_chain` now
compares only the grouped maps with the chain.
##########
datafusion/physical-plan/src/joins/array_map.rs:
##########
@@ -167,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;
Review Comment:
Thanks for tracking this down. A sparse `ArrayMap` now makes no pass over
all its slots: it keeps main's chain when it is built, and `group_rows` groups
its rows later, as for a hash map, counting and placing them in two passes over
the keys that touch only the slots of the keys that occur. Only maps with at
least two rows for each slot still sort their rows as they are built. A density
gate would have cost Q13 its speedup, since it has the same build side as Q14.
The numbers are in my comment on the PR.
##########
datafusion/physical-plan/src/joins/array_map.rs:
##########
@@ -261,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> {
+ let slots = match &self.rows {
+ KeyRows::Unique { data } => data.len(),
+ KeyRows::Grouped { starts, .. } => starts.len() - 1,
Review Comment:
`slot` is gone: after @comphead's review, `contain_keys` matches on the
layout once, with a loop for each. `test_array_map_layouts` probes the keys
just outside the range through `contain_keys` for each layout, and fails with
your mutation's equivalent there, `starts.len() / 1`.
--
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]
