viirya commented on code in PR #5778:
URL: https://github.com/apache/datafusion-comet/pull/5778#discussion_r3997202194
##########
native/spark-expr/src/hash_funcs/utils.rs:
##########
@@ -526,34 +526,111 @@ macro_rules! hash_list_array {
let values = list_array.values();
let offsets = list_array.offsets();
- if list_array.null_count() == 0 {
- // Fast path: no nulls, skip null checks
- for (row_idx, hash) in $hashes.iter_mut().enumerate() {
- let start = offsets[row_idx] as usize;
- let end = offsets[row_idx + 1] as usize;
- let len = end - start;
- // Hash each element in sequence, chaining the hash values
- for elem_idx in 0..len {
- let elem_array = values.slice(start + elem_idx, 1);
- let mut single_hash = [*hash];
- $recursive_hash_method(&[elem_array], &mut single_hash)?;
- *hash = single_hash[0];
+ // Spark chains the element hashes in order, so the elements of one
row have to be hashed
+ // in sequence. What does not have to happen per element is the
allocation and dispatch:
+ // slicing a one-element array and re-entering the hash dispatch for
it costs an Arrow
+ // array plus a full type match every time, and for a struct element
the dispatch also
+ // copies the field vector on every call.
+ //
+ // Instead, hash one element per row at a time in a single batched
call, seeding each
+ // slot with the running hash of the row it belongs to. That is
exactly what the
+ // per-element call did, so the result is bit-identical.
+ let total_elements = offsets[$hashes.len()] as usize - offsets[0] as
usize;
+ if total_elements == 0 {
+ // Every list is empty or null; the seeds already hold the answer.
+ } else {
+ let first_offset = offsets[0] as usize;
+ let elements = values.slice(first_offset, total_elements);
+
+ // Chaining means element k of a row can only be hashed once
element k-1 is known, so
+ // batch by position: all the first elements together, then all
the second, and so on.
+ // Rows are independent, so one pass per position is enough.
+ //
+ // Only rows that still have an element at the current position
take part, and a row
+ // never becomes alive again once exhausted, so carry the
surviving rows forward instead
+ // of rescanning all of them each pass. Rescanning would cost rows
x longest-list, which
+ // for one long list among short ones is almost all wasted: 8192
rows with one list of
+ // 1024 scans 8.4M slots for 9215 elements. Carrying the survivors
makes the scheduling
+ // work proportional to the elements actually hashed.
+ //
+ // Index the gather by the list's own offset width. A `LargeList`
can hold more than
+ // `u32::MAX` elements, so narrowing the positions to `u32` would
silently wrap and
+ // hash the wrong elements.
+ let mut active: Vec<usize> = Vec::with_capacity($hashes.len());
+ // The same pass records whether every row is non-null with the
same length. When it
+ // is, no row ever drops out early, so the survivor bookkeeping is
pure overhead and
+ // the rows can simply be walked directly.
+ let mut uniform_len: Option<usize> = None;
+ let mut all_same = true;
+ for row_idx in 0..$hashes.len() {
+ if list_array.is_null(row_idx) {
+ all_same = false;
+ continue;
+ }
+ let len = offsets[row_idx + 1] as usize - offsets[row_idx] as
usize;
+ if len > 0 {
+ active.push(row_idx);
+ }
+ match uniform_len {
+ None => uniform_len = Some(len),
+ Some(seen) if seen == len => {}
+ Some(_) => all_same = false,
}
}
- } else {
- // Slow path: array has nulls, check each row
- for (row_idx, hash) in $hashes.iter_mut().enumerate() {
- if !list_array.is_null(row_idx) {
- let start = offsets[row_idx] as usize;
- let end = offsets[row_idx + 1] as usize;
- let len = end - start;
- // Hash each element in sequence, chaining the hash values
- for elem_idx in 0..len {
- let elem_array = values.slice(start + elem_idx, 1);
- let mut single_hash = [*hash];
- $recursive_hash_method(&[elem_array], &mut
single_hash)?;
- *hash = single_hash[0];
+ let uniform = all_same && uniform_len.unwrap_or(0) > 0;
+
+ let mut positions: Vec<$offset_type> =
Vec::with_capacity($hashes.len());
+ let mut rows_at_position: Vec<usize> =
Vec::with_capacity($hashes.len());
+ let mut still_active: Vec<usize> =
Vec::with_capacity($hashes.len());
+ let mut position = 0usize;
+ let uniform_passes = if uniform { uniform_len.unwrap_or(0) } else
{ 0 };
+ while (uniform && position < uniform_passes) || (!uniform &&
!active.is_empty()) {
+ positions.clear();
+ rows_at_position.clear();
+ if uniform {
+ // Every row survives every pass, so skip the survivor
bookkeeping.
+ for row_idx in active.iter().copied() {
+ let start = offsets[row_idx] as usize;
+ positions.push((start + position - first_offset) as
$offset_type);
+ rows_at_position.push(row_idx);
}
+ } else {
+ still_active.clear();
+ for row_idx in active.iter().copied() {
+ let start = offsets[row_idx] as usize;
+ let end = offsets[row_idx + 1] as usize;
+ positions.push((start + position - first_offset) as
$offset_type);
+ rows_at_position.push(row_idx);
+ // Alive for the next pass only if it has an element
beyond this one.
+ if start + position + 1 < end {
+ still_active.push(row_idx);
+ }
+ }
+ std::mem::swap(&mut active, &mut still_active);
+ }
+ position += 1;
+ // `take` accepts any integer index type, so index by the
offset width: a
+ // `LargeList` can exceed `u32::MAX` elements.
+ let taken = if std::mem::size_of::<$offset_type>() > 4 {
+ let indices = arrow::array::Int64Array::from_iter_values(
+ positions.iter().map(|p| *p as i64),
+ );
+ arrow::compute::take(&elements, &indices, None)?
+ } else {
+ let indices = arrow::array::Int32Array::from_iter_values(
+ positions.iter().map(|p| *p as i32),
+ );
+ arrow::compute::take(&elements, &indices, None)?
Review Comment:
This was addressed by the narrow eligibility rule and the added
measurements, as confirmed in [sunchao’s follow-up
review](https://github.com/apache/datafusion-comet/pull/5778#pullrequestreview-5159508807).
--
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]
