HippoBaro commented on code in PR #11262:
URL: https://github.com/apache/arrow-rs/pull/11262#discussion_r4231155658
##########
parquet/src/column/reader.rs:
##########
@@ -639,6 +695,268 @@ mod tests {
use crate::util::test_common::page_util::InMemoryPageReader;
use crate::util::test_common::rand_gen::make_pages;
+ #[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;
+ }
+ let err = make_reader(vec![page])
+ .read_records(1, Some(&mut vec![]), Some(&mut vec![]),
&mut vec![])
+ .unwrap_err()
+ .to_string();
+ assert!(err.contains("definition levels"), "{err}");
+ }
Review Comment:
Sorry I had missed that one! You’re right: the first-row failure assertion
imposed unnecessary eager validation. The latest version changes this test to
detect the mismatch during page consumption and adds an explicit
first-read-succeeds/later-read-fails regression.
I updated the PR description to match the current behavior 🙇
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
