Rich-T-kid commented on code in PR #11423:
URL: https://github.com/apache/arrow-rs/pull/11423#discussion_r4214299811
##########
parquet/src/arrow/array_reader/byte_array_dictionary.rs:
##########
@@ -414,18 +441,24 @@ where
return Ok(0); // All data must be NULL
}
- match out.as_keys(dict) {
+ let decoded_len = match out.as_keys_with_bit_width(dict,
Some(*bit_width)) {
Some(keys) => {
- // Happy path - can just copy keys
- // Keys will be validated on conversion to arrow
-
+ // Happy path: decode directly into the vec's spare
capacity to skip the zero-fill.
// TODO: Push vec into decoder (#5177)
let start = keys.len();
- keys.resize(start + len, K::default());
- let len = decoder.get_batch(&mut keys[start..])?;
- keys.truncate(start + len);
- *max_remaining_values -= len;
- Ok(len)
+ keys.reserve(len);
+ let decoded_len = {
+ let spare = keys.spare_capacity_mut();
+ // SAFETY: `MaybeUninit<K>` matches `K`'s layout;
`get_batch` is write-only.
+ let dst: &mut [K] = unsafe {
+
std::slice::from_raw_parts_mut(spare.as_mut_ptr().cast::<K>(), len)
+ };
+ decoder.get_batch(dst)?
+ };
+ // SAFETY: `get_batch` initialised the first
`decoded_len` spare slots.
+ unsafe { keys.set_len(start + decoded_len) };
+ *max_remaining_values -= decoded_len;
+ decoded_len
Review Comment:
I need to take a closer look at this
##########
parquet/src/arrow/buffer/dictionary_buffer.rs:
##########
@@ -172,9 +237,13 @@ impl<K: ArrowNativeType + Ord, V: OffsetSizeTrait>
DictionaryBuffer<K, V> {
values
};
+ let keys_len = keys.len();
+ let keys_buffer = Buffer::from_vec(keys);
+ *keys_buffer_sink = Some(keys_buffer.clone());
+
let builder = ArrayDataBuilder::new(data_type.clone())
- .len(keys.len())
- .add_buffer(Buffer::from_vec(keys))
+ .len(keys_len)
+ .add_buffer(keys_buffer)
.add_child_data(values.into_data())
.null_bit_buffer(null_buffer);
Review Comment:
use unsafe builder here?
##########
parquet/src/arrow/buffer/dictionary_buffer.rs:
##########
@@ -172,9 +237,13 @@ impl<K: ArrowNativeType + Ord, V: OffsetSizeTrait>
DictionaryBuffer<K, V> {
values
};
+ let keys_len = keys.len();
+ let keys_buffer = Buffer::from_vec(keys);
Review Comment:
99% sure this is zero copy but worth another check as theres a similar
method that allocates
##########
parquet/src/arrow/buffer/dictionary_buffer.rs:
##########
@@ -117,40 +170,52 @@ impl<K: ArrowNativeType + Ord, V: OffsetSizeTrait>
DictionaryBuffer<K, V> {
spilled.extend_from_dictionary(keys.as_slice(),
dict_offsets, dict_values)?;
}
- *self = Self::Values { values: spilled };
+ *self = Self::Values {
+ values: spilled,
+ recycled_keys: None,
+ };
match self {
- Self::Values { values } => Ok(values),
+ Self::Values { values, .. } => Ok(values),
Self::Dict { .. } => unreachable!(),
}
}
}
}
- /// Converts this into an [`ArrayRef`] with the provided `data_type` and
`null_buffer`
+ /// Converts this into an [`ArrayRef`]. `keys_buffer_sink` receives a clone
+ /// of the keys buffer for next-batch reclamation; pass `&mut None` to
skip.
pub fn into_array(
self,
null_buffer: Option<Buffer>,
data_type: &ArrowType,
hash_scratch: &mut MutableBuffer,
+ keys_buffer_sink: &mut Option<Buffer>,
) -> Result<ArrayRef> {
assert!(matches!(data_type, ArrowType::Dictionary(_, _)));
match self {
- Self::Dict { keys, values } => {
- // Validate keys unless dictionary is empty
- if !values.is_empty() {
+ Self::Dict {
+ keys,
+ values,
+ key_upper_bound,
+ } => {
+ // Skip the key-bounds scan when the RLE bit width guarantees
every key is in range.
+ let keys_provably_in_range = key_upper_bound
+ .map(|bound| bound <= values.len() as u64)
+ .unwrap_or(false);
+
+ if !values.is_empty() && !keys_provably_in_range {
let min = K::from_usize(0).unwrap();
let max = K::from_usize(values.len()).unwrap();
- // using copied and fold gets auto-vectorized since rust
1.70
- // all/any would allow early exit on invalid values
- // but in the happy case all values have to be checked
anyway
- if !keys
- .as_slice()
- .iter()
- .copied()
- .fold(true, |a, x| a && x >= min && x < max)
- {
+ let (observed_min, observed_max) =
keys.as_slice().iter().copied().fold(
+ (max, min),
+ |(current_min, current_max), key| {
+ (current_min.min(key), current_max.max(key))
+ },
+ );
+
+ if observed_min < min || observed_max >= max {
Review Comment:
worth looking into speeding this up
--
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]
