diff --git a/parquet/src/arrow/array_reader/byte_view_array.rs b/parquet/src/arrow/array_reader/byte_view_array.rs index 1933654118f3..7c1cf145100f 100644 --- a/parquet/src/arrow/array_reader/byte_view_array.rs +++ b/parquet/src/arrow/array_reader/byte_view_array.rs @@ -30,7 +30,6 @@ use crate::schema::types::ColumnDescPtr; use crate::util::utf8::check_valid_utf8; use arrow_array::{ArrayRef, builder::make_view}; use arrow_buffer::Buffer; -use arrow_data::ByteView; use arrow_schema::DataType as ArrowType; use bytes::Bytes; use std::any::Any; @@ -496,56 +495,14 @@ impl ByteViewArrayDecoderDictionary { } // Calculate the offset of the dictionary buffers in the output buffers - // For example if the 2nd buffer in the dictionary is the 5th buffer in the output buffers, - // then the base_buffer_idx is 5 - 2 = 3 let base_buffer_idx = output.buffers.len() as u32 - dict.buffers.len() as u32; // Pre-reserve output capacity to avoid per-chunk reallocation in extend output.views.reserve(len); - let mut error = None; - let read = self.decoder.read(len, |keys| { - if base_buffer_idx == 0 { - // the dictionary buffers are the last buffers in output, we can directly use the views - output - .views - .extend(keys.iter().map(|k| match dict.views.get(*k as usize) { - Some(&view) => view, - None => { - if error.is_none() { - error = Some(general_err!("invalid key={} for dictionary", *k)); - } - 0 - } - })); - Ok(()) - } else { - output - .views - .extend(keys.iter().map(|k| match dict.views.get(*k as usize) { - Some(&view) => { - let len = view as u32; - if len <= 12 { - view - } else { - let mut view = ByteView::from(view); - view.buffer_index += base_buffer_idx; - view.into() - } - } - None => { - if error.is_none() { - error = Some(general_err!("invalid key={} for dictionary", *k)); - } - 0 - } - })); - Ok(()) - } - })?; - if let Some(e) = error { - return Err(e); - } + let read = + self.decoder + .read_gather_views(len, &dict.views, &mut output.views, base_buffer_idx)?; Ok(read) } diff --git a/parquet/src/arrow/buffer/dict_view_buffer.rs b/parquet/src/arrow/buffer/dict_view_buffer.rs new file mode 100644 index 000000000000..0810ae6d79af --- /dev/null +++ b/parquet/src/arrow/buffer/dict_view_buffer.rs @@ -0,0 +1,387 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::buffer::view_buffer::ViewBuffer; +use crate::arrow::record_reader::buffer::ValuesBuffer; +use crate::errors::{ParquetError, Result}; +use arrow_array::{ArrayRef, StringViewArray}; +use arrow_buffer::Buffer; +use arrow_data::ArrayDataBuilder; +use arrow_schema::DataType as ArrowType; +use std::sync::Arc; + +/// A buffer that can store variable-length byte view data either as dictionary-encoded +/// (keys + dictionary) or as expanded views. +/// +/// This is analogous to [`DictionaryBuffer`] but for view-based types (`Utf8View`, `BinaryView`). +/// When dictionary encoding is preserved, the output is a `DictionaryArray`. +/// When dictionary encoding cannot be preserved (e.g., non-dictionary pages are encountered), +/// the buffer falls back to expanded views. +/// +/// [`DictionaryBuffer`]: crate::arrow::buffer::dictionary_buffer::DictionaryBuffer +pub enum DictViewBuffer { + /// Dictionary-encoded: stores keys and a reference to the dictionary + Dict { + keys: Vec, + /// The dictionary as a StringViewArray (cheap Arc clone) + dict: ArrayRef, + }, + /// Fallback: expanded views (same as ViewBuffer) + Values { + values: ViewBuffer, + }, +} + +impl Default for DictViewBuffer { + fn default() -> Self { + Self::Values { + values: ViewBuffer::default(), + } + } +} + +impl std::fmt::Debug for DictViewBuffer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Dict { keys, .. } => f + .debug_struct("DictViewBuffer::Dict") + .field("num_keys", &keys.len()) + .finish(), + Self::Values { values } => f + .debug_struct("DictViewBuffer::Values") + .field("values", values) + .finish(), + } + } +} + +impl DictViewBuffer { + #[allow(unused)] + pub fn len(&self) -> usize { + match self { + Self::Dict { keys, .. } => keys.len(), + Self::Values { values } => values.views.len(), + } + } + + /// Returns a mutable reference to the keys vec if we can preserve the dictionary. + /// + /// Returns `None` if: + /// - The dictionary changed (different pointer) and keys are non-empty + /// - The buffer has already fallen back to Values mode with data + /// + /// When `None` is returned, the caller should use `spill_values()` instead. + pub fn as_keys(&mut self, dictionary: &ArrayRef) -> Option<&mut Vec> { + match self { + Self::Dict { keys, dict } => { + // Check if the dictionary is the same object + let dict_ptr = dict.as_ref() as *const _ as *const (); + let new_ptr = dictionary.as_ref() as *const _ as *const (); + if dict_ptr == new_ptr { + Some(keys) + } else if keys.is_empty() { + *dict = Arc::clone(dictionary); + Some(keys) + } else { + // Dictionary changed and we have existing keys - can't preserve + None + } + } + Self::Values { values } if values.views.is_empty() => { + *self = Self::Dict { + keys: Vec::new(), + dict: Arc::clone(dictionary), + }; + match self { + Self::Dict { keys, .. } => Some(keys), + _ => unreachable!(), + } + } + _ => None, + } + } + + /// Spills dictionary-encoded data into expanded views. + /// + /// If already in `Values` mode, returns the existing `ViewBuffer`. + /// If in `Dict` mode, expands keys through the dictionary into views, + /// then switches to `Values` mode. + pub fn spill_values(&mut self) -> Result<&mut ViewBuffer> { + match self { + Self::Values { values } => Ok(values), + Self::Dict { keys, dict } => { + let mut spilled = ViewBuffer::default(); + + if !dict.is_empty() && !keys.is_empty() { + // Expand the dictionary keys into views + let dict_view = dict + .as_any() + .downcast_ref::() + .ok_or_else(|| { + general_err!("DictViewBuffer dictionary is not StringViewArray") + })?; + + // Copy dictionary buffers to the output + for buf in dict_view.data_buffers() { + spilled.buffers.push(buf.clone()); + } + + let dict_views = dict_view.views(); + spilled.views.reserve(keys.len()); + + for &key in keys.iter() { + if key < 0 || (key as usize) >= dict_views.len() { + return Err(general_err!( + "dictionary key {} beyond bounds of dictionary: 0..{}", + key, + dict_views.len() + )); + } + let view = dict_views[key as usize]; + // Adjust buffer indices since dict buffers are the first buffers + // in the output (base_buffer_idx = 0) + spilled.views.push(view); + } + } else if !keys.is_empty() { + // Empty dictionary but non-empty keys - zero-length views + spilled.views.resize(keys.len(), 0u128); + } + + *self = Self::Values { values: spilled }; + match self { + Self::Values { values } => Ok(values), + _ => unreachable!(), + } + } + } + } + + /// Converts this buffer into an [`ArrayRef`]. + /// + /// When in `Dict` mode, produces a `DictionaryArray`. + /// When in `Values` mode, packs the values into a fresh `DictionaryArray`. + pub fn into_array(self, null_buffer: Option, data_type: &ArrowType) -> Result { + assert!( + matches!(data_type, ArrowType::Dictionary(_, _)), + "DictViewBuffer requires Dictionary data type, got {data_type}" + ); + + match self { + Self::Dict { keys, dict } => { + // Validate keys + if !dict.is_empty() { + let max = dict.len(); + if !keys + .iter() + .copied() + .fold(true, |a, x| a && x >= 0 && (x as usize) < max) + { + return Err(general_err!( + "dictionary key beyond bounds of dictionary: 0..{}", + max + )); + } + } + + let builder = ArrayDataBuilder::new(data_type.clone()) + .len(keys.len()) + .add_buffer(Buffer::from_vec(keys)) + .add_child_data(dict.to_data()) + .null_bit_buffer(null_buffer); + + let data = match cfg!(debug_assertions) { + true => builder.build().unwrap(), + false => unsafe { builder.build_unchecked() }, + }; + + Ok(arrow_array::make_array(data)) + } + Self::Values { values } => { + // Need to create a dictionary from expanded values + let ArrowType::Dictionary(_, value_type) = data_type else { + unreachable!() + }; + let values_array = values.into_array(null_buffer, value_type); + pack_string_view_values(&values_array, data_type) + } + } + } +} + +/// Pack a `StringViewArray` into a `DictionaryArray`. +/// +/// This creates an identity-mapping dictionary where each row maps to its +/// own dictionary entry. This is a fallback for when dictionary encoding +/// couldn't be preserved, and is suboptimal but correct. +fn pack_string_view_values(array: &ArrayRef, dict_type: &ArrowType) -> Result { + let len = array.len(); + + // Create identity keys [0, 1, 2, ..., len-1] + let keys: Vec = (0..len as i32).collect(); + + let builder = ArrayDataBuilder::new(dict_type.clone()) + .len(len) + .add_buffer(Buffer::from_vec(keys)) + .add_child_data(array.to_data()) + .null_bit_buffer(array.nulls().map(|n| n.buffer().clone())); + + let data = match cfg!(debug_assertions) { + true => builder.build().unwrap(), + false => unsafe { builder.build_unchecked() }, + }; + + Ok(arrow_array::make_array(data)) +} + +impl ValuesBuffer for DictViewBuffer { + fn pad_nulls( + &mut self, + read_offset: usize, + values_read: usize, + levels_read: usize, + valid_mask: &[u8], + ) { + match self { + Self::Dict { keys, .. } => { + keys.resize(read_offset + levels_read, 0i32); + keys.pad_nulls(read_offset, values_read, levels_read, valid_mask); + } + Self::Values { values } => { + values.pad_nulls(read_offset, values_read, levels_read, valid_mask); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::compute::cast; + use arrow_array::{Array, StringViewArray}; + use arrow_buffer::Buffer; + + fn utf8view_dictionary() -> ArrowType { + ArrowType::Dictionary(Box::new(ArrowType::Int32), Box::new(ArrowType::Utf8View)) + } + + #[test] + fn test_dict_view_buffer_dict_mode() { + let data_type = utf8view_dictionary(); + + // Create a dictionary + let dict: ArrayRef = Arc::new(StringViewArray::from(vec!["hello", "world", ""])); + + let mut buffer = DictViewBuffer::default(); + + // Write keys preserving dictionary + let keys = buffer.as_keys(&dict).unwrap(); + keys.extend_from_slice(&[1, 0, 2, 0]); + + assert!(matches!(buffer, DictViewBuffer::Dict { .. })); + + let array = buffer.into_array(None, &data_type).unwrap(); + assert_eq!(array.data_type(), &data_type); + + let strings = cast(&array, &ArrowType::Utf8View).unwrap(); + let strings = strings + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + strings.iter().collect::>(), + vec![Some("world"), Some("hello"), Some(""), Some("hello")] + ); + } + + #[test] + fn test_dict_view_buffer_with_nulls() { + let data_type = utf8view_dictionary(); + + let dict: ArrayRef = Arc::new(StringViewArray::from(vec!["a", "b"])); + + let mut buffer = DictViewBuffer::default(); + let keys = buffer.as_keys(&dict).unwrap(); + keys.extend_from_slice(&[0, 1]); + + let valid = [false, true, true, false]; + let valid_buffer = Buffer::from_iter(valid.iter().copied()); + buffer.pad_nulls(0, 2, valid.len(), valid_buffer.as_slice()); + + let array = buffer + .into_array(Some(valid_buffer), &data_type) + .unwrap(); + assert_eq!(array.len(), 4); + assert_eq!(array.null_count(), 2); + + let strings = cast(&array, &ArrowType::Utf8View).unwrap(); + let strings = strings + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + strings.iter().collect::>(), + vec![None, Some("a"), Some("b"), None] + ); + } + + #[test] + fn test_dict_view_buffer_spill() { + let data_type = utf8view_dictionary(); + + let dict: ArrayRef = Arc::new(StringViewArray::from(vec!["hello", "world"])); + + let mut buffer = DictViewBuffer::default(); + let keys = buffer.as_keys(&dict).unwrap(); + keys.extend_from_slice(&[0, 1]); + + // Spill to values mode + let _values = buffer.spill_values().unwrap(); + + assert!(matches!(buffer, DictViewBuffer::Values { .. })); + + let array = buffer.into_array(None, &data_type).unwrap(); + assert_eq!(array.data_type(), &data_type); + + let strings = cast(&array, &ArrowType::Utf8View).unwrap(); + let strings = strings + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + strings.iter().collect::>(), + vec![Some("hello"), Some("world")] + ); + } + + #[test] + fn test_dict_view_buffer_validates_keys() { + let data_type = utf8view_dictionary(); + + let dict: ArrayRef = Arc::new(StringViewArray::from(vec!["a"])); + + let mut buffer = DictViewBuffer::default(); + let keys = buffer.as_keys(&dict).unwrap(); + keys.extend_from_slice(&[0, 2, 0]); // key 2 is out of bounds + + let err = buffer.into_array(None, &data_type).unwrap_err().to_string(); + assert!( + err.contains("dictionary key beyond bounds"), + "{}", + err + ); + } +} diff --git a/parquet/src/arrow/decoder/dictionary_index.rs b/parquet/src/arrow/decoder/dictionary_index.rs index 7a4b77f89d59..c17f1cc8dd07 100644 --- a/parquet/src/arrow/decoder/dictionary_index.rs +++ b/parquet/src/arrow/decoder/dictionary_index.rs @@ -17,8 +17,8 @@ use bytes::Bytes; -use crate::encodings::rle::RleDecoder; -use crate::errors::Result; +use crate::encodings::rle::{RleDecodedBatch, RleDecoder}; +use crate::errors::{ParquetError, Result}; /// Decoder for `Encoding::RLE_DICTIONARY` indices pub struct DictIndexDecoder { @@ -27,7 +27,7 @@ pub struct DictIndexDecoder { /// We want to decode the offsets in chunks so we will maintain an internal buffer of decoded /// offsets - index_buf: Box<[i32; 1024]>, + index_buf: Vec, /// Current length of `index_buf` index_buf_len: usize, /// Current offset into `index_buf`. If `index_buf_offset` == `index_buf_len` then we've consumed @@ -46,13 +46,14 @@ impl DictIndexDecoder { let bit_width = data[0]; let mut decoder = RleDecoder::new(bit_width); decoder.set_data(data.slice(1..))?; + let max_remaining = num_values.unwrap_or(num_levels); Ok(Self { decoder, - index_buf: Box::new([0; 1024]), + index_buf: vec![0; max_remaining.min(1024)], index_buf_len: 0, index_offset: 0, - max_remaining_values: num_values.unwrap_or(num_levels), + max_remaining_values: max_remaining, }) } @@ -91,6 +92,115 @@ impl DictIndexDecoder { Ok(values_read) } + /// Decode indices and gather views directly into `output`. + /// + /// For RLE runs, fills output with the repeated view directly (no index buffer needed). + /// For bit-packed runs, decodes indices to a stack-local buffer and gathers immediately. + pub fn read_gather_views( + &mut self, + len: usize, + dict_views: &[u128], + output: &mut Vec, + base_buffer_idx: u32, + ) -> Result { + let to_read = len.min(self.max_remaining_values); + let mut values_read = 0; + let mut error = None; + + // Flush any buffered indices from previous reads + if self.index_offset < self.index_buf_len { + let n = (self.index_buf_len - self.index_offset).min(to_read); + let keys = &self.index_buf[self.index_offset..self.index_offset + n]; + Self::gather_views(keys, dict_views, output, base_buffer_idx, &mut error); + self.index_offset += n; + self.max_remaining_values -= n; + values_read += n; + } + + if values_read < to_read { + let read = + self.decoder + .get_batch_direct(to_read - values_read, |batch| match batch { + RleDecodedBatch::Rle { index, count } => { + match dict_views.get(index as usize) { + Some(&view) => { + let view = if base_buffer_idx == 0 || (view as u32) <= 12 { + view + } else { + let mut bv = arrow_data::ByteView::from(view); + bv.buffer_index += base_buffer_idx; + bv.into() + }; + output.extend(std::iter::repeat_n(view, count)); + } + None => { + if error.is_none() { + error = Some(general_err!( + "invalid key={} for dictionary", + index + )); + } + } + } + } + RleDecodedBatch::BitPacked(keys) => { + Self::gather_views( + keys, + dict_views, + output, + base_buffer_idx, + &mut error, + ); + } + })?; + self.max_remaining_values -= read; + values_read += read; + } + + if let Some(e) = error { + return Err(e); + } + Ok(values_read) + } + + fn gather_views( + keys: &[i32], + dict_views: &[u128], + output: &mut Vec, + base_buffer_idx: u32, + error: &mut Option, + ) { + if base_buffer_idx == 0 { + output.extend(keys.iter().map(|k| match dict_views.get(*k as usize) { + Some(&view) => view, + None => { + if error.is_none() { + *error = Some(general_err!("invalid key={} for dictionary", *k)); + } + 0 + } + })); + } else { + output.extend(keys.iter().map(|k| match dict_views.get(*k as usize) { + Some(&view) => { + if (view as u32) <= 12 { + view + } else { + let mut bv = arrow_data::ByteView::from(view); + bv.buffer_index += base_buffer_idx; + bv.into() + } + } + None => { + if error.is_none() { + *error = Some(general_err!("invalid key={} for dictionary", *k)); + } + 0 + } + })); + } + } + /// Skip up to `to_skip` values, returning the number of values skipped pub fn skip(&mut self, to_skip: usize) -> Result { let to_skip = to_skip.min(self.max_remaining_values); diff --git a/parquet/src/encodings/rle.rs b/parquet/src/encodings/rle.rs index c95a46c634d2..ff79402969cf 100644 --- a/parquet/src/encodings/rle.rs +++ b/parquet/src/encodings/rle.rs @@ -287,6 +287,15 @@ impl RleEncoder { /// Size, in number of `i32s` of buffer to use for RLE batch reading const RLE_DECODER_INDEX_BUFFER_SIZE: usize = 1024; +/// A decoded batch from [`RleDecoder::get_batch_direct`]. +#[cfg(feature = "arrow")] +pub enum RleDecodedBatch<'a> { + /// An RLE run: all values are the same index, repeated `count` times + Rle { index: i32, count: usize }, + /// A batch of bit-packed indices + BitPacked(&'a [i32]), +} + /// A RLE/Bit-Packing hybrid decoder. pub struct RleDecoder { // Number of bits used to encode the value. Must be between [0, 64]. @@ -414,6 +423,52 @@ impl RleDecoder { Ok(values_read) } + /// Decode up to `max_values` indices and call `f` with each decoded batch. + /// + /// For RLE runs, provides [`RleDecodedBatch::Rle`] so callers can fill output directly. + /// For bit-packed runs, provides [`RleDecodedBatch::BitPacked`] with decoded indices. + #[cfg(feature = "arrow")] + pub fn get_batch_direct(&mut self, max_values: usize, mut f: F) -> Result + where + F: FnMut(RleDecodedBatch<'_>), + { + let mut values_read = 0; + let mut index_buf = [0i32; 1024]; + while values_read < max_values { + if self.rle_left > 0 { + let num_values = cmp::min(max_values - values_read, self.rle_left as usize); + let idx = self.current_value.unwrap() as i32; + f(RleDecodedBatch::Rle { + index: idx, + count: num_values, + }); + self.rle_left -= num_values as u32; + values_read += num_values; + } else if self.bit_packed_left > 0 { + let to_read = (max_values - values_read) + .min(self.bit_packed_left as usize) + .min(index_buf.len()); + let bit_reader = self + .bit_reader + .as_mut() + .ok_or_else(|| general_err!("bit_reader should be set"))?; + + let num_values = + bit_reader.get_batch::(&mut index_buf[..to_read], self.bit_width as usize); + if num_values == 0 { + self.bit_packed_left = 0; + continue; + } + f(RleDecodedBatch::BitPacked(&index_buf[..num_values])); + self.bit_packed_left -= num_values as u32; + values_read += num_values; + } else if !self.reload()? { + break; + } + } + Ok(values_read) + } + #[inline(never)] pub fn skip(&mut self, num_values: usize) -> Result { let mut values_skipped = 0; @@ -458,6 +513,9 @@ impl RleDecoder { { assert!(buffer.len() >= max_values); + if dict.is_empty() { + return Ok(0); + } let mut values_read = 0; while values_read < max_values { let index_buf = self.index_buf.get_or_insert_with(|| Box::new([0; 1024])); @@ -497,7 +555,9 @@ impl RleDecoder { buffer[values_read..values_read + num_values] .iter_mut() .zip(index_buf[..num_values].iter()) - .for_each(|(b, i)| b.clone_from(&dict[*i as usize])); + .for_each(|(b, i)| { + b.clone_from(&dict[*i as usize]); + }); self.bit_packed_left -= num_values as u32; values_read += num_values; if num_values < to_read {