etseidl commented on code in PR #11262:
URL: https://github.com/apache/arrow-rs/pull/11262#discussion_r4200564217
##########
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:
So this part of the test is adding phantom levels to the count and then
expecting the legacy checks to catch it...but this pretty much requires walking
the level data, negating any speed up from avoiding that walk. Wouldn't this
mismatch be detected later if more than a single record was read? Just
wondering if this is strictly necessary.
--
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]
