From d63765091025bc5f8b61387f7f275026f68c3bc5 Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Wed, 29 Jul 2026 19:31:08 +0800 Subject: [PATCH 1/2] perf(aggregate): consecutive-keys (last-group) cache for single-column GROUP BY MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Real-world data frequently arrives with runs of identical group keys (time-ordered logs, clustered writes). Today `GroupValues::intern` pays the full hash + hash-table-probe cost for every row, even when a row's key equals the previous row's — and for high-cardinality keys the table is far too large to stay cache-resident, so each probe is typically an L3/DRAM round trip. Remember the previous row's key and its group/payload; when the current row matches, reuse the answer and skip the probe (and for primitives, the hash as well). A hit replaces a random memory access with one or two well-predicted compares against L1-resident state; a miss costs a single compare (<1ns). Break-even hit rate is under 1%. This is the same optimization ClickHouse ships as its "consecutive keys optimization" (`LastElementCache` in ColumnsHashing, refined in ClickHouse#57872). Measured run lengths on ClickBench hits.parquet (file order, which is what the Partial aggregation phase sees): UserID avg run 10.1 (90.1% hit rate), SearchPhrase 11.7 (91.5%), RegionID 11.0 (90.9%), URL 1.4 (30.2%). Three implementations, one per single-column path: - `GroupValuesPrimitive` (primitive keys): cache `(last_key, last_group)` fields; the guard compares the canonicalized value itself, so a hit skips the hash too. Invalidated in `emit` / `clear_shrink`, since both reassign group indices. - `ArrowBytesMap::insert_if_new_inner` (Utf8 / LargeUtf8 / Binary via `GroupValuesBytes`, also used by distinct-aggregate accumulators): compare the current row's bytes against the previous input row's — adjacent rows in the same values buffer, so the comparison is cache-hot. Per-call state, no invalidation concerns. - `ArrowBytesViewMap::insert_if_new_inner` (Utf8View / BinaryView via `GroupValuesBytesView`): for inline views (len <= 12) the u128 view comparison is a complete equality check; longer values get length + 4-byte-prefix guards from the views before touching bytes. Intervening nulls do not invalidate the cached mapping (a value's group/payload never changes within an accumulation), and the make-payload-once / observe-per-row-in-order contract of the maps is preserved by the fast path. ClickBench (hits.parquet, 5 iterations, release-nonlto, both binaries built from this commit's base): Q33 (URL) 3621ms -> 2969ms (-18.0%) Q15 (UserID) 304ms -> 271ms (-10.7%) Q12 (SearchPhrase) 389ms -> 357ms ( -8.2%) Q7 (AdvEngineID) 20ms -> 18ms ( -6.2%) Q27 (CounterID) 739ms -> 704ms ( -4.8%) The win scales with (hit rate x probe cost): Q33's URL table is the largest string hash table, so even a 30% hit rate pays the most per skipped probe; Q27's ~6K-entry table is L1/L2-resident, so gains are small even at ~100% hit rate. Tests: consecutive-run + interleaved-null payload-assignment tests for both maps, and a primitive test covering the emit(First(n)) cache invalidation (a stale cache would assign shifted group ids). --- .../physical-expr-common/src/binary_map.rs | 61 +++++++++++++++ .../src/binary_view_map.rs | 71 +++++++++++++++++ .../group_values/single_group_by/primitive.rs | 78 ++++++++++++++++++- 3 files changed, 208 insertions(+), 2 deletions(-) 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/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. /// From 4c053537d2a305534d6d6a4076de200a53dc7294 Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Thu, 30 Jul 2026 09:32:47 +0800 Subject: [PATCH 2/2] perf(aggregate): extend the consecutive-keys fast path to multi-column GROUP BY MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends the last-group cache to `GroupValuesColumn` (both the vectorized and the streaming/scalarized intern), covering multi-column group keys such as ClickBench's `GROUP BY "UserID", "SearchPhrase"`. Multi-column runs need every column to match its predecessor, so instead of a per-row cached slot the batch computes an adjacent-equality mask up front, guarded by a two-tier gate: - A single sequential scan over the (already computed) hashes counts adjacent hash matches. Hash equality is a necessary condition for row equality, so a shuffled batch — e.g. the Final phase after hash repartitioning, where runs cannot survive — fails the gate and pays nothing beyond this one scan. Correctness never rests on the hashes: they only decide whether the mask is worth building. - Run-heavy batches build the precise mask: one vectorized `not_distinct` pass per column, AND-ed. `not_distinct` treats `null == null` as equal, matching GROUP BY semantics, and never returns nulls. Column types it does not support (e.g. nested types handled by `RowsGroupColumn`) make the helper return `None`, falling back to the normal path. Integration is deliberately non-invasive to the vectorized machinery: masked rows are skipped in `collect_vectorized_process_context` (they enter none of the append / equal-to / remaining lists), and a final ascending fill pass assigns each one its predecessor's group index — forward propagation resolves multi-row runs from the run head, which took the normal path. The streaming path checks the mask at the top of its row loop and reuses the previous row's group index directly; sorted input has perfect runs, so it benefits the most. ClickBench multi-column queries (hits.parquet, 5-8 iterations, release-nonlto, same-commit baseline): Q18 (UserID, m, SearchPhrase) 3714ms -> 3312ms (-10.8%) Q32 (WatchID, ClientIP) 3478ms -> 3122ms (-10.2%, 8 iters) Q17 (UserID, SearchPhrase) 740ms -> 700ms ( -5.4%) Q16 (UserID, SearchPhrase) 819ms -> 795ms ( -3.0%) Tests: a multi-column runs test asserting identical group assignments on both intern paths — including the "one column runs, the other breaks the run" case and null==null runs — plus a no-runs batch exercising the gate fallback. Probe-verified (temporary eprintln) that the fast path fires for exactly the designed rows on both paths and that the gate rejects the no-runs batch. --- .../group_values/multi_group_by/mod.rs | 197 +++++++++++++++++- 1 file changed, 196 insertions(+), 1 deletion(-) 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();