adriangb commented on code in PR #26447:
URL: https://github.com/apache/datafusion/pull/26447#discussion_r4428231300
##########
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:
Test gap: mutating `starts.len() - 1` → `starts.len() / 1` survives both
unit tests and a fuzz test. `slot()` is only reached via `contain_keys`, and
nothing probes the key `max + 1` there. That key would index `starts[size +
1]`. A one-row case in `test_array_map_grouped_rows_keep_build_order` would
cover it.
##########
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:
This is where `hj` Q14 (`density=0.01`, 100K build rows → ~10M slots) loses
~20%. Grouping makes 3 passes over all `size` slots (`fill`, prefix sum,
placement), on top of the aborted unique pass. `try_new` goes from 6.7% to
22.5% of samples.
I tried allocating a fresh `vec![0; size + 1]` here instead of `fill` +
`reserve_exact` + `push`: no difference, so the cost is the O(range) passes,
not the realloc.
Worth a density gate? E.g. keep the chain when `num_rows` is small relative
to `size`, similar to the avg-≥2-rows-per-hash heuristic in `grouped_rows_size`.
##########
datafusion/physical-plan/src/joins/join_hash_map.rs:
##########
@@ -383,6 +641,33 @@ where
let zero = T::try_from(0).unwrap();
let one = T::try_from(1).unwrap();
+ if !groups.is_empty() {
+ for (row_idx, hash_value) in iter {
+ if let Some((_, pos)) =
+ map.find(*hash_value, |(hash, _)| *hash_value == *hash)
+ {
+ let pos: u64 = (*pos).into();
+ let pos = pos as usize;
+ let len: u64 = groups[pos].into();
+ for &row in &groups[pos + 1..=pos + len as usize] {
+ let match_row_idx = if let Some(offset) = deleted_offset {
+ let offset = T::try_from(offset).unwrap();
+ if row < offset {
Review Comment:
Test gap: `row < offset` → `==` / `>` / `<=` all survive, and so does
`PruningJoinHashMap::get_matched_indices_with_limit_offset → None` in
`stream_join_utils.rs`. The grouped + `deleted_offset` (pruning) path doesn't
seem to be covered by any test. Is a SHJ test with pruning plus repeated keys
worth adding?
--
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]
