This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11524-193b9d30a65a791b1b2d5e7f2b3df215195dfade in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit cacb42eaf39646c3ca52459c954bba348bc343b1 Author: Raz Luvaton <[email protected]> AuthorDate: Thu Oct 29 06:06:00 2026 +0000 perf(parquet): optimize dictionary computation for byte arrays (#11524) # Which issue does this PR close? None # Rationale for this change I want dictionary encoding to be faster for strings. The old implementation have a few things that are slow: 1. whenever a resize is needed we rehash the entire inner map - so this version saves the hash in the map to avoid that 2. for small strings we need to do random access to memory that is not in the cache - this version fixes that by inlining the small strings in the hash map Later we can add to bulk hash the values These changes are heavily inspired by DataFusion `ArrowBytesMap` and `GroupIndexView` # What changes are included in this PR? Add custom implementation for `Interner` that is for byte array. in that implementation, the hash map contain the value hashes to avoid rehashing and contain inline values for short strings to avoid random memory access # Are these changes tested? Existing tests # Are there any significant user-facing changes? I'm not sure if there are breaking changes or not since I removed some pub functions but not sure if they were reachable in the first place. --- parquet/src/arrow/arrow_writer/byte_array.rs | 314 ++++++++++++++++++++++++--- parquet/src/util/interner.rs | 7 +- parquet/tests/arrow_writer/mod.rs | 47 +++- 3 files changed, 325 insertions(+), 43 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 70687874b8..7f3120b71e 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -32,7 +32,6 @@ use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accu use crate::geospatial::statistics::GeospatialStatistics; use crate::schema::types::ColumnDescPtr; use crate::util::bit_util::num_required_bits; -use crate::util::interner::{Interner, Storage}; use crate::util::prefix::common_prefix_length; use arrow_array::types::ByteArrayType; use arrow_array::{ @@ -302,53 +301,287 @@ impl FallbackEncoder { } } -/// [`Storage`] for the [`Interner`] used by [`DictEncoder`] +/// Maximum size of a value that can be inlined in the hash table +const SHORT_VALUE_LEN: usize = size_of::<u64>() - 1; + +/// Entry in the hash table +#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)] +struct Entry { + /// hash of the value (stored to avoid recomputing it in hash table check) + hash: u64, + /// if len =< [`SHORT_VALUE_LEN`]: the data inlined + /// if len > [`SHORT_VALUE_LEN`], the offset of where the data starts + offset_or_inline: OffsetOrInlineView, + index_in_dict: u64, +} + +/// Either offset to the dictionary page bytes or inline bytes when the bytes are short enough +/// to avoid doing the indirection +/// +/// For inlined view: +/// +------------+---------------+------------------------------+ +/// | 1 bit set | len {7 bits} | value {56 bits} | +/// +------------+---------------+------------------------------+ +/// +/// For non-inlined view: +/// +--------------+--------------------------------------------+ +/// | 1 bit unset | offset in dict {63 bits} | +/// +--------------+--------------------------------------------+ +/// +/// the inline view contain it's length to avoid having another field for length in the entry +/// since if the hashes are the same, most probably the length will be the same, so we can do the indirection check +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +struct OffsetOrInlineView(u64); + +impl OffsetOrInlineView { + const INLINE_FLAG: u64 = 1 << 63; + const LEN_SHIFT: u32 = 56; + + #[inline] + fn is_non_inlined(&self) -> bool { + (self.0 & Self::INLINE_FLAG) == 0 + } + + #[inline] + fn is_inlined(&self) -> bool { + !self.is_non_inlined() + } + + #[inline] + fn new_inlined(value: &[u8]) -> Self { + debug_assert!(value.len() <= 7); + + // Value bytes occupy the low 56 bits, byte 0 in the least significant position + let mut buf = [0u8; 8]; + buf[..value.len()].copy_from_slice(value); + let packed = u64::from_le_bytes(buf); + + Self(Self::INLINE_FLAG | ((value.len() as u64) << Self::LEN_SHIFT) | packed) + } + + #[inline] + fn new_non_inlined(offset: u64) -> Self { + // `offset` is a position in the dictionary page `Vec`, whose length is at + // most `isize::MAX`, so it always fits in 63 bits and can't collide with + // `INLINE_FLAG` + // also, having the inline flag set for non-inlined values would mean the dictionary page is at least 9 exabytes. + debug_assert_eq!( + offset & Self::INLINE_FLAG, + 0, + "dictionary page offset must fit in 63 bits" + ); + let non_inlined_value = offset; + Self(non_inlined_value) + } + + fn maybe_non_inlined_value(&self) -> Option<u64> { + if self.is_non_inlined() { + // Not inlined is when the first bit is unset, so the value + Some(self.0) + } else { + None + } + } +} + #[derive(Debug, Default)] -struct ByteArrayStorage { - /// Encoded dictionary data +struct ByteArrayInterner { + state: ahash::RandomState, + /// Underlying hash set for each distinct value + map: hashbrown::hash_table::HashTable<Entry>, + /// Encoded dictionary page: each value prefixed by its length as a `u32` page: Vec<u8>, - - values: Vec<std::ops::Range<usize>>, } -impl Storage for ByteArrayStorage { - type Key = u64; - type Value = [u8]; +impl ByteArrayInterner { + fn intern(&mut self, value: &[u8]) -> Result<u64> { + let hash = self.state.hash_one(value); + + // value is "small" - can be inlined + let (found_entry, inline) = if let Some(inline) = maybe_inline_value(value) { + // is value is already present in the set? + let entry = self.map.find(hash, |header| { + // value is stored inline so no need to consult buffer + // (this is the "small string optimization") + // + // no need to compare the hashes since the hash is a function of the value which is compared here + let is_same = inline == header.offset_or_inline; + + if is_same { + debug_assert_eq!( + header.hash, hash, + "hashes must match if the values are the same" + ); + } + + is_same + }); + + (entry, Some(inline)) + } else { + // value is not "small" + // Check if the value is already present in the set + let entry = self.map.find(hash, |header| { + // compare hash + if header.hash != hash { + return false; + } + + // If the value is not inlined compare it + let Some(offset) = header.offset_or_inline.maybe_non_inlined_value() else { + // that value is inline, so it must be different + return false; + }; + + self.is_at(offset, value) + }); + + (entry, None) + }; + + if let Some(entry) = found_entry { + return Ok(entry.index_in_dict); + } - fn get(&self, idx: Self::Key) -> &Self::Value { - &self.page[self.values[idx as usize].clone()] + // if no existing entry, make a new one + self.insert(value, hash, inline) } - fn push(&mut self, value: &Self::Value) -> Self::Key { - let key = self.values.len(); + /// Returns whether `value` is in the dictionary page at `offset` + #[inline(always)] + fn is_at(&self, offset: u64, value: &[u8]) -> bool { + let offset = offset as usize; + // SAFETY: every value in the page is preceded by its length + let stored_len = unsafe { self.page.get_unchecked(offset - 4..offset) }; + // Rarely false, as the hashes matched, so the branch predicts well and + // the comparison's reads issue before the length is known + if u32::from_le_bytes(stored_len.try_into().unwrap()) as usize != value.len() { + return false; + } + // SAFETY: the value at `offset` has the length just checked + let existing = unsafe { self.page.get_unchecked(offset..offset + value.len()) }; + existing == value + } - self.page.reserve(4 + value.len()); - self.page.extend_from_slice((value.len() as u32).as_bytes()); + fn num_unique_values(&self) -> usize { + self.map.len() + } - let start = self.page.len(); + /// Appends `value` to the dictionary, in the empty slot at `pos` + fn insert( + &mut self, + value: &[u8], + hash: u64, + inline: Option<OffsetOrInlineView>, + ) -> Result<u64> { + // If inline provided it should inlined + debug_assert!(inline.is_none_or(|i| i.is_inlined())); + + // The dictionary page is PLAIN encoded, where each value is prefixed by + // its length as a 4-byte little endian integer so we need to reject that since we use unsafe access later + // + // From the parquet-format spec: + // > Encoding PLAIN: + // > [...] For the byte array type, it encodes the length as a 4-byte little + // > endian integer, followed by the bytes. + // > + // > Dictionary Encoding: + // > [...] Dictionary page format: the entries in the dictionary using the plain encoding. + // > + // + // Reference: + // - Plain encoding: https://github.com/apache/parquet-format/blob/53b8bd8576827080d4382fa457003f45b41cda6e/Encodings.md?plain=1#L72-L73 + // - Plain dictionary encoding: https://github.com/apache/parquet-format/blob/53b8bd8576827080d4382fa457003f45b41cda6e/Encodings.md?plain=1#L83 + if value.len() > u32::MAX as usize { + return Err(ParquetError::General(format!( + "Byte array value is too large to be written in dictionary: {} bytes", + value.len() + ))); + } + + let key = self.num_unique_values() as u64; + let len = value.len(); + + self.page.reserve(4 + value.len()); + self.page.extend_from_slice(&(len as u32).to_le_bytes()); + let offset = self.page.len() as u64; self.page.extend_from_slice(value); - self.values.push(start..self.page.len()); - key as u64 + // Make the new entry + let new_header = Entry { + hash, + offset_or_inline: inline.unwrap_or_else(|| OffsetOrInlineView::new_non_inlined(offset)), + index_in_dict: key, + }; + self.map + .insert_unique(hash, new_header, |header| header.hash); + + Ok(key) } fn estimated_memory_size(&self) -> usize { - self.page.capacity() * std::mem::size_of::<u8>() - + self.values.capacity() * std::mem::size_of::<std::ops::Range<usize>>() + self.page.capacity() * std::mem::size_of::<u8>() + self.map.allocation_size() + } + + fn iter(&self) -> impl Iterator<Item = &[u8]> { + ByteArrayInternerIterator { + values: &self.page, + current_index: 0, + } + } +} + +/// Iterates the values of a [`ByteArrayInterner`]'s dictionary page +/// +/// `values` must be a page built by [`ByteArrayInterner::insert`], which is a +/// sequence of values each preceded by its exact length as a `u32`, and +/// `current_index` must start at 0, so it is always at the start of a value +struct ByteArrayInternerIterator<'a> { + values: &'a [u8], + current_index: usize, +} + +impl<'a> Iterator for ByteArrayInternerIterator<'a> { + type Item = &'a [u8]; + + fn next(&mut self) -> Option<Self::Item> { + if self.current_index >= self.values.len() { + return None; + } + + // SAFETY: `current_index` is at the start of a value and is not past + // the end of the page, so the value's 4-byte length prefix is in bounds + let len_bytes = unsafe { + self.values + .get_unchecked(self.current_index..self.current_index + 4) + }; + let len = u32::from_le_bytes(len_bytes.try_into().unwrap()) as usize; + + let start = self.current_index + 4; + let end = start + len; + + // SAFETY: `insert` writes exactly `len` bytes after the prefix, and + // rejects values whose length does not fit in a `u32`, so the prefix is + // never truncated and `start..end` is within the page + let value = unsafe { self.values.get_unchecked(start..end) }; + + self.current_index = end; + + Some(value) } } /// A dictionary encoder for byte array data #[derive(Debug, Default)] struct DictEncoder { - interner: Interner<ByteArrayStorage>, + interner: ByteArrayInterner, indices: Vec<u64>, variable_length_bytes: i64, } impl DictEncoder { /// Encode `values` to the in-progress page - fn encode<T>(&mut self, values: T, indices: impl ExactSizeIterator<Item = usize>) + fn encode<T>(&mut self, values: T, indices: impl ExactSizeIterator<Item = usize>) -> Result<()> where T: ArrayAccessor + Copy, T::Item: AsRef<[u8]>, @@ -357,14 +590,16 @@ impl DictEncoder { for idx in indices { let value = values.value(idx); - let interned = self.interner.intern(value.as_ref()); + let interned = self.interner.intern(value.as_ref())?; self.indices.push(interned); self.variable_length_bytes += value.as_ref().len() as i64; } + + Ok(()) } fn bit_width(&self) -> u8 { - let length = self.interner.storage().values.len(); + let length = self.interner.num_unique_values(); num_required_bits(length.saturating_sub(1) as u64) } @@ -378,15 +613,16 @@ impl DictEncoder { } fn estimated_dict_page_size(&self) -> usize { - self.interner.storage().page.len() + self.interner.page.len() } fn flush_dict_page(self) -> DictionaryPage { - let storage = self.interner.into_inner(); + let unique_values = self.interner.num_unique_values(); + let bytes = self.interner.page; DictionaryPage { - buf: storage.page.into(), - num_values: storage.values.len(), + buf: bytes.into(), + num_values: unique_values, is_sorted: false, } } @@ -422,6 +658,14 @@ impl DictEncoder { } } +fn maybe_inline_value(value: &[u8]) -> Option<OffsetOrInlineView> { + if value.len() <= SHORT_VALUE_LEN { + Some(OffsetOrInlineView::new_inlined(value)) + } else { + None + } +} + pub struct ByteArrayEncoder { fallback: FallbackEncoder, dict_encoder: Option<DictEncoder>, @@ -487,8 +731,7 @@ impl ColumnValueEncoder for ByteArrayEncoder { encode, indices.iter().copied(), self - ); - Ok(()) + ) } fn count_values_within_byte_budget_gather( @@ -633,9 +876,8 @@ impl ColumnValueEncoder for ByteArrayEncoder { } if let Some(bloom_filter) = &mut self.bloom_filter { - let storage = encoder.interner.storage(); - for range in &storage.values { - bloom_filter.insert(&storage.page[range.clone()]); + for value in encoder.interner.iter() { + bloom_filter.insert(value); } } @@ -663,7 +905,7 @@ impl ColumnValueEncoder for ByteArrayEncoder { /// Encodes the provided `values` and `indices` to `encoder` /// /// This is a free function so it can be used with `downcast_op!` -fn encode<T, I>(values: T, indices: I, encoder: &mut ByteArrayEncoder) +fn encode<T, I>(values: T, indices: I, encoder: &mut ByteArrayEncoder) -> Result<()> where T: ArrayAccessor + Copy, T::Item: Copy + Ord + AsRef<[u8]>, @@ -709,7 +951,9 @@ where bloom_filter.insert(values.value(idx).as_ref()); } } - encoder.fallback.encode(values, indices) + encoder.fallback.encode(values, indices); + + Ok(()) } } } diff --git a/parquet/src/util/interner.rs b/parquet/src/util/interner.rs index f4e5e170c9..96b8b5225c 100644 --- a/parquet/src/util/interner.rs +++ b/parquet/src/util/interner.rs @@ -17,6 +17,7 @@ use crate::data_type::AsBytes; use hashbrown::HashTable; +use std::fmt::Debug; const DEFAULT_DEDUP_CAPACITY: usize = 4096; @@ -89,10 +90,4 @@ impl<S: Storage> Interner<S> { pub fn storage(&self) -> &S { &self.storage } - - /// Unwraps the inner storage - #[cfg(feature = "arrow")] - pub fn into_inner(self) -> S { - self.storage - } } diff --git a/parquet/tests/arrow_writer/mod.rs b/parquet/tests/arrow_writer/mod.rs index caf89d3641..791205dfe5 100644 --- a/parquet/tests/arrow_writer/mod.rs +++ b/parquet/tests/arrow_writer/mod.rs @@ -29,7 +29,10 @@ use std::fs::File; use std::io::{Read as _, Seek, SeekFrom, Write as _}; use std::sync::Arc; -use arrow::array::{ArrayRef, BinaryArray, Float64Array, Int32Array, RecordBatch, StringArray}; +use arrow::array::{ + ArrayRef, BinaryArray, Float64Array, Int32Array, LargeBinaryArray, RecordBatch, StringArray, +}; +use arrow::buffer::{Buffer, OffsetBuffer}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use bytes::Bytes; use parquet::arrow::arrow_writer::{ @@ -39,7 +42,7 @@ use parquet::arrow::arrow_writer::{ use parquet::arrow::{ArrowSchemaConverter, ArrowWriter}; use parquet::basic::Encoding; use parquet::errors::Result; -use parquet::file::properties::{WriterProperties, WriterVersion}; +use parquet::file::properties::{EnabledStatistics, WriterProperties, WriterVersion}; use parquet::file::writer::SerializedFileWriter; #[test] @@ -141,6 +144,14 @@ unsafe impl GlobalAlloc for TrackingAllocator { } new_ptr } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { self.inner.alloc_zeroed(layout) }; + if !ptr.is_null() { + add_live_bytes(layout.size()); + } + ptr + } } /// Run `f` and return the peak *additional* live heap (bytes) observed on this @@ -593,3 +604,35 @@ fn unused_delta_fallback_does_not_preallocate() { assert!(peak < COLUMNS * 64 * 1024, "peak: {peak}"); } + +#[test] +// The value is 2^32 bytes, which does not fit in a 32-bit address space +#[cfg(target_pointer_width = "64")] +fn test_dictionary_rejects_value_longer_than_u32() { + // The dictionary page prefixes each value with its length as a u32, so a + // 2^32 byte value would have its length truncated to 0 + let len = 1usize << 32; + let mut data = vec![0u8; len]; + // Read as the next value's length if the truncated prefix were trusted + data[..4].copy_from_slice(&[0xFF; 4]); + let array = LargeBinaryArray::new( + OffsetBuffer::new(vec![0i64, len as i64].into()), + Buffer::from_vec(data), + None, + ); + let batch = RecordBatch::try_from_iter([("col", Arc::new(array) as ArrayRef)]).unwrap(); + + let props = WriterProperties::builder() + // statistics are disabled to avoid copying the value for min/max and using a lot of extra memory in this test + .set_statistics_enabled(EnabledStatistics::None) + // The bloom filter is built from the dictionary page, + .set_bloom_filter_enabled(true) + .build(); + let mut writer = ArrowWriter::try_new(std::io::sink(), batch.schema(), Some(props)).unwrap(); + + let err = writer.write(&batch).unwrap_err().to_string(); + assert_eq!( + err, + "Parquet error: Byte array value is too large to be written in dictionary: 4294967296 bytes" + ); +}
