Jefffrey commented on code in PR #10614:
URL: https://github.com/apache/arrow-rs/pull/10614#discussion_r3801020253
##########
parquet/src/arrow/array_reader/byte_array_dictionary.rs:
##########
@@ -128,6 +128,7 @@ struct ByteArrayDictionaryReader<K: ArrowNativeType, V:
OffsetSizeTrait> {
def_levels_buffer: Option<Vec<i16>>,
rep_levels_buffer: Option<Vec<i16>>,
record_reader: GenericRecordReader<DictionaryBuffer<K, V>,
DictionaryDecoder<K, V>>,
+ hash_scratch: MutableBuffer,
Review Comment:
worth leaving a comment here explaining what its used for, since its quite
detached in space from `DictionaryBuffer`
##########
parquet/src/arrow/buffer/dictionary_buffer.rs:
##########
@@ -127,6 +131,7 @@ impl<K: ArrowNativeType + Ord, V: OffsetSizeTrait>
DictionaryBuffer<K, V> {
self,
null_buffer: Option<Buffer>,
data_type: &ArrowType,
+ hash_scratch: &mut MutableBuffer,
Review Comment:
not an api change since `DictionaryBuffer` isnt part of our public api 👍
##########
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
+ }
+ }
+ 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 raw = builder.finish();
- Ok(Arc::new(raw))
+
+ let arrow_value_type = if V::IS_LARGE {
Review Comment:
we need to consider binary types too; considering our test suite didnt break
maybe we need a test to lock it in too
##########
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 suppose this should be rare, but if we have multiple of the collided
values theyll keep getting appended as new values
--
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]
