diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index ad184d6500d56..42979a0897499 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -360,6 +360,19 @@ where // Ensure lengths are equivalent assert_eq!(values.len(), batch_hashes.len()); + // Consecutive-keys fast path: real-world data frequently arrives + // with runs of identical values (time-ordered logs, clustered + // writes). Remembering the previous non-null row's value and + // payload lets a run reuse the payload with one (adjacent, + // cache-hot) byte comparison instead of a hash-table probe into + // a table that may not be cache-resident. A miss costs a single + // well-predicted compare. Same idea as ClickHouse's "consecutive + // keys optimization". Intervening nulls don't invalidate the + // cached mapping (the payload for a value never changes within + // one `insert_if_new` call or across calls). + let mut prev_value: Option<&[u8]> = None; + let mut prev_payload = V::default(); + for (value, &hash) in values.iter().zip(batch_hashes.iter()) { // handle null value let Some(value) = value else { @@ -380,6 +393,17 @@ where // get the value as bytes let value: &[u8] = value.as_ref(); + + // Consecutive-keys fast path: the comparison reads bytes that + // were just scanned (adjacent rows in the same values buffer), + // so it is cache-hot; a hit skips the hash-table probe. + if let Some(prev) = prev_value + && prev == value + { + observe_payload_fn(prev_payload); + continue; + } + let value_len = O::usize_as(value.len()); // value is "small" @@ -464,6 +488,8 @@ where payload } }; + prev_value = Some(value); + prev_payload = payload; observe_payload_fn(payload); } // Check for overflow in offsets (if more data was sent than can be represented) @@ -646,6 +672,41 @@ mod tests { assert_set(set, &[None]); } + /// Runs of identical values (with interleaved nulls) must map to the + /// same payload as scattered occurrences — exercises the + /// consecutive-keys fast path in `insert_if_new_inner`. + #[test] + fn string_map_consecutive_keys_runs() { + let mut map: ArrowBytesMap = ArrowBytesMap::new(OutputType::Utf8); + // runs: aaa aaa | null | aaa bbb bbb | long(>8B) x2 | aaa + let values: ArrayRef = Arc::new(StringArray::from(vec![ + Some("aaa"), + Some("aaa"), + None, + Some("aaa"), + Some("bbb"), + Some("bbb"), + Some("longer_than_short_value_len"), + Some("longer_than_short_value_len"), + Some("aaa"), + ])); + let mut next = 0; + let mut seen = vec![]; + map.insert_if_new( + &values, + |_| { + let id = next; + next += 1; + id + }, + |id| seen.push(id), + ); + // aaa=0, null=1, bbb=2, long=3; run members share ids, the "aaa" + // after the null and at the end still resolve to 0. + assert_eq!(seen, vec![0, 0, 1, 0, 2, 2, 3, 3, 0]); + assert_eq!(next, 4, "make_payload_fn must run once per distinct"); + } + #[test] fn string_set_basic_i32() { test_string_set_basic::(); diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 9d4b556393a24..8f95eeae1c616 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -270,6 +270,20 @@ where // Ensure lengths are equivalent assert_eq!(values.len(), self.hashes_buffer.len()); + // Consecutive-keys fast path state: real-world data frequently + // arrives with runs of identical values (time-ordered logs, + // clustered writes), so remembering the previous non-null row's + // view and payload lets a run reuse the payload without probing + // the hash table. For inline views (len <= 12) the u128 view + // comparison is a complete equality check; for longer values the + // view still provides the length + 4-byte-prefix guards before a + // byte comparison of two adjacent (cache-hot) rows. A miss costs + // a couple of well-predicted compares. Same idea as ClickHouse's + // "consecutive keys optimization". Intervening nulls don't + // invalidate the cached mapping. + let mut prev: Option<(u128, usize)> = None; + let mut prev_payload = V::default(); + for i in 0..values.len() { let view_u128 = input_views[i]; let hash = self.hashes_buffer[i]; @@ -293,6 +307,29 @@ where // Extract length from the view (first 4 bytes of u128 in little-endian) let len = view_u128 as u32; + // Consecutive-keys fast path (see comment above the loop). + if let Some((prev_view, prev_i)) = prev { + let equal = if len <= 12 { + // Inline views embed the full value: u128 equality is + // a complete equality check. + view_u128 == prev_view + } else if prev_view as u32 == len + && (prev_view >> 32) as u32 == (view_u128 >> 32) as u32 + { + // Length and 4-byte prefix (read from the views) matched; + // compare the actual bytes of two adjacent, cache-hot rows. + let cur: &[u8] = values.value(i).as_ref(); + let prv: &[u8] = values.value(prev_i).as_ref(); + cur == prv + } else { + false + }; + if equal { + observe_payload_fn(prev_payload); + continue; + } + } + // Check if value already exists let maybe_payload = { // Borrow completed and in_progress for comparison @@ -371,6 +408,8 @@ where .insert_accounted(new_header, |h| h.hash, &mut self.map_size); payload }; + prev = Some((view_u128, i)); + prev_payload = payload; observe_payload_fn(payload); } } @@ -528,6 +567,38 @@ mod tests { use arrow::array::{GenericByteViewArray, StringViewArray}; use datafusion_common::HashMap; + /// Consecutive-keys fast path: runs of inline (<=12B) and non-inline + /// (>12B) values, with interleaved nulls, must produce the same + /// payload assignment as scattered occurrences. + #[test] + fn view_map_consecutive_keys_runs() { + let mut map: ArrowBytesViewMap = + ArrowBytesViewMap::new(OutputType::Utf8View); + let values: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("inline"), + Some("inline"), + None, + Some("inline"), + Some("this value is longer than twelve bytes"), + Some("this value is longer than twelve bytes"), + Some("this value is also longer, but different"), + Some("inline"), + ])); + let mut next = 0; + let mut seen = vec![]; + map.insert_if_new( + &values, + |_| { + let id = next; + next += 1; + id + }, + |id| seen.push(id), + ); + assert_eq!(seen, vec![0, 0, 1, 0, 2, 2, 3, 0]); + assert_eq!(next, 4, "make_payload_fn must run once per distinct"); + } + use super::*; // asserts that the set contains the expected strings, in the same order diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 5163948bd594a..c4cf5a54ae657 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -32,6 +32,7 @@ use crate::aggregates::group_values::multi_group_by::{ row_backed::RowsGroupColumn, }; use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; +use arrow::buffer::BooleanBuffer; use arrow::compute::cast; use arrow::datatypes::{ BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Field, Float32Type, @@ -41,6 +42,7 @@ use arrow::datatypes::{ TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; +use arrow_ord::cmp::not_distinct; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; @@ -53,6 +55,66 @@ use hashbrown::hash_table::HashTable; const NON_INLINED_FLAG: u64 = 0x8000000000000000; const VALUE_MASK: u64 = 0x7FFFFFFFFFFFFFFF; +/// Minimum fraction of adjacent-hash-equal row pairs (1/16 ≈ 6%) for the +/// consecutive-keys run mask to be worth computing. Below it, the mask's +/// per-column comparison passes cost more than the probes they could skip. +const RUN_MASK_MIN_HINT_DENOMINATOR: usize = 16; + +/// Returns a null-safe "row equals the previous row on every group column" +/// mask when the batch looks run-heavy, `None` otherwise. +/// +/// `mask[i] == true` means row `i + 1` equals row `i` on all columns, with +/// `null == null` counting as equal (matching GROUP BY semantics, via +/// [`not_distinct`]). Rows covered by the mask can reuse the previous row's +/// group index and skip both the hash-table probe and (for high-cardinality +/// keys) its likely cache miss — the same "consecutive keys optimization" +/// ClickHouse ships, and the multi-column analogue of the single-column +/// caches in `GroupValuesPrimitive` / `ArrowBytesMap`. +/// +/// Cost control is a two-tier gate: +/// - A single sequential scan counts adjacent *hash* matches (hash equality +/// is a necessary condition for row equality). Shuffled inputs — e.g. the +/// Final aggregation phase after hash repartitioning, where runs cannot +/// survive — fail this gate and pay nothing beyond the scan itself. +/// Correctness never rests on the hashes: they only decide whether the +/// precise mask below is worth building. +/// - Only run-heavy batches pay for the per-column vectorized +/// [`not_distinct`] passes that build the precise mask. +/// +/// Returns `None` (falling back to the normal per-row path) when the batch +/// is too small, fails the gate, or a column type is not supported by +/// [`not_distinct`] (e.g. nested types handled by `RowsGroupColumn`). +fn adjacent_equal_run_mask(cols: &[ArrayRef], hashes: &[u64]) -> Option { + let n = hashes.len(); + if n < 2 { + return None; + } + + let run_hint = hashes.windows(2).filter(|w| w[0] == w[1]).count(); + if run_hint * RUN_MASK_MIN_HINT_DENOMINATOR < n { + return None; + } + + let mut mask: Option = None; + for col in cols { + let head = col.slice(0, n - 1); + let tail = col.slice(1, n - 1); + let eq = match not_distinct(&head, &tail) { + Ok(eq) => eq, + // e.g. nested types: fall back to the normal path + Err(_) => return None, + }; + // `not_distinct` never returns nulls, so the values buffer is the + // complete answer. + let (buf, _nulls) = eq.into_parts(); + mask = Some(match mask { + None => buf, + Some(acc) => &acc & &buf, + }); + } + mask +} + /// Trait for storing a single column of group values in [`GroupValuesColumn`] /// /// Implementations of this trait store an in-progress collection of group values @@ -362,7 +424,25 @@ impl GroupValuesColumn { batch_hashes.resize(n_rows, 0); create_hashes(cols, &self.random_state, batch_hashes)?; + // Consecutive-keys fast path (see `adjacent_equal_run_mask`). The + // streaming path benefits the most: sorted-on-group-keys input has + // perfect runs, so all but the first row of every group skip the + // probe below. + let run_mask = adjacent_equal_run_mask(cols, batch_hashes); + for (row, &target_hash) in batch_hashes.iter().enumerate() { + if let Some(mask) = &run_mask + && row > 0 + && mask.value(row - 1) + { + // Same key as the previous row: reuse its group index. + let prev_group = *groups + .last() + .expect("row > 0 implies a previous group index"); + groups.push(prev_group); + continue; + } + let entry = self .map .find_mut(target_hash, |(exist_hash, group_idx_view)| { @@ -463,6 +543,11 @@ impl GroupValuesColumn { batch_hashes.resize(n_rows, 0); create_hashes(cols, &self.random_state, &mut batch_hashes)?; + // Consecutive-keys fast path: rows equal to their predecessor reuse + // its group index (filled in a final pass below) and skip the map + // probe entirely. See `adjacent_equal_run_mask` for the gating. + let run_mask = adjacent_equal_run_mask(cols, &batch_hashes); + // General steps for one round `vectorized equal_to & append`: // 1. Collect vectorized context by checking hash values of `cols` in `map`, // mainly fill `vectorized_append_row_indices`, `vectorized_equal_to_row_indices` @@ -484,7 +569,7 @@ impl GroupValuesColumn { // // 1. Collect vectorized context by checking hash values of `cols` in `map` - self.collect_vectorized_process_context(&batch_hashes, groups); + self.collect_vectorized_process_context(&batch_hashes, groups, run_mask.as_ref()); // 2. Perform `vectorized_append` self.vectorized_append(cols)?; @@ -496,6 +581,20 @@ impl GroupValuesColumn { // (about remaining rows, can see comments for `remaining_row_indices`) self.scalarized_intern_remaining(cols, &batch_hashes, groups)?; + // 5. Fill the rows skipped by the consecutive-keys fast path with + // their predecessor's group index. Ascending order makes multi-row + // runs propagate forward from the run head, which was resolved by + // the normal path above. + if let Some(mask) = run_mask { + for row in 1..n_rows { + if mask.value(row - 1) { + debug_assert_eq!(groups[row], usize::MAX); + debug_assert_ne!(groups[row - 1], usize::MAX); + groups[row] = groups[row - 1]; + } + } + } + self.hashes_buffer = batch_hashes; Ok(()) @@ -518,6 +617,7 @@ impl GroupValuesColumn { &mut self, batch_hashes: &[u64], groups: &mut [usize], + run_mask: Option<&BooleanBuffer>, ) { self.vectorized_operation_buffers.append_row_indices.clear(); self.vectorized_operation_buffers @@ -528,6 +628,15 @@ impl GroupValuesColumn { .clear(); for (row, &target_hash) in batch_hashes.iter().enumerate() { + // Consecutive-keys fast path: this row equals its predecessor on + // every group column; it reuses the predecessor's group index in + // the final fill pass of `vectorized_intern` and needs no probe. + if let Some(mask) = run_mask + && row > 0 + && mask.value(row - 1) + { + continue; + } let entry = self .map .find(target_hash, |(exist_hash, _)| target_hash == *exist_hash); @@ -1660,6 +1769,92 @@ mod tests { } } + /// Consecutive-keys fast path on the multi-column paths: a run only + /// counts when EVERY column matches its predecessor (with null == null + /// being equal, per GROUP BY semantics). Covers both the vectorized + /// (`STREAMING = false`) and scalarized (`STREAMING = true`) interns — + /// both must produce identical group assignments to a run-free shuffle + /// of the same values. + #[test] + fn test_intern_consecutive_keys_multi_column() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8View, true), + ])); + + // Runs designed to hit every interesting case: + // rows 0-1: identical on both cols (run) + // row 2: `a` continues the run but `b` changes → NOT a run + // row 3: identical to row 2 (run) + // row 4: null in both cols + // row 5: null in both cols again → null == null run + // row 6: same `a` as 5 is null vs 1 → not a run + let col_a: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + Some(1), + Some(1), + Some(1), + None, + None, + Some(1), + ])); + let col_b: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("x"), + Some("x"), + Some("y"), + Some("y"), + None, + None, + Some("x"), + ])); + let cols = vec![col_a, col_b]; + + // Expected groups: (1,x)=0, (1,y)=1, (null,null)=2, (1,x)=0 again + let expected = vec![0, 0, 1, 1, 2, 2, 0]; + + // Vectorized path + let mut vectorized = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + let mut groups = vec![]; + // Repeat the batch: the second pass exercises "run head found in + // existing map" (equal_to path) rather than append. + for _ in 0..2 { + vectorized.intern(&cols, &mut groups).unwrap(); + assert_eq!(groups, expected, "vectorized path"); + } + assert_eq!(vectorized.len(), 3); + + // Scalarized (streaming) path + let mut scalarized = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + for _ in 0..2 { + scalarized.intern(&cols, &mut groups).unwrap(); + assert_eq!(groups, expected, "scalarized path"); + } + assert_eq!(scalarized.len(), 3); + } + + /// The run mask is gated on an adjacent-hash heuristic; a batch with no + /// runs at all must take the normal path and still be correct. + #[test] + fn test_intern_consecutive_keys_no_runs() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8View, true), + ])); + let col_a: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3, 4, 1, 2, 3, 4])); + let col_b: ArrayRef = Arc::new(StringViewArray::from(vec![ + "p", "q", "r", "s", "p", "q", "r", "s", + ])); + let cols = vec![col_a, col_b]; + + let mut gv = GroupValuesColumn::::try_new(schema).unwrap(); + let mut groups = vec![]; + gv.intern(&cols, &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2, 3, 0, 1, 2, 3]); + assert_eq!(gv.len(), 4); + } + #[test] fn test_intern_for_vectorized_group_values() { let data_set = VectorizedTestDataSet::new(); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index e254aebcfd7ce..6666cabf79b9c 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -116,6 +116,23 @@ pub struct GroupValuesPrimitive { values: Vec, /// The random state used to generate hashes random_state: RandomState, + /// The key interned by the previous non-null row (canonicalized). + /// Valid only when `last_group != usize::MAX`. + /// + /// Real-world data frequently arrives with runs of identical keys + /// (time-ordered logs, clustered writes), so remembering the previous + /// row's `(key, group)` lets `intern` skip both the hash computation + /// and the hash-table probe for every row of a run after the first — + /// one register compare against state that stays in L1 instead of a + /// random access into a table that does not. Same idea as + /// ClickHouse's "consecutive keys optimization". A miss costs a + /// single well-predicted compare, so the break-even hit rate is + /// far below 1%. + last_key: T::Native, + /// Group index of `last_key`; `usize::MAX` means "no cached key". + /// Invalidated on `emit` / `clear_shrink` because both reassign + /// group indices. + last_group: usize, } impl GroupValuesPrimitive { @@ -127,6 +144,8 @@ impl GroupValuesPrimitive { values: Vec::with_capacity(128), null_group: None, random_state: crate::aggregates::AGGREGATION_HASH_SEED, + last_key: Default::default(), + last_group: usize::MAX, } } } @@ -151,6 +170,15 @@ where // so the bit-equal `is_eq` matches and the stored value is // the canonical representative. let key = key.canonicalize(); + + // Consecutive-keys fast path: if this row's key equals + // the previous row's, reuse its group index and skip + // both the hash and the table probe entirely. + if self.last_group != usize::MAX && key.is_eq(self.last_key) { + groups.push(self.last_group); + continue; + } + let state = &self.random_state; let hash = key.hash(state); let insert = self.map.entry( @@ -161,7 +189,7 @@ where |&(_, h)| h, ); - match insert { + let g = match insert { hashbrown::hash_table::Entry::Occupied(o) => o.get().0, hashbrown::hash_table::Entry::Vacant(v) => { let g = self.values.len(); @@ -169,7 +197,10 @@ where self.values.push(key); g } - } + }; + self.last_key = key; + self.last_group = g; + g } }; groups.push(group_id) @@ -190,6 +221,11 @@ where } fn emit(&mut self, emit_to: EmitTo) -> Result> { + // Both emit modes reassign group indices (All clears them, + // First(n) shifts them down by n), so the consecutive-keys cache + // must be invalidated. + self.last_group = usize::MAX; + fn build_primitive( values: Vec, null_idx: Option, @@ -244,6 +280,7 @@ where self.values.shrink_to(num_rows); self.map.clear(); self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared + self.last_group = usize::MAX; } } @@ -256,6 +293,43 @@ mod tests { use datafusion_expr::EmitTo; use std::sync::Arc; + /// Consecutive-keys fast path: runs (with interleaved nulls) map to the + /// same group ids as scattered occurrences, and `emit(First(n))` must + /// invalidate the cache (group indices shift down by `n`, so a stale + /// `last_group` would assign wrong ids). + #[test] + fn consecutive_keys_runs_and_emit_invalidation() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Int32); + + let arr: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(7), + Some(7), + None, + Some(7), + Some(9), + Some(9), + ])); + let mut groups = vec![]; + gv.intern(&[arr], &mut groups)?; + // 7=0, null=1, 9=2; run members and the post-null 7 share ids + assert_eq!(groups, vec![0, 0, 1, 0, 2, 2]); + + // Emit the first 2 groups: "7"->gone, "null"->gone, "9" shifts 2->0. + gv.emit(EmitTo::First(2))?; + + // Re-intern a run of 9s followed by 7s. A stale cache would claim + // "7 is still group 0" — but group 0 is now "9". + let arr2: ArrayRef = Arc::new(Int32Array::from(vec![9, 9, 7, 7])); + gv.intern(&[arr2], &mut groups)?; + assert_eq!( + groups, + vec![0, 0, 1, 1], + "9 is group 0 after shift; 7 is new group 1" + ); + + Ok(()) + } + /// Mirror of the `EmitTo::take_needed` regression test, applied to the /// concrete `GroupValuesPrimitive` accumulator. ///