andygrove commented on code in PR #26447:
URL: https://github.com/apache/datafusion/pull/26447#discussion_r4432072272
##########
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:
Done. Grouping now runs after every reservation the build needs, just before
the map is shared, and keeps the chain if the pool has no room for the groups.
`grouping_needs_no_more_memory_than_the_build` checks that a build that fits
without its groups still succeeds; with the old order it fails with
`ResourcesExhausted`. `group_rows(&reservation, &mem_used)` now makes the size
check, grows the reservation and gives back the chain's memory, and both maps
call one generic function for it.
##########
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:
The NULL keys: the threshold and the size of the groups now count the rows
inserted into the map, from the nonzero entries of `next`, so the rows left out
of it count towards neither. `test_which_build_rows_are_grouped` covers both.
The skew: in `Partitioned` mode the hot key ends up in one partition, whose
hash values have many rows each, so its rows are grouped. A single map with
fewer than two rows per hash value on average, as in `CollectLeft` mode, isn't
grouped. With your shape at 15M build rows, 5M of them on one key at scattered
positions, and 100 probe rows on that key, the join takes 1.56 s on main and
0.97 s with this PR when `Partitioned`, and 1.6 s with both in `CollectLeft`
mode.
I added a skewed case to `hj`: Q26 has 6M build rows, 1.5M of them on one
key, and 100 probe rows on it. Like Q23 to Q25 it runs `Partitioned`, and it
takes 0.53 to 0.60 s on main and 0.36 s with this PR. Catching the skew in a
single map would take more than the average, such as sampling the chain
lengths, which I'd leave for a follow-up.
##########
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:
Done: both public functions are as on main, and grouped rows go through
`grouped_matched_indices` instead, which takes `limit` by value and which
`ArrayMap` now shares. The `stream_join_utils.rs` hunk and those test lines are
gone. With the enum from your next comment, `get_matched_indices` on grouped
rows is `unreachable!`.
##########
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:
Done: the rows are `KeyRows::Chain { next, chained }` or `KeyRows::Grouped {
groups }`, which replaces the `groups.is_empty()` checks, and inserting rows
once they are grouped is `unreachable!`.
##########
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:
Done, thanks.
--
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]
