|
18 | 18 | use crate::aggregates::group_values::GroupValues; |
19 | 19 | use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; |
20 | 20 | use arrow::array::{ |
21 | | - ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, |
22 | | - cast::AsArray, |
| 21 | + Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, |
| 22 | + PrimitiveArray, cast::AsArray, |
23 | 23 | }; |
24 | 24 | use arrow::datatypes::{DataType, i256}; |
25 | 25 | use datafusion_common::Result; |
26 | | -use datafusion_common::hash_utils::RandomState; |
| 26 | +use datafusion_common::hash_utils::{RandomState, with_hashes}; |
27 | 27 | use datafusion_execution::memory_pool::proxy::VecAllocExt; |
28 | 28 | use datafusion_expr::EmitTo; |
29 | 29 | use half::f16; |
@@ -81,13 +81,12 @@ hash_float!(f16, f32, f64); |
81 | 81 | pub struct GroupValuesPrimitive<T: ArrowPrimitiveType> { |
82 | 82 | /// The data type of the output array |
83 | 83 | 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 |
85 | 85 | /// |
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)>, |
91 | 90 | /// The group index of the null value if any |
92 | 91 | null_group: Option<usize>, |
93 | 92 | /// The values for each group index |
@@ -117,42 +116,79 @@ where |
117 | 116 | assert_eq!(cols.len(), 1); |
118 | 117 | groups.clear(); |
119 | 118 |
|
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( |
131 | 138 | 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), |
136 | 141 | ); |
137 | | - |
138 | | - match insert { |
| 142 | + let group_id = match insert { |
139 | 143 | hashbrown::hash_table::Entry::Occupied(o) => o.get().0, |
140 | 144 | 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); |
144 | 148 | g |
145 | 149 | } |
146 | | - } |
| 150 | + }; |
| 151 | + groups.push(group_id); |
147 | 152 | } |
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 | + })?; |
151 | 186 | Ok(()) |
152 | 187 | } |
153 | 188 |
|
154 | 189 | 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() |
156 | 192 | } |
157 | 193 |
|
158 | 194 | fn is_empty(&self) -> bool { |
@@ -219,6 +255,6 @@ where |
219 | 255 | self.values.clear(); |
220 | 256 | self.values.shrink_to(num_rows); |
221 | 257 | 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)); |
223 | 259 | } |
224 | 260 | } |
0 commit comments