make ready for review

Author: Rich-T-kid

datafusion/common/src/hash_utils.rs

@@ -521,9 +521,8 @@ fn hash_generic_byte_view_array<T: ByteViewType>(
 #[cfg(not(feature = "force_hash_collisions"))]
 const DICT_SCATTER_LANES: usize = 8;
 
-/// Fold `dict_hash` into `prev` when `MULTI_COL` (secondary column), else
-/// overwrite with `dict_hash` (first column). Const generic +
-/// `#[inline(always)]` makes it branch-free per instantiation.
+/// Fold `dict_hash` into `prev` when `MULTI_COL`, else
+/// overwrite with `dict_hash`.
 #[cfg(not(feature = "force_hash_collisions"))]
 #[inline(always)]
 fn maybe_combine_dict_hash<const MULTI_COL: bool>(prev: u64, dict_hash: u64) -> u64 {
@@ -540,9 +539,8 @@ fn maybe_combine_dict_hash<const MULTI_COL: bool>(prev: u64, dict_hash: u64) ->
 /// The 8-wide unroll is for performance, splitting the loop into
 /// `std::array::from_fn` (gather 8 dict hashes by key index) followed by a
 /// fixed-length inner loop encourages the compiler to overlap 8
-/// independent load latencies, saturating the gather pipeline on
-/// ARM64/x86 instead of serializing one lookup at a time. The scalar tail
-/// handles the non-multiple-of-8 remainder.
+/// independent load latencies, saturating the gather pipeline instead of
+/// serializing one lookup at a time.
 ///
 /// Invariants: `keys.len() == hashes.len()`; no key or dict value referenced
 /// here is null.
@@ -571,6 +569,7 @@ fn scatter_dict_hashes_no_null_checks<
                 maybe_combine_dict_hash::<MULTI_COL>(hash_group[lane], lane_hashes[lane]);
         }
     }
+    // need to handle the remainder for non 64 byte multiples
     for (key, hash) in keys_tail.iter().zip(hashes_tail.iter_mut()) {
         let dict_hash = unsafe { *dict_hashes.get_unchecked(key.as_usize()) };
         *hash = maybe_combine_dict_hash::<MULTI_COL>(*hash, dict_hash);
@@ -582,7 +581,7 @@ fn scatter_dict_hashes_no_null_checks<
 /// nullable-values case; the per-lane `is_valid` check is unavoidable.
 #[cfg(not(feature = "force_hash_collisions"))]
 #[inline(always)]
-fn scatter_dict_hashes_skip_null_values<
+fn scatter_dict_hashes_skip_value_nulls<
     K: ArrowDictionaryKeyType,
     const MULTI_COL: bool,
 >(
@@ -600,10 +599,7 @@ fn scatter_dict_hashes_skip_null_values<
     }
 }
 
-/// Scatter dict-value hashes inside a mixed-validity 64-key chunk, with
-/// per-lane key-validity given by `mask`. All-valid-values path: branchless
-/// 64-lane gather + csel-blend by `mask`. Possibly-null-values path:
-/// bit-walk set lanes via `trailing_zeros`.
+/// Scatter inside a mixed-validity 64-key chunk governed by `mask`.
 #[cfg(not(feature = "force_hash_collisions"))]
 #[inline(always)]
 fn scatter_dict_hashes_in_masked_chunk<
@@ -633,8 +629,10 @@ fn scatter_dict_hashes_in_masked_chunk<
         }
         return;
     }
+    // null-key lanes carry garbage indices; we gather them anyway and mask the result out below
     let max_valid_idx = dict_hashes.len().saturating_sub(1);
     for lane_base in (0..64).step_by(DICT_SCATTER_LANES) {
+        // gather 8 dict hashes in parallel; clamp keeps garbage indices in-bounds
         let lane_hashes: [u64; DICT_SCATTER_LANES] = std::array::from_fn(|lane_offset| {
             let raw_idx =
                 unsafe { keys.get_unchecked(lane_base + lane_offset).as_usize() };
@@ -653,10 +651,10 @@ fn scatter_dict_hashes_in_masked_chunk<
 
 /// Scatter precomputed dictionary value hashes to key positions.
 ///
-/// Const generics eliminate runtime branching:
-/// - `HAS_NULL_KEYS`: dictionary has null keys to skip
-/// - `HAS_NULL_VALUES`: dictionary values may be null (per-index check)
-/// - `MULTI_COL`: combine with existing hash vs. overwrite
+/// Uses const generics to eliminate runtime branching in the hot loop:
+/// - `HAS_NULL_KEYS`: Whether to check for null dictionary keys
+/// - `HAS_NULL_VALUES`: Whether to check for null dictionary values
+/// - `MULTI_COL`: Whether to combine with existing hash (true) or initialize (false)
 #[cfg(not(feature = "force_hash_collisions"))]
 #[inline(never)]
 fn hash_dictionary_scatter<
@@ -674,7 +672,7 @@ fn hash_dictionary_scatter<
 
     if !HAS_NULL_KEYS {
         if HAS_NULL_VALUES {
-            scatter_dict_hashes_skip_null_values::<K, MULTI_COL>(
+            scatter_dict_hashes_skip_value_nulls::<K, MULTI_COL>(
                 keys,
                 dict_values.as_ref(),
                 dict_hashes,
@@ -694,13 +692,15 @@ fn hash_dictionary_scatter<
     let bit_chunks = null_buffer.inner().bit_chunks();
     let chunk_len = bit_chunks.chunk_len();
 
+    // Walk the key null buffer 64 bits at a time; bit i set = lane i's key is non-null.
     for (chunk_idx, mask) in bit_chunks.iter().enumerate() {
         let base = chunk_idx * 64;
         let keys_chunk = &keys[base..base + 64];
         let hashes_chunk = &mut hashes_buffer[base..base + 64];
         if mask == u64::MAX {
+            // all 64 keys valid: dense helper
             if HAS_NULL_VALUES {
-                scatter_dict_hashes_skip_null_values::<K, MULTI_COL>(
+                scatter_dict_hashes_skip_value_nulls::<K, MULTI_COL>(
                     keys_chunk,
                     dict_values.as_ref(),
                     dict_hashes,
@@ -714,6 +714,7 @@ fn hash_dictionary_scatter<
                 );
             }
         } else if mask != 0 {
+            // mixed: masked-chunk helper (mask == 0 falls through, nothing to write)
             scatter_dict_hashes_in_masked_chunk::<K, HAS_NULL_VALUES, MULTI_COL>(
                 keys_chunk,
                 dict_values.as_ref(),
@@ -724,6 +725,7 @@ fn hash_dictionary_scatter<
         }
     }
 
+    // non-64 multiple: scalar path for the tail
     let remainder_len = bit_chunks.remainder_len();
     if remainder_len > 0 {
         let remainder = bit_chunks.remainder_bits();