This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new a2bb4c90dc feat(parquet): selective null padding for list child
readers (#9848)
a2bb4c90dc is described below
commit a2bb4c90dc2c28e261fffc1ff777bbbe2df572cc
Author: Hippolyte Barraud <[email protected]>
AuthorDate: Thu Jul 2 06:26:39 2026 -0400
feat(parquet): selective null padding for list child readers (#9848)
# 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.
-->
- Contributes to #9731
- Depend on #9847
- Depend on #9846
# 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.
-->
Parquet list decoding currently materializes padding for null or empty
parent lists and then copies the child array to filter that padding back
out. This is expensive for nested list columns, especially sparse lists
and fixed-width children, where memory can scale with decoded levels
instead of actual emitted child values.
# 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.
-->
This PR makes list child readers emit compact child arrays directly by
pushing selective null padding into the leaf `RecordReader`. It also
builds definition-level validity bitmaps word-at-a-time, sizes child
buffers after levels are decoded, and adds list runtime and peak-memory
benchmarks across element types and null densities.
# Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
Extensive test coverage for the new logic through existing list reader
tests, which now exercise the production `PrimitiveArrayReader` path
with in-memory Parquet pages (see #9846), and
`BooleanBufferBuilder::append_word` has targeted unit coverage.
Benchmarks results:
```
Name Before After
Delta
ListArray/StringList/no NULLs 7.3395 ms 5.8001 ms
(-21.0%)
ListArray/StringList/half NULLs 4.1255 ms 3.4274 ms
(-16.9%)
ListArray/Int32List/90pct NULLs 1.0366 ms 975.63 us
(-5.9%)
ListArray/Fixed32List/no NULLs 4.9298 ms 3.3137 ms
(-32.8%)
ListArray/Fixed32List/half NULLs 2.8998 ms 2.7258 ms
(-6.0%)
ListArray/Fixed32List/90pct NULLs 1.0556 ms 995.82 us
(-5.7%)
ListArray/Fixed32List/99pct NULLs 508.65 us 467.16 us
(-8.2%)
Name Before After
Delta
ListArray_peak_memory/Int32List/no NULLs 836.51 KiB 574.79 KiB
(-31.3%)
ListArray_peak_memory/Int32List/half NULLs 482.01 KiB 336.29 KiB
(-30.2%)
ListArray_peak_memory/Int32List/90pct NULLs 271.95 KiB 175.32 KiB
(-35.5%)
ListArray_peak_memory/Int32List/99pct NULLs 217.96 KiB 120.04 KiB
(-44.9%)
ListArray_peak_memory/DoubleList/no NULLs 1.2399 MiB 715.31 KiB
(-43.7%)
ListArray_peak_memory/DoubleList/half NULLs 753.66 KiB 400.39 KiB
(-46.9%)
ListArray_peak_memory/DoubleList/90pct NULLs 380.62 KiB 190.89 KiB
(-49.8%)
ListArray_peak_memory/DoubleList/99pct NULLs 315.21 KiB 121.61 KiB
(-61.4%)
ListArray_peak_memory/Fixed32List/no NULLs 3.8031 MiB 1.5760 MiB
(-58.6%)
ListArray_peak_memory/Fixed32List/half NULLs 2.1710 MiB 849.94 KiB
(-61.8%)
ListArray_peak_memory/Fixed32List/90pct NULLs 1.0017 MiB 277.35 KiB
(-73.0%)
ListArray_peak_memory/Fixed32List/99pct NULLs 898.69 KiB 130.93 KiB
(-85.4%)
ListArray_peak_memory/StringList/no NULLs 3.7925 MiB 2.4715 MiB
(-34.8%)
ListArray_peak_memory/StringList/half NULLs 1.2541 MiB 772.94 KiB
(-39.8%)
ListArray_peak_memory/StringList/90pct NULLs 296.63 KiB 188.96 KiB
(-36.3%)
ListArray_peak_memory/StringList/99pct NULLs 226.75 KiB 120.37 KiB
(-46.9%)
Name Before After
Delta
ListArray_allocated_bytes/Int32List/no NULLs 10.458 MiB 6.8018 MiB
(-35.0%)
ListArray_allocated_bytes/Int32List/half NULLs 5.9797 MiB 4.0127 MiB
(-32.9%)
ListArray_allocated_bytes/Int32List/90pct NULLs 2.9985 MiB 1.8210 MiB
(-39.3%)
ListArray_allocated_bytes/Int32List/99pct NULLs 2.5579 MiB 1.3733 MiB
(-46.3%)
ListArray_allocated_bytes/DoubleList/no NULLs 16.083 MiB 8.6546 MiB
(-46.2%)
ListArray_allocated_bytes/DoubleList/half NULLs 8.9134 MiB 4.8497 MiB
(-45.6%)
ListArray_allocated_bytes/DoubleList/90pct NULLs 4.3656 MiB 2.0179 MiB
(-53.8%)
ListArray_allocated_bytes/DoubleList/99pct NULLs 3.7482 MiB 1.3903 MiB
(-62.9%)
ListArray_allocated_bytes/Fixed32List/no NULLs 49.441 MiB 19.505 MiB
(-60.5%)
ListArray_allocated_bytes/Fixed32List/half NULLs 26.846 MiB 10.459 MiB
(-61.0%)
ListArray_allocated_bytes/Fixed32List/90pct NULLs 12.483 MiB 3.1127 MiB
(-75.1%)
ListArray_allocated_bytes/Fixed32List/99pct NULLs 10.895 MiB 1.4980 MiB
(-86.3%)
ListArray_allocated_bytes/StringList/no NULLs 47.519 MiB 21.743 MiB
(-54.2%)
ListArray_allocated_bytes/StringList/half NULLs 19.097 MiB 10.478 MiB
(-45.1%)
ListArray_allocated_bytes/StringList/90pct NULLs 3.4203 MiB 2.1165 MiB
(-38.1%)
ListArray_allocated_bytes/StringList/99pct NULLs 2.6424 MiB 1.3777 MiB
(-47.9%)
```
# Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
None.
---------
Signed-off-by: Hippolyte Barraud <[email protected]>
---
arrow-buffer/src/builder/boolean.rs | 136 ++++++++
parquet/benches/arrow_reader.rs | 34 +-
parquet/benches/arrow_reader_peak_memory.rs | 42 ++-
parquet/src/arrow/array_reader/builder.rs | 86 ++++-
parquet/src/arrow/array_reader/byte_array.rs | 20 +-
.../arrow/array_reader/byte_array_dictionary.rs | 19 +-
parquet/src/arrow/array_reader/byte_view_array.rs | 18 +-
.../src/arrow/array_reader/cached_array_reader.rs | 4 +
.../src/arrow/array_reader/fixed_len_byte_array.rs | 56 +++-
.../arrow/array_reader/fixed_size_list_array.rs | 266 ++++++++-------
parquet/src/arrow/array_reader/list_array.rs | 367 +++++++++++++++------
parquet/src/arrow/array_reader/list_view_array.rs | 21 +-
parquet/src/arrow/array_reader/map_array.rs | 15 +-
parquet/src/arrow/array_reader/mod.rs | 13 +
parquet/src/arrow/array_reader/null_array.rs | 16 +-
parquet/src/arrow/array_reader/primitive_array.rs | 27 +-
parquet/src/arrow/array_reader/struct_array.rs | 142 ++++----
parquet/src/arrow/array_reader/test_util.rs | 19 +-
parquet/src/arrow/arrow_reader/mod.rs | 8 +-
parquet/src/arrow/arrow_reader/read_plan.rs | 4 +-
parquet/src/arrow/buffer/dictionary_buffer.rs | 7 +
parquet/src/arrow/buffer/offset_buffer.rs | 4 +
parquet/src/arrow/buffer/view_buffer.rs | 4 +
parquet/src/arrow/record_reader/buffer.rs | 15 +-
.../src/arrow/record_reader/definition_levels.rs | 99 ++++++
parquet/src/arrow/record_reader/mod.rs | 300 +++++++++++++++--
parquet/src/column/reader.rs | 27 +-
parquet/src/util/bit_util.rs | 76 +++++
28 files changed, 1440 insertions(+), 405 deletions(-)
diff --git a/arrow-buffer/src/builder/boolean.rs
b/arrow-buffer/src/builder/boolean.rs
index 956f4a3e22..4f84b0102b 100644
--- a/arrow-buffer/src/builder/boolean.rs
+++ b/arrow-buffer/src/builder/boolean.rs
@@ -186,6 +186,46 @@ impl BooleanBufferBuilder {
}
}
+ /// Appends the low `count` bits from `word` into the buffer.
+ ///
+ /// `word` is treated as a packed LSB-first bitmap. Only the lowest
+ /// `count` bits are appended; higher bits are ignored.
+ ///
+ /// This is significantly faster than calling [`Self::append`] in a
+ /// loop when the caller already has bits packed into a `u64`.
+ #[inline]
+ pub fn append_word(&mut self, word: u64, count: usize) {
+ debug_assert!(count <= 64);
+ let mask = (u64::MAX >> ((64 - count) & 63)) * ((count != 0) as u64);
+ let word = word & mask;
+
+ let new_len = self.len + count;
+ let new_len_bytes = bit_util::ceil(new_len, 8);
+ if new_len_bytes > self.buffer.len() {
+ self.buffer.resize(new_len_bytes, 0);
+ }
+
+ let bit_offset = self.len & 7;
+ let byte_start = self.len / 8;
+ let buf = self.buffer.as_slice_mut();
+
+ // Shift the word to align with the bit offset within the
+ // current byte, then OR it into the buffer. The shift merges
+ // correctly with any existing bits in the partial trailing byte.
+ let shifted = word << bit_offset;
+ let shifted_bytes = shifted.to_le_bytes();
+ let bytes_to_write = bit_util::ceil(count + bit_offset, 8).min(8);
+ for i in 0..bytes_to_write {
+ buf[byte_start + i] |= shifted_bytes[i];
+ }
+ // When bit_offset > 0 the shift can overflow into a 9th byte.
+ if bit_offset > 0 && count + bit_offset > 64 {
+ buf[byte_start + 8] |= (word >> (64 - bit_offset)) as u8;
+ }
+
+ self.len = new_len;
+ }
+
/// Appends n `additional` bits of value `v` into the buffer
#[inline]
pub fn append_n(&mut self, additional: usize, v: bool) {
@@ -619,6 +659,102 @@ mod tests {
}
}
+ /// Helper: build via append_word and verify against bit-by-bit append.
+ fn check_append_word(initial_bits: usize, word: u64, count: usize) {
+ let mut got = BooleanBufferBuilder::new(0);
+ let mut expected = BooleanBufferBuilder::new(0);
+ got.append_n(initial_bits, true);
+ expected.append_n(initial_bits, true);
+ got.append_word(word, count);
+ for i in 0..count {
+ expected.append(word & (1 << i) != 0);
+ }
+ assert_eq!(got.len(), expected.len());
+ assert_eq!(got.finish(), expected.finish());
+ }
+
+ #[test]
+ fn test_append_word_zero_count() {
+ check_append_word(0, u64::MAX, 0);
+ check_append_word(3, u64::MAX, 0);
+ }
+
+ #[test]
+ fn test_append_word_aligned() {
+ for count in [1, 5, 8, 17, 64] {
+ check_append_word(0, 0xDEAD_BEEF_CAFE_BABE, count);
+ }
+ }
+
+ #[test]
+ fn test_append_word_unaligned() {
+ for offset in 1..=7 {
+ check_append_word(offset, 0xDEAD_BEEF_CAFE_BABE, 13);
+ }
+ }
+
+ #[test]
+ fn test_append_word_overflow_9th_byte() {
+ check_append_word(3, u64::MAX, 64);
+ check_append_word(7, 0xA5A5_A5A5_A5A5_A5A5, 64);
+ }
+
+ #[test]
+ fn test_append_word_small_counts() {
+ check_append_word(0, 0b1, 1);
+ check_append_word(0, 0b1010101, 7);
+ check_append_word(3, 0b1, 1);
+ check_append_word(3, 0b1010101, 7);
+ }
+
+ #[test]
+ fn test_append_word_ignores_high_bits_before_later_appends() {
+ let mut builder = BooleanBufferBuilder::new(0);
+ builder.append_word(0b10, 1);
+ builder.append(false);
+
+ let finished = builder.finish();
+ assert_eq!(finished.len(), 2);
+ assert!(!finished.value(0));
+ assert!(!finished.value(1));
+ }
+
+ #[test]
+ fn test_append_word_full_word() {
+ check_append_word(0, u64::MAX, 64);
+ check_append_word(0, 0, 64);
+ }
+
+ #[test]
+ fn test_append_word_sequential() {
+ let mut got = BooleanBufferBuilder::new(0);
+ let mut expected = BooleanBufferBuilder::new(0);
+ for (word, count) in [(0b1010u64, 4), (0b111u64, 3), (0u64, 5),
(u64::MAX, 64)] {
+ got.append_word(word, count);
+ for i in 0..count {
+ expected.append(word & (1 << i) != 0);
+ }
+ }
+ assert_eq!(got.finish(), expected.finish());
+ }
+
+ #[test]
+ fn test_append_word_mixed_with_append() {
+ let mut got = BooleanBufferBuilder::new(0);
+ let mut expected = BooleanBufferBuilder::new(0);
+ got.append(true);
+ expected.append(true);
+ got.append_word(0b1100, 4);
+ for i in 0..4 {
+ expected.append(0b1100u64 & (1 << i) != 0);
+ }
+ got.append(false);
+ expected.append(false);
+ got.append_word(0xFF, 8);
+ expected.append_n(8, true);
+ assert_eq!(got.finish(), expected.finish());
+ }
+
#[test]
fn test_extend_misaligned_end() {
for len in 1..130 {
diff --git a/parquet/benches/arrow_reader.rs b/parquet/benches/arrow_reader.rs
index 7a6ed4ec9f..6e91a04f4e 100644
--- a/parquet/benches/arrow_reader.rs
+++ b/parquet/benches/arrow_reader.rs
@@ -785,12 +785,15 @@ fn create_int32_list_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ Some(2),
)
.unwrap(),
) as Box<dyn ArrayReader>;
let field = Field::new_list_field(DataType::Int32, true);
let data_type = DataType::List(Arc::new(field));
- Box::new(ListArrayReader::<i32>::new(items, data_type, 2, 1, true))
+ Box::new(ListArrayReader::<i32>::new(
+ items, data_type, 2, 1, true, None,
+ ))
}
const FIXED_BYTE_LEN: usize = 32;
@@ -858,11 +861,14 @@ fn create_fixed32_list_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ Some(2),
)
.unwrap();
let field = Field::new_list_field(DataType::FixedSizeBinary(FIXED_BYTE_LEN
as i32), true);
let data_type = DataType::List(Arc::new(field));
- Box::new(ListArrayReader::<i32>::new(items, data_type, 2, 1, true))
+ Box::new(ListArrayReader::<i32>::new(
+ items, data_type, 2, 1, true, None,
+ ))
}
fn bench_array_reader(mut array_reader: Box<dyn ArrayReader>) -> usize {
@@ -912,6 +918,7 @@ fn create_primitive_array_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap();
Box::new(reader)
@@ -922,6 +929,7 @@ fn create_primitive_array_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap();
Box::new(reader)
@@ -932,6 +940,7 @@ fn create_primitive_array_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap();
Box::new(reader)
@@ -951,6 +960,7 @@ fn create_f16_by_bytes_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap(),
_ => unimplemented!(),
@@ -968,6 +978,7 @@ fn create_decimal_by_bytes_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap(),
Type::FIXED_LEN_BYTE_ARRAY => make_fixed_len_byte_array_reader(
@@ -975,6 +986,7 @@ fn create_decimal_by_bytes_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap(),
_ => unimplemented!(),
@@ -990,6 +1002,7 @@ fn create_fixed_len_byte_array_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap()
}
@@ -1003,6 +1016,7 @@ fn create_byte_array_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap()
}
@@ -1016,6 +1030,7 @@ fn create_byte_view_array_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap()
}
@@ -1029,6 +1044,7 @@ fn create_string_view_byte_array_reader(
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap()
}
@@ -1045,6 +1061,7 @@ fn create_string_byte_array_dictionary_reader(
column_desc,
Some(arrow_type),
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap()
}
@@ -1053,10 +1070,19 @@ fn create_string_list_reader(
page_iterator: impl PageIterator + 'static,
column_desc: ColumnDescPtr,
) -> Box<dyn ArrayReader> {
- let items = create_byte_array_reader(page_iterator, column_desc);
+ let items = make_byte_array_reader(
+ Box::new(page_iterator),
+ column_desc,
+ None,
+ DEFAULT_BATCH_SIZE,
+ Some(2),
+ )
+ .unwrap();
let field = Field::new_list_field(DataType::Utf8, true);
let data_type = DataType::List(Arc::new(field));
- Box::new(ListArrayReader::<i32>::new(items, data_type, 2, 1, true))
+ Box::new(ListArrayReader::<i32>::new(
+ items, data_type, 2, 1, true, None,
+ ))
}
fn bench_byte_decimal<T>(
diff --git a/parquet/benches/arrow_reader_peak_memory.rs
b/parquet/benches/arrow_reader_peak_memory.rs
index 6486bf2da0..2a717ec4a1 100644
--- a/parquet/benches/arrow_reader_peak_memory.rs
+++ b/parquet/benches/arrow_reader_peak_memory.rs
@@ -443,12 +443,20 @@ fn int32_list_wrapper(
column_desc: ColumnDescPtr,
) -> Box<dyn ArrayReader> {
let child: Box<dyn ArrayReader> = Box::new(
- PrimitiveArrayReader::<Int32Type>::new(pages, column_desc, None,
DEFAULT_BATCH_SIZE)
- .unwrap(),
+ PrimitiveArrayReader::<Int32Type>::new(
+ pages,
+ column_desc,
+ None,
+ DEFAULT_BATCH_SIZE,
+ Some(2),
+ )
+ .unwrap(),
);
let field = Field::new_list_field(DataType::Int32, true);
let data_type = DataType::List(Arc::new(field));
- Box::new(ListArrayReader::<i32>::new(child, data_type, 2, 1, true))
+ Box::new(ListArrayReader::<i32>::new(
+ child, data_type, 2, 1, true, None,
+ ))
}
fn double_list_wrapper(
@@ -456,12 +464,20 @@ fn double_list_wrapper(
column_desc: ColumnDescPtr,
) -> Box<dyn ArrayReader> {
let child: Box<dyn ArrayReader> = Box::new(
- PrimitiveArrayReader::<DoubleType>::new(pages, column_desc, None,
DEFAULT_BATCH_SIZE)
- .unwrap(),
+ PrimitiveArrayReader::<DoubleType>::new(
+ pages,
+ column_desc,
+ None,
+ DEFAULT_BATCH_SIZE,
+ Some(2),
+ )
+ .unwrap(),
);
let field = Field::new_list_field(DataType::Float64, true);
let data_type = DataType::List(Arc::new(field));
- Box::new(ListArrayReader::<i32>::new(child, data_type, 2, 1, true))
+ Box::new(ListArrayReader::<i32>::new(
+ child, data_type, 2, 1, true, None,
+ ))
}
fn fixed32_list_wrapper(
@@ -469,20 +485,26 @@ fn fixed32_list_wrapper(
column_desc: ColumnDescPtr,
) -> Box<dyn ArrayReader> {
let child =
- make_fixed_len_byte_array_reader(pages, column_desc, None,
DEFAULT_BATCH_SIZE).unwrap();
+ make_fixed_len_byte_array_reader(pages, column_desc, None,
DEFAULT_BATCH_SIZE, Some(2))
+ .unwrap();
let field = Field::new_list_field(DataType::FixedSizeBinary(FIXED_BYTE_LEN
as i32), true);
let data_type = DataType::List(Arc::new(field));
- Box::new(ListArrayReader::<i32>::new(child, data_type, 2, 1, true))
+ Box::new(ListArrayReader::<i32>::new(
+ child, data_type, 2, 1, true, None,
+ ))
}
fn string_list_wrapper(
pages: Box<dyn PageIterator>,
column_desc: ColumnDescPtr,
) -> Box<dyn ArrayReader> {
- let child = make_byte_array_reader(pages, column_desc, None,
DEFAULT_BATCH_SIZE).unwrap();
+ let child =
+ make_byte_array_reader(pages, column_desc, None, DEFAULT_BATCH_SIZE,
Some(2)).unwrap();
let field = Field::new_list_field(DataType::Utf8, true);
let data_type = DataType::List(Arc::new(field));
- Box::new(ListArrayReader::<i32>::new(child, data_type, 2, 1, true))
+ Box::new(ListArrayReader::<i32>::new(
+ child, data_type, 2, 1, true, None,
+ ))
}
fn add_benches<M: Measurement>(c: &mut Criterion<M>, measurement_name: &str) {
diff --git a/parquet/src/arrow/array_reader/builder.rs
b/parquet/src/arrow/array_reader/builder.rs
index 827ebe39ce..f89a789647 100644
--- a/parquet/src/arrow/array_reader/builder.rs
+++ b/parquet/src/arrow/array_reader/builder.rs
@@ -97,7 +97,7 @@ pub struct ArrayReaderBuilder<'a> {
parquet_metadata: Option<&'a ParquetMetaData>,
/// metrics
metrics: &'a ArrowReaderMetrics,
- /// Batch size for pre-allocating internal buffers
+ /// Batch size hint for pre-allocating internal buffers (see
[`Self::with_batch_size`])
batch_size: usize,
}
@@ -112,6 +112,11 @@ struct ReaderArgs<'a> {
field: &'a ParquetField,
/// Which leaf columns are being read.
mask: &'a ProjectionMask,
+ /// Definition-level threshold at or above which the nearest enclosing
+ /// list/map keeps a leaf value. Used to push selective null padding down
+ /// into leaf readers so they emit compact child arrays. `None` at the top
+ /// level (no enclosing list/map), which disables selective padding.
+ padding_threshold: Option<i16>,
}
impl<'a> ReaderArgs<'a> {
@@ -121,6 +126,17 @@ impl<'a> ReaderArgs<'a> {
fn with_field(self, field: &'a ParquetField) -> Self {
Self { field, ..self }
}
+
+ /// Returns a copy of these arguments with a different `padding_threshold`.
+ ///
+ /// Used by list/map readers to set the threshold their children should
+ /// use, without disturbing the threshold applied to the list/map itself.
+ fn with_padding_threshold(self, padding_threshold: Option<i16>) -> Self {
+ Self {
+ padding_threshold,
+ ..self
+ }
+ }
}
impl<'a> ArrayReaderBuilder<'a> {
@@ -135,9 +151,14 @@ impl<'a> ArrayReaderBuilder<'a> {
}
}
- /// Set the batch size used to pre-allocate internal buffers.
+ /// Set the batch size hint used to pre-allocate internal buffers.
///
- /// This avoids reallocations when reading the first batch of data.
+ /// This is a hint, not a hard capacity contract. Readers without
+ /// selective null padding reserve their value buffers up front from this
+ /// size to avoid reallocations when reading the first batch. Readers that
+ /// emit compact (selectively padded) child arrays instead size their
+ /// buffers from the number of values actually decoded, and may allocate
+ /// less than the batch size.
pub fn with_batch_size(mut self, batch_size: usize) -> Self {
self.batch_size = batch_size;
self
@@ -162,7 +183,14 @@ impl<'a> ArrayReaderBuilder<'a> {
mask: &ProjectionMask,
) -> Result<Box<dyn ArrayReader>> {
let reader = field
- .and_then(|field| self.build_reader(ReaderArgs { field, mask
}).transpose())
+ .and_then(|field| {
+ self.build_reader(ReaderArgs {
+ field,
+ mask,
+ padding_threshold: None,
+ })
+ .transpose()
+ })
.transpose()?
.unwrap_or_else(|| make_empty_array_reader(self.num_rows()));
@@ -246,11 +274,16 @@ impl<'a> ArrayReaderBuilder<'a> {
/// Build array reader for map type.
fn build_map_reader(&self, args: ReaderArgs<'_>) -> Result<Option<Box<dyn
ArrayReader>>> {
let field = args.field;
+ let padding_threshold = args.padding_threshold;
let children = field.children().unwrap();
assert_eq!(children.len(), 2);
- let key_reader = self.build_reader(args.with_field(&children[0]))?;
- let value_reader = self.build_reader(args.with_field(&children[1]))?;
+ // The map is internally List(Struct(key, value)). The key/value
+ // readers are children of the struct inside the list, so their
+ // padding threshold is the map's (= list's) def_level.
+ let child_args = args.with_padding_threshold(Some(field.def_level));
+ let key_reader =
self.build_reader(child_args.with_field(&children[0]))?;
+ let value_reader =
self.build_reader(child_args.with_field(&children[1]))?;
match (key_reader, value_reader) {
(Some(key_reader), Some(value_reader)) => {
@@ -282,6 +315,7 @@ impl<'a> ArrayReaderBuilder<'a> {
field.def_level,
field.rep_level,
field.nullable,
+ padding_threshold,
))))
}
(None, None) => Ok(None),
@@ -294,10 +328,14 @@ impl<'a> ArrayReaderBuilder<'a> {
/// Build array reader for list type.
fn build_list_reader(&self, args: ReaderArgs<'_>) -> Result<Option<Box<dyn
ArrayReader>>> {
let field = args.field;
+ let padding_threshold = args.padding_threshold;
let children = field.children().unwrap();
assert_eq!(children.len(), 1);
- let reader = match self.build_reader(args.with_field(&children[0]))? {
+ // Build the child with this list's def_level as its padding threshold.
+ // Pass the received padding_threshold as this list's parent_threshold.
+ let child_args = args.with_padding_threshold(Some(field.def_level));
+ let reader = match
self.build_reader(child_args.with_field(&children[0]))? {
Some(item_reader) => {
// Need to retrieve underlying data type to handle projection
let item_type = item_reader.get_data_type().clone();
@@ -311,6 +349,7 @@ impl<'a> ArrayReaderBuilder<'a> {
field.def_level,
field.rep_level,
field.nullable,
+ padding_threshold,
))
}
DataType::LargeList(f) => {
@@ -323,6 +362,7 @@ impl<'a> ArrayReaderBuilder<'a> {
field.def_level,
field.rep_level,
field.nullable,
+ padding_threshold,
))
}
DataType::ListView(f) => {
@@ -335,6 +375,7 @@ impl<'a> ArrayReaderBuilder<'a> {
field.def_level,
field.rep_level,
field.nullable,
+ padding_threshold,
))
}
DataType::LargeListView(f) => {
@@ -347,6 +388,7 @@ impl<'a> ArrayReaderBuilder<'a> {
field.def_level,
field.rep_level,
field.nullable,
+ padding_threshold,
))
}
_ => unreachable!(),
@@ -364,10 +406,13 @@ impl<'a> ArrayReaderBuilder<'a> {
args: ReaderArgs<'_>,
) -> Result<Option<Box<dyn ArrayReader>>> {
let field = args.field;
+ let padding_threshold = args.padding_threshold;
let children = field.children().unwrap();
assert_eq!(children.len(), 1);
- let reader = match self.build_reader(args.with_field(&children[0]))? {
+ // Build the child with this list's def_level as its padding threshold.
+ let child_args = args.with_padding_threshold(Some(field.def_level));
+ let reader = match
self.build_reader(child_args.with_field(&children[0]))? {
Some(item_reader) => {
let item_type = item_reader.get_data_type().clone();
let reader = match &field.arrow_type {
@@ -384,6 +429,7 @@ impl<'a> ArrayReaderBuilder<'a> {
field.def_level,
field.rep_level,
field.nullable,
+ padding_threshold,
)) as _
}
_ => unimplemented!(),
@@ -398,6 +444,7 @@ impl<'a> ArrayReaderBuilder<'a> {
/// Creates primitive array reader for each primitive type.
fn build_primitive_reader(&self, args: ReaderArgs<'_>) ->
Result<Option<Box<dyn ArrayReader>>> {
let field = args.field;
+ let padding_threshold = args.padding_threshold;
let (col_idx, primitive_type) = match &field.field_type {
ParquetFieldType::Primitive {
col_idx,
@@ -441,6 +488,7 @@ impl<'a> ArrayReaderBuilder<'a> {
page_iterator,
column_desc,
self.batch_size,
+ padding_threshold,
)?) as _;
return Ok(Some(reader));
}
@@ -451,36 +499,42 @@ impl<'a> ArrayReaderBuilder<'a> {
column_desc,
arrow_type,
self.batch_size,
+ padding_threshold,
)?) as _,
PhysicalType::INT32 =>
Box::new(PrimitiveArrayReader::<Int32Type>::new(
page_iterator,
column_desc,
arrow_type,
self.batch_size,
+ padding_threshold,
)?) as _,
PhysicalType::INT64 =>
Box::new(PrimitiveArrayReader::<Int64Type>::new(
page_iterator,
column_desc,
arrow_type,
self.batch_size,
+ padding_threshold,
)?) as _,
PhysicalType::INT96 =>
Box::new(PrimitiveArrayReader::<Int96Type>::new(
page_iterator,
column_desc,
arrow_type,
self.batch_size,
+ padding_threshold,
)?) as _,
PhysicalType::FLOAT =>
Box::new(PrimitiveArrayReader::<FloatType>::new(
page_iterator,
column_desc,
arrow_type,
self.batch_size,
+ padding_threshold,
)?) as _,
PhysicalType::DOUBLE =>
Box::new(PrimitiveArrayReader::<DoubleType>::new(
page_iterator,
column_desc,
arrow_type,
self.batch_size,
+ padding_threshold,
)?) as _,
PhysicalType::BYTE_ARRAY => match arrow_type {
Some(DataType::Dictionary(_, _)) =>
make_byte_array_dictionary_reader(
@@ -488,16 +542,22 @@ impl<'a> ArrayReaderBuilder<'a> {
column_desc,
arrow_type,
self.batch_size,
+ padding_threshold,
)?,
Some(DataType::Utf8View | DataType::BinaryView) =>
make_byte_view_array_reader(
page_iterator,
column_desc,
arrow_type,
self.batch_size,
+ padding_threshold,
+ )?,
+ _ => make_byte_array_reader(
+ page_iterator,
+ column_desc,
+ arrow_type,
+ self.batch_size,
+ padding_threshold,
)?,
- _ => {
- make_byte_array_reader(page_iterator, column_desc,
arrow_type, self.batch_size)?
- }
},
PhysicalType::FIXED_LEN_BYTE_ARRAY => match arrow_type {
Some(DataType::Dictionary(_, _)) =>
make_byte_array_dictionary_reader(
@@ -505,12 +565,14 @@ impl<'a> ArrayReaderBuilder<'a> {
column_desc,
arrow_type,
self.batch_size,
+ padding_threshold,
)?,
_ => make_fixed_len_byte_array_reader(
page_iterator,
column_desc,
arrow_type,
self.batch_size,
+ padding_threshold,
)?,
},
};
@@ -519,6 +581,7 @@ impl<'a> ArrayReaderBuilder<'a> {
fn build_struct_reader(&self, args: ReaderArgs<'_>) ->
Result<Option<Box<dyn ArrayReader>>> {
let field = args.field;
+ let padding_threshold = args.padding_threshold;
let arrow_fields = match &field.arrow_type {
DataType::Struct(children) => children,
_ => unreachable!(),
@@ -548,6 +611,7 @@ impl<'a> ArrayReaderBuilder<'a> {
field.def_level,
field.rep_level,
field.nullable,
+ padding_threshold,
))))
}
}
diff --git a/parquet/src/arrow/array_reader/byte_array.rs
b/parquet/src/arrow/array_reader/byte_array.rs
index db50313046..4fc32dd803 100644
--- a/parquet/src/arrow/array_reader/byte_array.rs
+++ b/parquet/src/arrow/array_reader/byte_array.rs
@@ -38,14 +38,12 @@ use std::any::Any;
use std::sync::Arc;
/// Returns an [`ArrayReader`] that decodes the provided byte array column
-///
-/// `batch_size` is used to pre-allocate internal buffers,
-/// avoiding reallocations when reading the first batch of data.
pub fn make_byte_array_reader(
pages: Box<dyn PageIterator>,
column_desc: ColumnDescPtr,
arrow_type: Option<ArrowType>,
batch_size: usize,
+ padding_threshold: Option<i16>,
) -> Result<Box<dyn ArrayReader>> {
// Check if Arrow type is specified, else create it from Parquet type
let data_type = match arrow_type {
@@ -60,13 +58,19 @@ pub fn make_byte_array_reader(
| ArrowType::Utf8
| ArrowType::Decimal128(_, _)
| ArrowType::Decimal256(_, _) => {
- let reader = GenericRecordReader::new(column_desc, batch_size);
+ let mut reader = GenericRecordReader::new(column_desc, batch_size);
+ if let Some(threshold) = padding_threshold {
+ reader.set_padding_threshold(threshold);
+ }
Ok(Box::new(ByteArrayReader::<i32>::new(
pages, data_type, reader,
)))
}
ArrowType::LargeUtf8 | ArrowType::LargeBinary => {
- let reader = GenericRecordReader::new(column_desc, batch_size);
+ let mut reader = GenericRecordReader::new(column_desc, batch_size);
+ if let Some(threshold) = padding_threshold {
+ reader.set_padding_threshold(threshold);
+ }
Ok(Box::new(ByteArrayReader::<i64>::new(
pages, data_type, reader,
)))
@@ -118,7 +122,7 @@ impl<I: OffsetSizeTrait> ArrayReader for ByteArrayReader<I>
{
fn consume_batch(&mut self) -> Result<ArrayRef> {
let buffer = self.record_reader.consume_record_data();
- let null_buffer = self.record_reader.consume_bitmap_buffer();
+ let null_buffer = self.record_reader.consume_compact_bitmap();
self.def_levels_buffer = self.record_reader.consume_def_levels();
self.rep_levels_buffer = self.record_reader.consume_rep_levels();
self.record_reader.reset();
@@ -168,6 +172,10 @@ impl<I: OffsetSizeTrait> ArrayReader for
ByteArrayReader<I> {
fn get_rep_levels(&self) -> Option<&[i16]> {
self.rep_levels_buffer.as_deref()
}
+
+ fn max_def_level(&self) -> i16 {
+ self.record_reader.max_def_level()
+ }
}
/// A [`ColumnValueDecoder`] for variable length byte arrays
diff --git a/parquet/src/arrow/array_reader/byte_array_dictionary.rs
b/parquet/src/arrow/array_reader/byte_array_dictionary.rs
index ba14b4d6b0..dfd44a8efa 100644
--- a/parquet/src/arrow/array_reader/byte_array_dictionary.rs
+++ b/parquet/src/arrow/array_reader/byte_array_dictionary.rs
@@ -39,14 +39,17 @@ use crate::util::bit_util::FromBitpacked;
/// A macro to reduce verbosity of [`make_byte_array_dictionary_reader`]
macro_rules! make_reader {
(
- ($pages:expr, $column_desc:expr, $data_type:expr, $batch_size:expr) =>
match ($k:expr, $v:expr) {
+ ($pages:expr, $column_desc:expr, $data_type:expr, $batch_size:expr,
$padding_threshold:expr) => match ($k:expr, $v:expr) {
$(($key_arrow:pat, $value_arrow:pat) => ($key_type:ty,
$value_type:ty),)+
}
) => {
match (($k, $v)) {
$(
($key_arrow, $value_arrow) => {
- let reader = GenericRecordReader::new($column_desc,
$batch_size);
+ let mut reader = GenericRecordReader::new($column_desc,
$batch_size);
+ if let Some(threshold) = $padding_threshold {
+ reader.set_padding_threshold(threshold);
+ }
Ok(Box::new(ByteArrayDictionaryReader::<$key_type,
$value_type>::new(
$pages, $data_type, reader,
)))
@@ -71,14 +74,12 @@ macro_rules! make_reader {
///
/// It is therefore recommended that if `pages` contains data from multiple
column chunks,
/// that the read batch size used is a divisor of the row group size
-///
-/// `batch_size` is used to pre-allocate internal buffers,
-/// avoiding reallocations when reading the first batch of data.
pub fn make_byte_array_dictionary_reader(
pages: Box<dyn PageIterator>,
column_desc: ColumnDescPtr,
arrow_type: Option<ArrowType>,
batch_size: usize,
+ padding_threshold: Option<i16>,
) -> Result<Box<dyn ArrayReader>> {
// Check if Arrow type is specified, else create it from Parquet type
let data_type = match arrow_type {
@@ -91,7 +92,7 @@ pub fn make_byte_array_dictionary_reader(
match &data_type {
ArrowType::Dictionary(key_type, value_type) => {
make_reader! {
- (pages, column_desc, data_type, batch_size) => match
(key_type.as_ref(), value_type.as_ref()) {
+ (pages, column_desc, data_type, batch_size, padding_threshold)
=> match (key_type.as_ref(), value_type.as_ref()) {
(ArrowType::UInt8, ArrowType::Binary | ArrowType::Utf8 |
ArrowType::FixedSizeBinary(_)) => (u8, i32),
(ArrowType::UInt8, ArrowType::LargeBinary |
ArrowType::LargeUtf8) => (u8, i64),
(ArrowType::Int8, ArrowType::Binary | ArrowType::Utf8 |
ArrowType::FixedSizeBinary(_)) => (i8, i32),
@@ -179,7 +180,7 @@ where
}
let buffer = self.record_reader.consume_record_data();
- let null_buffer = self.record_reader.consume_bitmap_buffer();
+ let null_buffer = self.record_reader.consume_compact_bitmap();
let array = buffer.into_array(null_buffer, &self.data_type)?;
self.record_reader.reset();
@@ -197,6 +198,10 @@ where
fn get_rep_levels(&self) -> Option<&[i16]> {
self.rep_levels_buffer.as_deref()
}
+
+ fn max_def_level(&self) -> i16 {
+ self.record_reader.max_def_level()
+ }
}
/// If the data is dictionary encoded decode the key data directly, so that
the dictionary
diff --git a/parquet/src/arrow/array_reader/byte_view_array.rs
b/parquet/src/arrow/array_reader/byte_view_array.rs
index ba78502145..cc1abf7ff6 100644
--- a/parquet/src/arrow/array_reader/byte_view_array.rs
+++ b/parquet/src/arrow/array_reader/byte_view_array.rs
@@ -37,13 +37,16 @@ use std::any::Any;
/// Returns an [`ArrayReader`] that decodes the provided byte array column to
view types.
///
-/// `batch_size` is used to pre-allocate internal buffers,
-/// avoiding reallocations when reading the first batch of data.
+/// `batch_size` is a hint used to pre-allocate internal buffers and avoid
+/// reallocations when reading the first batch. It is not a hard capacity
+/// contract: with selective null padding (`padding_threshold`), value buffers
+/// are sized from the number of values actually decoded.
pub fn make_byte_view_array_reader(
pages: Box<dyn PageIterator>,
column_desc: ColumnDescPtr,
arrow_type: Option<ArrowType>,
batch_size: usize,
+ padding_threshold: Option<i16>,
) -> Result<Box<dyn ArrayReader>> {
// Check if Arrow type is specified, else create it from Parquet type
let data_type = match arrow_type {
@@ -56,7 +59,10 @@ pub fn make_byte_view_array_reader(
match data_type {
ArrowType::BinaryView | ArrowType::Utf8View => {
- let reader = GenericRecordReader::new(column_desc, batch_size);
+ let mut reader = GenericRecordReader::new(column_desc, batch_size);
+ if let Some(threshold) = padding_threshold {
+ reader.set_padding_threshold(threshold);
+ }
Ok(Box::new(ByteViewArrayReader::new(pages, data_type, reader)))
}
@@ -107,7 +113,7 @@ impl ArrayReader for ByteViewArrayReader {
fn consume_batch(&mut self) -> Result<ArrayRef> {
let buffer = self.record_reader.consume_record_data();
- let null_buffer = self.record_reader.consume_bitmap_buffer();
+ let null_buffer = self.record_reader.consume_compact_bitmap();
self.def_levels_buffer = self.record_reader.consume_def_levels();
self.rep_levels_buffer = self.record_reader.consume_rep_levels();
self.record_reader.reset();
@@ -128,6 +134,10 @@ impl ArrayReader for ByteViewArrayReader {
fn get_rep_levels(&self) -> Option<&[i16]> {
self.rep_levels_buffer.as_deref()
}
+
+ fn max_def_level(&self) -> i16 {
+ self.record_reader.max_def_level()
+ }
}
/// A [`ColumnValueDecoder`] for variable length byte arrays
diff --git a/parquet/src/arrow/array_reader/cached_array_reader.rs
b/parquet/src/arrow/array_reader/cached_array_reader.rs
index 21f0c2afa4..73f3ba6c8f 100644
--- a/parquet/src/arrow/array_reader/cached_array_reader.rs
+++ b/parquet/src/arrow/array_reader/cached_array_reader.rs
@@ -346,6 +346,10 @@ impl ArrayReader for CachedArrayReader {
fn get_rep_levels(&self) -> Option<&[i16]> {
None
}
+
+ fn max_def_level(&self) -> i16 {
+ self.inner.max_def_level()
+ }
}
#[cfg(test)]
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 f48a219160..e895f79921 100644
--- a/parquet/src/arrow/array_reader/fixed_len_byte_array.rs
+++ b/parquet/src/arrow/array_reader/fixed_len_byte_array.rs
@@ -40,14 +40,12 @@ use std::ops::Range;
use std::sync::Arc;
/// Returns an [`ArrayReader`] that decodes the provided fixed length byte
array column
-///
-/// `batch_size` is used to pre-allocate internal buffers,
-/// avoiding reallocations when reading the first batch of data.
pub fn make_fixed_len_byte_array_reader(
pages: Box<dyn PageIterator>,
column_desc: ColumnDescPtr,
arrow_type: Option<ArrowType>,
batch_size: usize,
+ padding_threshold: Option<i16>,
) -> Result<Box<dyn ArrayReader>> {
// Check if Arrow type is specified, else create it from Parquet type
let data_type = match arrow_type {
@@ -131,6 +129,7 @@ pub fn make_fixed_len_byte_array_reader(
data_type,
byte_length,
batch_size,
+ padding_threshold,
)))
}
@@ -150,8 +149,12 @@ impl FixedLenByteArrayReader {
data_type: ArrowType,
byte_length: usize,
batch_size: usize,
+ padding_threshold: Option<i16>,
) -> Self {
- let record_reader = GenericRecordReader::new(column_desc, batch_size);
+ let mut record_reader = GenericRecordReader::new(column_desc,
batch_size);
+ if let Some(threshold) = padding_threshold {
+ record_reader.set_padding_threshold(threshold);
+ }
Self {
data_type,
byte_length,
@@ -177,12 +180,23 @@ impl ArrayReader for FixedLenByteArrayReader {
}
fn consume_batch(&mut self) -> Result<ArrayRef> {
+ let len = self.record_reader.values_written();
+
let record_data = self.record_reader.consume_record_data();
+ debug_assert_eq!(
+ record_data.buffer.len(),
+ len * self.byte_length,
+ "fixed-len byte array buffer size mismatch: {} bytes for {}
elements of size {}",
+ record_data.buffer.len(),
+ len,
+ self.byte_length,
+ );
+ let null_bit_buffer = self.record_reader.consume_compact_bitmap();
let array_data =
ArrayDataBuilder::new(ArrowType::FixedSizeBinary(self.byte_length as i32))
- .len(self.record_reader.num_values())
+ .len(len)
.add_buffer(Buffer::from_vec(record_data.buffer))
- .null_bit_buffer(self.record_reader.consume_bitmap_buffer());
+ .null_bit_buffer(null_bit_buffer);
let binary = FixedSizeBinaryArray::from(unsafe {
array_data.build_unchecked() });
@@ -258,6 +272,10 @@ impl ArrayReader for FixedLenByteArrayReader {
fn get_rep_levels(&self) -> Option<&[i16]> {
self.rep_levels_buffer.as_deref()
}
+
+ fn max_def_level(&self) -> i16 {
+ self.record_reader.max_def_level()
+ }
}
#[derive(Default)]
@@ -296,7 +314,7 @@ fn move_values<F>(
impl ValuesBuffer for FixedLenByteArrayBuffer {
fn with_capacity(capacity: usize) -> Self {
// `byte_length` is not known initially, so preserve the value-count
- // hint so the first decode can allocate the exact byte capacity.
+ // hint so the first decode can allocate the initial byte capacity.
Self {
buffer: Vec::new(),
byte_length: None,
@@ -304,6 +322,18 @@ impl ValuesBuffer for FixedLenByteArrayBuffer {
}
}
+ fn reserve_exact(&mut self, additional: usize) {
+ match self.byte_length {
+ Some(byte_length) => self
+ .buffer
+ .reserve_exact(additional.saturating_mul(byte_length)),
+ None => {
+ let capacity = self.values_capacity.get_or_insert(0);
+ *capacity = (*capacity).max(additional);
+ }
+ }
+ }
+
fn pad_nulls(
&mut self,
read_offset: usize,
@@ -566,6 +596,18 @@ mod tests {
use bytes::Bytes;
use std::sync::Arc;
+ #[test]
+ fn test_pending_reservation_tracks_max_capacity() {
+ let mut buffer = FixedLenByteArrayBuffer::with_capacity(0);
+
+ buffer.reserve_exact(8);
+ buffer.reserve_exact(4);
+ assert_eq!(buffer.values_capacity, Some(8));
+
+ buffer.reserve_exact(12);
+ assert_eq!(buffer.values_capacity, Some(12));
+ }
+
#[test]
fn test_decimal_list() {
let decimals = Decimal256Array::from_iter_values(
diff --git a/parquet/src/arrow/array_reader/fixed_size_list_array.rs
b/parquet/src/arrow/array_reader/fixed_size_list_array.rs
index 8eb08daad0..976aeb98ba 100644
--- a/parquet/src/arrow/array_reader/fixed_size_list_array.rs
+++ b/parquet/src/arrow/array_reader/fixed_size_list_array.rs
@@ -15,10 +15,10 @@
// specific language governing permissions and limitations
// under the License.
-use std::cmp::Ordering;
use std::sync::Arc;
use crate::arrow::array_reader::ArrayReader;
+use
crate::arrow::record_reader::definition_levels::build_filtered_validity_bitmap;
use crate::errors::ParquetError;
use crate::errors::Result;
use arrow_array::FixedSizeListArray;
@@ -38,10 +38,17 @@ pub struct FixedSizeListArrayReader {
rep_level: i16,
/// If the list is nullable
nullable: bool,
+ /// When set, this reader is a child of another list and should
+ /// exclude entries where def < threshold from its output.
+ parent_threshold: Option<i16>,
}
impl FixedSizeListArrayReader {
/// Construct fixed-size list array reader.
+ ///
+ /// The child reader must already be constructed with a padding
+ /// threshold of `def_level`. `parent_threshold` is set when this
+ /// reader is itself a child of another list.
pub fn new(
item_reader: Box<dyn ArrayReader>,
fixed_size: usize,
@@ -49,6 +56,7 @@ impl FixedSizeListArrayReader {
def_level: i16,
rep_level: i16,
nullable: bool,
+ parent_threshold: Option<i16>,
) -> Self {
Self {
item_reader,
@@ -57,6 +65,7 @@ impl FixedSizeListArrayReader {
def_level,
rep_level,
nullable,
+ parent_threshold,
}
}
}
@@ -75,11 +84,12 @@ impl ArrayReader for FixedSizeListArrayReader {
Ok(size)
}
+ /// Consume batch. The compact child array contains entries only for
+ /// items within non-null fixed-size lists. Null lists get `fixed_size`
+ /// null entries inserted via MutableArrayData. When all lists are
+ /// non-null the compact child is used directly without any copying.
fn consume_batch(&mut self) -> Result<ArrayRef> {
- let next_batch_array = self.item_reader.consume_batch()?;
- if next_batch_array.is_empty() {
- return Ok(new_empty_array(&self.data_type));
- }
+ let child_array = self.item_reader.consume_batch()?;
let def_levels = self
.get_def_levels()
@@ -88,130 +98,112 @@ impl ArrayReader for FixedSizeListArrayReader {
.get_rep_levels()
.ok_or_else(|| general_err!("item_reader rep levels are None"))?;
- if !rep_levels.is_empty() && rep_levels[0] != 0 {
- // This implies either the source data was invalid, or the leaf
column
- // reader did not correctly delimit semantic records
+ if def_levels.is_empty() {
+ return Ok(new_empty_array(&self.data_type));
+ }
+
+ if rep_levels[0] != 0 {
return Err(general_err!("first repetition level of batch must be
0"));
}
- let mut validity = self
- .nullable
- .then(|| BooleanBufferBuilder::new(next_batch_array.len()));
-
- let data = next_batch_array.to_data();
- let mut child_data_builder =
- MutableArrayData::new(vec![&data], true, next_batch_array.len());
-
- // The current index into the child array entries
- let mut child_idx = 0;
- // The total number of rows (valid and invalid) in the list array
- let mut list_len = 0;
- // Start of the current run of valid values
- let mut start_idx = None;
- let mut row_len = 0;
-
- def_levels.iter().zip(rep_levels).try_for_each(|(d, r)| {
- match r.cmp(&self.rep_level) {
- Ordering::Greater => {
- // Repetition level greater than current => already
handled by inner array
- if *d < self.def_level {
- return Err(general_err!(
- "Encountered repetition level too large for
definition level"
- ));
- }
+ // Single pass: collect the def level at each list boundary and
+ // validate that each non-null row contains exactly fixed_size items.
+ let mut boundary_defs: Vec<i16> = Vec::with_capacity(def_levels.len());
+ let mut current_row_len = None;
+ for (&d, &r) in def_levels.iter().zip(rep_levels) {
+ if let Some(threshold) = self.parent_threshold {
+ if d < threshold {
+ continue;
}
- Ordering::Equal => {
- // Item inside of the current list
- child_idx += 1;
- row_len += 1;
+ }
+ if r < self.rep_level {
+ if let Some(row_len) = current_row_len.take().filter(|&len|
len != self.fixed_size)
+ {
+ return Err(general_err!(
+ "Encountered misaligned row with length {} (expected
length {})",
+ row_len,
+ self.fixed_size
+ ));
}
- Ordering::Less => {
- // Start of new list row
- list_len += 1;
-
- // Length of the previous row should be equal to:
- // - the list's fixed size (valid entries)
- // - zero (null entries, start of array)
- // Any other length indicates invalid data
- if start_idx.is_some() && row_len != self.fixed_size {
- return Err(general_err!(
- "Encountered misaligned row with length {}
(expected length {})",
- row_len,
- self.fixed_size
- ));
- }
- row_len = 0;
-
- if *d >= self.def_level {
- row_len += 1;
-
- // Valid list entry
- if let Some(validity) = validity.as_mut() {
- validity.append(true);
- }
- // Start a run of valid rows if not already inside of
one
- start_idx.get_or_insert(child_idx);
- } else {
- // Null list entry
-
- if let Some(start) = start_idx.take() {
- // Flush pending child items
- child_data_builder
- .try_extend(0, start, child_idx)
- .map_err(|e| general_err!("{}", e))?;
- }
- // Pad list with nulls
- child_data_builder
- .try_extend_nulls(self.fixed_size)
- .map_err(|e| general_err!("{}", e))?;
-
- if let Some(validity) = validity.as_mut() {
- // Valid if empty list
- validity.append(*d + 1 == self.def_level);
- }
- }
- child_idx += 1;
+
+ boundary_defs.push(d);
+ current_row_len = (d >= self.def_level).then_some(1);
+ } else if r == self.rep_level {
+ if let Some(row_len) = current_row_len.as_mut() {
+ *row_len += 1;
}
+ } else if d < self.def_level {
+ return Err(general_err!(
+ "Encountered repetition level too large for definition
level"
+ ));
}
- Ok(())
- })?;
+ }
- let child_data = match start_idx {
- Some(0) => {
- // No null entries - can reuse original array
- next_batch_array.to_data()
- }
- Some(start) => {
- // Flush pending child items
- child_data_builder
- .try_extend(0, start, child_idx)
- .map_err(|e| general_err!("{}", e))?;
- child_data_builder.freeze()
- }
- None => child_data_builder.freeze(),
- };
+ if let Some(row_len) = current_row_len.filter(|&len| len !=
self.fixed_size) {
+ return Err(general_err!(
+ "Encountered misaligned row with length {} (expected length
{})",
+ row_len,
+ self.fixed_size
+ ));
+ }
- // Verify total number of elements is aligned with fixed list size
- if list_len * self.fixed_size != child_data.len() {
+ let list_len = boundary_defs.len();
+ let items_count = boundary_defs
+ .iter()
+ .filter(|&&d| d >= self.def_level)
+ .count();
+
+ if child_array.len() != items_count * self.fixed_size {
return Err(general_err!(
- "fixed-size list length must be a multiple of {} but array
contains {} elements",
- self.fixed_size,
- child_data.len()
+ "Fixed-size list child mismatch: expected {} entries but got
{}",
+ items_count * self.fixed_size,
+ child_array.len()
));
}
+ // Build child data. If all lists are non-null, reuse the compact
+ // child directly. Otherwise insert fixed_size null blocks.
+ let child_data = if items_count == list_len {
+ child_array.to_data()
+ } else {
+ let data = child_array.to_data();
+ let mut builder = MutableArrayData::new(vec![&data], true,
list_len * self.fixed_size);
+ let mut src_idx = 0;
+ for &d in &boundary_defs {
+ if d >= self.def_level {
+ builder
+ .try_extend(0, src_idx, src_idx + self.fixed_size)
+ .map_err(|e| general_err!("{}", e))?;
+ src_idx += self.fixed_size;
+ } else {
+ builder
+ .try_extend_nulls(self.fixed_size)
+ .map_err(|e| general_err!("{}", e))?;
+ }
+ }
+ builder.freeze()
+ };
+
+ debug_assert_eq!(child_data.len(), list_len * self.fixed_size);
+
let mut list_builder = ArrayData::builder(self.get_data_type().clone())
.len(list_len)
.add_child_data(child_data);
- if let Some(builder) = validity {
- list_builder = list_builder.null_bit_buffer(Some(builder.into()));
+ if self.nullable {
+ let mut validity = BooleanBufferBuilder::new(list_len);
+ build_filtered_validity_bitmap(
+ &boundary_defs,
+ None,
+ None,
+ self.def_level - 1,
+ &mut validity,
+ );
+ list_builder = list_builder.null_bit_buffer(Some(validity.into()));
}
let list_data = unsafe { list_builder.build_unchecked() };
-
- let result_array = FixedSizeListArray::from(list_data);
- Ok(Arc::new(result_array))
+ Ok(Arc::new(FixedSizeListArray::from(list_data)))
}
fn skip_records(&mut self, num_records: usize) -> Result<usize> {
@@ -225,6 +217,10 @@ impl ArrayReader for FixedSizeListArrayReader {
fn get_rep_levels(&self) -> Option<&[i16]> {
self.item_reader.get_rep_levels()
}
+
+ fn max_def_level(&self) -> i16 {
+ self.def_level
+ }
}
#[cfg(test)]
@@ -266,6 +262,7 @@ mod tests {
&[0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1],
3,
1,
+ Some(2),
);
let mut list_array_reader = FixedSizeListArrayReader::new(
@@ -275,6 +272,7 @@ mod tests {
2,
1,
true,
+ None,
);
let actual = list_array_reader.next_batch(1024).unwrap();
let actual = actual
@@ -303,6 +301,7 @@ mod tests {
&[0, 1, 0, 1, 0, 1, 0, 1],
2,
1,
+ Some(1),
);
let mut list_array_reader = FixedSizeListArrayReader::new(
@@ -312,6 +311,7 @@ mod tests {
1,
1,
false,
+ None,
);
let actual = list_array_reader.next_batch(1024).unwrap();
let actual = actual
@@ -321,6 +321,31 @@ mod tests {
assert_eq!(&expected, actual)
}
+ #[test]
+ fn test_misaligned_rows_error() {
+ // Two rows with fixed_size == 2 are encoded with row lengths 1 and 3.
+ // The total child count is still 4, so aggregate count validation
alone
+ // would silently repartition the values into two valid-looking rows.
+ let item_array_reader =
+ make_int32_page_reader(&[1, 2, 3, 4], &[2, 2, 2, 2], &[0, 0, 1,
1], 2, 1, Some(1));
+
+ let mut list_array_reader = FixedSizeListArrayReader::new(
+ item_array_reader,
+ 2,
+
ArrowType::FixedSizeList(Arc::new(Field::new_list_field(ArrowType::Int32,
true)), 2),
+ 1,
+ 1,
+ false,
+ None,
+ );
+
+ let err = list_array_reader.next_batch(1024).unwrap_err();
+ assert_eq!(
+ err.to_string(),
+ "Parquet error: Encountered misaligned row with length 1 (expected
length 2)"
+ );
+ }
+
#[test]
fn test_nested_list() {
// [
@@ -372,10 +397,11 @@ mod tests {
&[0, 0, 2, 0, 2, 0, 0, 2, 0, 2],
5,
2,
+ Some(4),
);
- let l2 = FixedSizeListArrayReader::new(item_array_reader, 2, l2_type,
4, 2, false);
- let mut l1 = FixedSizeListArrayReader::new(Box::new(l2), 1, l1_type,
3, 1, true);
+ let l2 = FixedSizeListArrayReader::new(item_array_reader, 2, l2_type,
4, 2, false, Some(3));
+ let mut l1 = FixedSizeListArrayReader::new(Box::new(l2), 1, l1_type,
3, 1, true, None);
let expected_1 = expected.slice(0, 2);
let expected_2 = expected.slice(2, 4);
@@ -395,7 +421,8 @@ mod tests {
0,
);
- let item_array_reader = make_int32_page_reader(&[], &[0, 1, 0, 1],
&[0, 0, 0, 0], 2, 1);
+ let item_array_reader =
+ make_int32_page_reader(&[], &[0, 1, 0, 1], &[0, 0, 0, 0], 2, 1,
Some(2));
let mut list_array_reader = FixedSizeListArrayReader::new(
item_array_reader,
@@ -404,6 +431,7 @@ mod tests {
2,
1,
true,
+ None,
);
let actual = list_array_reader.next_batch(1024).unwrap();
let actual = actual
@@ -441,13 +469,21 @@ mod tests {
&[0, 2, 2, 1, 0, 1, 0, 2, 1, 2, 0],
5,
2,
+ Some(4),
);
let inner_array_reader =
- ListArrayReader::<i32>::new(item_array_reader, inner_type, 4, 2,
true);
+ ListArrayReader::<i32>::new(item_array_reader, inner_type, 4, 2,
true, Some(2));
- let mut list_array_reader =
- FixedSizeListArrayReader::new(Box::new(inner_array_reader), 2,
list_type, 2, 1, true);
+ let mut list_array_reader = FixedSizeListArrayReader::new(
+ Box::new(inner_array_reader),
+ 2,
+ list_type,
+ 2,
+ 1,
+ true,
+ None,
+ );
let actual = list_array_reader.next_batch(1024).unwrap();
let actual = actual
.as_any()
diff --git a/parquet/src/arrow/array_reader/list_array.rs
b/parquet/src/arrow/array_reader/list_array.rs
index bc579a543a..66c93b7e1f 100644
--- a/parquet/src/arrow/array_reader/list_array.rs
+++ b/parquet/src/arrow/array_reader/list_array.rs
@@ -24,10 +24,9 @@ use arrow_array::{
};
use arrow_buffer::Buffer;
use arrow_buffer::ToByteSlice;
-use arrow_data::{ArrayData, transform::MutableArrayData};
+use arrow_data::ArrayData;
use arrow_schema::DataType as ArrowType;
use std::any::Any;
-use std::cmp::Ordering;
use std::marker::PhantomData;
use std::sync::Arc;
@@ -41,17 +40,31 @@ pub struct ListArrayReader<OffsetSize: OffsetSizeTrait> {
rep_level: i16,
/// If this list is nullable
nullable: bool,
+ /// When set, this list is itself a child of another list and should
+ /// exclude entries where def < threshold from its output.
+ parent_threshold: Option<i16>,
_marker: PhantomData<OffsetSize>,
}
impl<OffsetSize: OffsetSizeTrait> ListArrayReader<OffsetSize> {
/// Construct list array reader.
+ ///
+ /// The child reader must already be constructed with a padding
+ /// threshold of `def_level`, so it produces item-level padded arrays
+ /// (values + null items, no list-level padding). `consume_batch`
+ /// then computes offsets directly from def/rep levels without
+ /// MutableArrayData.
+ ///
+ /// `parent_threshold` is set when this list is itself a child of
+ /// another list — entries where `def < threshold` are excluded from
+ /// this list's output.
pub fn new(
item_reader: Box<dyn ArrayReader>,
data_type: ArrowType,
def_level: i16,
rep_level: i16,
nullable: bool,
+ parent_threshold: Option<i16>,
) -> Self {
Self {
item_reader,
@@ -59,12 +72,13 @@ impl<OffsetSize: OffsetSizeTrait>
ListArrayReader<OffsetSize> {
def_level,
rep_level,
nullable,
+ parent_threshold,
_marker: PhantomData,
}
}
}
-/// Implementation of ListArrayReader. Nested lists and lists of structs are
not yet supported.
+/// Implementation of ListArrayReader.
impl<OffsetSize: OffsetSizeTrait> ArrayReader for ListArrayReader<OffsetSize> {
fn as_any(&self) -> &dyn Any {
self
@@ -82,10 +96,7 @@ impl<OffsetSize: OffsetSizeTrait> ArrayReader for
ListArrayReader<OffsetSize> {
}
fn consume_batch(&mut self) -> Result<ArrayRef> {
- let next_batch_array = self.item_reader.consume_batch()?;
- if next_batch_array.is_empty() {
- return Ok(new_empty_array(&self.data_type));
- }
+ let child_array = self.item_reader.consume_batch()?;
let def_levels = self
.item_reader
@@ -97,125 +108,131 @@ impl<OffsetSize: OffsetSizeTrait> ArrayReader for
ListArrayReader<OffsetSize> {
.get_rep_levels()
.ok_or_else(|| general_err!("item_reader rep levels are None."))?;
- if OffsetSize::from_usize(next_batch_array.len()).is_none() {
- return Err(general_err!(
- "offset of {} would overflow list array",
- next_batch_array.len()
- ));
+ if def_levels.is_empty() {
+ return Ok(new_empty_array(&self.data_type));
}
- if !rep_levels.is_empty() && rep_levels[0] != 0 {
- // This implies either the source data was invalid, or the leaf
column
- // reader did not correctly delimit semantic records
+ if rep_levels[0] != 0 {
return Err(general_err!("first repetition level of batch must be
0"));
}
- // A non-nullable list has a single definition level indicating if the
list is empty
+ if OffsetSize::from_usize(child_array.len()).is_none() {
+ return Err(general_err!(
+ "offset of {} would overflow list array",
+ child_array.len()
+ ));
+ }
+
+ // Definition levels identify whether each list boundary contributes a
+ // child slot. For a nullable list, the states are:
//
- // A nullable list has two definition levels associated with it:
+ // d >= self.def_level : list is present with a child item
+ // d == self.def_level - 1 : list is present but empty
+ // d <= self.def_level - 2 : list is null
//
- // The first identifies if the list is null
- // The second identifies if the list is empty
+ // Required lists do not have the null state, but still use the same
+ // `d >= self.def_level` test to distinguish child items from empty
+ // lists. Repetition levels identify list boundaries and whether a
+ // child item belongs directly to this list or to a nested child list.
+ let levels_len = def_levels.len();
+ let mut list_offsets: Vec<OffsetSize> = Vec::with_capacity(levels_len
+ 1);
+ let mut validity = self.nullable.then(||
BooleanBufferBuilder::new(levels_len));
+
+ // Count direct child values of this list.
//
- // The child data returned above is padded with a value for each
not-fully defined level.
- // Therefore null and empty lists will correspond to a value in the
child array.
+ // `d >= self.def_level` excludes null/empty parent lists: an empty
+ // list is encoded as `d == self.def_level - 1` (a null list lower
+ // still), so it carries no child value and must not advance the
+ // offset. Emptiness is handled by this definition-level guard, not
+ // by the repetition level.
//
- // Whilst nulls may have a non-zero slice in the offsets array, empty
lists must
- // be of zero length. As a result we MUST filter out values
corresponding to empty
- // lists, and for consistency we do the same for nulls.
-
- // The output offsets for the computed ListArray
- let mut list_offsets: Vec<OffsetSize> =
Vec::with_capacity(next_batch_array.len() + 1);
-
- // The validity mask of the computed ListArray if nullable
- let mut validity = self
- .nullable
- .then(|| BooleanBufferBuilder::new(next_batch_array.len()));
-
- // The offset into the filtered child data of the current level being
considered
- let mut cur_offset = 0;
-
- // Identifies the start of a run of values to copy from the source
child data
- let mut filter_start = None;
-
- // The number of child values skipped due to empty lists or nulls
- let mut skipped = 0;
-
- // Builder used to construct the filtered child data, skipping empty
lists and nulls
- let data = next_batch_array.to_data();
- let mut child_data_builder =
- MutableArrayData::new(vec![&data], false, next_batch_array.len());
-
- def_levels.iter().zip(rep_levels).try_for_each(|(d, r)| {
- match r.cmp(&self.rep_level) {
- Ordering::Greater => {
- // Repetition level greater than current => already
handled by inner array
- if *d < self.def_level {
- return Err(general_err!(
- "Encountered repetition level too large for
definition level"
- ));
+ // `r <= self.rep_level` (not `==`) keeps only direct children: the
+ // first element of each list has `r < self.rep_level` (it coincides
+ // with an outer/row boundary), so `==` would drop it; continuation
+ // entries of a nested child's own list have `r > self.rep_level` and
+ // are excluded.
+ //
+ // Split the scan by nullable/parent-threshold mode to avoid checking
+ // these options once per decoded level in the hot loop.
+ let mut cur_offset: usize = 0;
+ let def_level = self.def_level;
+ let rep_level = self.rep_level;
+
+ match (self.parent_threshold, validity.as_mut()) {
+ (None, Some(validity)) => {
+ for (&d, &r) in def_levels.iter().zip(rep_levels) {
+ if r < rep_level {
+
list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap());
+ validity.append(d + 1 >= def_level);
+ }
+
+ if d >= def_level && r <= rep_level {
+ cur_offset += 1;
}
}
- Ordering::Equal => {
- // New value in the current list
- cur_offset += 1;
- }
- Ordering::Less => {
- // Create new array slice
- // Already checked that this cannot overflow
-
list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap());
+ }
+ (None, None) => {
+ for (&d, &r) in def_levels.iter().zip(rep_levels) {
+ if r < rep_level {
+
list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap());
+ }
- if *d >= self.def_level {
- // Fully defined value
+ if d >= def_level && r <= rep_level {
+ cur_offset += 1;
+ }
+ }
+ }
+ (Some(threshold), Some(validity)) => {
+ for (&d, &r) in def_levels.iter().zip(rep_levels) {
+ // When this list is a child of another list, skip entries
+ // belonging to null/empty parent lists.
+ if d < threshold {
+ continue;
+ }
- // Record current offset if it is None
- filter_start.get_or_insert(cur_offset + skipped);
+ if r < rep_level {
+
list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap());
+ validity.append(d + 1 >= def_level);
+ }
+ if d >= def_level && r <= rep_level {
cur_offset += 1;
-
- if let Some(validity) = validity.as_mut() {
- validity.append(true)
- }
- } else {
- // Flush the current slice of child values if any
- if let Some(start) = filter_start.take() {
- child_data_builder
- .try_extend(0, start, cur_offset + skipped)
- .map_err(|e| general_err!("{}", e))?;
- }
-
- if let Some(validity) = validity.as_mut() {
- // Valid if empty list
- validity.append(*d + 1 == self.def_level)
- }
-
- skipped += 1;
}
}
}
- Ok(())
- })?;
+ (Some(threshold), None) => {
+ for (&d, &r) in def_levels.iter().zip(rep_levels) {
+ // When this list is a child of another list, skip entries
+ // belonging to null/empty parent lists.
+ if d < threshold {
+ continue;
+ }
- list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap());
+ if r < rep_level {
+
list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap());
+ }
- let child_data = if skipped == 0 {
- // No filtered values - can reuse original array
- next_batch_array.to_data()
- } else {
- // One or more filtered values - must build new array
- if let Some(start) = filter_start.take() {
- child_data_builder
- .try_extend(0, start, cur_offset + skipped)
- .map_err(|e| general_err!("{}", e))?;
+ if d >= def_level && r <= rep_level {
+ cur_offset += 1;
+ }
+ }
}
+ }
- child_data_builder.freeze()
- };
+ list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap());
- if cur_offset != child_data.len() {
- return Err(general_err!("Failed to reconstruct list from level
data"));
+ if cur_offset != child_array.len() {
+ return Err(general_err!(
+ "Failed to reconstruct list from level data: \
+ expected {} child values but got {}",
+ cur_offset,
+ child_array.len()
+ ));
}
+ debug_assert!(list_offsets.windows(2).all(|w| w[0] <= w[1]));
+
+ let child_data = child_array.to_data();
let value_offsets = Buffer::from(list_offsets.to_byte_slice());
let mut data_builder = ArrayData::builder(self.get_data_type().clone())
@@ -224,14 +241,19 @@ impl<OffsetSize: OffsetSizeTrait> ArrayReader for
ListArrayReader<OffsetSize> {
.add_child_data(child_data);
if let Some(builder) = validity {
- assert_eq!(builder.len(), list_offsets.len() - 1);
- data_builder = data_builder.null_bit_buffer(Some(builder.into()))
+ if builder.len() != list_offsets.len() - 1 {
+ return Err(general_err!(
+ "Failed to decode level data for list array: \
+ expected {} validity bits but got {}",
+ list_offsets.len() - 1,
+ builder.len()
+ ));
+ }
+ data_builder = data_builder.null_bit_buffer(Some(builder.into()));
}
let list_data = unsafe { data_builder.build_unchecked() };
-
- let result_array = GenericListArray::<OffsetSize>::from(list_data);
- Ok(Arc::new(result_array))
+ Ok(Arc::new(GenericListArray::<OffsetSize>::from(list_data)))
}
fn skip_records(&mut self, num_records: usize) -> Result<usize> {
@@ -245,6 +267,10 @@ impl<OffsetSize: OffsetSizeTrait> ArrayReader for
ListArrayReader<OffsetSize> {
fn get_rep_levels(&self) -> Option<&[i16]> {
self.item_reader.get_rep_levels()
}
+
+ fn max_def_level(&self) -> i16 {
+ self.def_level
+ }
}
#[cfg(test)]
@@ -364,13 +390,15 @@ mod tests {
&[0, 3, 2, 2, 2, 1, 1, 1, 1, 3, 3, 2, 3, 3, 2, 0, 0, 0],
6,
3,
+ Some(5),
);
- let l3 = ListArrayReader::<OffsetSize>::new(item_array_reader,
l3_type, 5, 3, true);
+ let l3 =
+ ListArrayReader::<OffsetSize>::new(item_array_reader, l3_type, 5,
3, true, Some(3));
- let l2 = ListArrayReader::<OffsetSize>::new(Box::new(l3), l2_type, 3,
2, false);
+ let l2 = ListArrayReader::<OffsetSize>::new(Box::new(l3), l2_type, 3,
2, false, Some(2));
- let mut l1 = ListArrayReader::<OffsetSize>::new(Box::new(l2), l1_type,
2, 1, true);
+ let mut l1 = ListArrayReader::<OffsetSize>::new(Box::new(l2), l1_type,
2, 1, true, None);
let expected_1 = expected.slice(0, 2);
let expected_2 = expected.slice(2, 2);
@@ -400,6 +428,7 @@ mod tests {
&[0, 1, 1, 0, 0, 1, 0, 0, 0, 1],
2,
1,
+ Some(1),
);
let mut list_array_reader = ListArrayReader::<OffsetSize>::new(
@@ -408,6 +437,7 @@ mod tests {
1,
1,
false,
+ None,
);
let actual = list_array_reader.next_batch(1024).unwrap();
@@ -437,6 +467,7 @@ mod tests {
&[0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1],
3,
1,
+ Some(2),
);
let mut list_array_reader = ListArrayReader::<OffsetSize>::new(
@@ -445,6 +476,7 @@ mod tests {
2,
1,
true,
+ None,
);
let actual = list_array_reader.next_batch(1024).unwrap();
@@ -469,6 +501,123 @@ mod tests {
test_list_array::<i64>()
}
+ #[test]
+ fn test_list_multi_row_group_selective_padding() {
+ use crate::arrow::array_reader::PrimitiveArrayReader;
+ use crate::basic::{Encoding, Type as PhysicalType};
+ use crate::data_type::Int32Type as ParquetInt32;
+ use crate::schema::types::{ColumnDescriptor, ColumnPath, Type};
+ use crate::util::InMemoryPageIterator;
+ use crate::util::test_common::page_util::{DataPageBuilder,
DataPageBuilderImpl};
+
+ // Schema: OPTIONAL GROUP my_list (LIST) {
+ // repeated group list { optional INT32 element; }
+ // }
+ // max_def_level = 3, max_rep_level = 1
+ let leaf_type = Type::primitive_type_builder("element",
PhysicalType::INT32)
+ .build()
+ .unwrap();
+
+ let desc = Arc::new(ColumnDescriptor::new(
+ Arc::new(leaf_type),
+ 3, // max_def_level
+ 1, // max_rep_level
+ ColumnPath::new(vec![]),
+ ));
+
+ // Row group 0 — 3 records:
+ // Record 1: [10, null, 20]
+ // Record 2: null list
+ // Record 3: [null, 30]
+ let mut pb1 = DataPageBuilderImpl::new(desc.clone(), 6, true);
+ pb1.add_rep_levels(1, &[0, 1, 1, 0, 0, 1]);
+ pb1.add_def_levels(3, &[3, 2, 3, 0, 2, 3]);
+ pb1.add_values::<ParquetInt32>(Encoding::PLAIN, &[10, 20, 30]);
+
+ // Row group 1 — 2 records:
+ // Record 4: [40, null]
+ // Record 5: [null, null, 50]
+ let mut pb2 = DataPageBuilderImpl::new(desc.clone(), 5, true);
+ pb2.add_rep_levels(1, &[0, 1, 0, 1, 1]);
+ pb2.add_def_levels(3, &[3, 2, 2, 2, 3]);
+ pb2.add_values::<ParquetInt32>(Encoding::PLAIN, &[40, 50]);
+
+ // Two row groups, one page each → forces two read_one_batch calls
+ let pages = vec![vec![pb1.consume()], vec![pb2.consume()]];
+ let page_iter = InMemoryPageIterator::new(pages);
+
+ let item_reader = Box::new(
+ PrimitiveArrayReader::<ParquetInt32>::new(
+ Box::new(page_iter),
+ desc,
+ None,
+ DEFAULT_BATCH_SIZE,
+ Some(2), // padding_threshold = list's def_level
+ )
+ .unwrap(),
+ );
+
+ let list_type =
ArrowType::List(Arc::new(Field::new_list_field(ArrowType::Int32, true)));
+ let mut list_reader = ListArrayReader::<i32>::new(item_reader,
list_type, 2, 1, true, None);
+
+ let result = list_reader.next_batch(5).unwrap();
+ let list = result
+ .as_any()
+ .downcast_ref::<GenericListArray<i32>>()
+ .unwrap();
+
+ assert_eq!(list.len(), 5);
+
+ // Record 1: [10, null, 20]
+ assert!(list.is_valid(0));
+ let r1 = list.value(0);
+ let r1 = r1
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap();
+ assert_eq!(r1.len(), 3);
+ assert_eq!(r1.value(0), 10);
+ assert!(r1.is_null(1));
+ assert_eq!(r1.value(2), 20);
+
+ // Record 2: null
+ assert!(list.is_null(1));
+
+ // Record 3: [null, 30]
+ assert!(list.is_valid(2));
+ let r3 = list.value(2);
+ let r3 = r3
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap();
+ assert_eq!(r3.len(), 2);
+ assert!(r3.is_null(0));
+ assert_eq!(r3.value(1), 30);
+
+ // Record 4: [40, null]
+ assert!(list.is_valid(3));
+ let r4 = list.value(3);
+ let r4 = r4
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap();
+ assert_eq!(r4.len(), 2);
+ assert_eq!(r4.value(0), 40);
+ assert!(r4.is_null(1));
+
+ // Record 5: [null, null, 50]
+ assert!(list.is_valid(4));
+ let r5 = list.value(4);
+ let r5 = r5
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap();
+ assert_eq!(r5.len(), 3);
+ assert!(r5.is_null(0));
+ assert!(r5.is_null(1));
+ assert_eq!(r5.value(2), 50);
+ }
+
#[test]
fn test_nested_lists() {
// Construct column schema
diff --git a/parquet/src/arrow/array_reader/list_view_array.rs
b/parquet/src/arrow/array_reader/list_view_array.rs
index 357ab9dc14..0a7f58c87e 100644
--- a/parquet/src/arrow/array_reader/list_view_array.rs
+++ b/parquet/src/arrow/array_reader/list_view_array.rs
@@ -38,6 +38,7 @@ impl<OffsetSize: OffsetSizeTrait>
ListViewArrayReader<OffsetSize> {
def_level: i16,
rep_level: i16,
nullable: bool,
+ parent_threshold: Option<i16>,
) -> Self {
// Create the underlying ListArrayReader with the corresponding List
type
let list_data_type = match &data_type {
@@ -46,8 +47,14 @@ impl<OffsetSize: OffsetSizeTrait>
ListViewArrayReader<OffsetSize> {
_ => unreachable!(),
};
- let inner =
- ListArrayReader::new(item_reader, list_data_type, def_level,
rep_level, nullable);
+ let inner = ListArrayReader::new(
+ item_reader,
+ list_data_type,
+ def_level,
+ rep_level,
+ nullable,
+ parent_threshold,
+ );
Self { inner, data_type }
}
@@ -101,6 +108,10 @@ impl<OffsetSize: OffsetSizeTrait> ArrayReader for
ListViewArrayReader<OffsetSize
fn get_rep_levels(&self) -> Option<&[i16]> {
self.inner.get_rep_levels()
}
+
+ fn max_def_level(&self) -> i16 {
+ self.inner.max_def_level()
+ }
}
#[cfg(test)]
@@ -130,6 +141,7 @@ mod tests {
&[0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1],
3,
1,
+ Some(2),
);
let field =
Arc::new(arrow_schema::Field::new_list_field(ArrowType::Int32, true));
@@ -140,7 +152,7 @@ mod tests {
};
let mut list_view_array_reader =
- ListViewArrayReader::<OffsetSize>::new(item_array_reader,
data_type, 2, 1, true);
+ ListViewArrayReader::<OffsetSize>::new(item_array_reader,
data_type, 2, 1, true, None);
let actual = list_view_array_reader.next_batch(1024).unwrap();
let actual = actual
@@ -169,6 +181,7 @@ mod tests {
&[0, 1, 1, 0, 0, 1, 0, 0, 0, 1],
2,
1,
+ Some(1),
);
let field =
Arc::new(arrow_schema::Field::new_list_field(ArrowType::Int32, true));
@@ -179,7 +192,7 @@ mod tests {
};
let mut list_view_array_reader =
- ListViewArrayReader::<OffsetSize>::new(item_array_reader,
data_type, 1, 1, false);
+ ListViewArrayReader::<OffsetSize>::new(item_array_reader,
data_type, 1, 1, false, None);
let actual = list_view_array_reader.next_batch(1024).unwrap();
let actual = actual
diff --git a/parquet/src/arrow/array_reader/map_array.rs
b/parquet/src/arrow/array_reader/map_array.rs
index 1639aca629..5d88df6cea 100644
--- a/parquet/src/arrow/array_reader/map_array.rs
+++ b/parquet/src/arrow/array_reader/map_array.rs
@@ -39,11 +39,12 @@ impl MapArrayReader {
def_level: i16,
rep_level: i16,
nullable: bool,
+ parent_threshold: Option<i16>,
) -> Self {
- let struct_def_level = match nullable {
- true => def_level + 2,
- false => def_level + 1,
- };
+ // The struct exists when the key (always required in maps) exists.
+ // Derive struct_def_level from the key's max def level rather than
+ // a fixed formula, so it matches the schema exactly.
+ let struct_def_level = key_reader.max_def_level();
let struct_rep_level = rep_level + 1;
let element = match &data_type {
@@ -66,6 +67,7 @@ impl MapArrayReader {
struct_def_level,
struct_rep_level,
false,
+ Some(def_level),
);
let reader = ListArrayReader::new(
@@ -74,6 +76,7 @@ impl MapArrayReader {
def_level,
rep_level,
nullable,
+ parent_threshold,
);
Self { data_type, reader }
@@ -119,6 +122,10 @@ impl ArrayReader for MapArrayReader {
fn get_rep_levels(&self) -> Option<&[i16]> {
self.reader.get_rep_levels()
}
+
+ fn max_def_level(&self) -> i16 {
+ self.reader.max_def_level()
+ }
}
#[cfg(test)]
diff --git a/parquet/src/arrow/array_reader/mod.rs
b/parquet/src/arrow/array_reader/mod.rs
index 989c934809..347058329f 100644
--- a/parquet/src/arrow/array_reader/mod.rs
+++ b/parquet/src/arrow/array_reader/mod.rs
@@ -135,6 +135,19 @@ pub trait ArrayReader: Send {
///
/// This is used by parent [`ArrayReader`] to compute their array offsets
fn get_rep_levels(&self) -> Option<&[i16]>;
+
+ /// Returns the maximum definition level for this reader as defined by
+ /// the Parquet schema. For leaf readers this is the column's max def
level;
+ /// for composite readers it is the level at which the composite itself is
+ /// fully defined.
+ ///
+ /// The default panics. Synthetic readers that are not backed by a Parquet
+ /// column (e.g. row-number or empty-struct readers) may rely on this
+ /// default because they are never used as children of schema-driven
+ /// composite readers (list, map, struct), which are the only callers.
+ fn max_def_level(&self) -> i16 {
+ panic!("max_def_level called on a reader that does not track
definition levels")
+ }
}
/// Interface for reading data pages from the columns of one or more RowGroups.
diff --git a/parquet/src/arrow/array_reader/null_array.rs
b/parquet/src/arrow/array_reader/null_array.rs
index 97bc68d6d2..d3740b5e82 100644
--- a/parquet/src/arrow/array_reader/null_array.rs
+++ b/parquet/src/arrow/array_reader/null_array.rs
@@ -47,14 +47,16 @@ where
T::T: ArrowNativeType,
{
/// Construct null array reader.
- ///
- /// `batch_size` is used to pre-allocate internal buffers.
pub fn new(
pages: Box<dyn PageIterator>,
column_desc: ColumnDescPtr,
batch_size: usize,
+ padding_threshold: Option<i16>,
) -> Result<Self> {
- let record_reader = RecordReader::<T>::new(column_desc, batch_size);
+ let mut record_reader = RecordReader::<T>::new(column_desc,
batch_size);
+ if let Some(threshold) = padding_threshold {
+ record_reader.set_padding_threshold(threshold);
+ }
Ok(Self {
data_type: ArrowType::Null,
@@ -87,14 +89,14 @@ where
fn consume_batch(&mut self) -> Result<ArrayRef> {
// convert to arrays
- let array =
arrow_array::NullArray::new(self.record_reader.num_values());
+ let array =
arrow_array::NullArray::new(self.record_reader.values_written());
// save definition and repetition buffers
self.def_levels_buffer = self.record_reader.consume_def_levels();
self.rep_levels_buffer = self.record_reader.consume_rep_levels();
// Must consume bitmap buffer
- self.record_reader.consume_bitmap_buffer();
+ self.record_reader.consume_compact_bitmap();
self.record_reader.reset();
Ok(Arc::new(array))
@@ -111,4 +113,8 @@ where
fn get_rep_levels(&self) -> Option<&[i16]> {
self.rep_levels_buffer.as_deref()
}
+
+ fn max_def_level(&self) -> i16 {
+ self.record_reader.max_def_level()
+ }
}
diff --git a/parquet/src/arrow/array_reader/primitive_array.rs
b/parquet/src/arrow/array_reader/primitive_array.rs
index 037f0af81f..a6581aea93 100644
--- a/parquet/src/arrow/array_reader/primitive_array.rs
+++ b/parquet/src/arrow/array_reader/primitive_array.rs
@@ -105,12 +105,16 @@ where
{
/// Construct primitive array reader.
///
- /// `batch_size` is used to pre-allocate internal buffers.
+ /// `padding_threshold` controls how null padding is applied. When
+ /// `None`, the reader pads all null positions (full padding). When
+ /// `Some(threshold)`, entries with `def < threshold` are excluded
+ /// from the value buffer (selective padding for list children).
pub fn new(
pages: Box<dyn PageIterator>,
column_desc: ColumnDescPtr,
arrow_type: Option<ArrowType>,
batch_size: usize,
+ padding_threshold: Option<i16>,
) -> Result<Self> {
// Check if Arrow type is specified, else create it from Parquet type
let data_type = match arrow_type {
@@ -120,7 +124,10 @@ where
.clone(),
};
- let record_reader = RecordReader::<T>::new(column_desc, batch_size);
+ let mut record_reader = RecordReader::<T>::new(column_desc,
batch_size);
+ if let Some(threshold) = padding_threshold {
+ record_reader.set_padding_threshold(threshold);
+ }
Ok(Self {
data_type,
@@ -157,15 +164,16 @@ where
// Convert physical data to equivalent arrow type, and then perform
// coercion as needed
+ let len = self.record_reader.values_written();
+
let record_data = self
.record_reader
.consume_record_data()
.into_buffer(target_type);
- let len = self.record_reader.num_values();
let nulls = self
.record_reader
- .consume_bitmap_buffer()
+ .consume_compact_bitmap()
.and_then(|b| NullBuffer::from_unsliced_buffer(b, len));
let array: ArrayRef = match T::get_physical_type() {
@@ -219,6 +227,10 @@ where
fn get_rep_levels(&self) -> Option<&[i16]> {
self.rep_levels_buffer.as_deref()
}
+
+ fn max_def_level(&self) -> i16 {
+ self.record_reader.max_def_level()
+ }
}
/// Coerce the parquet physical type array to the target type
@@ -515,6 +527,7 @@ mod tests {
schema.column(0),
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap();
@@ -562,6 +575,7 @@ mod tests {
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap();
@@ -633,6 +647,7 @@ mod tests {
column_desc.clone(),
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.expect("Unable to get array reader");
@@ -773,6 +788,7 @@ mod tests {
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap();
@@ -853,6 +869,7 @@ mod tests {
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap();
@@ -916,6 +933,7 @@ mod tests {
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap();
@@ -982,6 +1000,7 @@ mod tests {
column_desc,
None,
DEFAULT_BATCH_SIZE,
+ None,
)
.unwrap();
diff --git a/parquet/src/arrow/array_reader/struct_array.rs
b/parquet/src/arrow/array_reader/struct_array.rs
index da92d410f3..2032768976 100644
--- a/parquet/src/arrow/array_reader/struct_array.rs
+++ b/parquet/src/arrow/array_reader/struct_array.rs
@@ -16,6 +16,7 @@
// under the License.
use crate::arrow::array_reader::ArrayReader;
+use
crate::arrow::record_reader::definition_levels::build_filtered_validity_bitmap;
use crate::errors::{ParquetError, Result};
use arrow_array::{Array, ArrayRef, StructArray, builder::BooleanBufferBuilder};
use arrow_buffer::NullBuffer;
@@ -30,6 +31,10 @@ pub struct StructArrayReader {
struct_def_level: i16,
struct_rep_level: i16,
nullable: bool,
+ /// When set, entries with def < threshold are excluded from child
+ /// arrays. Set when this struct is inside a list (to the parent
+ /// list's def_level).
+ padding_threshold: Option<i16>,
}
impl StructArrayReader {
@@ -40,6 +45,7 @@ impl StructArrayReader {
def_level: i16,
rep_level: i16,
nullable: bool,
+ padding_threshold: Option<i16>,
) -> Self {
Self {
data_type,
@@ -47,6 +53,7 @@ impl StructArrayReader {
struct_def_level: def_level,
struct_rep_level: rep_level,
nullable,
+ padding_threshold,
}
}
}
@@ -82,107 +89,68 @@ impl ArrayReader for StructArrayReader {
Ok(read.unwrap_or(0))
}
- /// Consume struct records.
- ///
- /// Definition levels of struct array is calculated as following:
- /// ```ignore
- /// def_levels[i] = min(child1_def_levels[i], child2_def_levels[i], ...,
- /// childn_def_levels[i]);
- /// ```
- ///
- /// Repetition levels of struct array is calculated as following:
- /// ```ignore
- /// rep_levels[i] = child1_rep_levels[i];
- /// ```
- ///
- /// The null bitmap of struct array is calculated from def_levels:
- /// ```ignore
- /// null_bitmap[i] = (def_levels[i] >= self.def_level);
- /// ```
- ///
fn consume_batch(&mut self) -> Result<ArrayRef> {
if self.children.is_empty() {
return Ok(Arc::new(StructArray::from(Vec::new())));
}
- let children_array = self
+ let children_arrays = self
.children
.iter_mut()
.map(|reader| reader.consume_batch())
.collect::<Result<Vec<_>>>()?;
- // check that array child data has same size
- let children_array_len = children_array
- .first()
- .map(|arr| arr.len())
- .ok_or_else(|| general_err!("Struct array reader should have at
least one child!"))?;
-
- let all_children_len_eq = children_array
- .iter()
- .all(|arr| arr.len() == children_array_len);
- if !all_children_len_eq {
- return Err(general_err!("Not all children array length are the
same!"));
- }
-
let DataType::Struct(fields) = &self.data_type else {
return Err(general_err!(
"Internal: StructArrayReader must have struct data type, got
{:?}",
self.data_type
));
};
- let fields = fields.clone(); // cloning Fields is cheap (Arc
internally)
+ let fields = fields.clone();
- let mut nulls = None;
- if self.nullable {
- // calculate struct def level data
+ let item_count = children_arrays.first().map(|a| a.len()).unwrap_or(0);
- // children should have consistent view of parent, only need to
inspect first child
+ if !children_arrays.windows(2).all(|w| w[0].len() == w[1].len()) {
+ return Err(general_err!("Not all children array length are the
same!"));
+ }
+
+ // Build struct null bitmap if the struct is nullable.
+ // We iterate def/rep levels and select entries that correspond to
+ // struct rows: skip parent-level padding (d < threshold) and
+ // inner-list continuations (r > struct_rep_level).
+ let nulls = if self.nullable {
let def_levels = self.children[0]
.get_def_levels()
- .expect("child with nullable parents must have definition
level");
-
- // calculate bitmap for current array
- let mut bitmap_builder =
BooleanBufferBuilder::new(children_array_len);
-
- match self.children[0].get_rep_levels() {
- Some(rep_levels) => {
- // Sanity check
- assert_eq!(rep_levels.len(), def_levels.len());
-
- for (rep_level, def_level) in
rep_levels.iter().zip(def_levels) {
- if rep_level > &self.struct_rep_level {
- // Already handled by inner list - SKIP
- continue;
- }
- bitmap_builder.append(*def_level >=
self.struct_def_level)
- }
- }
- None => {
- // Safety: slice iterator has a trusted length
- unsafe {
- bitmap_builder.extend_trusted_len(
- def_levels
- .iter()
- .map(|level| *level >= self.struct_def_level),
- )
- }
- }
+ .ok_or_else(|| general_err!("child def levels are None"))?;
+ let rep_levels = self.children[0].get_rep_levels();
+
+ let mut bitmap = BooleanBufferBuilder::new(item_count);
+ build_filtered_validity_bitmap(
+ def_levels,
+ rep_levels.map(|r| (r, self.struct_rep_level)),
+ self.padding_threshold,
+ self.struct_def_level,
+ &mut bitmap,
+ );
+ if bitmap.len() != item_count {
+ return Err(general_err!(
+ "Failed to decode level data for struct array: \
+ expected {} validity bits but got {}",
+ item_count,
+ bitmap.len()
+ ));
}
+ Some(NullBuffer::from(bitmap))
+ } else {
+ None
+ };
- if bitmap_builder.len() != children_array_len {
- return Err(general_err!("Failed to decode level data for
struct array"));
- }
- nulls = Some(NullBuffer::from(bitmap_builder));
- }
-
- // Safety: checked above that all children array data have same
- // length and correct type
unsafe {
Ok(Arc::new(StructArray::new_unchecked_with_length(
fields,
- children_array,
+ children_arrays,
nulls,
- children_array_len,
+ item_count,
)))
}
}
@@ -218,6 +186,10 @@ impl ArrayReader for StructArrayReader {
// parent structure, so return first child's
self.children.first().and_then(|l| l.get_rep_levels())
}
+
+ fn max_def_level(&self) -> i16 {
+ self.struct_def_level
+ }
}
#[cfg(test)]
@@ -233,9 +205,11 @@ mod tests {
#[test]
fn test_struct_array_reader() {
- let array_reader_1 = make_int32_page_reader(&[4], &[0, 1, 2, 3, 1],
&[0, 1, 1, 1, 1], 3, 1);
+ let array_reader_1 =
+ make_int32_page_reader(&[4], &[0, 1, 2, 3, 1], &[0, 1, 1, 1, 1],
3, 1, None);
- let array_reader_2 = make_int32_page_reader(&[3], &[0, 1, 3, 1, 2],
&[0, 1, 1, 1, 1], 3, 1);
+ let array_reader_2 =
+ make_int32_page_reader(&[3], &[0, 1, 3, 1, 2], &[0, 1, 1, 1, 1],
3, 1, None);
let struct_type = ArrowType::Struct(Fields::from(vec![
Field::new("f1", ArrowType::Int32, true),
@@ -248,6 +222,7 @@ mod tests {
1,
1,
true,
+ None,
);
let struct_array = struct_array_reader.next_batch(5).unwrap();
@@ -294,11 +269,17 @@ mod tests {
)];
let expected = StructArray::from((struct_fields, validity));
- let reader =
- make_int32_page_reader(&[1, 2], &[4, 4, 3, 2, 1, 0], &[0, 1, 1, 0,
0, 0], 4, 1);
+ let reader = make_int32_page_reader(
+ &[1, 2],
+ &[4, 4, 3, 2, 1, 0],
+ &[0, 1, 1, 0, 0, 0],
+ 4,
+ 1,
+ Some(3),
+ );
let list_reader =
- ListArrayReader::<i32>::new(reader,
expected_l.data_type().clone(), 3, 1, true);
+ ListArrayReader::<i32>::new(reader,
expected_l.data_type().clone(), 3, 1, true, None);
let mut struct_reader = StructArrayReader::new(
expected.data_type().clone(),
@@ -306,6 +287,7 @@ mod tests {
1,
0,
true,
+ None,
);
let actual = struct_reader.next_batch(1024).unwrap();
diff --git a/parquet/src/arrow/array_reader/test_util.rs
b/parquet/src/arrow/array_reader/test_util.rs
index afba7c2c30..4cbf9d8f03 100644
--- a/parquet/src/arrow/array_reader/test_util.rs
+++ b/parquet/src/arrow/array_reader/test_util.rs
@@ -90,17 +90,22 @@ pub fn byte_array_all_encodings(
/// Build a real `PrimitiveArrayReader<Int32Type>` from raw non-null values
/// and definition/repetition levels. This exercises the full production
-/// `RecordReader` code path (including selective padding when a parent
-/// `ListArrayReader` calls `enable_selective_padding`).
+/// `RecordReader` code path, including selective padding when
+/// `padding_threshold` is set.
///
/// `values` must contain only the non-null values (entries where
/// `def_levels[i] == max_def_level`), in order, as Parquet encodes them.
+///
+/// `padding_threshold` controls null filtering: `None` for full padding
+/// (top-level columns), `Some(threshold)` for selective padding (list
+/// children — typically the parent list's def_level).
pub fn make_int32_page_reader(
values: &[i32],
def_levels: &[i16],
rep_levels: &[i16],
max_def_level: i16,
max_rep_level: i16,
+ padding_threshold: Option<i16>,
) -> Box<dyn ArrayReader> {
use crate::arrow::array_reader::PrimitiveArrayReader;
use crate::arrow::arrow_reader::DEFAULT_BATCH_SIZE;
@@ -130,8 +135,14 @@ pub fn make_int32_page_reader(
let pages = vec![vec![pb.consume()]];
let page_iter = InMemoryPageIterator::new(pages);
Box::new(
- PrimitiveArrayReader::<Int32Type>::new(Box::new(page_iter), desc,
None, DEFAULT_BATCH_SIZE)
- .unwrap(),
+ PrimitiveArrayReader::<Int32Type>::new(
+ Box::new(page_iter),
+ desc,
+ None,
+ DEFAULT_BATCH_SIZE,
+ padding_threshold,
+ )
+ .unwrap(),
)
}
diff --git a/parquet/src/arrow/arrow_reader/mod.rs
b/parquet/src/arrow/arrow_reader/mod.rs
index 12c3e192cd..814b825050 100644
--- a/parquet/src/arrow/arrow_reader/mod.rs
+++ b/parquet/src/arrow/arrow_reader/mod.rs
@@ -200,8 +200,12 @@ impl<T> ArrowReaderBuilder<T> {
&self.schema
}
- /// Set the size of [`RecordBatch`] to produce. Defaults to
[`DEFAULT_BATCH_SIZE`]
- /// If the batch_size more than the file row count, use the file row count.
+ /// Set the size of [`RecordBatch`] to produce. Defaults to
[`DEFAULT_BATCH_SIZE`].
+ ///
+ /// This may be used as a hint for internal allocations, but does not
+ /// guarantee exact internal buffer capacities.
+ ///
+ /// If `batch_size` is more than the file row count, use the file row
count.
pub fn with_batch_size(self, batch_size: usize) -> Self {
// Try to avoid allocate large buffer
let batch_size =
batch_size.min(self.metadata.file_metadata().num_rows() as usize);
diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs
b/parquet/src/arrow/arrow_reader/read_plan.rs
index 71b64802b0..7ca54f961e 100644
--- a/parquet/src/arrow/arrow_reader/read_plan.rs
+++ b/parquet/src/arrow/arrow_reader/read_plan.rs
@@ -493,13 +493,13 @@ mod tests {
let data: Vec<i32> = (0..TOTAL_ROWS as i32).collect();
let levels = vec![0; TOTAL_ROWS];
- let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0);
+ let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None);
let struct_type = ArrowType::Struct(Fields::from(vec![Field::new(
"c0",
ArrowType::Int32,
false,
)]));
- let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0,
0, false);
+ let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0,
0, false, None);
let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(),
|batch| {
Ok(BooleanArray::from(vec![true; batch.num_rows()]))
diff --git a/parquet/src/arrow/buffer/dictionary_buffer.rs
b/parquet/src/arrow/buffer/dictionary_buffer.rs
index abf7653029..33537b40f8 100644
--- a/parquet/src/arrow/buffer/dictionary_buffer.rs
+++ b/parquet/src/arrow/buffer/dictionary_buffer.rs
@@ -200,6 +200,13 @@ impl<K: ArrowNativeType, V: OffsetSizeTrait> ValuesBuffer
for DictionaryBuffer<K
}
}
+ fn reserve_exact(&mut self, additional: usize) {
+ match self {
+ Self::Dict { keys, .. } => keys.reserve_exact(additional),
+ Self::Values { values, .. } => values.reserve_exact(additional),
+ }
+ }
+
fn pad_nulls(
&mut self,
read_offset: usize,
diff --git a/parquet/src/arrow/buffer/offset_buffer.rs
b/parquet/src/arrow/buffer/offset_buffer.rs
index 3090842679..ab67694c29 100644
--- a/parquet/src/arrow/buffer/offset_buffer.rs
+++ b/parquet/src/arrow/buffer/offset_buffer.rs
@@ -151,6 +151,10 @@ impl<I: OffsetSizeTrait> ValuesBuffer for OffsetBuffer<I> {
Self::with_capacity(capacity)
}
+ fn reserve_exact(&mut self, additional: usize) {
+ self.offsets.reserve_exact(additional);
+ }
+
fn pad_nulls(
&mut self,
read_offset: usize,
diff --git a/parquet/src/arrow/buffer/view_buffer.rs
b/parquet/src/arrow/buffer/view_buffer.rs
index f20ca116e9..9670d0e9bb 100644
--- a/parquet/src/arrow/buffer/view_buffer.rs
+++ b/parquet/src/arrow/buffer/view_buffer.rs
@@ -76,6 +76,10 @@ impl ValuesBuffer for ViewBuffer {
Self::with_capacity(capacity)
}
+ fn reserve_exact(&mut self, additional: usize) {
+ self.views.reserve_exact(additional);
+ }
+
fn pad_nulls(
&mut self,
read_offset: usize,
diff --git a/parquet/src/arrow/record_reader/buffer.rs
b/parquet/src/arrow/record_reader/buffer.rs
index 40fb598e59..70605866b7 100644
--- a/parquet/src/arrow/record_reader/buffer.rs
+++ b/parquet/src/arrow/record_reader/buffer.rs
@@ -20,12 +20,15 @@ use crate::errors::Result;
/// A buffer that supports padding with nulls
pub trait ValuesBuffer {
- /// Create a new buffer with capacity for at least `capacity` elements
- ///
- /// This allows pre-allocating buffers to avoid reallocations during
reading,
- /// improving performance when the number of values is known in advance.
+ /// Create a new buffer with capacity for at least `capacity` logical
values
fn with_capacity(capacity: usize) -> Self;
+ /// Reserve capacity for `additional` more logical values.
+ ///
+ /// Callers should use this once the level decoders have determined how
+ /// many output slots are actually needed for the next read.
+ fn reserve_exact(&mut self, additional: usize);
+
/// If a column contains nulls, more level data may be read than value
data, as null
/// values are not encoded. Therefore, first the levels data is read, the
null count
/// determined, and then the corresponding number of values read to a
[`ValuesBuffer`].
@@ -56,6 +59,10 @@ impl<T: Copy + Default> ValuesBuffer for Vec<T> {
Vec::with_capacity(capacity)
}
+ fn reserve_exact(&mut self, additional: usize) {
+ Vec::reserve_exact(self, additional)
+ }
+
fn pad_nulls(
&mut self,
read_offset: usize,
diff --git a/parquet/src/arrow/record_reader/definition_levels.rs
b/parquet/src/arrow/record_reader/definition_levels.rs
index 9e7345e2d7..0720c6cdbe 100644
--- a/parquet/src/arrow/record_reader/definition_levels.rs
+++ b/parquet/src/arrow/record_reader/definition_levels.rs
@@ -115,8 +115,107 @@ impl DefinitionLevelBuffer {
BufferInner::Mask { nulls } => nulls,
}
}
+
+ /// Returns the raw definition levels accumulated so far, if available.
+ /// Only available when the buffer is in Full mode (nested columns).
+ pub fn levels(&self) -> Option<&[i16]> {
+ match &self.inner {
+ BufferInner::Full { levels, .. } => Some(levels.as_slice()),
+ BufferInner::Mask { .. } => None,
+ }
+ }
}
+/// Build a filtered validity bitmap from definition/repetition levels.
+///
+/// For each level entry where `d >= include_threshold` (when set) and
+/// `r <= max_rep` (when `rep_filter` is provided), appends one bit to
+/// `bitmap`: set when `d >= value_level`, unset otherwise.
+///
+/// Returns the number of bits appended.
+///
+/// This is the shared implementation behind the compact bitmap (leaf
+/// readers with selective padding) and the struct validity bitmap.
+/// Processing uses 64-level chunks with [`compress`] for word-at-a-time
+/// packing.
+pub(crate) fn build_filtered_validity_bitmap(
+ def_levels: &[i16],
+ rep_filter: Option<(&[i16], i16)>,
+ include_threshold: Option<i16>,
+ value_level: i16,
+ bitmap: &mut BooleanBufferBuilder,
+) -> usize {
+ // Fast path: no filtering — every def level produces a bit.
+ if include_threshold.is_none() && rep_filter.is_none() {
+ let chunks = def_levels.chunks_exact(u64::BITS as usize);
+ let remainder = chunks.remainder();
+ for chunk in chunks {
+ let mut word: u64 = 0;
+ for (i, &d) in chunk.iter().enumerate() {
+ word |= ((d >= value_level) as u64) << i;
+ }
+ bitmap.append_word(word, u64::BITS as usize);
+ }
+ for &d in remainder {
+ bitmap.append(d >= value_level);
+ }
+ return def_levels.len();
+ }
+
+ // Filtered path: build include mask + value mask per chunk, compress.
+ let mut item_count: usize = 0;
+ let chunks = def_levels.chunks_exact(u64::BITS as usize);
+ let remainder_offset = def_levels.len() - chunks.remainder().len();
+
+ for (chunk_idx, chunk) in chunks.enumerate() {
+ let base = chunk_idx * u64::BITS as usize;
+ let mut include_mask: u64 = 0;
+ let mut value_mask: u64 = 0;
+ for (i, &d) in chunk.iter().enumerate() {
+ let mut include = true;
+ if let Some(threshold) = include_threshold {
+ if d < threshold {
+ include = false;
+ }
+ }
+ if include {
+ if let Some((reps, max_rep)) = rep_filter {
+ if reps[base + i] > max_rep {
+ include = false;
+ }
+ }
+ }
+ include_mask |= (include as u64) << i;
+ value_mask |= ((d >= value_level) as u64) << i;
+ }
+
+ let count = include_mask.count_ones() as usize;
+ let compact = compress(value_mask, include_mask);
+ bitmap.append_word(compact, count);
+ item_count += count;
+ }
+
+ for idx in remainder_offset..def_levels.len() {
+ let d = def_levels[idx];
+ if let Some(threshold) = include_threshold {
+ if d < threshold {
+ continue;
+ }
+ }
+ if let Some((reps, max_rep)) = rep_filter {
+ if reps[idx] > max_rep {
+ continue;
+ }
+ }
+ bitmap.append(d >= value_level);
+ item_count += 1;
+ }
+
+ item_count
+}
+
+use crate::util::bit_util::compress;
+
enum MaybePacked {
Packed(PackedDecoder),
Fallback(DefinitionLevelDecoderImpl),
diff --git a/parquet/src/arrow/record_reader/mod.rs
b/parquet/src/arrow/record_reader/mod.rs
index 625aa674b2..cdb8084eb4 100644
--- a/parquet/src/arrow/record_reader/mod.rs
+++ b/parquet/src/arrow/record_reader/mod.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use arrow_buffer::Buffer;
+use arrow_buffer::{BooleanBufferBuilder, Buffer};
use crate::arrow::record_reader::{
buffer::ValuesBuffer,
@@ -34,7 +34,7 @@ use crate::errors::{ParquetError, Result};
use crate::schema::types::ColumnDescPtr;
pub(crate) mod buffer;
-mod definition_levels;
+pub(crate) mod definition_levels;
/// A `RecordReader` is a stateful column reader that delimits semantic
records.
pub type RecordReader<T> = GenericRecordReader<Vec<<T as DataType>::T>,
ColumnValueDecoderImpl<T>>;
@@ -62,6 +62,29 @@ pub struct GenericRecordReader<V, CV> {
num_records: usize,
/// Capacity hint for pre-allocating buffers based on batch size
capacity_hint: usize,
+ /// Number of values in the values buffer (may differ from num_values when
+ /// padding_threshold is set, since parent-level padding is excluded).
+ values_written: usize,
+ /// Definition-level threshold used for selective null padding.
+ ///
+ /// With full padding (`None`), the leaf values buffer has one slot for
each
+ /// decoded definition level. This includes placeholders for null or empty
+ /// parent lists, which parent `ListArrayReader`s later have to filter out
+ /// before computing offsets.
+ ///
+ /// With selective padding (`Some(threshold)`), the threshold is the
nearest
+ /// enclosing list/map definition level. Entries with `def < threshold`
+ /// describe a null/empty parent and are skipped entirely. Entries with
+ /// `def >= threshold` belong to an actual child item slot: real values are
+ /// copied, and item-level nulls are padded. The companion `compact_bitmap`
+ /// has the same compact length and becomes the leaf null bitmap.
+ padding_threshold: Option<i16>,
+ /// Compact bitmap accumulated during selective padding. Each bit
+ /// corresponds to an item-level entry (def >= threshold): set when the
+ /// value is real (def >= max_def), unset for item-level nulls. Used both
+ /// as the valid_mask for `pad_nulls` (via `as_slice()`) and as the null
+ /// bitmap consumed by the leaf reader (via `consume_compact_bitmap`).
+ compact_bitmap: Option<BooleanBufferBuilder>,
}
impl<V, CV> GenericRecordReader<V, CV>
@@ -71,9 +94,8 @@ where
{
/// Create a new [`GenericRecordReader`]
///
- /// The capacity is used to pre-allocate internal buffers, avoiding
reallocations
- /// when reading the first batch of data. For optimal performance, set
this to
- /// the expected batch size.
+ /// The capacity is used to pre-allocate internal buffers for full-padding
+ /// reads, avoiding reallocations when reading fragmented row selections.
pub fn new(desc: ColumnDescPtr, capacity: usize) -> Self {
let def_levels = (desc.max_def_level() > 0)
.then(|| DefinitionLevelBuffer::new(&desc,
packed_null_mask(&desc)));
@@ -89,6 +111,9 @@ where
num_values: 0,
num_records: 0,
capacity_hint: capacity,
+ values_written: 0,
+ padding_threshold: None,
+ compact_bitmap: None,
}
}
@@ -183,19 +208,52 @@ where
self.values.take().unwrap_or_else(|| V::with_capacity(0))
}
- /// Returns currently stored null bitmap data for nullable columns.
- /// For non-nullable columns, the bitmap is discarded.
- /// The side effect is similar to `consume_def_levels`.
- pub fn consume_bitmap_buffer(&mut self) -> Option<Buffer> {
- self.consume_bitmap()
- }
-
/// Reset state of record reader.
/// Should be called after consuming data, e.g. `consume_rep_levels`,
- /// `consume_rep_levels`, `consume_record_data` and
`consume_bitmap_buffer`.
+ /// `consume_rep_levels`, `consume_record_data` and
`consume_compact_bitmap`.
pub fn reset(&mut self) {
self.num_values = 0;
self.num_records = 0;
+ self.values_written = 0;
+ self.compact_bitmap = None;
+ }
+
+ /// Returns the maximum definition level for the column being read.
+ pub fn max_def_level(&self) -> i16 {
+ self.column_desc.max_def_level()
+ }
+
+ /// Set the padding threshold. When set, `pad_nulls` only pads entries
+ /// where `def >= threshold` (item-level nulls within non-null lists),
+ /// skipping list-level padding entries (def < threshold).
+ pub fn set_padding_threshold(&mut self, threshold: i16) {
+ self.padding_threshold = Some(threshold);
+ }
+
+ /// Returns the number of values in the values buffer.
+ /// When padding_threshold is None, this equals `num_values` (full
padding).
+ /// When padding_threshold is set, this is the item_count (selective
padding).
+ pub fn values_written(&self) -> usize {
+ if self.padding_threshold.is_some() {
+ self.values_written
+ } else {
+ self.num_values
+ }
+ }
+
+ /// Consume the compact null bitmap built during selective padding.
+ /// Returns the full bitmap when not using selective padding.
+ pub fn consume_compact_bitmap(&mut self) -> Option<Buffer> {
+ if self.padding_threshold.is_some() {
+ if let Some(levels) = self.def_levels.as_mut() {
+ levels.consume_bitmask();
+ }
+ self.compact_bitmap
+ .as_mut()
+ .map(|b| b.finish().into_inner())
+ } else {
+ self.consume_bitmap()
+ }
}
/// Returns bitmap data for nullable columns.
@@ -222,30 +280,106 @@ where
if batch_size == 0 {
return Ok(0);
}
- // Update capacity hint to the largest batch size seen
+ // Update capacity hint to the largest batch size seen.
if batch_size > self.capacity_hint {
self.capacity_hint = batch_size;
}
- // Lazily initialize buffer on first read
- let capacity_hint = self.capacity_hint;
- let values = self
- .values
- .get_or_insert_with(|| V::with_capacity(capacity_hint));
+ let padding_threshold = self.padding_threshold;
+ let max_def = self.column_desc.max_def_level();
+ let mut physical_values_to_read = 0usize;
+ let mut output_slots_to_read = 0usize;
+
+ let values = self.values.get_or_insert_with(|| V::with_capacity(0));
+ let compact_bitmap = &mut self.compact_bitmap;
+
+ if padding_threshold.is_none() {
+
values.reserve_exact(self.capacity_hint.saturating_sub(self.num_values));
+ }
- let (records_read, values_read, levels_read) =
- self.column_reader.as_mut().unwrap().read_records(
+ let (records_read, values_read, levels_read) = self
+ .column_reader
+ .as_mut()
+ .unwrap()
+ .read_records_with_reservation(
batch_size,
self.def_levels.as_mut(),
self.rep_levels.as_mut(),
values,
+ |values, values_to_read, levels_to_read, def_levels| {
+ let output_slots = if let Some(threshold) =
padding_threshold {
+ let def_levels = def_levels.ok_or_else(|| {
+ general_err!(
+ "Definition levels should exist when data is
less than levels!"
+ )
+ })?;
+ let all_levels = def_levels.levels().ok_or_else(|| {
+ general_err!(
+ "Raw definition levels must be available for
selective padding"
+ )
+ })?;
+ let batch_levels = &all_levels[all_levels.len() -
levels_to_read..];
+ let bitmap =
+ compact_bitmap.get_or_insert_with(||
BooleanBufferBuilder::new(0));
+
+ definition_levels::build_filtered_validity_bitmap(
+ batch_levels,
+ None,
+ Some(threshold),
+ max_def,
+ bitmap,
+ )
+ } else {
+ levels_to_read
+ };
+
+ let additional = output_slots_to_read
+ .saturating_add(output_slots)
+ .saturating_sub(physical_values_to_read);
+ values.reserve_exact(additional);
+
+ output_slots_to_read += output_slots;
+ physical_values_to_read += values_to_read;
+ Ok(())
+ },
)?;
- if values_read < levels_read {
+ if self.padding_threshold.is_some() {
+ debug_assert_eq!(
+ physical_values_to_read, values_read,
+ "reservation accounting must match decoded values"
+ );
+ let item_count = output_slots_to_read;
+ let bitmap = compact_bitmap.get_or_insert_with(||
BooleanBufferBuilder::new(0));
+
+ // Pad values to item_count positions (only if there are gaps).
+ if values_read < item_count {
+ values.reserve_exact(item_count - values_read);
+ values.pad_nulls(
+ self.values_written,
+ values_read,
+ item_count,
+ bitmap.as_slice(),
+ )?;
+ }
+
+ debug_assert_eq!(
+ bitmap.len(),
+ self.values_written + item_count,
+ "compact bitmap length must equal total items written"
+ );
+ self.values_written += item_count;
+ } else if values_read < levels_read {
+ debug_assert_eq!(
+ output_slots_to_read, levels_read,
+ "reservation accounting must match decoded levels"
+ );
+ // Full padding: pad all null positions.
let def_levels = self.def_levels.as_ref().ok_or_else(|| {
general_err!("Definition levels should exist when data is less
than levels!")
})?;
+ values.reserve_exact(levels_read - values_read);
values.pad_nulls(
self.num_values,
values_read,
@@ -354,6 +488,39 @@ mod tests {
assert_eq!(None, record_reader.consume_bitmap());
}
+ #[test]
+ fn test_capacity_hint_preserved_across_fragmented_reads() {
+ let message_type = "
+ message test_schema {
+ REQUIRED INT32 leaf;
+ }
+ ";
+ let desc = parse_message_type(message_type)
+ .map(|t| SchemaDescriptor::new(Arc::new(t)))
+ .map(|s| s.column(0))
+ .unwrap();
+
+ let mut record_reader = RecordReader::<Int32Type>::new(desc.clone(),
8);
+ let values = [1, 2, 3, 4];
+ let mut pb = DataPageBuilderImpl::new(desc, 4, true);
+ pb.add_values::<Int32Type>(Encoding::PLAIN, &values);
+ let page = pb.consume();
+
+ let page_reader = Box::new(InMemoryPageReader::new(vec![page]));
+ record_reader.set_page_reader(page_reader).unwrap();
+
+ assert_eq!(1, record_reader.read_records(1).unwrap());
+ assert_eq!(1, record_reader.read_records(1).unwrap());
+
+ let record_data = record_reader.consume_record_data();
+ assert_eq!(record_data, &[1, 2]);
+ assert!(
+ record_data.capacity() >= 8,
+ "capacity hint should survive fragmented reads, got {}",
+ record_data.capacity()
+ );
+ }
+
#[test]
fn test_read_optional_records() {
// Construct column schema
@@ -561,6 +728,95 @@ mod tests {
}
}
+ #[test]
+ fn test_selective_padding_consumes_full_bitmap() {
+ let message_type = "
+ message test_schema {
+ OPTIONAL GROUP my_list (LIST) {
+ REPEATED GROUP list {
+ OPTIONAL INT32 element;
+ }
+ }
+ }
+ ";
+
+ let desc = parse_message_type(message_type)
+ .map(|t| SchemaDescriptor::new(Arc::new(t)))
+ .map(|s| s.column(0))
+ .unwrap();
+
+ let mut record_reader = RecordReader::<Int32Type>::new(desc.clone(),
DEFAULT_BATCH_SIZE);
+ record_reader.set_padding_threshold(2);
+
+ let values = [10, 20];
+ let def_levels = [3i16, 2i16, 0i16, 2i16, 3i16];
+ let rep_levels = [0i16, 1i16, 0i16, 0i16, 1i16];
+ let mut pb = DataPageBuilderImpl::new(desc, 5, true);
+ pb.add_rep_levels(1, &rep_levels);
+ pb.add_def_levels(3, &def_levels);
+ pb.add_values::<Int32Type>(Encoding::PLAIN, &values);
+
+ let page_reader =
Box::new(InMemoryPageReader::new(vec![pb.consume()]));
+ record_reader.set_page_reader(page_reader).unwrap();
+ assert_eq!(3, record_reader.read_records(3).unwrap());
+
+ let expected_compact = Buffer::from_iter([true, false, false, true]);
+ assert_eq!(
+ Some(expected_compact),
+ record_reader.consume_compact_bitmap()
+ );
+ assert_eq!(None, record_reader.consume_bitmap());
+ }
+
+ #[test]
+ fn test_selective_padding_reserves_compact_capacity() {
+ let message_type = "
+ message test_schema {
+ OPTIONAL GROUP my_list (LIST) {
+ REPEATED GROUP list {
+ OPTIONAL INT32 element;
+ }
+ }
+ }
+ ";
+
+ let desc = parse_message_type(message_type)
+ .map(|t| SchemaDescriptor::new(Arc::new(t)))
+ .map(|s| s.column(0))
+ .unwrap();
+
+ let mut record_reader = RecordReader::<Int32Type>::new(desc.clone(),
DEFAULT_BATCH_SIZE);
+ record_reader.set_padding_threshold(2);
+
+ let values = [10];
+ let mut def_levels = vec![0i16; 30];
+ def_levels.extend([3i16, 2i16]);
+ let rep_levels = vec![0i16; def_levels.len()];
+
+ let mut pb = DataPageBuilderImpl::new(desc, def_levels.len() as u32,
true);
+ pb.add_rep_levels(1, &rep_levels);
+ pb.add_def_levels(3, &def_levels);
+ pb.add_values::<Int32Type>(Encoding::PLAIN, &values);
+
+ let page_reader =
Box::new(InMemoryPageReader::new(vec![pb.consume()]));
+ record_reader.set_page_reader(page_reader).unwrap();
+ assert_eq!(def_levels.len(), record_reader.read_records(128).unwrap());
+
+ let expected_compact = Buffer::from_iter([true, false]);
+ assert_eq!(
+ Some(expected_compact),
+ record_reader.consume_compact_bitmap()
+ );
+
+ let actual = record_reader.consume_record_data();
+ assert_eq!(actual.len(), 2);
+ assert_eq!(actual[0], 10);
+ assert!(
+ actual.capacity() < def_levels.len(),
+ "selective padding should not reserve one child slot per
list-level NULL"
+ );
+ }
+
#[test]
fn test_read_more_than_one_batch() {
// Construct column schema
diff --git a/parquet/src/column/reader.rs b/parquet/src/column/reader.rs
index 5ab78038f4..81960abbe4 100644
--- a/parquet/src/column/reader.rs
+++ b/parquet/src/column/reader.rs
@@ -200,12 +200,34 @@ where
/// `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
pub fn read_records(
+ &mut self,
+ max_records: usize,
+ def_levels: Option<&mut D::Buffer>,
+ rep_levels: Option<&mut R::Buffer>,
+ values: &mut V::Buffer,
+ ) -> Result<(usize, usize, usize)> {
+ self.read_records_with_reservation(
+ max_records,
+ def_levels,
+ rep_levels,
+ values,
+ |_, _, _, _| Ok(()),
+ )
+ }
+
+ /// Like [`Self::read_records`], but invokes `reserve_values` after levels
+ /// are decoded and before values are read.
+ pub(crate) fn read_records_with_reservation<F>(
&mut self,
max_records: usize,
mut def_levels: Option<&mut D::Buffer>,
mut rep_levels: Option<&mut R::Buffer>,
values: &mut V::Buffer,
- ) -> Result<(usize, usize, usize)> {
+ mut reserve_values: F,
+ ) -> Result<(usize, usize, usize)>
+ where
+ F: FnMut(&mut V::Buffer, usize, usize, Option<&D::Buffer>) ->
Result<()>,
+ {
let mut total_records_read = 0;
let mut total_levels_read = 0;
let mut total_values_read = 0;
@@ -262,6 +284,9 @@ where
None => levels_to_read,
};
+ let def_levels = def_levels.as_deref();
+ reserve_values(values, values_to_read, levels_to_read,
def_levels)?;
+
let values_read = self.values_decoder.read(values,
values_to_read)?;
if values_read != values_to_read {
diff --git a/parquet/src/util/bit_util.rs b/parquet/src/util/bit_util.rs
index 326f4a61e2..d5f5945a12 100644
--- a/parquet/src/util/bit_util.rs
+++ b/parquet/src/util/bit_util.rs
@@ -921,6 +921,44 @@ impl From<Vec<u8>> for BitReader {
}
}
+/// Parallel bit extract: for each set bit in `mask`, extract the
+/// corresponding bit from `value` and pack them contiguously into the low
+/// bits of the return value.
+///
+/// Equivalent to the x86 BMI2 `PEXT` instruction. When compiled with the
+/// `bmi2` target feature enabled (for example `-C target-cpu=x86-64-v3`)
+/// this lowers to the hardware `pext` instruction; otherwise it falls back
+/// to a portable scalar loop.
+///
+/// Replace with `value.compress(mask)` when `uint_gather_scatter_bits`
+/// is stabilised: <https://github.com/rust-lang/rust/issues/149069>
+#[allow(dead_code)]
+#[inline]
+pub(crate) fn compress(value: u64, mask: u64) -> u64 {
+ #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
+ {
+ // SAFETY: the `bmi2` target feature is statically enabled for this
+ // build, so the `pext` instruction is guaranteed to be available.
+ unsafe { std::arch::x86_64::_pext_u64(value, mask) }
+ }
+
+ #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))]
+ {
+ let mut mask = mask;
+ let mut result: u64 = 0;
+ let mut dest_bit: u64 = 1;
+ while mask != 0 {
+ let lowest = mask & mask.wrapping_neg();
+ if value & lowest != 0 {
+ result |= dest_bit;
+ }
+ dest_bit <<= 1;
+ mask ^= lowest;
+ }
+ result
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -929,6 +967,44 @@ mod tests {
use rand::distr::{Distribution, StandardUniform};
use std::fmt::Debug;
+ #[test]
+ fn test_compress() {
+ // Reference: gather the `mask`-selected bits of `value` into
+ // contiguous low bits, least-significant first.
+ fn reference(value: u64, mut mask: u64) -> u64 {
+ let mut result = 0u64;
+ let mut dest = 0u32;
+ while mask != 0 {
+ let lowest = mask & mask.wrapping_neg();
+ result |= (((value & lowest) != 0) as u64) << dest;
+ dest += 1;
+ mask ^= lowest;
+ }
+ result
+ }
+
+ // Hand-picked edge cases.
+ assert_eq!(compress(0b1010, 0b1111), 0b1010);
+ assert_eq!(compress(0b1010, 0b1010), 0b11);
+ assert_eq!(compress(0b1010, 0b0101), 0);
+ assert_eq!(compress(u64::MAX, 0), 0);
+ assert_eq!(compress(0, u64::MAX), 0);
+ assert_eq!(compress(u64::MAX, u64::MAX), u64::MAX);
+
+ // Randomised cross-check against the reference. On a `bmi2` build
+ // this validates the hardware `pext` path; otherwise it exercises
+ // the portable fallback.
+ let values = random_numbers::<u64>(1024);
+ let masks = random_numbers::<u64>(1024);
+ for (&value, &mask) in values.iter().zip(masks.iter()) {
+ assert_eq!(
+ compress(value, mask),
+ reference(value, mask),
+ "value={value:#x} mask={mask:#x}"
+ );
+ }
+ }
+
#[test]
fn test_ceil() {
assert_eq!(ceil(0, 1), 0);
