Rich-T-kid commented on code in PR #10614:
URL: https://github.com/apache/arrow-rs/pull/10614#discussion_r3788475851
##########
parquet/src/arrow/buffer/dictionary_buffer.rs:
##########
@@ -226,58 +233,169 @@ 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
+ }
+ }
+ Entry::Vacant(entry) => {
+ let 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);
+ entry.insert((input_idx, output_idx));
+ output_idx
+ }
+ };
+
+ let key = K::Native::from_usize(output_idx)
+ .ok_or_else(|| general_err!("dictionary key overflow"))?;
+ keys.push(key);
+ }
+
+ let arrow_value_type = if V::IS_LARGE {
+ ArrowType::LargeUtf8
+ } else {
+ ArrowType::Utf8
+ };
+ let num_unique = unique_offsets.len() - 1;
+
+ // SAFETY: buffers are constructed directly from typed Vecs above; offsets
are
+ // monotonically non-decreasing and bounded by unique_bytes.len(), and all
+ // key values are within 0..num_unique, so the invariants Arrow requires
hold.
+ let value_data = unsafe {
+ arrow_data::ArrayData::builder(arrow_value_type)
+ .len(num_unique)
+ .add_buffer(Buffer::from_vec(unique_offsets))
+ .add_buffer(Buffer::from_vec(unique_bytes))
+ .build_unchecked()
+ };
+
+ // SAFETY: keys are within 0..num_unique and value_data is valid.
+ let dict_array: DictionaryArray<K> = unsafe {
+ arrow_data::ArrayData::builder(dict_type)
+ .len(keys.len())
+ .add_buffer(Buffer::from_vec(keys))
+ .add_child_data(value_data)
+ .null_bit_buffer(null_buffer)
+ .build_unchecked()
+ .into()
+ };
+
+ Ok(Arc::new(dict_array))
+}
+
+fn hash_byte_slices<I: ArrowNativeType>(offsets: &[I], values: &[u8], scratch:
&mut Vec<u8>) {
+ let count = offsets.len().saturating_sub(1);
+ scratch.clear();
+ scratch.resize(count * size_of::<u64>(), 0u8);
+
+ let state = RandomState::new();
+ // SAFETY: scratch is sized to exactly count * size_of::<u64>() above
+ let hash_slots = unsafe { from_raw_parts_mut(scratch.as_mut_ptr() as *mut
u64, count) };
Review Comment:
ah 🤔 , commit
https://github.com/apache/arrow-rs/pull/10614/changes/4c4bcadff8e1a6611291377e69fc206b5b181136
replaces the vector scratch space with a
https://docs.rs/arrow/latest/arrow/buffer/struct.MutableBuffer.html, which is
guaranteed to 64 byte aligned.
--
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]
