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-11262-2f33b89afe0e0a361554d0db4daafd5018ab493d in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit bece8ca54ed1f776b2b3a1c4cb663ac4f8dd83dc Author: Hippolyte Barraud <[email protected]> AuthorDate: Wed Oct 14 01:34:59 2026 +0000 fix(parquet): encode `FixedSizeBinary` dictionaries as FLBA (#11262) # Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. --> - Closes #11261 # Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> The writer can encode `Dictionary<K, FixedSizeBinary>` values using the `BYTE_ARRAY` wire representation even though the Parquet schema declares `FIXED_LEN_BYTE_ARRAY` (FLBA). This distinction matters for PLAIN encoding, which is also used for dictionary entries: `BYTE_ARRAY` prefixes each value with a four-byte length, whereas FLBA stores only the value bytes because their width is defined by the schema. # What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> - Route Arrow `Dictionary<K, FixedSizeBinary>` values through the FLBA writer rather than the byte-array writer. - Align dictionary decoding, fallback decoding, and Arrow dictionary spill conversion with the declared physical type. - Stop emitting `BYTE_ARRAY`-specific size statistics for FLBA columns. - Require decompressed FLBA dictionary payloads to contain exactly `num_values * type_length` bytes, using checked arithmetic, in the typed column reader and both Arrow dictionary-decoding paths. - Validate PLAIN data-page value sections independently of repetition and definition levels, counting physical non-null values. - Validate ordinary nullable V1 FLBA pages incrementally during reads and in-page skips. Count-dependent corruption may be reported only as the page is consumed; partial reads can return values before encountering an error. Potentially legacy PLAIN pages require an independent physical-value count before choosing canonical decoding or legacy normalization. The reader follows the schema rather than attempting to infer which historical writer produced a payload. # Are these changes tested? Yes. Regression coverage exercises FLBA dictionary handling and malformed dictionary/data-page payloads, including partial reads, skips, and dictionary-to-PLAIN transitions. # Are there any user-facing changes? Yes. New writes from the affected dictionary-writing path now use the `FIXED_LEN_BYTE_ARRAY` representation declared by the schema. Readers also accept the known malformed historical layouts through narrowly scoped validation and normalization, allowing those existing files to remain readable on a best-effort basis. Already-conforming files retain their existing interpretation, including files produced by the dense `FixedSizeBinary` writer path. There are no public Rust API signature changes. --------- Signed-off-by: Hippolyte Barraud <[email protected]> Co-authored-by: Andrew Lamb <[email protected]> --- parquet-testing | 2 +- .../arrow/array_reader/byte_array_dictionary.rs | 171 ++++++- .../src/arrow/array_reader/fixed_len_byte_array.rs | 158 +++--- parquet/src/arrow/arrow_writer/mod.rs | 1 - parquet/src/arrow/buffer/dictionary_buffer.rs | 85 +++- parquet/src/column/reader.rs | 554 ++++++++++++++++++++- parquet/src/column/reader/decoder.rs | 350 ++++++++++++- .../column/reader/legacy_fixed_len_byte_array.rs | 466 +++++++++++++++++ parquet/src/encodings/rle.rs | 138 +++++ parquet/src/util/bit_util.rs | 190 +++++++ parquet/src/util/test_common/page_util.rs | 83 ++- parquet/tests/arrow_reader/bad_data.rs | 8 +- .../arrow_reader/legacy_fixed_len_byte_array.rs | 262 ++++++++++ parquet/tests/arrow_reader/mod.rs | 1 + parquet/tests/arrow_writer/roundtrip.rs | 60 ++- 15 files changed, 2414 insertions(+), 115 deletions(-) diff --git a/parquet-testing b/parquet-testing index 09f3cdbde4..1ba58c2971 160000 --- a/parquet-testing +++ b/parquet-testing @@ -1 +1 @@ -Subproject commit 09f3cdbde45302f0f0c689c950e465e98a9df960 +Subproject commit 1ba58c2971cab87e2337b6185c82a21c219210e3 diff --git a/parquet/src/arrow/array_reader/byte_array_dictionary.rs b/parquet/src/arrow/array_reader/byte_array_dictionary.rs index 4b4e1ad852..01138abfab 100644 --- a/parquet/src/arrow/array_reader/byte_array_dictionary.rs +++ b/parquet/src/arrow/array_reader/byte_array_dictionary.rs @@ -28,13 +28,14 @@ use arrow_schema::DataType as ArrowType; use bytes::Bytes; use crate::arrow::array_reader::byte_array::{ByteArrayDecoder, ByteArrayDecoderPlain}; +use crate::arrow::array_reader::fixed_len_byte_array::ValueDecoder as FixedByteArrayDecoder; use crate::arrow::array_reader::{ArrayReader, read_records, skip_records}; use crate::arrow::buffer::{dictionary_buffer::DictionaryBuffer, offset_buffer::OffsetBuffer}; use crate::arrow::record_reader::GenericRecordReader; use crate::arrow::schema::parquet_to_arrow_field; -use crate::basic::{ConvertedType, Encoding}; +use crate::basic::{ConvertedType, Encoding, Type}; use crate::column::page::PageIterator; -use crate::column::reader::decoder::ColumnValueDecoder; +use crate::column::reader::decoder::{ColumnValueDecoder, normalize_fixed_len_byte_array_payload}; use crate::encodings::rle::{MAX_RLE_DICTIONARY_BIT_WIDTH, RleDecoder}; use crate::errors::{ParquetError, Result}; use crate::schema::types::ColumnDescPtr; @@ -151,7 +152,7 @@ fn convert_values_to_view(array: ArrayRef, to_type: &ArrowType) -> Result<ArrayR Ok(array.with_values(new_values)) } -/// An [`ArrayReader`] for dictionary encoded variable length byte arrays +/// An [`ArrayReader`] for dictionary encoded byte arrays /// /// Will attempt to preserve any dictionary encoding present in the parquet data struct ByteArrayDictionaryReader<K: ArrowNativeType, V: OffsetSizeTrait> { @@ -281,9 +282,10 @@ enum MaybeDictionaryDecoder { max_remaining_values: usize, }, Fallback(ByteArrayDecoder), + Fixed(FixedByteArrayDecoder), } -/// A [`ColumnValueDecoder`] for dictionary encoded variable length byte arrays +/// A [`ColumnValueDecoder`] for dictionary encoded byte arrays struct DictionaryDecoder<K, V> { /// The current dictionary dict: Option<ArrayRef>, @@ -295,6 +297,10 @@ struct DictionaryDecoder<K, V> { value_type: ArrowType, + /// `Some` only for `FIXED_LEN_BYTE_ARRAY` columns, selecting fixed-width + /// decoding; `None` for other physical types. + flba_column_descr: Option<ColumnDescPtr>, + phantom: PhantomData<(K, V)>, } @@ -320,6 +326,8 @@ where decoder: None, validate_utf8, value_type, + flba_column_descr: (col.physical_type() == Type::FIXED_LEN_BYTE_ARRAY) + .then(|| col.clone()), phantom: Default::default(), } } @@ -347,8 +355,24 @@ where let len = num_values as usize; let mut buffer = OffsetBuffer::<V>::with_capacity(0); - let mut decoder = ByteArrayDecoderPlain::new(buf, len, Some(len), self.validate_utf8); - decoder.read(&mut buffer, usize::MAX)?; + if let Some(desc) = &self.flba_column_descr { + let buf = normalize_fixed_len_byte_array_payload( + buf, + len, + desc.type_length() as usize, + "dictionary page", + )?; + let mut decoder = FixedByteArrayDecoder::new(desc); + decoder.set_data(Encoding::PLAIN, buf, len, Some(len))?; + if decoder.read_offsets(&mut buffer, len)? != len { + return Err(general_err!( + "too few bytes in fixed-length dictionary page" + )); + } + } else { + ByteArrayDecoderPlain::new(buf, len, Some(len), self.validate_utf8) + .read(&mut buffer, usize::MAX)?; + } let array = buffer.into_array(None, self.value_type.clone()); self.dict = Some(array); @@ -379,6 +403,12 @@ where max_remaining_values: num_values.unwrap_or(num_levels), } } + _ if self.flba_column_descr.is_some() => { + let mut decoder = + FixedByteArrayDecoder::new(self.flba_column_descr.as_ref().unwrap()); + decoder.set_data(encoding, data, num_levels, num_values)?; + MaybeDictionaryDecoder::Fixed(decoder) + } _ => MaybeDictionaryDecoder::Fallback(ByteArrayDecoder::new( encoding, data, @@ -397,6 +427,9 @@ where MaybeDictionaryDecoder::Fallback(decoder) => { decoder.read(out.spill_values()?, num_values, None) } + MaybeDictionaryDecoder::Fixed(decoder) => { + decoder.read_offsets(out.spill_values()?, num_values) + } MaybeDictionaryDecoder::Dict { decoder, max_remaining_values, @@ -455,6 +488,7 @@ where fn skip_values(&mut self, num_values: usize) -> Result<usize> { match self.decoder.as_mut().expect("decoder set") { MaybeDictionaryDecoder::Fallback(decoder) => decoder.skip::<V>(num_values, None), + MaybeDictionaryDecoder::Fixed(decoder) => decoder.skip_values(num_values), MaybeDictionaryDecoder::Dict { decoder, max_remaining_values, @@ -481,6 +515,131 @@ mod tests { use super::*; + #[test] + fn fixed_dictionary_canonical_wire_read_skip_and_truncation() { + use crate::schema::types::{ColumnDescriptor, ColumnPath, Type as SchemaType}; + use arrow_array::cast::AsArray; + use std::sync::Arc; + let primitive = Arc::new( + SchemaType::primitive_type_builder("fixed", Type::FIXED_LEN_BYTE_ARRAY) + .with_length(4) + .build() + .unwrap(), + ); + let desc = Arc::new(ColumnDescriptor::new( + primitive, + 0, + 0, + ColumnPath::from("fixed"), + )); + let dtype = ArrowType::Dictionary( + Box::new(ArrowType::Int32), + Box::new(ArrowType::FixedSizeBinary(4)), + ); + // Canonical PLAIN has no length prefixes, independent of ArrowWriter. + let raw = Bytes::from_static(b"abcdEFGHijkl"); + let mut decoder = DictionaryDecoder::<i32, i32>::new(&desc); + decoder + .set_dict(raw.clone(), 3, Encoding::PLAIN, false) + .unwrap(); + let dict = decoder.dict.as_ref().unwrap().as_binary::<i32>(); + assert_eq!(dict.value(0), b"abcd"); + assert_eq!(dict.value(1), b"EFGH"); + assert_eq!(dict.value(2), b"ijkl"); + assert!( + decoder + .set_dict(raw.slice(..11), 3, Encoding::PLAIN, false) + .unwrap_err() + .to_string() + .contains("Invalid FIXED_LEN_BYTE_ARRAY dictionary page payload length") + ); + + decoder.set_data(Encoding::PLAIN, raw, 3, Some(3)).unwrap(); + let mut output = DictionaryBuffer::<i32, i32>::with_capacity(0); + assert_eq!(decoder.read(&mut output, 1).unwrap(), 1); + assert_eq!(decoder.skip_values(1).unwrap(), 1); + assert_eq!(decoder.read(&mut output, 10).unwrap(), 1); + assert_eq!(decoder.skip_values(10).unwrap(), 0); + let array = output + .into_array(None, &dtype, &mut MutableBuffer::new(0)) + .unwrap(); + let array = cast(&array, &ArrowType::FixedSizeBinary(4)).unwrap(); + assert_eq!(array.as_fixed_size_binary().value(0), b"abcd"); + assert_eq!(array.as_fixed_size_binary().value(1), b"ijkl"); + + // Truncated PLAIN data yields only complete values; it cannot manufacture a row. + decoder + .set_data(Encoding::PLAIN, Bytes::from_static(b"abcdEFG"), 2, Some(2)) + .unwrap(); + let mut output = DictionaryBuffer::<i32, i32>::with_capacity(0); + assert_eq!(decoder.read(&mut output, 2).unwrap(), 1); + assert_eq!(decoder.read(&mut output, 1).unwrap(), 0); + assert_eq!(decoder.skip_values(1).unwrap(), 0); + } + + #[test] + fn fixed_dictionary_payload_length_validation() { + use crate::column::reader::decoder::ColumnValueDecoderImpl; + use crate::data_type::FixedLenByteArrayType; + use crate::schema::types::{ColumnDescriptor, ColumnPath, Type as SchemaType}; + use std::sync::Arc; + + fn check<D: ColumnValueDecoder>(desc: &ColumnDescPtr, encoding: Encoding) { + // Length-looking and high-bit bytes are valid values, not legacy markers. + let raw = Bytes::from_static(b"\x04\0\0\0\xff\xff\xff\xffabcd"); + D::new(desc) + .set_dict(raw.clone(), 3, encoding, false) + .unwrap(); + D::new(desc) + .set_dict(Bytes::new(), 0, encoding, false) + .unwrap(); + + // A noncanonical payload with a mismatched prefix must be rejected. + // Legacy acceptance is covered separately by the compatibility tests. + let bad_prefix = + Bytes::from_static(b"\x04\0\0\0\x04\0\0\0\x03\0\0\0\xff\xff\xff\xff\x04\0\0\0abcd"); + for (buf, count) in [ + (bad_prefix, 3), + (raw.slice(..11), 3), + (raw.clone(), 2), + (raw, 0), + (Bytes::new(), 3), + ] { + let err = D::new(desc) + .set_dict(buf, count, encoding, false) + .unwrap_err() + .to_string(); + assert!( + err.contains("Invalid FIXED_LEN_BYTE_ARRAY dictionary page payload length"), + "{} {encoding}: {err}", + std::any::type_name::<D>() + ); + } + } + + let primitive = Arc::new( + SchemaType::primitive_type_builder("fixed", Type::FIXED_LEN_BYTE_ARRAY) + .with_length(4) + .build() + .unwrap(), + ); + let desc = Arc::new(ColumnDescriptor::new( + primitive, + 0, + 0, + ColumnPath::from("fixed"), + )); + for encoding in [ + Encoding::PLAIN, + Encoding::PLAIN_DICTIONARY, + Encoding::RLE_DICTIONARY, + ] { + check::<DictionaryDecoder<i32, i32>>(&desc, encoding); + check::<FixedByteArrayDecoder>(&desc, encoding); + check::<ColumnValueDecoderImpl<FixedLenByteArrayType>>(&desc, encoding); + } + } + fn utf8_dictionary() -> ArrowType { ArrowType::Dictionary(Box::new(ArrowType::Int32), Box::new(ArrowType::Utf8)) } diff --git a/parquet/src/arrow/array_reader/fixed_len_byte_array.rs b/parquet/src/arrow/array_reader/fixed_len_byte_array.rs index 436e1db343..c2efb1c5d6 100644 --- a/parquet/src/arrow/array_reader/fixed_len_byte_array.rs +++ b/parquet/src/arrow/array_reader/fixed_len_byte_array.rs @@ -17,18 +17,20 @@ use crate::arrow::array_reader::{ArrayReader, read_records, skip_records}; use crate::arrow::buffer::bit_util::{iter_set_bits_rev, sign_extend_be}; +use crate::arrow::buffer::offset_buffer::OffsetBuffer; use crate::arrow::decoder::{DeltaByteArrayDecoder, DictIndexDecoder}; use crate::arrow::record_reader::GenericRecordReader; use crate::arrow::record_reader::buffer::ValuesBuffer; use crate::arrow::schema::parquet_to_arrow_field; use crate::basic::{Encoding, Type}; use crate::column::page::PageIterator; -use crate::column::reader::decoder::ColumnValueDecoder; +use crate::column::reader::decoder::{ColumnValueDecoder, normalize_fixed_len_byte_array_payload}; use crate::errors::{ParquetError, Result}; use crate::schema::types::ColumnDescPtr; use arrow_array::{ ArrayRef, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, FixedSizeBinaryArray, Float16Array, IntervalDayTimeArray, IntervalYearMonthArray, + OffsetSizeTrait, }; use arrow_buffer::{Buffer, IntervalDayTime, i256}; use arrow_data::ArrayDataBuilder; @@ -279,7 +281,7 @@ impl ArrayReader for FixedLenByteArrayReader { } #[derive(Default)] -struct FixedLenByteArrayBuffer { +pub(super) struct FixedLenByteArrayBuffer { buffer: Vec<u8>, /// The length of each element in bytes byte_length: Option<usize>, @@ -379,12 +381,86 @@ impl ValuesBuffer for FixedLenByteArrayBuffer { } } -struct ValueDecoder { +pub(super) struct ValueDecoder { byte_length: usize, dict_page: Option<Bytes>, decoder: Option<Decoder>, } +impl ValueDecoder { + pub(super) fn read_offsets<I: OffsetSizeTrait>( + &mut self, + out: &mut OffsetBuffer<I>, + num_values: usize, + ) -> Result<usize> { + let start = out.values.len(); + let read = self.read_values(&mut out.values, num_values)?; + out.offsets.reserve(read); + for value in 1..=read { + out.offsets.push( + I::from_usize(start + value * self.byte_length) + .ok_or_else(|| general_err!("index overflow decoding fixed byte array"))?, + ); + } + Ok(read) + } + + fn read_values(&mut self, out: &mut Vec<u8>, num_values: usize) -> Result<usize> { + if self.byte_length == 0 { + return Ok(num_values); + } + + match self.decoder.as_mut().unwrap() { + Decoder::Plain { offset, buf } => { + let to_read = + (num_values * self.byte_length).min(buf.len() - *offset) / self.byte_length; + let end_offset = *offset + to_read * self.byte_length; + out.extend_from_slice(&buf.as_ref()[*offset..end_offset]); + *offset = end_offset; + Ok(to_read) + } + Decoder::Dict { decoder } => { + let dict = self.dict_page.as_ref().unwrap(); + if dict.is_empty() { + return Ok(0); + } + + decoder.read(num_values, |keys| { + out.reserve(keys.len() * self.byte_length); + for key in keys { + let offset = *key as usize * self.byte_length; + out.extend_from_slice(&dict.as_ref()[offset..offset + self.byte_length]); + } + Ok(()) + }) + } + Decoder::Delta { decoder } => { + let to_read = num_values.min(decoder.remaining()); + out.reserve(to_read * self.byte_length); + + decoder.read(to_read, |slice| { + if slice.len() != self.byte_length { + return Err(general_err!( + "encountered array with incorrect length, got {} expected {}", + slice.len(), + self.byte_length + )); + } + out.extend_from_slice(slice); + Ok(()) + }) + } + Decoder::ByteStreamSplit { buf, offset } => { + let total_values = buf.len() / self.byte_length; + let to_read = num_values.min(total_values - *offset); + read_byte_stream_split(out, buf, *offset, to_read, self.byte_length); + *offset += to_read; + Ok(to_read) + } + } + } +} + impl ColumnValueDecoder for ValueDecoder { type Buffer = FixedLenByteArrayBuffer; @@ -412,14 +488,12 @@ impl ColumnValueDecoder for ValueDecoder { encoding )); } - let expected_len = num_values as usize * self.byte_length; - if expected_len > buf.len() { - return Err(general_err!( - "too few bytes in dictionary page, expected {} got {}", - expected_len, - buf.len() - )); - } + let buf = normalize_fixed_len_byte_array_payload( + buf, + num_values as usize, + self.byte_length, + "dictionary page", + )?; self.dict_page = Some(buf); Ok(()) @@ -473,66 +547,14 @@ impl ColumnValueDecoder for ValueDecoder { } } - match self.decoder.as_mut().unwrap() { - Decoder::Plain { offset, buf } => { - let to_read = - (num_values * self.byte_length).min(buf.len() - *offset) / self.byte_length; - let end_offset = *offset + to_read * self.byte_length; - out.buffer - .extend_from_slice(&buf.as_ref()[*offset..end_offset]); - *offset = end_offset; - Ok(to_read) - } - Decoder::Dict { decoder } => { - let dict = self.dict_page.as_ref().unwrap(); - // All data must be NULL - if dict.is_empty() { - return Ok(0); - } - - decoder.read(num_values, |keys| { - out.buffer.reserve(keys.len() * self.byte_length); - for key in keys { - let offset = *key as usize * self.byte_length; - let val = &dict.as_ref()[offset..offset + self.byte_length]; - out.buffer.extend_from_slice(val); - } - Ok(()) - }) - } - Decoder::Delta { decoder } => { - let to_read = num_values.min(decoder.remaining()); - out.buffer.reserve(to_read * self.byte_length); - - decoder.read(to_read, |slice| { - if slice.len() != self.byte_length { - return Err(general_err!( - "encountered array with incorrect length, got {} expected {}", - slice.len(), - self.byte_length - )); - } - out.buffer.extend_from_slice(slice); - Ok(()) - }) - } - Decoder::ByteStreamSplit { buf, offset } => { - // we have n=`byte_length` streams of length `buf.len/byte_length` - // to read value i, we need the i'th byte from each of the streams - // so `offset` should be the value offset, not the byte offset - let total_values = buf.len() / self.byte_length; - let to_read = num_values.min(total_values - *offset); - - // now read the n streams and reassemble values into the output buffer - read_byte_stream_split(&mut out.buffer, buf, *offset, to_read, self.byte_length); - - *offset += to_read; - Ok(to_read) - } - } + self.read_values(&mut out.buffer, num_values) } fn skip_values(&mut self, num_values: usize) -> Result<usize> { + if self.byte_length == 0 { + return Ok(num_values); + } + match self.decoder.as_mut().unwrap() { Decoder::Plain { offset, buf } => { let to_read = num_values.min((buf.len() - *offset) / self.byte_length); diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index b8258b8de0..180d20625a 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -1533,7 +1533,6 @@ impl ArrowColumnWriterFactory { ArrowDataType::Utf8View | ArrowDataType::BinaryView => { out.push(bytes(leaves.next().unwrap())?) } - ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?), _ => out.push(col(leaves.next().unwrap())?), }, ArrowDataType::RunEndEncoded(_, value_field) => { diff --git a/parquet/src/arrow/buffer/dictionary_buffer.rs b/parquet/src/arrow/buffer/dictionary_buffer.rs index 441da3a501..eb81058726 100644 --- a/parquet/src/arrow/buffer/dictionary_buffer.rs +++ b/parquet/src/arrow/buffer/dictionary_buffer.rs @@ -21,7 +21,8 @@ use crate::errors::{ParquetError, Result}; use ahash::RandomState; use arrow_array::{Array, DictionaryArray, downcast_integer}; use arrow_array::{ - ArrayRef, FixedSizeBinaryArray, OffsetSizeTrait, cast::AsArray, make_array, + ArrayRef, FixedSizeBinaryArray, GenericBinaryArray, OffsetSizeTrait, + builder::FixedSizeBinaryDictionaryBuilder, cast::AsArray, make_array, types::ArrowDictionaryKeyType, }; use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer}; @@ -191,6 +192,17 @@ impl<K: ArrowNativeType + Ord, V: OffsetSizeTrait> DictionaryBuffer<K, V> { _ => unreachable!(), }; + // Plain byte-array pages are decoded into an offset buffer. Convert + // this representation explicitly before rebuilding an FSB dictionary, + // as FixedSizeBinary arrays have no offsets buffer. + if let ArrowType::FixedSizeBinary(size) = value_type { + let binary = values.into_array(null_buffer, GenericBinaryArray::<V>::DATA_TYPE); + let binary = binary.as_binary::<V>(); + let array = + FixedSizeBinaryArray::try_from_sparse_iter_with_size(binary.iter(), size)?; + return pack_fixed_values(key_type, &array); + } + hash_byte_slices(&values.offsets, &values.values, hash_scratch); let hashes = hashes_as_u64(hash_scratch); @@ -233,6 +245,38 @@ impl<K: ArrowNativeType, V: OffsetSizeTrait> ValuesBuffer for DictionaryBuffer<K } } +macro_rules! fixed_dict_helper { + ($k:ty, $array:ident) => { + pack_fixed_values_impl::<$k>($array) + }; +} + +fn pack_fixed_values(key_type: &ArrowType, array: &FixedSizeBinaryArray) -> Result<ArrayRef> { + downcast_integer! { + key_type => (fixed_dict_helper, array), + _ => unreachable!(), + } +} + +fn pack_fixed_values_impl<K: ArrowDictionaryKeyType>( + array: &FixedSizeBinaryArray, +) -> 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(x)?; + } + None => builder.append_null(), + } + } + Ok(Arc::new(builder.finish())) +} + 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, _>( @@ -622,6 +666,45 @@ mod tests { } } + #[test] + fn test_fixed_size_binary_key_overflow_boundary() { + for (key_type, capacity) in [(ArrowType::Int8, 128u32), (ArrowType::UInt8, 256)] { + for unique in [capacity, capacity + 1] { + let values = (0..unique) + .map(|value| Some(value.to_le_bytes())) + .chain([None, Some(0u32.to_le_bytes())]); + let values = + FixedSizeBinaryArray::try_from_sparse_iter_with_size(values, 4).unwrap(); + let result = pack_fixed_values(&key_type, &values); + if unique > capacity { + let error = result.unwrap_err(); + let ParquetError::External(source) = &error else { + panic!("expected an Arrow dictionary overflow error, got {error}"); + }; + assert!( + matches!( + source.downcast_ref::<arrow_schema::ArrowError>(), + Some(arrow_schema::ArrowError::DictionaryKeyOverflowError) + ), + "{error}" + ); + } else { + let array = result.unwrap(); + assert_eq!(array.len(), unique as usize + 2); + assert_eq!(array.null_count(), 1); + assert!(array.is_null(unique as usize)); + assert!(array.is_valid(unique as usize + 1)); + let dictionary = array.as_any_dictionary(); + assert_eq!(dictionary.values().len(), unique as usize); + assert_eq!( + dictionary.values().data_type(), + &ArrowType::FixedSizeBinary(4) + ); + } + } + } + } + /// A colliding string repeated after its initial collision must reuse its key, /// not be inserted again as a new dictionary entry. #[test] diff --git a/parquet/src/column/reader.rs b/parquet/src/column/reader.rs index 9c2345b097..e6fe0b6a27 100644 --- a/parquet/src/column/reader.rs +++ b/parquet/src/column/reader.rs @@ -23,7 +23,8 @@ use super::page::{Page, PageReader}; use crate::basic::*; use crate::column::reader::decoder::{ ColumnValueDecoder, ColumnValueDecoderImpl, DefinitionLevelDecoder, DefinitionLevelDecoderImpl, - RepetitionLevelDecoder, RepetitionLevelDecoderImpl, + RepetitionLevelDecoder, RepetitionLevelDecoderImpl, normalize_fixed_len_byte_array_data, + prepare_v1_fixed_len_byte_array, }; use crate::data_type::*; use crate::errors::{ParquetError, Result}; @@ -31,6 +32,7 @@ use crate::schema::types::ColumnDescPtr; use crate::util::bit_util::{ceil, num_required_bits, read_num_bytes}; pub(crate) mod decoder; +mod legacy_fixed_len_byte_array; /// Column reader for a Parquet type. pub enum ColumnReader { @@ -136,6 +138,10 @@ pub struct GenericColumnReader<R, D, V> { /// The decoder for the values values_decoder: V, + + /// Width and remaining payload bytes for a nullable V1 FLBA page whose + /// physical count is checked during normal level decoding, not a pre-scan. + remaining_fixed_len_bytes: Option<(usize, usize)>, } impl<V> GenericColumnReader<RepetitionLevelDecoderImpl, DefinitionLevelDecoderImpl, V> @@ -184,6 +190,7 @@ where num_decoded_values: 0, values_decoder, has_record_delimiter: false, + remaining_fixed_len_bytes: None, } } @@ -200,7 +207,11 @@ where /// the number of levels read, with an error returned if it is `None`. /// /// `values` will be contiguously populated with the non-null values. Note that if the column - /// is not required, this may be less than either `max_records` or the number of levels read + /// is not required, this may be less than either `max_records` or the number of levels read. + /// + /// Nullable V1 FLBA payload sizes may be validated incrementally. Extra values or + /// truncated levels can therefore be reported only when the rest of a page is consumed. + /// Pages that could require legacy PLAIN repair are still validated eagerly. pub fn read_records( &mut self, max_records: usize, @@ -284,6 +295,8 @@ where None => levels_to_read, }; + self.check_fixed_len_bytes(values_to_read, levels_to_read == remaining_levels)?; + let def_levels = def_levels.as_deref(); reserve_values(values, values_to_read, levels_to_read, def_levels)?; @@ -373,7 +386,8 @@ where self.num_decoded_values += rep_levels_read; remaining_records -= records_read; - if self.num_buffered_values == self.num_decoded_values { + let page_finished = self.num_buffered_values == self.num_decoded_values; + if page_finished && self.remaining_fixed_len_bytes.is_none() { // Exhausted buffered page - no need to advance other decoders continue; } @@ -391,6 +405,11 @@ where )); } + self.check_fixed_len_bytes(values_read, page_finished)?; + if page_finished { + continue; + } + let values = self.values_decoder.skip_values(values_read)?; if values != values_read { return Err(general_err!( @@ -403,6 +422,33 @@ where Ok(num_records - remaining_records) } + /// Account for actual physical values decoded or skipped from definition levels. + /// The payload length is a byte budget, not evidence of the expected value count. + fn check_fixed_len_bytes(&mut self, values: usize, page_finished: bool) -> Result<()> { + if let Some((width, remaining)) = self.remaining_fixed_len_bytes.as_mut() { + let bytes = values + .checked_mul(*width) + .ok_or_else(|| general_err!("FIXED_LEN_BYTE_ARRAY payload size overflow"))?; + let left = remaining.checked_sub(bytes).ok_or_else(|| { + general_err!( + "Invalid FIXED_LEN_BYTE_ARRAY PLAIN data page payload length: \ + need {bytes} value bytes, only {remaining} remain" + ) + })?; + if page_finished && left != 0 { + return Err(general_err!( + "Invalid FIXED_LEN_BYTE_ARRAY PLAIN data page payload length: \ + {left} trailing value bytes" + )); + } + *remaining = left; + if page_finished { + self.remaining_fixed_len_bytes = None; + } + } + Ok(()) + } + /// Read the next page as a dictionary page. If the next page is not a dictionary page, /// this will return an error. fn read_dictionary_page(&mut self) -> Result<()> { @@ -454,6 +500,7 @@ where let max_rep_level = self.descr.max_rep_level(); let max_def_level = self.descr.max_def_level(); + self.remaining_fixed_len_bytes = None; let mut offset = 0; @@ -475,7 +522,7 @@ where .set_data(rep_level_encoding, level_data)?; } - if max_def_level > 0 { + let definition_levels = if max_def_level > 0 { let (bytes_read, level_data) = parse_v1_level( max_def_level, num_values, @@ -483,16 +530,39 @@ where buf.slice(offset..), )?; offset += bytes_read; - + Some(level_data) + } else { + None + }; + + let data = buf.slice(offset..); + let (encoding, data) = + if self.descr.physical_type() == Type::FIXED_LEN_BYTE_ARRAY { + let width = self.descr.type_length() as usize; + let page = prepare_v1_fixed_len_byte_array( + data, + num_values as usize, + width, + encoding, + definition_levels.as_ref().map(|levels| { + (max_def_level, def_level_encoding, levels) + }), + )?; + self.remaining_fixed_len_bytes = + page.validate_count.then_some((width, page.data.len())); + (page.encoding, page.data) + } else { + (encoding, data) + }; + if let Some(level_data) = definition_levels { self.def_level_decoder .as_mut() .unwrap() .set_data(def_level_encoding, level_data)?; } - self.values_decoder.set_data( encoding, - buf.slice(offset..), + data, num_values as usize, None, )?; @@ -521,6 +591,8 @@ where self.num_buffered_values = num_values as _; self.num_decoded_values = 0; + self.remaining_fixed_len_bytes = None; + // DataPage v2 only supports RLE encoding for repetition // levels if self.descr.max_rep_level() > 0 { @@ -548,9 +620,22 @@ where )?; } + let data = + buf.slice((rep_levels_byte_len + def_levels_byte_len) as usize..); + let (encoding, data) = + if self.descr.physical_type() == Type::FIXED_LEN_BYTE_ARRAY { + normalize_fixed_len_byte_array_data( + data, + (num_values - num_nulls) as usize, + self.descr.type_length() as usize, + encoding, + )? + } else { + (encoding, data) + }; self.values_decoder.set_data( encoding, - buf.slice((rep_levels_byte_len + def_levels_byte_len) as usize..), + data, num_values as usize, Some((num_values - num_nulls) as usize), )?; @@ -639,6 +724,459 @@ mod tests { use crate::util::test_common::page_util::InMemoryPageReader; use crate::util::test_common::rand_gen::make_pages; + pub(super) fn nullable_fixed_page(def: &[i16], encoding: Encoding, payload: &[u8]) -> Page { + use crate::encodings::levels::LevelEncoder; + let mut levels = LevelEncoder::v1_streaming(1); + levels.put_with_observer(def, |_, _| {}); + let mut buf = levels.consume(); + buf.extend_from_slice(payload); + Page::DataPage { + buf: buf.into(), + num_values: def.len() as u32, + encoding, + def_level_encoding: Encoding::RLE, + rep_level_encoding: Encoding::RLE, + statistics: None, + } + } + + pub(super) fn nullable_fixed_reader( + pages: Vec<Page>, + ) -> ColumnReaderImpl<FixedLenByteArrayType> { + let descr = Arc::new(ColumnDescriptor::new( + Arc::new( + SchemaType::primitive_type_builder("fixed", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_length(4) + .build() + .unwrap(), + ), + 1, + 0, + ColumnPath::from("fixed"), + )); + ColumnReaderImpl::<FixedLenByteArrayType>::new( + descr, + Box::new(InMemoryPageReader::new(pages)), + ) + } + + /// These read/skip assertions apply to both canonical and normalized values. + pub(super) fn check_incremental_fixed_len_byte_array_pages( + encoding: Encoding, + payload: impl Fn(usize) -> Vec<u8>, + ) { + for count in [1, 2, 3] { + for finish_by_skipping in [false, true] { + let mut reader = nullable_fixed_reader(vec![nullable_fixed_page( + &[1, 0, 1, 0], + encoding, + &payload(count), + )]); + let mut values = vec![]; + assert_eq!( + reader + .read_records(1, Some(&mut vec![]), None, &mut values) + .unwrap(), + (1, 1, 1) + ); + // Count-dependent corruption can be detected after the first + // value, regardless of the original value representation. + assert_eq!(values[0].data(), b"abcd"); + assert_eq!(reader.remaining_fixed_len_bytes, Some((4, (count - 1) * 4))); + assert_eq!(reader.skip_records(1).unwrap(), 1); + let result = if finish_by_skipping { + reader.skip_records(2).map(|_| ()) + } else { + reader + .read_records(2, Some(&mut vec![]), None, &mut values) + .map(|_| ()) + }; + if count == 2 { + result.unwrap(); + assert!(reader.remaining_fixed_len_bytes.is_none()); + } else { + assert!(result.unwrap_err().to_string().contains("payload length")); + } + } + } + } + + #[test] + fn fixed_v1_incremental_count_validation() { + check_incremental_fixed_len_byte_array_pages(Encoding::PLAIN, |n| b"abcd".repeat(n)); + + // Finishing by skipping checks both missing and extra values, including + // an all-null page with an unexpected whole value in its payload. + for last in [0, 1] { + for payload in [b"".as_slice(), b"abcd".as_slice()] { + let mut reader = nullable_fixed_reader(vec![nullable_fixed_page( + &[0, 0, 0, last], + Encoding::PLAIN, + payload, + )]); + let mut values = vec![]; + assert_eq!( + reader + .read_records(1, Some(&mut vec![]), None, &mut values) + .unwrap(), + (1, 0, 1) + ); + assert_eq!( + reader.skip_records(3).is_ok(), + payload.len() == last as usize * 4 + ); + assert!(values.is_empty()); + } + } + } + + #[test] + fn fixed_v1_canonical_prefix_shaped_values() { + // Ten canonical values look exactly like five length-prefixed values. + let raw = 4_u32.to_le_bytes().repeat(10); + let mut reader = + nullable_fixed_reader(vec![nullable_fixed_page(&[1; 10], Encoding::PLAIN, &raw)]); + let mut values = vec![]; + assert_eq!( + reader + .read_records(1, Some(&mut vec![]), None, &mut values) + .unwrap(), + (1, 1, 1) + ); + reader + .read_records(usize::MAX, Some(&mut vec![]), None, &mut values) + .unwrap(); + assert_eq!(values.len(), 10); + assert!(values.iter().all(|v| v.data() == 4_u32.to_le_bytes())); + } + + #[test] + fn fixed_v1_incremental_page_transitions_and_short_levels() { + let first = nullable_fixed_page(&[1, 0, 1, 0], Encoding::PLAIN, b"abcdabcd"); + let second = nullable_fixed_page(&[1, 1], Encoding::PLAIN, b"efghefgh"); + let mut reader = nullable_fixed_reader(vec![first.clone(), second]); + let mut values = vec![]; + reader + .read_records(1, Some(&mut vec![]), None, &mut values) + .unwrap(); + assert_eq!(reader.skip_records(3).unwrap(), 3); + assert!(reader.remaining_fixed_len_bytes.is_none()); + assert_eq!( + reader + .read_records(usize::MAX, Some(&mut vec![]), None, &mut values) + .unwrap(), + (2, 2, 2) + ); + assert_eq!( + values.iter().map(|v| v.data()).collect::<Vec<_>>(), + [b"abcd", b"efgh", b"efgh"] + ); + + let mut short_levels = first; + if let Page::DataPage { num_values, .. } = &mut short_levels { + *num_values += 8; // exceed any packed padding + } + let mut reader = nullable_fixed_reader(vec![short_levels]); + reader + .read_records(1, Some(&mut vec![]), None, &mut vec![]) + .unwrap(); + let err = reader + .read_records(usize::MAX, Some(&mut vec![]), None, &mut vec![]) + .unwrap_err(); + assert!(err.to_string().contains("definition levels")); + } + + #[test] + fn fixed_plain_payload_length_validation() { + check_fixed_len_byte_array_pages( + |raw| vec![(Encoding::PLAIN, raw.to_vec())], + |raw| { + let mut extra = raw.to_vec(); + extra.push(0); + let mut malformed = vec![extra]; + if !raw.is_empty() { + malformed.push(raw[..raw.len() - 1].to_vec()); + } + malformed + }, + ); + } + + /// Shared page/level assertions, independent of the value representation. + /// Keeping the canonical test separate lets it survive removal of legacy support. + #[expect(deprecated, reason = "Cover legacy BIT_PACKED definition levels")] + pub(super) fn check_fixed_len_byte_array_pages( + payloads: impl Fn(&[u8]) -> Vec<(Encoding, Vec<u8>)>, + malformed_plain: impl Fn(&[u8]) -> Vec<Vec<u8>>, + ) { + use crate::encodings::levels::LevelEncoder; + use crate::encodings::rle::RleEncoder; + use crate::util::bit_util::BitWriter; + + fn levels(data: &[i16], max: i16, v2: bool, encoding: Encoding) -> Vec<u8> { + if max == 0 { + return vec![]; + } + if encoding == Encoding::BIT_PACKED { + let mut encoder = BitWriter::new(data.len()); + for &value in data { + encoder.put_value(value as u64, num_required_bits(max as u64) as usize); + } + return encoder.consume(); + } + let mut encoder = if v2 { + LevelEncoder::v2_streaming(max) + } else { + LevelEncoder::v1_streaming(max) + }; + encoder.put_with_observer(data, |_, _| {}); + encoder.consume() + } + + for (max_def, def, rep) in [ + (0, vec![], vec![]), + (1, vec![1, 0, 1, 0, 1], vec![]), + (1, vec![0; 5], vec![]), + // Null parents, empty containers, null elements and repeated values. + (3, vec![0, 1, 2, 3, 3, 1, 3], vec![0, 0, 0, 1, 1, 0, 1]), + // Exercise counting across multiple bounded definition-level batches. + ( + 1, + (0..2051).map(|i| i16::from(i % 3 == 0)).collect(), + vec![], + ), + ] { + let max_rep = i16::from(!rep.is_empty()); + let num_levels = if max_def == 0 { 3 } else { def.len() }; + let num_values = if max_def == 0 { + num_levels + } else { + def.iter().filter(|&&v| v == max_def).count() + }; + let num_rows = if max_rep == 0 { + num_levels + } else { + rep.iter().filter(|&&v| v == 0).count() + }; + let desc = Arc::new(ColumnDescriptor::new( + Arc::new( + SchemaType::primitive_type_builder("fixed", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_length(4) + .build() + .unwrap(), + ), + max_def, + max_rep, + ColumnPath::from("fixed"), + )); + let raw: Vec<u8> = b"\x04\0\0\0\xff\xff\xff\xffabcd" + .as_chunks::<4>() + .0 + .iter() + .cycle() + .take(num_values) + .flatten() + .copied() + .collect(); + let malformed = malformed_plain(&raw); + + for (v2, level_encoding) in [ + (false, Encoding::RLE), + (false, Encoding::BIT_PACKED), + (true, Encoding::RLE), + ] { + let rep_data = levels(&rep, max_rep, v2, level_encoding); + let def_data = levels(&def, max_def, v2, level_encoding); + let make_page = |encoding, payload: &[u8]| { + let mut buf = rep_data.clone(); + buf.extend_from_slice(&def_data); + buf.extend_from_slice(payload); + if v2 { + Page::DataPageV2 { + buf: Bytes::from(buf), + num_values: num_levels as u32, + encoding, + num_nulls: (num_levels - num_values) as u32, + num_rows: num_rows as u32, + def_levels_byte_len: def_data.len() as u32, + rep_levels_byte_len: rep_data.len() as u32, + is_compressed: false, + statistics: None, + } + } else { + Page::DataPage { + buf: Bytes::from(buf), + num_values: num_levels as u32, + encoding, + def_level_encoding: level_encoding, + rep_level_encoding: level_encoding, + statistics: None, + } + } + }; + let make_reader = |pages| { + ColumnReaderImpl::<FixedLenByteArrayType>::new( + desc.clone(), + Box::new(InMemoryPageReader::new(pages)), + ) + }; + for (encoding, payload) in payloads(&raw) { + let payload = payload.as_slice(); + let mut reader = make_reader(vec![make_page(encoding, payload)]); + let (mut values, mut actual_def, mut actual_rep) = (vec![], vec![], vec![]); + assert_eq!( + reader + .read_records( + usize::MAX, + Some(&mut actual_def), + Some(&mut actual_rep), + &mut values + ) + .unwrap(), + (num_rows, num_values, num_levels) + ); + assert_eq!(actual_def, def); + assert_eq!(actual_rep, rep); + assert_eq!( + values.iter().map(|v| v.data()).collect::<Vec<_>>(), + raw.as_chunks::<4>().0.iter().collect::<Vec<_>>() + ); + + let first_record_levels = if max_rep == 0 { + 1 + } else { + rep.iter().skip(1).position(|&v| v == 0).unwrap() + 1 + }; + let first_record_values = if max_def == 0 { + 1 + } else { + def[..first_record_levels] + .iter() + .filter(|&&v| v == max_def) + .count() + }; + let mut reader = make_reader(vec![make_page(encoding, payload)]); + assert_eq!(reader.skip_records(1).unwrap(), 1); + let mut values = vec![]; + reader + .read_records( + usize::MAX, + Some(&mut vec![]), + Some(&mut vec![]), + &mut values, + ) + .unwrap(); + assert_eq!( + values.iter().map(|v| v.data()).collect::<Vec<_>>(), + raw[first_record_values * 4..] + .as_chunks::<4>() + .0 + .iter() + .collect::<Vec<_>>() + ); + } + + if !v2 && max_def != 0 && level_encoding == Encoding::RLE { + let mut page = make_page(Encoding::PLAIN, &raw); + if let Page::DataPage { num_values, .. } = &mut page { + // Exceed even the final bit-packed run's padding. + *num_values += 8; + } + // Incremental V1 validation detects missing levels when + // the page is consumed, not necessarily on the first row. + let err = make_reader(vec![page]) + .read_records( + usize::MAX, + Some(&mut vec![]), + Some(&mut vec![]), + &mut vec![], + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("definition levels") + || (max_rep != 0 && err.contains("repetition levels")), + "{err}" + ); + } + + for payload in &malformed { + // Byte-alignment errors remain eager. Nullable V1 errors + // requiring the physical count can be delayed until page end. + for skip in [false, true] { + let mut reader = make_reader(vec![make_page(Encoding::PLAIN, payload)]); + let mut values = vec![]; + let mut result = if skip { + reader.skip_records(1).map(|_| ()) + } else { + reader + .read_records(1, Some(&mut vec![]), Some(&mut vec![]), &mut values) + .map(|_| ()) + }; + let provisional_values = values.len(); + if result.is_err() { + assert!(values.is_empty()); + } else { + assert!(!v2 && max_def != 0); + result = if skip { + reader.skip_records(usize::MAX).map(|_| ()) + } else { + reader + .read_records( + usize::MAX, + Some(&mut vec![]), + Some(&mut vec![]), + &mut values, + ) + .map(|_| ()) + }; + } + let err = result.unwrap_err().to_string(); + assert!( + err.contains( + "Invalid FIXED_LEN_BYTE_ARRAY PLAIN data page payload length" + ), + "v2={v2} {level_encoding} max_def={max_def} skip={skip}: {err}" + ); + assert_eq!(values.len(), provisional_values); + } + + // A valid dictionary and dictionary-encoded page must not + // conceal malformed PLAIN data after dictionary fallback. + let mut indices = RleEncoder::new(1, 64); + for _ in 0..num_values { + indices.put(0); + } + let mut encoded = vec![1]; + encoded.extend(indices.consume()); + let mut reader = make_reader(vec![ + Page::DictionaryPage { + buf: Bytes::from_static(b"abcd"), + num_values: 1, + encoding: Encoding::PLAIN, + is_sorted: false, + }, + make_page(Encoding::RLE_DICTIONARY, &encoded), + make_page(Encoding::PLAIN, payload), + ]); + let mut values = vec![]; + let err = reader + .read_records( + usize::MAX, + Some(&mut vec![]), + Some(&mut vec![]), + &mut values, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("PLAIN data page payload length"), "{err}"); + assert_eq!(values.len(), num_values); + assert!(values.iter().all(|v| v.data() == b"abcd")); + } + } + } + } + #[test] fn test_parse_v1_level_invalid_length() { // Say length is 10, but buffer is only 4 diff --git a/parquet/src/column/reader/decoder.rs b/parquet/src/column/reader/decoder.rs index ee24bb3837..4844bb4ee1 100644 --- a/parquet/src/column/reader/decoder.rs +++ b/parquet/src/column/reader/decoder.rs @@ -17,7 +17,7 @@ use bytes::Bytes; -use crate::basic::{Encoding, EncodingMask}; +use crate::basic::{Encoding, EncodingMask, Type}; use crate::data_type::DataType; use crate::encodings::{ decoding::{Decoder, DictDecoder, PlainDecoder, get_decoder}, @@ -124,6 +124,167 @@ pub trait ColumnValueDecoder { fn skip_values(&mut self, num_values: usize) -> Result<usize>; } +fn fixed_len_byte_array_payload_size( + num_values: usize, + type_length: usize, + page: &str, +) -> Result<usize> { + num_values.checked_mul(type_length).ok_or_else(|| { + general_err!( + "FIXED_LEN_BYTE_ARRAY {} payload size overflow: {} values of {} bytes", + page, + num_values, + type_length + ) + }) +} + +/// Validate a complete decompressed PLAIN FLBA value section, excluding levels. +/// `num_values` must count physical non-null values, not rows or requested values. +/// Canonical payloads are returned unchanged, even if values resemble prefixes. +pub(crate) fn normalize_fixed_len_byte_array_payload( + data: Bytes, + num_values: usize, + type_length: usize, + page: &str, +) -> Result<Bytes> { + let expected_len = fixed_len_byte_array_payload_size(num_values, type_length, page)?; + if data.len() == expected_len { + return Ok(data); + } + + // Temporary mitigation for arrow-rs #11261, reached only for a noncanonical + // size. Keep the repair separate from ordinary validation and value decoding. + super::legacy_fixed_len_byte_array::decode_length_prefixed_plain( + data, + num_values, + type_length, + expected_len, + page, + ) +} + +/// Prepare a complete FLBA data-page value section after extracting the levels. +/// PLAIN requires exact size validation; other compliant encodings are unchanged. +pub(crate) fn normalize_fixed_len_byte_array_data( + data: Bytes, + num_values: usize, + type_length: usize, + encoding: Encoding, +) -> Result<(Encoding, Bytes)> { + match encoding { + Encoding::PLAIN => Ok(( + encoding, + normalize_fixed_len_byte_array_payload( + data, + num_values, + type_length, + "PLAIN data page", + )?, + )), + Encoding::DELTA_LENGTH_BYTE_ARRAY => { + // Temporary mitigation only: this encoding is not defined for FLBA. + // Do not add support for it to the ordinary value decoders. + let expected_len = fixed_len_byte_array_payload_size( + num_values, + type_length, + "DELTA_LENGTH_BYTE_ARRAY data page", + )?; + let data = super::legacy_fixed_len_byte_array::decode_delta_length( + data, + num_values, + type_length, + expected_len, + )?; + Ok((Encoding::PLAIN, data)) + } + _ => Ok((encoding, data)), + } +} + +/// A value section ready for its ordinary decoder. When `validate_count` is set, +/// account for its physical values during normal level decoding. +pub(super) struct PreparedFixedLenByteArrayPage { + pub(super) encoding: Encoding, + pub(super) data: Bytes, + pub(super) validate_count: bool, +} + +/// Prepare V1 FLBA values without teaching the column reader historical layouts. +/// `definition_levels` contains the maximum level, encoding and encoded levels; +/// absent levels mean that the page's physical count is already known. +pub(super) fn prepare_v1_fixed_len_byte_array( + data: Bytes, + num_levels: usize, + type_length: usize, + encoding: Encoding, + definition_levels: Option<(i16, Encoding, &Bytes)>, +) -> Result<PreparedFixedLenByteArrayPage> { + let Some((max_level, level_encoding, levels)) = definition_levels else { + let (encoding, data) = + normalize_fixed_len_byte_array_data(data, num_levels, type_length, encoding)?; + return Ok(PreparedFixedLenByteArrayPage { + encoding, + data, + validate_count: false, + }); + }; + + // This callback is lazy: ordinary nullable pages do not create a second + // decoder or traverse their levels before normal reading begins. + let count_values = || { + let mut counter = DefinitionLevelDecoderImpl::new(max_level); + counter.set_data(level_encoding, levels.clone())?; + let (values, read) = counter.skip_def_levels(num_levels)?; + if read != num_levels { + return Err(general_err!( + "Invalid FIXED_LEN_BYTE_ARRAY {encoding} data page: \ + expected {num_levels} definition levels, got {read}" + )); + } + Ok(values) + }; + + let mut page = PreparedFixedLenByteArrayPage { + encoding, + data, + validate_count: false, + }; + // Temporary compatibility hook. Removing it leaves the canonical path and + // its physical-value validation intact. + if super::legacy_fixed_len_byte_array::prepare_v1_page( + &mut page, + num_levels, + type_length, + &count_values, + )? { + return Ok(page); + } + + if encoding == Encoding::PLAIN { + if num_levels == 0 || type_length == 0 { + // Preserve empty-page header validation and avoid a zero-width + // byte budget, which cannot establish the physical value count. + page.data = normalize_fixed_len_byte_array_payload( + page.data, + count_values()?, + type_length, + "PLAIN data page", + )?; + } else { + if !page.data.len().is_multiple_of(type_length) { + return Err(general_err!( + "Invalid FIXED_LEN_BYTE_ARRAY PLAIN data page payload length: \ + {} bytes is not a multiple of {type_length}", + page.data.len() + )); + } + page.validate_count = true; + } + } + Ok(page) +} + /// Bucket-based storage for decoder instances keyed by `Encoding`. /// /// This replaces `HashMap` lookups with direct indexing to avoid hashing overhead in the @@ -170,6 +331,16 @@ impl<T: DataType> ColumnValueDecoder for ColumnValueDecoderImpl<T> { } if encoding == Encoding::RLE_DICTIONARY { + let buf = if self.descr.physical_type() == Type::FIXED_LEN_BYTE_ARRAY { + normalize_fixed_len_byte_array_payload( + buf, + num_values as usize, + self.descr.type_length() as usize, + "dictionary page", + )? + } else { + buf + }; let mut dictionary = PlainDecoder::<T>::new(self.descr.type_length()); dictionary.set_data(buf, num_values as usize)?; @@ -248,8 +419,6 @@ impl<T: DataType> ColumnValueDecoder for ColumnValueDecoderImpl<T> { } } -const SKIP_BUFFER_SIZE: usize = 1024; - enum LevelDecoder { Packed(BitReader, u8), Rle(RleDecoder), @@ -323,25 +492,16 @@ impl DefinitionLevelDecoder for DefinitionLevelDecoderImpl { } fn skip_def_levels(&mut self, num_levels: usize) -> Result<(usize, usize)> { - let mut level_skip = 0; - let mut value_skip = 0; - let mut buf: Vec<i16> = vec![]; - while level_skip < num_levels { - let remaining_levels = num_levels - level_skip; - - let to_read = remaining_levels.min(SKIP_BUFFER_SIZE); - buf.resize(to_read, 0); - let (values_read, levels_read) = self.read_def_levels(&mut buf, to_read)?; - if levels_read == 0 { - // Reached end of page - break; + if num_levels == 0 { + return Ok((0, 0)); + } + match self.decoder.as_mut().unwrap() { + LevelDecoder::Rle(reader) => reader.skip_and_count(num_levels, self.max_level), + LevelDecoder::Packed(reader, 1) => Ok(reader.skip_and_count_ones(num_levels)), + LevelDecoder::Packed(reader, bit_width) => { + Ok(reader.skip_and_count(num_levels, *bit_width as usize, self.max_level)) } - - level_skip += levels_read; - value_skip += values_read; } - - Ok((value_skip, level_skip)) } } @@ -527,6 +687,156 @@ mod tests { use crate::encodings::rle::RleEncoder; use rand::{prelude::*, rng}; + #[test] + fn fixed_len_payload_validation() { + let mut rng = StdRng::seed_from_u64(11262); + for page in ["dictionary page", "PLAIN data page"] { + for width in [0, 1, 4, 8, 16, 32] { + for count in [0, 1, 2, 33] { + let raw: Bytes = (0..width * count).map(|_| rng.random::<u8>()).collect(); + let result = + normalize_fixed_len_byte_array_payload(raw.clone(), count, width, page) + .unwrap(); + assert_eq!(result, raw); + assert_eq!( + result.as_ptr(), + raw.as_ptr(), + "canonical data must stay zero-copy" + ); + } + } + // These are two canonical values, not one length-prefixed value. + let raw = Bytes::from_static(b"\x04\0\0\0abcd"); + assert_eq!( + normalize_fixed_len_byte_array_payload(raw.clone(), 2, 4, page).unwrap(), + raw + ); + // Truncated and trailing data must still be rejected. + for data in [b"abcdefg".as_slice(), b"abcdefghi"] { + assert!( + normalize_fixed_len_byte_array_payload( + Bytes::copy_from_slice(data), + 2, + 4, + page + ) + .is_err() + ); + } + assert!( + normalize_fixed_len_byte_array_payload(Bytes::from_static(b"x"), 0, 4, page) + .is_err() + ); + assert!( + normalize_fixed_len_byte_array_payload(Bytes::new(), 1, usize::MAX, page).is_err() + ); + let err = normalize_fixed_len_byte_array_payload(Bytes::new(), usize::MAX, 2, page) + .unwrap_err() + .to_string(); + assert!(err.contains("payload size overflow"), "{err}"); + assert!(err.contains(page), "{err}"); + } + } + + #[test] + fn fixed_len_other_encodings_unchanged() { + // No compatibility interpretation applies to other encodings. + for encoding in [ + Encoding::DELTA_BYTE_ARRAY, + Encoding::BYTE_STREAM_SPLIT, + Encoding::RLE_DICTIONARY, + ] { + let data = Bytes::from_static(b"\x04\0\0\0abcd"); + let (actual_encoding, actual) = + normalize_fixed_len_byte_array_data(data.clone(), 1, 4, encoding).unwrap(); + assert_eq!(actual_encoding, encoding); + assert_eq!(actual.as_ptr(), data.as_ptr()); + assert_eq!(actual, data); + } + } + + #[test] + #[expect(deprecated, reason = "Cover legacy BIT_PACKED definition levels")] + fn test_skip_def_levels() { + use crate::util::bit_util::BitWriter; + + // A multiple of eight avoids ambiguous padding in the final packed run. + let num_levels = 2056; + for max_level in [1, 2, 3, 7, 255, 32767] { + let bit_width = num_required_bits(max_level as u64); + for levels in [ + vec![0; num_levels], + vec![max_level; num_levels], + (0..num_levels) + .map(|i| (i % (max_level as usize + 1)) as i16) + .collect(), + ] { + for encoding in [Encoding::RLE, Encoding::BIT_PACKED] { + let data = match encoding { + Encoding::RLE => { + let mut encoder = RleEncoder::new(bit_width, 1024); + for &level in &levels { + encoder.put(level as u64); + } + Bytes::from(encoder.consume()) + } + Encoding::BIT_PACKED => { + let mut encoder = BitWriter::new(num_levels); + for &level in &levels { + encoder.put_value(level as u64, bit_width as usize); + } + Bytes::from(encoder.consume()) + } + _ => unreachable!(), + }; + for requested in [ + 0, + 1, + 63, + 64, + 65, + 1023, + 1024, + 1025, + num_levels, + num_levels + 17, + usize::MAX, + ] { + let mut decoder = DefinitionLevelDecoderImpl::new(max_level); + decoder.set_data(encoding, data.clone()).unwrap(); + let expected_levels = requested.min(num_levels); + let expected_values = levels[..expected_levels] + .iter() + .filter(|&&level| level == max_level) + .count(); + assert_eq!( + decoder.skip_def_levels(requested).unwrap(), + (expected_values, expected_levels), + "{encoding}, max_level={max_level}, requested={requested}" + ); + + let mut remaining = vec![]; + let expected_remaining = &levels[expected_levels..]; + assert_eq!( + decoder + .read_def_levels(&mut remaining, expected_remaining.len()) + .unwrap(), + ( + expected_remaining + .iter() + .filter(|&&level| level == max_level) + .count(), + expected_remaining.len(), + ) + ); + assert_eq!(remaining, expected_remaining); + assert_eq!(decoder.skip_def_levels(1).unwrap(), (0, 0)); + } + } + } + } + } + #[test] fn test_skip_padding() { let mut encoder = RleEncoder::new(1, 1024); diff --git a/parquet/src/column/reader/legacy_fixed_len_byte_array.rs b/parquet/src/column/reader/legacy_fixed_len_byte_array.rs new file mode 100644 index 0000000000..1f34e8dbb7 --- /dev/null +++ b/parquet/src/column/reader/legacy_fixed_len_byte_array.rs @@ -0,0 +1,466 @@ +// 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. + +//! Temporary reader mitigation for [arrow-rs #11261], fixed by #11262. +//! These historical layouts are not additional Parquet encodings. +//! +//! [arrow-rs #11261]: https://github.com/apache/arrow-rs/issues/11261 +//! +//! The ordinary decoder checks canonical PLAIN sizes before selecting a repair. +//! A bounded prefix probe can rule out the historical layout, but a match never +//! authorizes repair without an independent physical count. Repairs happen once +//! per page; normal value decoding does not need to know the historical layout. +//! Producer strings and embedded Arrow metadata never select a repair. +//! +//! Canonical PLAIN has `N * W` bytes; the historical representation has exactly +//! `N * (W + 4)` bytes with every little-endian length prefix equal to `W`. These +//! sizes are disjoint for nonempty pages with the same independently established +//! `N`. Without that count, the representations can overlap. +//! +//! To retire the mitigation, remove this module, the compatibility hook in +//! `decoder::prepare_v1_fixed_len_byte_array`, and the repair calls in the two +//! normalization helpers there. Replace noncanonical PLAIN repair with a length +//! error and reject FLBA/DELTA_LENGTH_BYTE_ARRAY instead of normalizing it. +//! No changes to the column reader's page setup, read/skip loops, byte-budget +//! validation or canonical tests are required. Keep physical-value validation: +//! it detects corruption independently of legacy support. Remove legacy fixture +//! tests separately. + +use bytes::Bytes; + +use super::decoder::{PreparedFixedLenByteArrayPage, normalize_fixed_len_byte_array_data}; +use crate::basic::Encoding; +use crate::data_type::Int32Type; +use crate::encodings::decoding::{Decoder, DeltaBitPackDecoder}; +use crate::errors::{ParquetError, Result}; + +/// Handle nullable V1 pages that may need compatibility processing. The caller +/// supplies a lazy, independent count; only ambiguous PLAIN or degenerate pages +/// invoke it. Returns `true` if preparation is complete; `false` leaves `page` +/// unchanged for the caller's ordinary canonical path. Work in place so that +/// ambiguous but canonical pages do not need additional buffer-reference clones. +pub(super) fn prepare_v1_page( + page: &mut PreparedFixedLenByteArrayPage, + num_levels: usize, + type_length: usize, + count_values: &impl Fn() -> Result<usize>, +) -> Result<bool> { + let (num_values, deferred) = match page.encoding { + Encoding::PLAIN + if num_levels != 0 + && type_length != 0 + && may_be_length_prefixed_plain(&page.data, type_length) => + { + // Even all prefixes matching is not permission to repair. The + // normalizer must first compare the independently counted canonical size. + (count_values()?, false) + } + Encoding::DELTA_LENGTH_BYTE_ARRAY => { + if num_levels == 0 || type_length == 0 { + (count_values()?, false) + } else { + let declared = delta_length_value_count(page.data.clone())?; + if declared > num_levels { + return Err(general_err!( + "DELTA_LENGTH_BYTE_ARRAY has more values than levels" + )); + } + (declared, true) + } + } + _ => return Ok(false), + }; + let (encoding, data) = normalize_fixed_len_byte_array_data( + std::mem::take(&mut page.data), + num_values, + type_length, + page.encoding, + )?; + page.encoding = encoding; + page.validate_count = deferred; + page.data = data; + Ok(true) +} + +/// A negative result rules out the supported historical PLAIN layout. A positive +/// result only requests an independent count: even every prefix matching could +/// be an ordinary canonical payload. Bound the probe independently of page size. +fn may_be_length_prefixed_plain(data: &[u8], type_length: usize) -> bool { + let Some(stride) = type_length.checked_add(4) else { + return false; + }; + !data.is_empty() + && data.len().is_multiple_of(stride) + && data + .chunks_exact(stride) + .take(5) + .all(|value| u32::from_le_bytes(value[..4].try_into().unwrap()) as usize == type_length) +} + +/// Read the declared number of lengths, not an independent physical count. +/// The caller must still check lengths, payload bytes and actual definition levels. +fn delta_length_value_count(data: Bytes) -> Result<usize> { + let mut decoder = DeltaBitPackDecoder::<Int32Type>::new(); + decoder.set_data(data, 0)?; + Ok(decoder.values_left()) +} + +/// Remove BYTE_ARRAY length prefixes from a noncanonical PLAIN FLBA section. +/// The caller has checked `expected_len = num_values * type_length` for overflow +/// and established that the input does not have that canonical length. +#[cold] +pub(super) fn decode_length_prefixed_plain( + data: Bytes, + num_values: usize, + type_length: usize, + expected_len: usize, + page: &str, +) -> Result<Bytes> { + let stride = type_length.checked_add(4); + if stride.and_then(|s| num_values.checked_mul(s)) != Some(data.len()) { + return Err(general_err!( + "Invalid FIXED_LEN_BYTE_ARRAY {} payload length: expected {} bytes \ + ({} values of {} bytes), got {}", + page, + expected_len, + num_values, + type_length, + data.len() + )); + } + // The checked size above guarantees a nonzero, representable stride. + let stride = stride.unwrap(); + for value in data.chunks_exact(stride) { + let length = u32::from_le_bytes(value[..4].try_into().unwrap()) as usize; + if length != type_length { + return Err(general_err!( + "Invalid FIXED_LEN_BYTE_ARRAY {} payload length prefix: expected {}, got {}", + page, + type_length, + length + )); + } + } + + // Validate the entire page before allocating or exposing any legacy values. + let mut values = Vec::with_capacity(expected_len); + for value in data.chunks_exact(stride) { + values.extend_from_slice(&value[4..]); + } + Ok(values.into()) +} + +/// Validate the BYTE_ARRAY-only delta-length representation and return its +/// contiguous fixed-width values. The caller supplies the overflow-checked +/// canonical payload size and passes the returned value section to PLAIN. +#[cold] +pub(super) fn decode_delta_length( + data: Bytes, + num_values: usize, + type_length: usize, + expected_len: usize, +) -> Result<Bytes> { + let mut decoder = DeltaBitPackDecoder::<Int32Type>::new(); + decoder.set_data(data.clone(), num_values)?; + if decoder.values_left() != num_values { + return Err(general_err!( + "Invalid FIXED_LEN_BYTE_ARRAY DELTA_LENGTH_BYTE_ARRAY value count: expected {}, got {}", + num_values, + decoder.values_left() + )); + } + // Bound scratch space independently of untrusted page counts. + let mut lengths = [0_i32; 128]; + let mut remaining = num_values; + while remaining != 0 { + let count = remaining.min(lengths.len()); + let read = decoder.get(&mut lengths[..count])?; + if read != count { + return Err(general_err!( + "Truncated FLBA DELTA_LENGTH_BYTE_ARRAY lengths" + )); + } + for &length in &lengths[..count] { + if usize::try_from(length).ok() != Some(type_length) { + return Err(general_err!( + "Invalid FIXED_LEN_BYTE_ARRAY DELTA_LENGTH_BYTE_ARRAY value length: expected {}, got {}", + type_length, + length + )); + } + } + remaining -= count; + } + let offset = decoder.get_offset(); + if data.len().checked_sub(offset) != Some(expected_len) { + return Err(general_err!( + "Invalid FIXED_LEN_BYTE_ARRAY DELTA_LENGTH_BYTE_ARRAY payload length: expected {} value bytes", + expected_len + )); + } + Ok(data.slice(offset..)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::basic::Encoding; + use crate::column::reader::decoder::{ + normalize_fixed_len_byte_array_data, normalize_fixed_len_byte_array_payload, + }; + use crate::column::reader::tests::{ + check_incremental_fixed_len_byte_array_pages, nullable_fixed_page, nullable_fixed_reader, + }; + use crate::encodings::encoding::{DeltaBitPackEncoder, Encoder}; + use rand::{prelude::*, rngs::StdRng}; + + #[test] + fn ordinary_pages_do_not_request_compatibility_counting() { + for encoding in [ + Encoding::PLAIN, + Encoding::RLE_DICTIONARY, + Encoding::DELTA_BYTE_ARRAY, + ] { + let mut page = PreparedFixedLenByteArrayPage { + encoding, + data: Bytes::from_static(b"abcdefghijkl"), + validate_count: false, + }; + assert!( + !prepare_v1_page(&mut page, 3, 4, &|| { + panic!("ordinary page must not be pre-scanned") + }) + .unwrap() + ); + assert_eq!(page.encoding, encoding); + assert_eq!(page.data.as_ref(), b"abcdefghijkl"); + assert!(!page.validate_count); + } + } + + #[test] + fn matching_prefixes_require_an_independent_count() { + let calls = std::cell::Cell::new(0); + let raw = Bytes::from(4_u32.to_le_bytes().repeat(10)); + let mut page = PreparedFixedLenByteArrayPage { + encoding: Encoding::PLAIN, + data: raw.clone(), + validate_count: false, + }; + assert!( + prepare_v1_page(&mut page, 20, 4, &|| { + calls.set(calls.get() + 1); + Ok(10) // ten present values, not twenty logical levels + }) + .unwrap() + ); + assert_eq!(calls.get(), 1); + assert_eq!(page.data, raw); + assert_eq!(page.encoding, Encoding::PLAIN); + assert!(!page.validate_count); + + // Five matching prefixes do not permit accepting a bad sixth prefix. + let mut malformed = 4_u32.to_le_bytes().repeat(12); + malformed[40..44].fill(0xff); + let mut reader = nullable_fixed_reader(vec![nullable_fixed_page( + &[1; 6], + Encoding::PLAIN, + &malformed, + )]); + let mut values = vec![]; + assert!( + reader + .read_records(1, Some(&mut vec![]), None, &mut values) + .is_err() + ); + assert!(values.is_empty()); + } + + #[test] + fn incremental_delta_length_count_validation() { + check_incremental_fixed_len_byte_array_pages(Encoding::DELTA_LENGTH_BYTE_ARRAY, |count| { + let mut lengths = DeltaBitPackEncoder::<Int32Type>::new(); + lengths.put(&vec![4; count]).unwrap(); + let mut payload = lengths.flush_buffer().unwrap().to_vec(); + payload.extend_from_slice(&b"abcd".repeat(count)); + payload + }); + + for lengths in [vec![], vec![4; 5], vec![3, 5]] { + let mut encoder = DeltaBitPackEncoder::<Int32Type>::new(); + encoder.put(&lengths).unwrap(); + let mut payload = encoder.flush_buffer().unwrap().to_vec(); + payload.extend(std::iter::repeat_n( + b'x', + lengths.iter().sum::<i32>() as usize, + )); + let def = if lengths.is_empty() { + [0; 4] + } else { + [1, 0, 1, 0] + }; + let mut reader = nullable_fixed_reader(vec![nullable_fixed_page( + &def, + Encoding::DELTA_LENGTH_BYTE_ARRAY, + &payload, + )]); + let result = reader.read_records(1, Some(&mut vec![]), None, &mut vec![]); + if lengths.is_empty() { + assert_eq!(result.unwrap(), (1, 0, 1)); + assert_eq!(reader.skip_records(3).unwrap(), 3); + } else { + // Invalid declared counts or individual lengths remain eager errors. + assert!(result.is_err()); + } + } + } + + #[test] + fn prefix_probe_only_rules_out_legacy() { + for width in [0, 1, 2, 3, 4, 16, 32] { + for count in [0, 1, 5, 6] { + let mut data = vec![]; + for _ in 0..count { + data.extend_from_slice(&(width as u32).to_le_bytes()); + data.extend(std::iter::repeat_n(0x55, width)); + } + assert_eq!(may_be_length_prefixed_plain(&data, width), count != 0); + if !data.is_empty() { + let mut truncated = data.clone(); + truncated.pop(); + assert!(!may_be_length_prefixed_plain(&truncated, width)); + data[0] ^= 1; + assert!(!may_be_length_prefixed_plain(&data, width)); + } + } + } + assert!(!may_be_length_prefixed_plain(&[], usize::MAX)); + let mut data = 4_u32.to_le_bytes().repeat(12); + data[40..44].fill(0xff); + assert!(may_be_length_prefixed_plain(&data, 4)); // probe is bounded + } + + #[test] + fn length_prefixed_plain() { + let mut rng = StdRng::seed_from_u64(11261); + for page in ["dictionary page", "PLAIN data page"] { + for width in [0, 1, 4, 8, 16, 32] { + for count in [0, 1, 2, 33] { + let raw: Bytes = (0..width * count).map(|_| rng.random::<u8>()).collect(); + let mut legacy = vec![]; + for i in 0..count { + legacy.extend_from_slice(&(width as u32).to_le_bytes()); + legacy.extend_from_slice(&raw[i * width..(i + 1) * width]); + } + assert_eq!( + normalize_fixed_len_byte_array_payload(legacy.into(), count, width, page) + .unwrap(), + raw + ); + } + } + // An exact aggregate size is insufficient: every prefix must match. + for data in [ + b"\x03\0\0\0abc\x05\0\0\0defgh".as_slice(), + b"\x04\0\0\0abcd\xff\xff\xff\xffefgh", + b"\x04\0\0\0abcd\x04\0\0\0efg", + b"\x04\0\0\0abcd\x04\0\0\0efghi", + ] { + assert!( + normalize_fixed_len_byte_array_payload( + Bytes::copy_from_slice(data), + 2, + 4, + page, + ) + .is_err() + ); + } + } + } + + #[test] + fn data_pages() { + let length_prefixed = |raw: &[u8]| { + let mut legacy = vec![]; + for value in raw.as_chunks::<4>().0 { + legacy.extend_from_slice(&4_u32.to_le_bytes()); + legacy.extend_from_slice(value); + } + legacy + }; + super::super::tests::check_fixed_len_byte_array_pages( + |raw| { + let mut lengths = DeltaBitPackEncoder::<Int32Type>::new(); + lengths.put(&vec![4; raw.len() / 4]).unwrap(); + let mut delta = lengths.flush_buffer().unwrap().to_vec(); + delta.extend_from_slice(raw); + vec![ + (Encoding::PLAIN, length_prefixed(raw)), + (Encoding::DELTA_LENGTH_BYTE_ARRAY, delta), + ] + }, + |raw| { + let mut bad_prefix = length_prefixed(raw); + if raw.is_empty() { + // All-null pages must have an empty PLAIN value section. + bad_prefix.extend_from_slice(&4_u32.to_le_bytes()); + } else { + bad_prefix[0] = 3; + } + vec![bad_prefix] + }, + ); + } + + #[test] + fn delta_length() { + let encode = |lengths: &[i32], values: &[u8]| { + let mut encoder = DeltaBitPackEncoder::<Int32Type>::new(); + encoder.put(lengths).unwrap(); + let mut data = encoder.flush_buffer().unwrap().to_vec(); + data.extend_from_slice(values); + Bytes::from(data) + }; + let normalize = |data, count| { + normalize_fixed_len_byte_array_data(data, count, 4, Encoding::DELTA_LENGTH_BYTE_ARRAY) + }; + for count in [0, 1, 2, 129, 2051] { + let values = b"abcd".repeat(count); + let data = encode(&vec![4; count], &values); + let (encoding, result) = normalize(data.clone(), count).unwrap(); + assert_eq!(encoding, Encoding::PLAIN); + assert_eq!(result.as_ref(), values); + assert_eq!( + result.as_ptr(), + data.as_ptr().wrapping_add(data.len() - result.len()) + ); + } + for (lengths, values, count) in [ + (vec![3, 5], b"abcdefgh".as_slice(), 2), + (vec![4, -1], b"abcdefgh", 2), + (vec![4, 4], b"abcdefgh", 1), + (vec![4, 4], b"abcdefg", 2), + (vec![4, 4], b"abcdefghi", 2), + (vec![], b"x", 0), + ] { + assert!(normalize(encode(&lengths, values), count).is_err()); + } + assert!(normalize(Bytes::new(), 1).is_err()); + assert!(normalize(Bytes::new(), usize::MAX).is_err()); + } +} diff --git a/parquet/src/encodings/rle.rs b/parquet/src/encodings/rle.rs index 2d56a30918..786197bf5b 100644 --- a/parquet/src/encodings/rle.rs +++ b/parquet/src/encodings/rle.rs @@ -548,6 +548,46 @@ impl RleDecoder { Ok(values_skipped) } + /// Skips up to `num_values`, returning the number equal to `value` and the + /// number skipped. Counts RLE runs directly, one-bit runs with popcount, and + /// wider packed values with the existing unpacker and bounded scratch space. + pub(crate) fn skip_and_count<T: BitPacking + Copy + PartialEq>( + &mut self, + num_values: usize, + value: T, + ) -> Result<(usize, usize)> { + let mut matching = 0; + let mut skipped = 0; + while skipped < num_values { + if self.rle_left > 0 { + let count = (num_values - skipped).min(self.rle_left as usize); + if T::from_u64(self.current_value.unwrap()) == value { + matching += count; + } + self.rle_left -= count as u32; + skipped += count; + } else if self.bit_packed_left > 0 { + let count = (num_values - skipped).min(self.bit_packed_left as usize); + let reader = self + .bit_reader + .as_mut() + .ok_or_else(|| general_err!("bit_reader should be set"))?; + let (matches, read) = reader.skip_and_count(count, self.bit_width as usize, value); + matching += matches; + if read == 0 { + // Match get_batch's handling of truncated final packed runs. + self.bit_packed_left = 0; + continue; + } + self.bit_packed_left -= read as u32; + skipped += read; + } else if !self.reload()? { + break; + } + } + Ok((matching, skipped)) + } + #[inline(never)] pub fn get_batch_with_dict<T>( &mut self, @@ -731,6 +771,104 @@ mod tests { assert_eq!(buffer, expected); } + #[test] + fn test_skip_and_count() { + for width in [0, 1, 2, 15, 64] { + let max = if width == 64 { + u64::MAX + } else { + (1_u64 << width) - 1 + }; + let values: Vec<u64> = std::iter::repeat_n(0, 1024) + .chain((0..128).map(|i| i & max)) + .chain(std::iter::repeat_n(max, 2056)) + .collect(); + let mut encoder = RleEncoder::new(width, 1024); + for &value in &values { + encoder.put(value); + } + let data = Bytes::from(encoder.consume()); + for value in [0, max] { + for prefix in [0, 1, 1023, 1024, 1027, 1152] { + for requested in [0, 1, 7, 1024, usize::MAX] { + let mut decoder = RleDecoder::new(width); + decoder.set_data(data.clone()).unwrap(); + let mut read = vec![0; prefix]; + assert_eq!(decoder.get_batch(&mut read).unwrap(), prefix); + assert_eq!(read, values[..prefix]); + let skipped = requested.min(values.len() - prefix); + let count = values[prefix..prefix + skipped] + .iter() + .filter(|&&v| v == value) + .count(); + assert_eq!( + decoder.skip_and_count(requested, value).unwrap(), + (count, skipped), + "width={width}, value={value}, prefix={prefix}, requested={requested}" + ); + let mut rest = vec![0; values.len() - prefix - skipped]; + assert_eq!(decoder.get_batch(&mut rest).unwrap(), rest.len()); + assert_eq!(rest, values[prefix + skipped..]); + assert_eq!(decoder.skip_and_count(1, value).unwrap(), (0, 0)); + } + } + } + } + } + + #[test] + fn test_skip_and_count_signed() { + let data = Bytes::from_static(&[3, 255, 0, 128, 255, 1, 2, 254, 255]); + let mut decoder = RleDecoder::new(8); + decoder.set_data(data.clone()).unwrap(); + assert_eq!(decoder.skip_and_count(8, -1_i8).unwrap(), (3, 8)); + decoder.set_data(data.clone()).unwrap(); + assert_eq!(decoder.skip_and_count(8, -1_i16).unwrap(), (0, 8)); + decoder.set_data(data).unwrap(); + assert_eq!(decoder.skip_and_count(8, 511_i16).unwrap(), (0, 8)); + } + + #[test] + fn test_skip_and_count_truncated() { + for width in [1, 8] { + // Header declares 24 values, but the final packed run is truncated. + let mut data = vec![7]; + data.extend(std::iter::repeat_n(0xff, if width == 1 { 2 } else { 20 })); + let mut decoder = RleDecoder::new(width); + decoder.set_data(data.into()).unwrap(); + let count = if width == 1 { 16 } else { 20 }; + assert_eq!( + decoder + .skip_and_count(usize::MAX, (1_u64 << width) - 1) + .unwrap(), + (count, count) + ); + assert_eq!(decoder.skip_and_count(1, 0_u64).unwrap(), (0, 0)); + } + let mut decoder = RleDecoder::new(1); + // A valid two-value run followed by a missing RLE value. + decoder.set_data(Bytes::from_static(&[4, 1, 4])).unwrap(); + assert!(decoder.skip_and_count(3, 1_u64).is_err()); + decoder.set_data(Bytes::from_static(&[4, 1, 4])).unwrap(); + assert_eq!(decoder.skip_and_count(2, 1_u64).unwrap(), (2, 2)); + assert!(decoder.skip_and_count(1, 1_u64).is_err()); + } + + #[test] + fn test_skip_and_count_long_run() { + let mut writer = BitWriter::new(8); + writer.put_vlq_int(u64::from(u32::MAX) * 2); + writer.put_aligned(1_u8, 1); + let mut decoder = RleDecoder::new(1); + decoder.set_data(writer.consume().into()).unwrap(); + // Must not expand this run or allocate scratch space proportional to it. + assert_eq!( + decoder.skip_and_count(usize::MAX, 1_u64).unwrap(), + (u32::MAX as usize, u32::MAX as usize) + ); + assert_eq!(decoder.skip_and_count(1, 1_u64).unwrap(), (0, 0)); + } + #[test] fn test_rle_consume_flush_buffer() { let data = vec![1, 1, 1, 2, 2, 3, 3, 3]; diff --git a/parquet/src/util/bit_util.rs b/parquet/src/util/bit_util.rs index b103655449..e0669e6fe7 100644 --- a/parquet/src/util/bit_util.rs +++ b/parquet/src/util/bit_util.rs @@ -806,6 +806,34 @@ impl BitWriter { /// MAX_VLQ_BYTE_LEN = 5 for i32, and MAX_VLQ_BYTE_LEN = 10 for i64 pub const MAX_VLQ_BYTE_LEN: usize = 10; +/// Counts set bits directly over byte slices, including unaligned endpoints. +fn count_set_bits(bytes: &[u8], range: std::ops::Range<usize>) -> usize { + assert!(range.start <= range.end); + if range.start == range.end { + return 0; + } + let data = &bytes[range.start / 8..range.end.div_ceil(8)]; + let (words, tail) = data.as_chunks::<8>(); + let mut count = words + .iter() + .map(|word| u64::from_le_bytes(*word).count_ones() as usize) + .sum::<usize>(); + if !tail.is_empty() { + let word = match data.last_chunk::<8>() { + Some(word) => u64::from_le_bytes(*word) >> ((8 - tail.len()) * 8), + None => read_num_bytes::<u64>(tail.len(), tail), + }; + count += word.count_ones() as usize; + } + if !range.start.is_multiple_of(8) { + count -= (data[0] & ((1 << (range.start % 8)) - 1)).count_ones() as usize; + } + if !range.end.is_multiple_of(8) { + count -= (data[data.len() - 1] >> (range.end % 8)).count_ones() as usize; + } + count +} + /// Reads bit packed values from an in-memory buffer. /// /// `BitReader` is the dual of [`BitWriter`] and reads values that are either @@ -1037,6 +1065,60 @@ impl BitReader { values_to_read } + /// Skips packed values, returning the number equal to `value` and the number + /// skipped. As with `get_batch`, `num_bits` must fit in `T`. + pub(crate) fn skip_and_count<T: BitPacking + Copy + PartialEq>( + &mut self, + num_values: usize, + num_bits: usize, + value: T, + ) -> (usize, usize) { + assert!(num_bits <= T::BIT_CAPACITY); + let zero = T::from_u64(0); + if num_bits == 0 { + return (if value == zero { num_values } else { 0 }, num_values); + } + if num_bits == 1 { + let (ones, skipped) = self.skip_and_count_ones(num_values); + let matching = if value == T::from_u64(1) { + ones + } else if value == zero { + skipped - ones + } else { + 0 + }; + return (matching, skipped); + } + // Reuse the existing unpacker for wider fields instead of maintaining + // a separate packed-field equality algorithm. + let mut buffer = [zero; 128]; + let mut skipped = 0; + let mut matching = 0; + while skipped < num_values { + let to_read = (num_values - skipped).min(buffer.len()); + let read = self.get_batch(&mut buffer[..to_read], num_bits); + if read == 0 { + break; + } + matching += buffer[..read].iter().filter(|&&v| v == value).count(); + skipped += read; + } + // Keep subsequent batch reads byte-aligned whenever possible. + self.skip(0, 1); + (matching, skipped) + } + + /// Skips up to `num_bits` bits, returning the number of set bits and bits skipped. + /// Counts packed bits directly without expanding them into individual values. + pub(crate) fn skip_and_count_ones(&mut self, num_bits: usize) -> (usize, usize) { + let start = self.byte_offset * 8 + self.bit_offset; + let skipped = self.skip(num_bits, 1); + ( + count_set_bits(&self.buffer, start..start + skipped), + skipped, + ) + } + /// Reads up to `num_bytes` bytes from the stream, appending them to `buf`, /// and returns the number of bytes actually appended. /// @@ -1223,6 +1305,114 @@ mod tests { assert_eq!(skipped, 0); } + #[test] + fn test_bit_reader_skip_and_count() { + for width in 1..=64 { + let max = trailing_bits(u64::MAX, width); + let values: Vec<u64> = (0..128) + .map(|i| match i % 3 { + 0 => 0, + 1 => max, + _ => i & max, + }) + .collect(); + let mut writer = BitWriter::new(1024); + for &value in &values { + writer.put_value(value, width); + } + let data = Bytes::from(writer.consume()); + for value in [0, max, max / 2] { + for prefix in [0, 1, 7, 31, 63, 64, 65, 127, 128] { + for requested in [0, 1, 7, 32, 63, 64, 65, usize::MAX] { + let mut reader = BitReader::new(data.clone()); + let mut before = vec![0_u64; prefix]; + assert_eq!(reader.get_batch(&mut before, width), prefix); + let skipped = requested.min(values.len() - prefix); + let count = values[prefix..prefix + skipped] + .iter() + .filter(|&&v| v == value) + .count(); + assert_eq!( + reader.skip_and_count(requested, width, value), + (count, skipped) + ); + let mut rest = vec![0_u64; values.len() - prefix - skipped]; + assert_eq!(reader.get_batch(&mut rest, width), rest.len()); + assert_eq!(rest, values[prefix + skipped..]); + } + } + } + } + let mut reader = BitReader::new(Bytes::new()); + assert_eq!( + reader.skip_and_count(usize::MAX, 0, 0), + (usize::MAX, usize::MAX) + ); + assert_eq!(reader.skip_and_count(123, 0, 1), (0, 123)); + } + + #[test] + fn test_bit_reader_count_unaligned() { + let data = Bytes::from((0..19).map(|i| (i * 73 + 19) as u8).collect::<Vec<_>>()); + for width in [1, 2, 3, 4, 5, 8, 16, 32, 64] { + for prefix in 0..=data.len() * 8 { + let mut reference = BitReader::new(data.clone()); + reference.skip(prefix, 1); + let mut values = vec![0_u64; (data.len() * 8 - prefix) / width]; + assert_eq!(reference.get_batch(&mut values, width), values.len()); + for value in [0, values.first().copied().unwrap_or(0)] { + for requested in [0, 1, 7, usize::MAX] { + let mut reader = BitReader::new(data.clone()); + reader.skip(prefix, 1); + let skipped = requested.min(values.len()); + let matching = values[..skipped].iter().filter(|&&v| v == value).count(); + assert_eq!( + reader.skip_and_count(requested, width, value), + (matching, skipped) + ); + let mut rest = vec![0_u64; values.len() - skipped]; + assert_eq!(reader.get_batch(&mut rest, width), rest.len()); + assert_eq!(rest, values[skipped..]); + } + } + } + } + } + + #[test] + fn test_bit_reader_skip_and_count_ones() { + let bits: Vec<u8> = (0..152).map(|i| u8::from(i % 3 == 1)).collect(); + let mut writer = BitWriter::new(19); + for &bit in &bits { + writer.put_value(bit as u64, 1); + } + let data = Bytes::from(writer.consume()); + for prefix in [0, 1, 7, 8, 31, 63, 64, 65, 151, 152] { + for requested in [0, 1, 7, 63, 64, 65, 130, usize::MAX] { + let mut reader = BitReader::new(data.clone()); + let first = prefix.min(64); + reader.get_value::<u64>(first).unwrap(); + if prefix > first { + assert_eq!(reader.skip(prefix - first, 1), prefix - first); + } + let skipped = requested.min(bits.len() - prefix); + let count = bits[prefix..prefix + skipped] + .iter() + .map(|&b| b as usize) + .sum(); + assert_eq!(reader.skip_and_count_ones(requested), (count, skipped)); + let mut remaining = vec![0; bits.len() - prefix - skipped]; + assert_eq!(reader.get_batch(&mut remaining, 1), remaining.len()); + assert_eq!(remaining, bits[prefix + skipped..]); + assert_eq!(reader.skip_and_count_ones(1), (0, 0)); + } + } + assert_eq!( + BitReader::new(Bytes::new()).skip_and_count_ones(usize::MAX), + (0, 0) + ); + } + #[test] fn test_bit_reader_get_value_boundary() { let buffer = vec![10, 0, 0, 0, 20, 0, 30, 0, 0, 0, 40, 0]; diff --git a/parquet/src/util/test_common/page_util.rs b/parquet/src/util/test_common/page_util.rs index a76d670159..1403a803e7 100644 --- a/parquet/src/util/test_common/page_util.rs +++ b/parquet/src/util/test_common/page_util.rs @@ -49,6 +49,7 @@ pub trait DataPageBuilder { pub struct DataPageBuilderImpl { encoding: Option<Encoding>, num_values: u32, + num_nulls: u32, buffer: Vec<u8>, rep_levels_byte_len: u32, def_levels_byte_len: u32, @@ -64,6 +65,7 @@ impl DataPageBuilderImpl { DataPageBuilderImpl { encoding: None, num_values, + num_nulls: 0, buffer: vec![], rep_levels_byte_len: 0, def_levels_byte_len: 0, @@ -102,6 +104,11 @@ impl DataPageBuilder for DataPageBuilderImpl { fn add_def_levels(&mut self, max_levels: i16, def_levels: &[i16]) { self.num_values = def_levels.len() as u32; + // Nulls and nested placeholders have no physical value in the data section. + self.num_nulls = def_levels + .iter() + .filter(|&&level| level < max_levels) + .count() as u32; self.def_levels_byte_len = self.add_levels(max_levels, def_levels); } @@ -146,8 +153,7 @@ impl DataPageBuilder for DataPageBuilderImpl { buf: Bytes::from(self.buffer), num_values: self.num_values, encoding: self.encoding.unwrap(), - num_nulls: 0, /* set to dummy value - don't need this when reading - * data page */ + num_nulls: self.num_nulls, num_rows: self.num_values, /* num_rows only needs in skip_records, now we not support skip REPEATED field, * so we can assume num_values == num_rows */ def_levels_byte_len: self.def_levels_byte_len, @@ -257,3 +263,76 @@ impl<I: Iterator<Item = Vec<Page>>> Iterator for InMemoryPageIterator<I> { #[cfg(feature = "arrow")] impl<I: Iterator<Item = Vec<Page>> + Send> PageIterator for InMemoryPageIterator<I> {} + +#[cfg(test)] +mod tests { + use super::*; + use crate::basic::Type; + use crate::data_type::{FixedLenByteArray, FixedLenByteArrayType}; + use crate::encodings::rle::RleEncoder; + + #[test] + fn test_data_page_null_counts() { + for (max_def_level, def_levels, rep_levels, expected_nulls) in [ + (0, vec![0; 3], vec![0; 3], 0), + (1, vec![1; 3], vec![0; 3], 0), + (1, vec![1, 0, 1], vec![0; 3], 1), + (1, vec![0; 3], vec![0; 3], 3), + // Null parents, empty lists, null elements and repeated physical values. + (3, vec![0, 1, 2, 3, 3], vec![0, 0, 0, 0, 1], 3), + (1, vec![], vec![], 0), + ] { + let max_rep_level = i16::from(rep_levels.contains(&1)); + let desc = Arc::new(ColumnDescriptor::new( + Arc::new( + SchemaType::primitive_type_builder("fixed", Type::FIXED_LEN_BYTE_ARRAY) + .with_length(2) + .build() + .unwrap(), + ), + max_def_level, + max_rep_level, + ColumnPath::from("fixed"), + )); + let num_levels = def_levels.len() as u32; + let num_values = num_levels - expected_nulls; + let values = vec![FixedLenByteArray::from(vec![1, 2]); num_values as usize]; + + for datapage_v2 in [false, true] { + for dictionary in [false, true] { + let mut builder = + DataPageBuilderImpl::new(desc.clone(), num_values, datapage_v2); + builder.add_rep_levels(max_rep_level, &rep_levels); + builder.add_def_levels(max_def_level, &def_levels); + if dictionary { + let mut encoder = RleEncoder::new(1, 64); + for _ in 0..num_values { + encoder.put(0); + } + let mut indices = vec![1]; + indices.extend(encoder.consume()); + builder.add_indices(indices.into()); + } else { + builder.add_values::<FixedLenByteArrayType>(Encoding::PLAIN, &values); + } + match builder.consume() { + Page::DataPageV2 { + num_values, + num_nulls, + .. + } => { + assert!(datapage_v2); + assert_eq!(num_values, num_levels); + assert_eq!(num_nulls, expected_nulls); + } + Page::DataPage { num_values, .. } => { + assert!(!datapage_v2); + assert_eq!(num_values, num_levels); + } + Page::DictionaryPage { .. } => unreachable!(), + } + } + } + } + } +} diff --git a/parquet/tests/arrow_reader/bad_data.rs b/parquet/tests/arrow_reader/bad_data.rs index b22e78a32b..cc7e26b214 100644 --- a/parquet/tests/arrow_reader/bad_data.rs +++ b/parquet/tests/arrow_reader/bad_data.rs @@ -15,7 +15,8 @@ // specific language governing permissions and limitations // under the License. -//! Tests that reading invalid parquet files returns an error +//! Regression tests for malformed Parquet files. +//! Historical FLBA compatibility is tested in `legacy_fixed_len_byte_array`. use arrow::compute::concat_batches; use arrow::util::test_util::parquet_test_data; @@ -34,6 +35,7 @@ static KNOWN_FILES: &[&str] = &[ "ARROW-GH-43605.parquet", "ARROW-RS-GH-6229-DICTHEADER.parquet", "ARROW-RS-GH-6229-LEVELS.parquet", + "ARROW-RS-GH-11261-FLBA-DICT.parquet", "ARROW-GH-45185.parquet", "ARROW-GH-47662.parquet", "README.md", @@ -41,7 +43,7 @@ static KNOWN_FILES: &[&str] = &[ ]; /// Returns the path to 'parquet-testing/bad_data' -fn bad_data_dir() -> PathBuf { +pub(super) fn bad_data_dir() -> PathBuf { // points to parquet-testing/data let parquet_testing_data = parquet_test_data(); PathBuf::from(parquet_testing_data) @@ -143,7 +145,7 @@ fn test_arrow_gh_47662() { let err = read_file("ARROW-GH-47662.parquet").unwrap_err(); assert_eq!( err.to_string(), - "External: Parquet argument error: Parquet error: insufficient values read from column - expected: 100, got: 91" + "External: Parquet argument error: Parquet error: Invalid FIXED_LEN_BYTE_ARRAY PLAIN data page payload length: expected 400 bytes (100 values of 4 bytes), got 364" ); } diff --git a/parquet/tests/arrow_reader/legacy_fixed_len_byte_array.rs b/parquet/tests/arrow_reader/legacy_fixed_len_byte_array.rs new file mode 100644 index 0000000000..efd451cb83 --- /dev/null +++ b/parquet/tests/arrow_reader/legacy_fixed_len_byte_array.rs @@ -0,0 +1,262 @@ +// 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. + +//! Compatibility with the malformed ARROW-RS-GH-11261-FLBA-DICT.parquet fixture. +//! See [arrow-rs #11261](https://github.com/apache/arrow-rs/issues/11261) for the defect. +//! The cases and expected values are defined by `ROW_GROUPS` and `expected` below. + +use std::sync::Arc; + +use super::bad_data::bad_data_dir; +use arrow::compute::{cast, concat}; +use arrow_array::{ + Array, ArrayRef, FixedSizeBinaryArray, ListArray, RecordBatch, builder::FixedSizeBinaryBuilder, +}; +use arrow_buffer::OffsetBuffer; +use arrow_schema::{DataType, Field}; +use bytes::Bytes; +use parquet::arrow::arrow_reader::{ + ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowSelection, RowSelector, +}; +use parquet::column::reader::ColumnReader; +use parquet::file::reader::{FileReader, SerializedFileReader}; + +// Every row group has four rows, with a nullable flat column and a list column +// containing nullable leaves. Column chunks have multiple data pages; spill +// cases mix dictionary-index and fallback pages. The final group uses Snappy. +const ROW_GROUPS: &[&str] = &[ + "v1 dictionary", + "v1 plain", + "v1 plain spill", + "v1 delta length", + "v1 delta length spill", + "v1 delta spill", + "v2 dictionary", + "v2 plain", + "v2 plain spill", + "v2 delta length", + "v2 delta length spill", + "v2 delta spill", + #[cfg(feature = "snap")] + "v1 snappy", +]; +const ROWS_PER_GROUP: usize = 4; + +fn fixture() -> Bytes { + let path = bad_data_dir().join("ARROW-RS-GH-11261-FLBA-DICT.parquet"); + // As with other corpus tests, missing data is a setup error, not a reason to skip. + std::fs::read(&path) + .unwrap_or_else(|e| { + panic!( + "failed to read {}: {e}; run `git submodule update --init`", + path.display() + ) + }) + .into() +} + +fn expected(nested: bool) -> ArrayRef { + // The first value deliberately resembles a BYTE_ARRAY length prefix. + let values = [b"\x04\0\0\0".as_slice(), b"ABCD", b"wxyz"]; + let keys = [ + Some(0), + Some(1), + None, + Some(2), + Some(0), + Some(1), + Some(2), + None, + Some(0), + ]; + let mut builder = FixedSizeBinaryBuilder::new(4); + for key in keys { + match key { + Some(key) => builder.append_value(values[key]).unwrap(), + None => builder.append_null(), + } + } + let leaf: ArrayRef = Arc::new(builder.finish()); + if nested { + Arc::new(ListArray::new( + Arc::new(Field::new("item", DataType::FixedSizeBinary(4), true)), + OffsetBuffer::new(vec![0_i32, 3, 3, 6, 9].into()), + leaf, + None, + )) + } else { + leaf.slice(0, ROWS_PER_GROUP) + } +} + +// Read each case independently, then together to exercise row-group transitions. +fn row_group_selections() -> impl Iterator<Item = Vec<usize>> { + (0..ROW_GROUPS.len()) + .map(|i| vec![i]) + .chain(std::iter::once((0..ROW_GROUPS.len()).collect())) +} + +fn selection(groups: usize, skip: bool) -> RowSelection { + let mut selectors = vec![]; + for _ in 0..groups { + if skip { + selectors.push(RowSelector::skip(1)); + } + selectors.push(RowSelector::select(ROWS_PER_GROUP - usize::from(skip))); + } + RowSelection::from(selectors) +} + +fn assert_batches(batches: &[RecordBatch], groups: &[usize], skip: bool) { + for (column, nested) in [false, true].into_iter().enumerate() { + let expected = expected(nested); + let start = usize::from(skip); + let expected = expected.slice(start, expected.len() - start); + let expected = concat(&vec![expected.as_ref(); groups.len()]).unwrap(); + let arrays: Vec<_> = batches.iter().map(|b| b.column(column).as_ref()).collect(); + let actual = concat(&arrays).unwrap(); + // Dictionary ordering is not prescribed; compare the decoded values. + let actual = cast(actual.as_ref(), expected.data_type()).unwrap(); + assert_eq!( + actual.to_data(), + expected.to_data(), + "row groups {groups:?}, column {column}, skip {skip}" + ); + } +} + +#[test] +fn read_legacy_fixed_len_byte_array() { + let data = fixture(); + for groups in row_group_selections() { + for skip_metadata in [false, true] { + for batch_size in [1, 3, 1024] { + for skip in [false, true] { + let options = ArrowReaderOptions::new().with_skip_arrow_metadata(skip_metadata); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + data.clone(), + options, + ) + .unwrap(); + let batches = builder + .with_row_groups(groups.clone()) + .with_batch_size(batch_size) + .with_row_selection(selection(groups.len(), skip)) + .build() + .unwrap() + .collect::<Result<Vec<_>, _>>() + .unwrap(); + assert_batches(&batches, &groups, skip); + } + } + } + } +} + +fn physical_values(array: &dyn Array, out: &mut Vec<Vec<u8>>) { + for i in 0..array.len() { + if array.is_null(i) { + continue; + } + if let Some(list) = array.as_any().downcast_ref::<ListArray>() { + physical_values(list.value(i).as_ref(), out); + } else { + let fixed = array + .as_any() + .downcast_ref::<FixedSizeBinaryArray>() + .unwrap(); + out.push(fixed.value(i).to_vec()); + } + } +} + +#[test] +fn read_legacy_fixed_len_byte_array_typed() { + let file = SerializedFileReader::new(fixture()).unwrap(); + assert_eq!(file.num_row_groups(), 13); + assert_eq!(file.metadata().file_metadata().num_rows(), 52); + for (index, name) in ROW_GROUPS.iter().enumerate() { + let row_group = file.get_row_group(index).unwrap(); + assert_eq!(row_group.metadata().num_rows(), ROWS_PER_GROUP as i64); + assert_eq!(row_group.num_columns(), 2); + for (column, nested) in [false, true].into_iter().enumerate() { + for skip in [false, true] { + let ColumnReader::FixedLenByteArrayColumnReader(mut reader) = + row_group.get_column_reader(column).unwrap() + else { + panic!("{name}: expected FLBA column"); + }; + if skip { + assert_eq!(reader.skip_records(1).unwrap(), 1); + } + let (mut values, mut def, mut rep) = (vec![], vec![], vec![]); + while reader + .read_records(1, Some(&mut def), Some(&mut rep), &mut values) + .unwrap() + .0 + != 0 + {} + let expected = expected(nested); + let start = usize::from(skip); + let mut physical = vec![]; + physical_values( + expected.slice(start, expected.len() - start).as_ref(), + &mut physical, + ); + assert_eq!( + values.iter().map(|v| v.data()).collect::<Vec<_>>(), + physical, + "{name}, column {column}, skip {skip}" + ); + } + } + } +} + +#[cfg(feature = "async")] +#[tokio::test] +async fn read_legacy_fixed_len_byte_array_async() { + use futures::TryStreamExt; + use parquet::arrow::async_reader::ParquetRecordBatchStreamBuilder; + + let data = fixture(); + for groups in row_group_selections() { + for skip_metadata in [false, true] { + for batch_size in [1, 3, 1024] { + for skip in [false, true] { + let options = ArrowReaderOptions::new().with_skip_arrow_metadata(skip_metadata); + let builder = ParquetRecordBatchStreamBuilder::new_with_options( + std::io::Cursor::new(data.clone()), + options, + ) + .await + .unwrap(); + let batches: Vec<_> = builder + .with_row_groups(groups.clone()) + .with_batch_size(batch_size) + .with_row_selection(selection(groups.len(), skip)) + .build() + .unwrap() + .try_collect() + .await + .unwrap(); + assert_batches(&batches, &groups, skip); + } + } + } + } +} diff --git a/parquet/tests/arrow_reader/mod.rs b/parquet/tests/arrow_reader/mod.rs index d4a4460c4a..428905a430 100644 --- a/parquet/tests/arrow_reader/mod.rs +++ b/parquet/tests/arrow_reader/mod.rs @@ -49,6 +49,7 @@ mod int96_stats_roundtrip; mod invalid_utf8; mod io; mod large_string_overflow; +mod legacy_fixed_len_byte_array; mod page_index; mod parquet_testing; mod partial_offset_index; diff --git a/parquet/tests/arrow_writer/roundtrip.rs b/parquet/tests/arrow_writer/roundtrip.rs index 9a122fabe9..fc64096900 100644 --- a/parquet/tests/arrow_writer/roundtrip.rs +++ b/parquet/tests/arrow_writer/roundtrip.rs @@ -18,7 +18,7 @@ //! Round-trip tests for Arrow data written to Parquet. use super::roundtrip_helpers::{ - RoundTripTest, SMALL_SIZE, required_and_optional, roundtrip, values_required, + RoundTripTest, SMALL_SIZE, required_and_optional, roundtrip, roundtrip_opts, values_required, }; use std::collections::HashMap; @@ -34,7 +34,7 @@ use arrow_array::cast::AsArray; use arrow_array::types::{ ArrowDictionaryKeyType, Date32Type, Date64Type, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType, Float16Type, Int8Type, Int16Type, Int32Type, Int64Type, - Time32MillisecondType, Time64MicrosecondType, UInt8Type, UInt16Type, UInt32Type, + Time32MillisecondType, Time64MicrosecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; use arrow_array::{ Array, ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array, @@ -59,9 +59,11 @@ use parquet::arrow::arrow_reader::{ ArrowReaderOptions, ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder, }; use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY}; -use parquet::basic::Type as PhysicalType; +use parquet::basic::{Encoding, Type as PhysicalType}; +use parquet::column::reader::ColumnReader; use parquet::errors::Result; use parquet::file::properties::WriterProperties; +use parquet::file::reader::{FileReader, SerializedFileReader}; use tempfile::tempfile; #[test] @@ -1023,19 +1025,67 @@ fn test_fixed_size_binary_in_dict() { let data = DictionaryArray::<K>::new(keys, Arc::new(values)); let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(data)]).unwrap(); - roundtrip(batch, None); + for file in roundtrip(batch, None) { + let parquet = SerializedFileReader::new(file).unwrap(); + let row_group = parquet.get_row_group(0).unwrap(); + let ColumnReader::FixedLenByteArrayColumnReader(mut reader) = + row_group.get_column_reader(0).unwrap() + else { + panic!("dictionary FSB was not written as FIXED_LEN_BYTE_ARRAY") + }; + let mut values = Vec::new(); + reader.read_records(3, None, None, &mut values).unwrap(); + assert_eq!( + values.iter().map(|value| value.data()).collect::<Vec<_>>(), + [b"\0\0\0\0".as_slice(), b"\0\0\0\0", b"\x01\x01\x01\x01"] + ); + } } test_fixed_size_binary_in_dict_inner::<UInt8Type>(); test_fixed_size_binary_in_dict_inner::<UInt16Type>(); test_fixed_size_binary_in_dict_inner::<UInt32Type>(); - test_fixed_size_binary_in_dict_inner::<UInt16Type>(); + test_fixed_size_binary_in_dict_inner::<UInt64Type>(); test_fixed_size_binary_in_dict_inner::<Int8Type>(); test_fixed_size_binary_in_dict_inner::<Int16Type>(); test_fixed_size_binary_in_dict_inner::<Int32Type>(); test_fixed_size_binary_in_dict_inner::<Int64Type>(); } +#[test] +fn test_fixed_size_binary_in_dict_dictionary_disabled() { + let field = Field::new( + "a", + ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt8), + Box::new(ArrowDataType::FixedSizeBinary(4)), + ), + true, + ); + let schema = Arc::new(Schema::new(vec![field])); + let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(0)]); + let values = + FixedSizeBinaryArray::try_from_iter(vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter()) + .unwrap(); + let data = DictionaryArray::<UInt8Type>::new(keys, Arc::new(values)); + let batch = RecordBatch::try_new(schema, vec![Arc::new(data)]).unwrap(); + for encoding in [ + Encoding::PLAIN, + Encoding::DELTA_BYTE_ARRAY, + Encoding::BYTE_STREAM_SPLIT, + ] { + for dictionary in [false, true] { + let props = WriterProperties::builder() + .set_dictionary_enabled(dictionary) + .set_dictionary_page_size_limit(1) + .set_write_batch_size(1) + .set_encoding(encoding) + .build(); + roundtrip_opts(&batch, props); + } + } +} + #[test] fn test_empty_dict() { let struct_fields = Fields::from(vec