etseidl commented on code in PR #11524:
URL: https://github.com/apache/arrow-rs/pull/11524#discussion_r4350045823
##########
parquet/src/arrow/arrow_writer/byte_array.rs:
##########
@@ -301,46 +301,234 @@ 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,
+ /// value stored by the entry
+ 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 {
+ assert_eq!(
+ offset & Self::INLINE_FLAG,
+ 0,
+ "offset must not have the inlined flag set, we don't support this
large of a byte array"
+ );
+ 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>>,
+ /// Number of distinct values
+ num_unique_values: usize,
}
-impl Storage for ByteArrayStorage {
- type Key = u64;
- type Value = [u8];
+impl ByteArrayInterner {
+ fn intern(&mut self, value: &[u8]) -> 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| {
+ // compare value if hashes match
+ if header.hash != hash {
+ return false;
+ }
+ // value is stored inline so no need to consult buffer
+ // (this is the "small string optimization")
+ inline == header.offset_or_inline
+ });
+
+ (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
+ if let Some(offset) =
header.offset_or_inline.maybe_non_inlined_value() {
+ self.is_at(offset, value)
+ } else {
+ // that value is inline, so it must be different
+ false
+ }
+ });
+
+ (entry, None)
+ };
+
+ if let Some(entry) = found_entry {
+ return 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());
+ /// Appends `value` to the dictionary, in the empty slot at `pos`
+ fn insert(&mut self, value: &[u8], hash: u64, inline:
Option<OffsetOrInlineView>) -> u64 {
+ // If inline provided it should inlined
+ debug_assert!(inline.is_none_or(|i| i.is_inlined()));
+ let key = self.num_unique_values as u64;
+ let len = value.len();
- let start = self.page.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);
Review Comment:
True, but now we're using the logic here to justify `unsafe` blocks in
`next()`. My LLM found a rare edge case where a single write of a large binary
with a length of 2^32 would output a 0 len here, followed by the data. Using
the iterator below would then see a 0 len, skip it, and then read the first 4
bytes of the payload as the next len, which could then cause an out-of-bounds
memory access in the unsafe block leading to UB. Without bloom filters enabled,
this would trigger an error when writing the page, but with bloom filters
enabled, we hit the UB in `next()` and then 🤷. I'm not personally deeply
steeped in Arrow, so I don't know how likely this scenario is in practice.
--
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]
