Rich-T-kid commented on code in PR #10614:
URL: https://github.com/apache/arrow-rs/pull/10614#discussion_r3805323152
##########
parquet/src/arrow/buffer/dictionary_buffer.rs:
##########
@@ -226,58 +233,170 @@ impl<K: ArrowNativeType, V: OffsetSizeTrait>
ValuesBuffer for DictionaryBuffer<K
}
}
-macro_rules! dict_helper {
- ($k:ty, $array:ident) => {
- match $array.data_type() {
- ArrowType::Utf8 => pack_values_impl::<$k,
_>($array.as_string::<i32>()),
- ArrowType::LargeUtf8 => pack_values_impl::<$k,
_>($array.as_string::<i64>()),
- ArrowType::Binary => pack_values_impl::<$k,
_>($array.as_binary::<i32>()),
- ArrowType::LargeBinary => pack_values_impl::<$k,
_>($array.as_binary::<i64>()),
- ArrowType::FixedSizeBinary(_) => {
- pack_fixed_values_impl::<$k>($array.as_fixed_size_binary())
- }
- _ => unreachable!(),
- }
+macro_rules! offsets_dict_helper {
+ ($k:ty, $key_type:ident, $value_type:ident, $values:ident, $hashes:ident,
$null_buffer:ident) => {
+ pack_values_from_offsets_impl::<$k, _>(
+ $values,
+ $hashes,
+ $null_buffer,
+ $key_type,
+ $value_type,
+ )
};
}
-fn pack_values(key_type: &ArrowType, values: &ArrayRef) -> Result<ArrayRef> {
+fn pack_values_from_offsets<V: OffsetSizeTrait>(
+ key_type: &ArrowType,
+ value_type: &ArrowType,
+ values: &OffsetBuffer<V>,
+ hashes: &[u64],
+ null_buffer: Option<Buffer>,
+) -> Result<ArrayRef> {
downcast_integer! {
- key_type => (dict_helper, values),
- _ => unreachable!(),
+ key_type => (offsets_dict_helper, key_type, value_type, values,
hashes, null_buffer),
+ _ => unreachable!(),
}
}
-fn pack_values_impl<K: ArrowDictionaryKeyType, T: ByteArrayType>(
- array: &GenericByteArray<T>,
-) -> Result<ArrayRef> {
- let mut builder = GenericByteDictionaryBuilder::<K,
T>::with_capacity(array.len(), 1024, 1024);
- for x in array {
- match x {
- Some(x) => builder.append_value(x),
- None => builder.append_null(),
- }
+// Avoids double-hashing: keys are already high-quality u64 hashes from ahash,
+// so we pass them through directly rather than re-hashing inside the HashMap.
+struct PassthroughHasher(u64);
+impl std::hash::Hasher for PassthroughHasher {
+ fn finish(&self) -> u64 {
+ self.0
+ }
+ fn write(&mut self, _: &[u8]) {
+ unreachable!()
+ }
+ fn write_u64(&mut self, value: u64) {
+ self.0 = value;
+ }
+}
+#[derive(Default)]
+struct BuildPassthroughHasher;
+impl std::hash::BuildHasher for BuildPassthroughHasher {
+ type Hasher = PassthroughHasher;
+ fn build_hasher(&self) -> PassthroughHasher {
+ PassthroughHasher(0)
}
- let raw = builder.finish();
- Ok(Arc::new(raw))
}
-fn pack_fixed_values_impl<K: ArrowDictionaryKeyType>(
- array: &FixedSizeBinaryArray,
+/// Builds a [`DictionaryArray`] directly from a flat [`OffsetBuffer`] using
pre-computed
+/// hashes to deduplicate values in a single pass, avoiding the intermediate
StringArray
+/// materialization
+fn pack_values_from_offsets_impl<K: ArrowDictionaryKeyType, V:
OffsetSizeTrait>(
+ offset_buffer: &OffsetBuffer<V>,
+ hashes: &[u64],
+ null_buffer: Option<Buffer>,
+ key_type: &ArrowType,
+ value_type: &ArrowType,
) -> Result<ArrayRef> {
- let mut builder = FixedSizeBinaryDictionaryBuilder::<K>::with_capacity(
- array.len(),
- 1024,
- array.value_length(),
- );
- for x in array {
- match x {
- Some(x) => builder.append_value(x),
- None => builder.append_null(),
- }
+ let dict_type = ArrowType::Dictionary(Box::new(key_type.clone()),
Box::new(value_type.clone()));
+ let num_values = offset_buffer.len();
+
+ let mut keys: Vec<K::Native> = Vec::with_capacity(num_values);
+ let mut unique_offsets: Vec<V> = Vec::with_capacity(num_values + 1);
+ unique_offsets.push(V::default());
+ let mut unique_bytes: Vec<u8> =
Vec::with_capacity(offset_buffer.values.len());
+
+ let mut dedup: HbHashMap<u64, (usize, usize), BuildPassthroughHasher> =
+ HbHashMap::with_capacity_and_hasher(num_values,
BuildPassthroughHasher);
+
+ for (input_idx, &hash) in hashes.iter().enumerate() {
+ let byte_start = offset_buffer.offsets[input_idx].as_usize();
+ let byte_end = offset_buffer.offsets[input_idx + 1].as_usize();
+ let bytes = &offset_buffer.values[byte_start..byte_end];
+
+ let output_idx = match dedup.entry(hash) {
+ Entry::Occupied(entry) => {
+ let (first_input_idx, existing_output_idx) = *entry.get();
+ let first_start =
offset_buffer.offsets[first_input_idx].as_usize();
+ let first_end = offset_buffer.offsets[first_input_idx +
1].as_usize();
+ if &offset_buffer.values[first_start..first_end] == bytes {
+ existing_output_idx
+ } else {
+ // True hash collision: same hash, different bytes —
insert as new unique value
+ let new_output_idx = unique_offsets.len() - 1;
+ unique_bytes.extend_from_slice(bytes);
+ let new_end = V::from_usize(unique_bytes.len())
+ .ok_or_else(|| general_err!("offset overflow building
dictionary"))?;
+ unique_offsets.push(new_end);
+ new_output_idx
Review Comment:
I just had claude iterate on this for a bit and the benchmarks regressed
from 15%.
collisions seem rare as you mentioned
```
arrow_array_reader/PlainStringToDictionary/low cardinality 1.41
1462.8±12.11µs ? ?/sec 1.00 1036.0±5.67µs ? ?/sec
```
and the low-card/high collision benchmark seems to be still be performant
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
