Skip to content

Commit 05b874d

Browse files
Dandandanclaude
andcommitted
Optimize GroupValuesPrimitive::intern with vectorized hashing and inline values
Two optimizations for the single-column primitive GROUP BY hot path: 1. Vectorized hashing: Split intern into two phases — batch hash computation via with_hashes (tight loop, better pipelining) followed by hash table probing with pre-computed hashes. 2. Inline values in hash table: Store the actual value in each hash table entry (usize, T::Native) instead of (usize, u64) with an indirect lookup into a separate values vec. This eliminates one cache miss per probe and removes the need to store the hash — the value can be rehashed from the inline copy when needed (rare, only during table growth). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 1e93a67 commit 05b874d

1 file changed

Lines changed: 71 additions & 35 deletions

File tree

  • datafusion/physical-plan/src/aggregates/group_values/single_group_by

datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs

Lines changed: 71 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,12 @@
1818
use crate::aggregates::group_values::GroupValues;
1919
use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano};
2020
use arrow::array::{
21-
ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray,
22-
cast::AsArray,
21+
Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder,
22+
PrimitiveArray, cast::AsArray,
2323
};
2424
use arrow::datatypes::{DataType, i256};
2525
use datafusion_common::Result;
26-
use datafusion_common::hash_utils::RandomState;
26+
use datafusion_common::hash_utils::{RandomState, with_hashes};
2727
use datafusion_execution::memory_pool::proxy::VecAllocExt;
2828
use datafusion_expr::EmitTo;
2929
use half::f16;
@@ -81,13 +81,12 @@ hash_float!(f16, f32, f64);
8181
pub struct GroupValuesPrimitive<T: ArrowPrimitiveType> {
8282
/// The data type of the output array
8383
data_type: DataType,
84-
/// Stores the `(group_index, hash)` based on the hash of its value
84+
/// Stores the `(group_index, value)` based on the hash of its value
8585
///
86-
/// We also store `hash` is for reducing cost of rehashing. Such cost
87-
/// is obvious in high cardinality group by situation.
88-
/// More details can see:
89-
/// <https://github.com/apache/datafusion/issues/15961>
90-
map: HashTable<(usize, u64)>,
86+
/// Storing the value inline avoids an indirection into `values` during
87+
/// hash-table probing (one fewer cache miss per probe) and allows
88+
/// direct equality comparison without a stored hash.
89+
map: HashTable<(usize, T::Native)>,
9190
/// The group index of the null value if any
9291
null_group: Option<usize>,
9392
/// The values for each group index
@@ -117,42 +116,79 @@ where
117116
assert_eq!(cols.len(), 1);
118117
groups.clear();
119118

120-
for v in cols[0].as_primitive::<T>() {
121-
let group_id = match v {
122-
None => *self.null_group.get_or_insert_with(|| {
123-
let group_id = self.values.len();
124-
self.values.push(Default::default());
125-
group_id
126-
}),
127-
Some(key) => {
128-
let state = &self.random_state;
129-
let hash = key.hash(state);
130-
let insert = self.map.entry(
119+
// Destructure to avoid borrow conflicts between random_state
120+
// (borrowed by with_hashes) and map/values/null_group (mutated in callback)
121+
let Self {
122+
map,
123+
null_group,
124+
values,
125+
random_state,
126+
..
127+
} = self;
128+
129+
// Phase 1: Compute all hashes in a tight loop (via with_hashes)
130+
// Phase 2: Probe hash table using pre-computed hashes
131+
with_hashes([&cols[0]], random_state, |hashes| {
132+
let arr = cols[0].as_primitive::<T>();
133+
134+
if arr.null_count() == 0 {
135+
for (idx, &key) in arr.values().iter().enumerate() {
136+
let hash = hashes[idx];
137+
let insert = map.entry(
131138
hash,
132-
|&(g, h)| unsafe {
133-
hash == h && self.values.get_unchecked(g).is_eq(key)
134-
},
135-
|&(_, h)| h,
139+
|&(_, v)| v.is_eq(key),
140+
|&(_, v)| v.hash(random_state),
136141
);
137-
138-
match insert {
142+
let group_id = match insert {
139143
hashbrown::hash_table::Entry::Occupied(o) => o.get().0,
140144
hashbrown::hash_table::Entry::Vacant(v) => {
141-
let g = self.values.len();
142-
v.insert((g, hash));
143-
self.values.push(key);
145+
let g = values.len();
146+
v.insert((g, key));
147+
values.push(key);
144148
g
145149
}
146-
}
150+
};
151+
groups.push(group_id);
147152
}
148-
};
149-
groups.push(group_id)
150-
}
153+
} else {
154+
for (idx, v) in arr.iter().enumerate() {
155+
let group_id = match v {
156+
None => *null_group.get_or_insert_with(|| {
157+
let group_id = values.len();
158+
values.push(Default::default());
159+
group_id
160+
}),
161+
Some(key) => {
162+
let hash = hashes[idx];
163+
let insert = map.entry(
164+
hash,
165+
|&(_, v)| v.is_eq(key),
166+
|&(_, v)| v.hash(random_state),
167+
);
168+
match insert {
169+
hashbrown::hash_table::Entry::Occupied(o) => {
170+
o.get().0
171+
}
172+
hashbrown::hash_table::Entry::Vacant(v) => {
173+
let g = values.len();
174+
v.insert((g, key));
175+
values.push(key);
176+
g
177+
}
178+
}
179+
}
180+
};
181+
groups.push(group_id);
182+
}
183+
}
184+
Ok(())
185+
})?;
151186
Ok(())
152187
}
153188

154189
fn size(&self) -> usize {
155-
self.map.capacity() * size_of::<(usize, u64)>() + self.values.allocated_size()
190+
self.map.capacity() * size_of::<(usize, T::Native)>()
191+
+ self.values.allocated_size()
156192
}
157193

158194
fn is_empty(&self) -> bool {
@@ -219,6 +255,6 @@ where
219255
self.values.clear();
220256
self.values.shrink_to(num_rows);
221257
self.map.clear();
222-
self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared
258+
self.map.shrink_to(num_rows, |&(_, v)| v.hash(&self.random_state));
223259
}
224260
}

0 commit comments

Comments
 (0)