adriangb commented on code in PR #11157:
URL: https://github.com/apache/arrow-rs/pull/11157#discussion_r4078078049
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -774,6 +776,32 @@ impl ArrowReaderOptions {
self
}
+ /// Sets the [`ColumnChunkMask`] for the Parquet [ColumnIndex] structure.
+ ///
+ /// The column index can be costly to decode and store, especially when it
is needed
+ /// only for a subset of row groups or columns (such as when filtering by
a predicate
+ /// on a single column). Providing a [`ColumnChunkMask`] can greatly
decrease
+ /// the time needed to decode this metadata.
+ ///
+ /// [ColumnIndex]:
https://github.com/apache/parquet-format/blob/master/PageIndex.md
+ pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self {
+ self.column_index_mask = mask;
+ self
+ }
+
+ /// Sets the [`ColumnChunkMask`] for the Parquet [OffsetIndex] structure.
+ ///
+ /// The offset index can be costly to decode and store, especially when it
is needed
+ /// only for a subset of row groups or columns (such as when projecting a
small subset
+ /// of columns). Providing a [`ColumnChunkMask`] can greatly decrease
+ /// the time needed to decode this metadata.
+ ///
+ /// [OffsetIndex]:
https://github.com/apache/parquet-format/blob/master/PageIndex.md
+ pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self {
+ self.offset_index_mask = mask;
+ self
+ }
Review Comment:
**Docs / Bug (C3).** These docs suggest a column index mask for the
predicate columns and an offset index mask for the projected columns. With that
setup, `StatisticsConverter` returns arrays of different lengths (see the
review body):
```text
CI columns([0]) (predicate), OI columns([1]) (projection): data_page_mins
len 30, data_page_row_counts len 0
```
Page pruning needs the offset index of the predicate columns too. The
suggestion also adds two other traps (C7): a mask without a policy does
nothing, and `ArrowReaderMetadata::try_new` ignores the masks.
```suggestion
/// Sets the [`ColumnChunkMask`] for the Parquet [ColumnIndex] structure.
///
/// The column index can be costly to decode and store, especially when
it is needed
/// only for a subset of row groups or columns (such as when filtering
by a predicate
/// on a single column). Providing a [`ColumnChunkMask`] can greatly
decrease
/// the time needed to decode this metadata.
///
/// The mask applies only if the column index policy is not
[`PageIndexPolicy::Skip`]
/// (the default), and only when the page index is loaded with these
options (for
/// example by [`ArrowReaderMetadata::load`]).
[`ArrowReaderMetadata::try_new`] ignores it.
///
/// [ColumnIndex]:
https://github.com/apache/parquet-format/blob/master/PageIndex.md
pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self {
self.column_index_mask = mask;
self
}
/// Sets the [`ColumnChunkMask`] for the Parquet [OffsetIndex] structure.
///
/// The offset index can be costly to decode and store, especially when
it is needed
/// only for a subset of row groups or columns (such as when projecting
a small subset
/// of columns). Providing a [`ColumnChunkMask`] can greatly decrease
/// the time needed to decode this metadata.
///
/// Page pruning with the column index also needs the offset index of
the same
/// columns. So include the predicate columns in this mask, not only the
projected
/// columns. The same notes as for [`Self::with_column_index_mask`]
apply.
///
/// [OffsetIndex]:
https://github.com/apache/parquet-format/blob/master/PageIndex.md
pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self {
self.offset_index_mask = mask;
self
}
```
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -488,6 +619,8 @@ impl ParquetMetaDataReader {
let push_decoder =
ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)?
.with_offset_index_policy(self.offset_index)
.with_column_index_policy(self.column_index)
+ .with_offset_index_mask(self.offset_index_mask.clone())
+ .with_column_index_mask(self.column_index_mask.clone())
Review Comment:
**Bug (C1).** With a prefetch hint and a mask, this call returns the page
index of another column chunk, with no error (also with `Required`). Without a
mask, the same call fails. Details, a diagram and a one-line fix are in the
review body.
```text
CI mask rg 0 x col 1, prefetch = file_len - 27: c1 page min/max = [2]/[11]
(true: [1]/[10], these are c2's stats)
OI mask rg 0 x col 1, prefetch = file_len - 11: c1 page offset = 232
(true: 145, 232 is c2's data_page_offset)
same prefetch, no mask: Err("Corrupted parquet
file: index data range (1372..1503) exceeds remainder length (1492)")
```
Cause: `load_metadata_via_suffix` (L826) returns the remainder with start
`0`, but its real file offset is `file_len - suffix_len`. The mask makes the
covering range short enough that the wrong slice is in bounds.
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -104,6 +107,114 @@ impl From<bool> for PageIndexPolicy {
}
}
+/// Struct to specify column chunks for which metadata is required.
+///
+/// Column chunks are identified by row group index and column index. This
struct
+/// allows for specifying vertical slices of column chunk data (via
[`Self::columns`]),
+/// horizontal slices (via [`Self::row_groups`]), or the intersection of the
two
+/// (via [`Self::row_groups_and_columns`]).
+///
+/// At present this is only used to select elements of the [Page Index] for
decoding.
+///
+/// # Examples
+///
+/// To select columns 0 and 1 from all row groups:
+/// ```rust
+/// # use parquet::file::metadata::ColumnChunkMask;
+/// let mask = ColumnChunkMask::columns([0, 1]);
+/// ```
+///
+/// To select all columns from row group 2:
+/// ```rust
+/// # use parquet::file::metadata::ColumnChunkMask;
+/// let mask = ColumnChunkMask::row_groups([2]);
+/// ```
+///
+/// To select columns 1 and 3 from row group 0:
+/// ```rust
+/// # use parquet::file::metadata::ColumnChunkMask;
+/// let mask = ColumnChunkMask::row_groups_and_columns([0], [1, 3]);
+/// ```
+///
+/// [Page Index]: https://parquet.apache.org/docs/file-format/pageindex/
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+pub struct ColumnChunkMask {
+ // using i32 because that's how thrift vectors are sized
+ row_groups: Option<Arc<BTreeSet<i32>>>,
+ columns: Option<Arc<BTreeSet<i32>>>,
+}
+
+impl ColumnChunkMask {
+ /// Select all row groups and columns.
+ pub fn all() -> Self {
+ Self::default()
+ }
+
+ /// Select only the listed columns.
+ ///
+ /// Any indices in `columns` that are less than zero will be ignored.
Passing an empty
+ /// set is treated the same as selecting all columns.
+ pub fn columns(columns: impl IntoIterator<Item = i32>) -> Self {
Review Comment:
**API, one-way door (C5).** The constructors take `i32`. The other APIs
around this one use `usize`: `includes_*` here, `ProjectionMask::leaves`,
`with_row_groups`, `PageIndexProvider::column_index`,
`ParquetStatisticsPolicy::skip_except`. So every caller writes `.map(|i| i as
i32)`. That cast wraps silently, and a negative result then means "ignore"
(and, with C4, "all"). A related detail:
`ColumnChunkMask::all().includes_column(u32::MAX as usize)` is `false`.
Suggestion: take `impl IntoIterator<Item = usize>` in the public API. Keep
`i32`/`u32` storage internal if you want it. A change after release breaks
every caller, so this is the time to decide.
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -104,6 +107,114 @@ impl From<bool> for PageIndexPolicy {
}
}
+/// Struct to specify column chunks for which metadata is required.
+///
+/// Column chunks are identified by row group index and column index. This
struct
+/// allows for specifying vertical slices of column chunk data (via
[`Self::columns`]),
+/// horizontal slices (via [`Self::row_groups`]), or the intersection of the
two
+/// (via [`Self::row_groups_and_columns`]).
Review Comment:
**Docs (C6).** The indexes are leaf (Parquet column) indexes, not Arrow
field indexes. For schema `a, s: {x, y}, b`, `columns([2])` selects `s.y`, not
`b`. A caller that passes a root field index loads the index of another column,
with no error.
```suggestion
/// Column chunks are identified by row group index and leaf column index
(the index of the
/// column in [`SchemaDescriptor::columns`], not the index of a root or
Arrow field). This struct
/// allows for specifying vertical slices of column chunk data (via
[`Self::columns`]),
/// horizontal slices (via [`Self::row_groups`]), or the intersection of the
two
/// (via [`Self::row_groups_and_columns`]).
```
##########
parquet/src/file/metadata/push_decoder.rs:
##########
@@ -481,25 +514,81 @@ enum DecodeState {
Intermediate,
}
-/// Returns the byte range needed to read the offset/page indexes, based on the
-/// specified policies
+/// Returns the minimum set of non-overlapping ranges needed to cover the
requested
+/// offset and column indexes.
+///
+/// If no page indexes are present in the file, or none are actually
+/// requested by the policies passed in, this will return an empty
+/// vector.
+///
+/// This may result in more ranges to fetch but may also result in fewer bytes
+/// being read (and stored). To get a single range for all indexes, sort
+/// the resultant vector by the range starts, and then create a single range
+/// using `start` from the head and `end` from the tail.
///
-/// Returns None if no page indexes are needed
-pub fn range_for_page_index(
+/// ```ignore
+/// # use core::ops::Range;
+/// # fn coalesce_page_index_ranges(ranges: &mut Vec<Range<u64>>) ->
Option<Range<u64>> {
+/// ranges.sort_by_key(|r| r.start);
+/// let range = (ranges.first()?.start..ranges.last()?.end);
+/// Some(range)
+/// # }
+/// ```
+pub fn ranges_for_page_index(
metadata: &ParquetMetaData,
column_index_policy: PageIndexPolicy,
offset_index_policy: PageIndexPolicy,
-) -> Option<Range<u64>> {
- let mut range = None;
- for c in metadata.row_groups().iter().flat_map(|r| r.columns()) {
- if column_index_policy != PageIndexPolicy::Skip {
- range = acc_range(range, c.column_index_range());
- }
- if offset_index_policy != PageIndexPolicy::Skip {
- range = acc_range(range, c.offset_index_range());
+ column_index_mask: &ColumnChunkMask,
+ offset_index_mask: &ColumnChunkMask,
+) -> Vec<Range<u64>> {
+ let mut result = Vec::new();
+
+ fn add_ranges<T>(
+ metadata: &ParquetMetaData,
+ mask: &ColumnChunkMask,
+ ranges: &mut Vec<Range<u64>>,
+ f: T,
+ ) where
+ T: Fn(&ColumnChunkMetaData) -> Option<Range<u64>>,
+ {
+ for (rg_idx, rg) in metadata.row_groups().iter().enumerate() {
+ if mask.includes_row_group(rg_idx) {
+ for (col_idx, col) in rg.columns().iter().enumerate() {
+ if mask.includes_column(col_idx)
+ && let Some(range) = f(col)
+ {
+ // ranges shouldn't overlap, so only check for
contiguous ranges
+ // [s1..e1], [s2..e2] where e1 == s2
+ if let Some(last) = ranges.last_mut()
+ && last.end == range.start
+ {
+ last.end = range.end;
+ } else {
+ ranges.push(range);
+ }
Review Comment:
**Perf / Robustness (C17, C8).** This merges a range only if it starts
exactly where the last emitted range ends. So:
- Ranges that are adjacent in the file, but not in emission order (CI rg x
col, then OI rg x col), stay separate.
- Duplicate or overlapping ranges stay separate. A crafted footer where 4
column indexes point at the same range makes the caller fetch that range 4
times: `requested 7 ranges, 46402 bytes total, file is 12460 bytes: [4..11582,
4..11582, 4..11582, 4..11582, 11446..11468, 11490..11534, 11558..11582]`. The
sync reader fetches the covering range once.
The doc comment (L517) says "minimum set of non-overlapping ranges", which
is not true in these cases. Fix: push every range, then sort and merge in place
(no extra allocation). With this change the crafted file gives `[4..11582]`,
and the existing tests pass (4 ranges for `columns([0])`, 1 range for `all()`):
```diff
- // ranges shouldn't overlap, so only check for
contiguous ranges
- // [s1..e1], [s2..e2] where e1 == s2
- if let Some(last) = ranges.last_mut()
- && last.end == range.start
- {
- last.end = range.end;
- } else {
- ranges.push(range);
- }
+ ranges.push(range);
...
+ // merge overlapping and adjacent ranges, whatever the emission order
+ result.sort_unstable_by_key(|r| r.start);
+ result.dedup_by(|next, prev| {
+ let merge = next.start <= prev.end;
+ if merge {
+ prev.end = prev.end.max(next.end);
+ }
+ merge
+ });
result
```
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -1392,4 +1525,32 @@ mod async_tests {
read_and_check(f.as_file(), PageIndexPolicy::Optional).unwrap();
read_and_check(f.as_file(), PageIndexPolicy::Skip).unwrap();
}
+
+ #[test]
+ fn test_chunk_mask() {
Review Comment:
**Test gap (N1).** `test_chunk_mask` is in `mod async_tests`, which has
`#[cfg(all(feature = "async", feature = "arrow", test))]`. The default features
of `parquet` do not include `async`, so `cargo test -p parquet` does not run
it. The test uses no async code. Move it to `mod tests` (L977).
##########
parquet/tests/arrow_reader/page_index.rs:
##########
@@ -0,0 +1,236 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Tests of selective page index population
+
+use parquet::file::metadata::{ColumnChunkMask, PageIndexPolicy,
ParquetMetaDataReader};
+
+use crate::custom_page_index_provider::create_test_file;
+
+#[test]
+fn test_parse_selected_columns() {
+ // test populating PageIndex with a subset of columns
+ let temp_file = create_test_file();
+
+ // populate column 0 for column index and columns 0 & 2 for the offset
index
+ let mut reader = ParquetMetaDataReader::new()
+ .with_page_index_policy(PageIndexPolicy::Optional)
+ .with_column_index_mask(ColumnChunkMask::columns([0]))
+ .with_offset_index_mask(ColumnChunkMask::columns([0, 2]));
+
+ // parse metadata
+ reader.try_parse(&temp_file).unwrap();
+ let metadata = reader.finish().unwrap();
+ let num_rg = metadata.num_row_groups();
+
+ // test page indexes
+ for rg in 0..num_rg {
+ let idx = metadata.page_index_for_row_group(rg);
+ // column 0 has both indexes
+ assert!(idx.column_index(0).is_some());
+ assert!(idx.offset_index(0).is_some());
+ // column 1 has no indexes
+ assert!(idx.column_index(1).is_none());
+ assert!(idx.offset_index(1).is_none());
+ // column 2 has offset index
+ assert!(idx.column_index(2).is_none());
+ assert!(idx.offset_index(2).is_some());
+ // column 3 has no indexes
+ assert!(idx.column_index(3).is_none());
+ assert!(idx.offset_index(3).is_none());
+ }
+}
+
+#[test]
+fn test_parse_selected_columns_mixed() {
+ // test populating PageIndex with a subset of columns
+ let temp_file = create_test_file();
+
+ // populate column 0 for column index and all columns for offset index
+ let mut reader = ParquetMetaDataReader::new()
+ .with_page_index_policy(PageIndexPolicy::Optional)
+ .with_column_index_mask(ColumnChunkMask::columns([0]))
+ .with_offset_index_mask(ColumnChunkMask::all());
+
+ // parse metadata
+ reader.try_parse(&temp_file).unwrap();
+ let metadata = reader.finish().unwrap();
+ let num_rg = metadata.num_row_groups();
+
+ // test page indexes
+ for rg in 0..num_rg {
+ let idx = metadata.page_index_for_row_group(rg);
+ // column 0 has both indexes
+ assert!(idx.column_index(0).is_some());
+ assert!(idx.offset_index(0).is_some());
+ // column 1 has no indexes
+ assert!(idx.column_index(1).is_none());
+ assert!(idx.offset_index(1).is_some());
+ // column 2 has offset index
+ assert!(idx.column_index(2).is_none());
+ assert!(idx.offset_index(2).is_some());
+ // column 3 has no indexes
+ assert!(idx.column_index(3).is_none());
+ assert!(idx.offset_index(3).is_some());
+ }
+}
+
+#[test]
+fn test_parse_selected_row_groups() {
+ // test populating PageIndex with a subset of row groups
+ let temp_file = create_test_file();
+
+ // populate indexes for row groups 0 and 2
+ let mut reader = ParquetMetaDataReader::new()
+ .with_page_index_policy(PageIndexPolicy::Optional)
+ .with_column_index_mask(ColumnChunkMask::row_groups([0, 2]))
+ .with_offset_index_mask(ColumnChunkMask::row_groups([0, 2]));
+
+ // parse metadata
+ reader.try_parse(&temp_file).unwrap();
+ let metadata = reader.finish().unwrap();
+ let num_cols = metadata.file_metadata().schema_descr().num_columns();
+ assert!(metadata.page_index().is_some());
+ let page_index = metadata.page_index().unwrap();
+
+ for col in 0..num_cols {
+ // row group 0 has both indexes
+ assert!(page_index.column_index(0, col).is_some());
+ assert!(page_index.offset_index(0, col).is_some());
+ // row group 1 has no indexes
+ assert!(page_index.column_index(1, col).is_none());
+ assert!(page_index.offset_index(1, col).is_none());
+ // row group 2 has both indexes
+ assert!(page_index.column_index(2, col).is_some());
+ assert!(page_index.offset_index(2, col).is_some());
+ }
+}
+
+#[test]
+fn test_parse_selected_row_groups_and_columns() {
+ // test populating PageIndex by row group and column
+ let temp_file = create_test_file();
+
+ // populate only row group 1, column index gets column 0, offset index gets
+ // columns 0 and 2.
+ let mut reader = ParquetMetaDataReader::new()
+ .with_page_index_policy(PageIndexPolicy::Optional)
+ .with_column_index_mask(ColumnChunkMask::row_groups_and_columns([1],
[0]))
+ .with_offset_index_mask(ColumnChunkMask::row_groups_and_columns([1],
[0, 2]));
+
+ // parse metadata
+ reader.try_parse(&temp_file).unwrap();
+ let metadata = reader.finish().unwrap();
+ assert!(metadata.page_index().is_some());
+ let num_cols = metadata.file_metadata().schema_descr().num_columns();
+
+ // row groups 0 and 2 should have no indexes
+ [0, 2].into_iter().for_each(|i| {
+ let rg_idx = metadata.page_index_for_row_group(i);
+ for col in 0..num_cols {
+ assert!(rg_idx.column_index(col).is_none());
+ assert!(rg_idx.offset_index(col).is_none());
+ }
+ });
+
+ // row group 1 should have column index for column 0 and offset index for
columns 0 & 2
+ let rg_idx = metadata.page_index_for_row_group(1);
+ assert!(rg_idx.column_index(0).is_some());
+ assert!(rg_idx.offset_index(0).is_some());
+ assert!(rg_idx.column_index(1).is_none());
+ assert!(rg_idx.offset_index(1).is_none());
+ assert!(rg_idx.column_index(2).is_none());
+ assert!(rg_idx.offset_index(2).is_some());
+ assert!(rg_idx.column_index(3).is_none());
+ assert!(rg_idx.offset_index(3).is_none());
+}
+
+#[test]
+fn test_page_index_sizes() {
+ // test populating PageIndex by row group and column
+ let temp_file = create_test_file();
+
+ // no index
+ let mut reader =
ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Skip);
+
+ // parse metadata
+ reader.try_parse(&temp_file).unwrap();
+ let metadata = reader.finish().unwrap();
+ assert!(metadata.page_index().is_none());
+ #[cfg(not(feature = "encryption"))]
+ assert_eq!(metadata.memory_size(), 7393);
+ #[cfg(feature = "encryption")]
+ assert_eq!(metadata.memory_size(), 7817);
+
+ // full index
+ let mut reader =
ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required);
+
+ // parse metadata
+ reader.try_parse(&temp_file).unwrap();
+ let metadata = reader.finish().unwrap();
+ assert!(metadata.page_index().is_some());
+ #[cfg(not(feature = "encryption"))]
+ assert_eq!(metadata.memory_size(), 13897);
+ #[cfg(feature = "encryption")]
+ assert_eq!(metadata.memory_size(), 14321);
+
+ // populate column 0 for column index and all columns for offset index
+ let mut reader = ParquetMetaDataReader::new()
+ .with_page_index_policy(PageIndexPolicy::Optional)
+ .with_column_index_mask(ColumnChunkMask::columns([0]))
+ .with_offset_index_mask(ColumnChunkMask::all());
+
+ // parse metadata
+ reader.try_parse(&temp_file).unwrap();
+ let metadata = reader.finish().unwrap();
+ assert!(metadata.page_index().is_some());
+ #[cfg(not(feature = "encryption"))]
+ assert_eq!(metadata.memory_size(), 12776);
+ #[cfg(feature = "encryption")]
+ assert_eq!(metadata.memory_size(), 13200);
+
+ // populate column 0 for column index and columns 0 & 2 for the offset
index
+ let mut reader = ParquetMetaDataReader::new()
+ .with_page_index_policy(PageIndexPolicy::Optional)
+ .with_column_index_mask(ColumnChunkMask::columns([0]))
+ .with_offset_index_mask(ColumnChunkMask::columns([0, 2]));
+
+ // parse metadata
+ reader.try_parse(&temp_file).unwrap();
+ let metadata = reader.finish().unwrap();
+ assert!(metadata.page_index().is_some());
+ #[cfg(not(feature = "encryption"))]
+ assert_eq!(metadata.memory_size(), 12056);
+ #[cfg(feature = "encryption")]
+ assert_eq!(metadata.memory_size(), 12480);
+
+ // populate only row group 1, column index gets column 0, offset index gets
+ // columns 0 and 2.
+ let mut reader = ParquetMetaDataReader::new()
+ .with_page_index_policy(PageIndexPolicy::Optional)
+ .with_column_index_mask(ColumnChunkMask::row_groups_and_columns([1],
[0]))
+ .with_offset_index_mask(ColumnChunkMask::row_groups_and_columns([1],
[0, 2]));
+
+ // parse metadata
+ reader.try_parse(&temp_file).unwrap();
+ let metadata = reader.finish().unwrap();
+ assert!(metadata.page_index().is_some());
+ #[cfg(not(feature = "encryption"))]
+ assert_eq!(metadata.memory_size(), 11326);
+ #[cfg(feature = "encryption")]
+ assert_eq!(metadata.memory_size(), 11750);
+}
Review Comment:
**Test gap (C14).** No test sets a mask through `ArrowReaderOptions` or
`ParquetMetaDataReader::with_page_index_mask`. Mutation testing (on the
[#11159](https://github.com/apache/arrow-rs/pull/11159) head, same code) shows
it: `ArrowReaderOptions::with_column_index_mask`, `with_offset_index_mask`,
both getters, and `ParquetMetaDataReader::with_page_index_mask` can each return
`Default::default()`, and all tests still pass. If a refactor drops the mask
plumbing, readers decode the full index again and no test fails. There, these
two tests kill all 5 mutants. They pass on this head.
Also missing: a data read (push decoder or async) with a mask. That test
would have found C2. Add it after the C2 fix.
```suggestion
}
/// Asserts which cells of the page index are populated
fn assert_page_index_cells(
metadata: &parquet::file::metadata::ParquetMetaData,
expect_ci: impl Fn(usize, usize) -> bool,
expect_oi: impl Fn(usize, usize) -> bool,
) {
let page_index = metadata.page_index().expect("page index should be
loaded");
let num_cols = metadata.file_metadata().schema_descr().num_columns();
for rg in 0..metadata.num_row_groups() {
for col in 0..num_cols {
let ci = page_index.column_index(rg, col).is_some();
let oi = page_index.offset_index(rg, col).is_some();
assert_eq!(ci, expect_ci(rg, col), "column index rg={rg}
col={col}");
assert_eq!(oi, expect_oi(rg, col), "offset index rg={rg}
col={col}");
}
}
}
#[test]
fn test_arrow_reader_options_page_index_masks() {
use parquet::arrow::arrow_reader::{ArrowReaderMetadata,
ArrowReaderOptions};
let file = create_test_file();
let column_mask = ColumnChunkMask::row_groups_and_columns([1], [0]);
let offset_mask = ColumnChunkMask::columns([2]);
let options = ArrowReaderOptions::new()
.with_page_index_policy(PageIndexPolicy::Required)
.with_column_index_mask(column_mask.clone())
.with_offset_index_mask(offset_mask.clone());
assert_eq!(options.column_index_mask(), &column_mask);
assert_eq!(options.offset_index_mask(), &offset_mask);
let expect_ci = |rg: usize, col: usize| rg == 1 && col == 0;
let expect_oi = |_rg: usize, col: usize| col == 2;
// sync reader builders
let arrow_metadata = ArrowReaderMetadata::load(&file,
options.clone()).unwrap();
assert_page_index_cells(arrow_metadata.metadata(), expect_ci, expect_oi);
// `AsyncFileReader::get_metadata` implementations
let metadata = ParquetMetaDataReader::new()
.with_arrow_reader_options(Some(&options))
.parse_and_finish(&file)
.unwrap();
assert_page_index_cells(&metadata, expect_ci, expect_oi);
}
#[test]
fn test_parse_with_page_index_mask() {
let file = create_test_file();
let metadata = ParquetMetaDataReader::new()
.with_page_index_policy(PageIndexPolicy::Required)
.with_page_index_mask(ColumnChunkMask::row_groups_and_columns([2],
[1, 3]))
.parse_and_finish(&file)
.unwrap();
let expect = |rg: usize, col: usize| rg == 2 && (col == 1 || col == 3);
assert_page_index_cells(&metadata, expect, expect);
}
```
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -104,6 +107,114 @@ impl From<bool> for PageIndexPolicy {
}
}
+/// Struct to specify column chunks for which metadata is required.
+///
+/// Column chunks are identified by row group index and column index. This
struct
+/// allows for specifying vertical slices of column chunk data (via
[`Self::columns`]),
+/// horizontal slices (via [`Self::row_groups`]), or the intersection of the
two
+/// (via [`Self::row_groups_and_columns`]).
+///
+/// At present this is only used to select elements of the [Page Index] for
decoding.
+///
+/// # Examples
+///
+/// To select columns 0 and 1 from all row groups:
+/// ```rust
+/// # use parquet::file::metadata::ColumnChunkMask;
+/// let mask = ColumnChunkMask::columns([0, 1]);
+/// ```
+///
+/// To select all columns from row group 2:
+/// ```rust
+/// # use parquet::file::metadata::ColumnChunkMask;
+/// let mask = ColumnChunkMask::row_groups([2]);
+/// ```
+///
+/// To select columns 1 and 3 from row group 0:
+/// ```rust
+/// # use parquet::file::metadata::ColumnChunkMask;
+/// let mask = ColumnChunkMask::row_groups_and_columns([0], [1, 3]);
+/// ```
+///
+/// [Page Index]: https://parquet.apache.org/docs/file-format/pageindex/
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+pub struct ColumnChunkMask {
+ // using i32 because that's how thrift vectors are sized
+ row_groups: Option<Arc<BTreeSet<i32>>>,
+ columns: Option<Arc<BTreeSet<i32>>>,
+}
+
+impl ColumnChunkMask {
+ /// Select all row groups and columns.
+ pub fn all() -> Self {
+ Self::default()
+ }
+
+ /// Select only the listed columns.
+ ///
+ /// Any indices in `columns` that are less than zero will be ignored.
Passing an empty
+ /// set is treated the same as selecting all columns.
+ pub fn columns(columns: impl IntoIterator<Item = i32>) -> Self {
+ Self {
+ row_groups: None,
+ columns: Self::iter_to_set(columns),
+ }
+ }
+
+ /// Select only the listed row groups.
+ ///
+ /// Any indices in `row_groups` that are less than zero will be ignored.
Passing an empty
+ /// set is treated the same as selecting all row groups.
+ pub fn row_groups(row_groups: impl IntoIterator<Item = i32>) -> Self {
+ Self {
+ row_groups: Self::iter_to_set(row_groups),
+ columns: None,
+ }
+ }
+
+ /// Select only the listed row groups and columns.
+ ///
+ /// Any indices in `row_groups` or `columns` that are less than zero will
be ignored.
+ /// Passing an empty set for `row_groups` is treated as selecting all row
groups, and
+ /// an empty set for `columns` as selectiong all columns.
Review Comment:
**Nit (C16).** Typo.
```suggestion
/// an empty set for `columns` as selecting all columns.
```
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -104,6 +107,114 @@ impl From<bool> for PageIndexPolicy {
}
}
+/// Struct to specify column chunks for which metadata is required.
+///
+/// Column chunks are identified by row group index and column index. This
struct
+/// allows for specifying vertical slices of column chunk data (via
[`Self::columns`]),
+/// horizontal slices (via [`Self::row_groups`]), or the intersection of the
two
+/// (via [`Self::row_groups_and_columns`]).
+///
+/// At present this is only used to select elements of the [Page Index] for
decoding.
+///
+/// # Examples
+///
+/// To select columns 0 and 1 from all row groups:
+/// ```rust
+/// # use parquet::file::metadata::ColumnChunkMask;
+/// let mask = ColumnChunkMask::columns([0, 1]);
+/// ```
+///
+/// To select all columns from row group 2:
+/// ```rust
+/// # use parquet::file::metadata::ColumnChunkMask;
+/// let mask = ColumnChunkMask::row_groups([2]);
+/// ```
+///
+/// To select columns 1 and 3 from row group 0:
+/// ```rust
+/// # use parquet::file::metadata::ColumnChunkMask;
+/// let mask = ColumnChunkMask::row_groups_and_columns([0], [1, 3]);
+/// ```
+///
+/// [Page Index]: https://parquet.apache.org/docs/file-format/pageindex/
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+pub struct ColumnChunkMask {
+ // using i32 because that's how thrift vectors are sized
+ row_groups: Option<Arc<BTreeSet<i32>>>,
+ columns: Option<Arc<BTreeSet<i32>>>,
+}
+
+impl ColumnChunkMask {
+ /// Select all row groups and columns.
+ pub fn all() -> Self {
+ Self::default()
+ }
+
+ /// Select only the listed columns.
+ ///
+ /// Any indices in `columns` that are less than zero will be ignored.
Passing an empty
+ /// set is treated the same as selecting all columns.
+ pub fn columns(columns: impl IntoIterator<Item = i32>) -> Self {
+ Self {
+ row_groups: None,
+ columns: Self::iter_to_set(columns),
+ }
+ }
+
+ /// Select only the listed row groups.
+ ///
+ /// Any indices in `row_groups` that are less than zero will be ignored.
Passing an empty
+ /// set is treated the same as selecting all row groups.
+ pub fn row_groups(row_groups: impl IntoIterator<Item = i32>) -> Self {
+ Self {
+ row_groups: Self::iter_to_set(row_groups),
+ columns: None,
+ }
+ }
+
+ /// Select only the listed row groups and columns.
+ ///
+ /// Any indices in `row_groups` or `columns` that are less than zero will
be ignored.
+ /// Passing an empty set for `row_groups` is treated as selecting all row
groups, and
+ /// an empty set for `columns` as selectiong all columns.
+ pub fn row_groups_and_columns(
+ row_groups: impl IntoIterator<Item = i32>,
+ columns: impl IntoIterator<Item = i32>,
+ ) -> Self {
+ Self {
+ row_groups: Self::iter_to_set(row_groups),
+ columns: Self::iter_to_set(columns),
+ }
+ }
+
+ /// Test if `idx` is in the row group set.
+ ///
+ /// Returns `false` if `idx > `[`i32::MAX`].
+ pub fn includes_row_group(&self, idx: usize) -> bool {
+ let Ok(idx) = i32::try_from(idx) else {
+ return false;
+ };
+ self.row_groups
+ .as_ref()
+ .is_none_or(|keep| keep.contains(&idx))
+ }
+
+ /// Test if `idx` is in the column set.
+ ///
+ /// Returns `false` if `idx > `[`i32::MAX`].
+ pub fn includes_column(&self, idx: usize) -> bool {
+ let Ok(idx) = i32::try_from(idx) else {
+ return false;
+ };
+ self.columns.as_ref().is_none_or(|keep| keep.contains(&idx))
+ }
+
+ fn iter_to_set(indices: impl IntoIterator<Item = i32>) ->
Option<Arc<BTreeSet<i32>>> {
+ let set: BTreeSet<i32> = indices.into_iter().filter(|&i| i >=
0).collect();
+ (!set.is_empty()).then_some(Arc::new(set))
+ }
Review Comment:
**API, one-way door (C4).** An empty set selects everything. The common
callers build the set from a list that can be empty:
```text
columns([]) -> loads 12 of 12 offset indexes (SELECT
count(*): no projected leaves)
row_groups(<no surviving rgs>) -> loads 12 of 12 (stats
pruned every row group)
columns([-1]) -> loads 12 of 12
columns([100]) -> loads 0 of 12, page_index() = None
columns([-1, 100]) -> loads 0 of 12
```
The neighbour APIs do the opposite: `ProjectionMask::leaves(schema, [])`
selects no leaves, `with_row_groups(vec![])` reads nothing, and
`ParquetStatisticsPolicy::skip_except(&[])` is `SkipAll`. After a release, a
change here silently changes what callers load. Suggestion: an empty set
selects nothing, and `all()` stays the only way to select everything.
`test_chunk_mask` (L1545-1549) and the constructor docs (L155-156, L166-167,
L177-179) need the same change.
```suggestion
fn iter_to_set(indices: impl IntoIterator<Item = i32>) ->
Option<Arc<BTreeSet<i32>>> {
// An empty set selects nothing. `Self::all()` selects everything.
Some(Arc::new(indices.into_iter().filter(|&i| i >= 0).collect()))
}
```
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -970,7 +1012,9 @@ impl ParquetMetaDataReader {
{
self = self
.with_column_index_policy(options.column_index_policy())
- .with_offset_index_policy(options.offset_index_policy());
+ .with_offset_index_policy(options.offset_index_policy())
+ .with_column_index_mask(options.column_index_mask().clone())
+ .with_offset_index_mask(options.offset_index_mask().clone());
}
Review Comment:
**Footgun (C7).** Masks are silently ignored in these paths (outputs from
this head, 4 columns x 3 row groups):
```text
options.with_offset_index_mask([0]) only (policy default Skip) ->
page_index() = None
preloading reader + with_arrow_reader_options(mask [0], policy Skip) ->
loads 12 of 12 offset indexes
ParquetObjectReader::with_preload_*(true) + options masks [0] ->
loads all columns (store.rs L239-248 has its own copy of this code)
ArrowReaderMetadata::try_new(md with mask [0], Required + mask [3]) -> Ok,
keeps mask [0], no error
```
This suggestion fixes the second line (then: 3 of 12, only column 0). The
same change is needed in `store.rs`. The docs change on
`ArrowReaderOptions::with_*_mask` covers the other lines.
```suggestion
{
self = self
.with_column_index_policy(options.column_index_policy())
.with_offset_index_policy(options.offset_index_policy());
}
// Apply the masks also when a preload setting on this reader
enables the page index
if options.column_index_mask() != &ColumnChunkMask::all() {
self =
self.with_column_index_mask(options.column_index_mask().clone());
}
if options.offset_index_mask() != &ColumnChunkMask::all() {
self =
self.with_offset_index_mask(options.offset_index_mask().clone());
}
```
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -137,6 +248,24 @@ impl ParquetMetaDataReader {
self
}
+ /// Selects the row groups and columns for which both page index
structures are read.
+ pub fn with_page_index_mask(self, mask: ColumnChunkMask) -> Self {
+ self.with_column_index_mask(mask.clone())
+ .with_offset_index_mask(mask)
+ }
+
+ /// Selects the row groups and columns for which column indexes are read.
+ pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self {
+ self.column_index_mask = mask;
+ self
+ }
+
+ /// Selects the row groups and columns for which offset indexes are read.
+ pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self {
+ self.offset_index_mask = mask;
+ self
+ }
Review Comment:
**Perf / Docs (C8, C12).** This reader merges all requested ranges into one
covering range (`needs_index_data`, L959-966), and the async path does the
same. So a column mask cuts decode work, but almost no I/O:
```text
4 columns x 10 row groups ParquetMetaDataReader push decoder
all() 1625 B in 1 request 1 range,
1625 B
columns([0]) 1586 B in 1 request 20 ranges,
406 B
row_groups_and_columns([9],[3]) 533 B in 1 request 2 ranges,
41 B
```
Options: fetch the ranges with a gap threshold (like
`object_store::coalesce_ranges`), or document it. This suggestion documents it,
and also the `Skip` and `Required` behavior (C7, C12).
```suggestion
/// Selects the row groups and columns for which both page index
structures are read.
///
/// See [`Self::with_column_index_mask`] and
[`Self::with_offset_index_mask`].
pub fn with_page_index_mask(self, mask: ColumnChunkMask) -> Self {
self.with_column_index_mask(mask.clone())
.with_offset_index_mask(mask)
}
/// Selects the row groups and columns for which column indexes are read.
///
/// The mask applies only if the column index policy is not
[`PageIndexPolicy::Skip`]
/// (the default). The mask reduces decode time, but not I/O: this
reader fetches one
/// byte range that covers all selected column indexes.
pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self {
self.column_index_mask = mask;
self
}
/// Selects the row groups and columns for which offset indexes are read.
///
/// The mask applies only if the offset index policy is not
[`PageIndexPolicy::Skip`]
/// (the default). [`PageIndexPolicy::Required`] checks only the
selected column chunks.
/// The mask reduces decode time, but not I/O: this reader fetches one
byte range that
/// covers all selected offset indexes.
pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self {
self.offset_index_mask = mask;
self
}
```
##########
parquet/src/file/metadata/push_decoder.rs:
##########
@@ -408,30 +435,36 @@ impl ParquetMetaDataPushDecoder {
DecodeState::ReadingPageIndex(mut metadata) => {
// First determine if any page indexes are needed based on
// the specified policies
- let range = range_for_page_index(
+ let ranges = ranges_for_page_index(
&metadata,
self.column_index_policy,
self.offset_index_policy,
+ &self.column_index_mask,
+ &self.offset_index_mask,
);
- let Some(page_index_range) = range else {
+ if ranges.is_empty() {
self.state = DecodeState::Finished;
return Ok(DecodeResult::Data(*metadata));
- };
+ }
- if !self.buffers.has_range(&page_index_range) {
+ let needed_ranges = ranges
+ .into_iter()
+ .filter(|r| !self.buffers.has_range(r))
+ .collect::<Vec<_>>();
Review Comment:
**Perf (C9).** `PushBuffers::has_range` (here) and `PushBuffers::get_bytes`
(once per chunk in `parse_*_index`) are linear scans over all pushed buffers.
With a sparse mask, a caller that pushes the requested ranges as they are (as
the docs show) gets O(ranges x buffers). Release build, Int32 columns, mask =
every other column, 3 runs on a shared machine:
| row groups x columns | ranges | one buffer per range | one covering buffer
|
|---|---|---|---|
| 100 x 100 | 10,000 | 49-203 ms | 3-16 ms |
| 300 x 100 | 30,000 | 0.97-1.6 s | 24-40 ms |
Fix: keep the buffers sorted by start and use a binary search in `has_range`
/ `get_bytes`, and cache the computed ranges in
`DecodeState::ReadingPageIndex`. Smaller change: reuse the `Vec` here (no new
allocation).
```suggestion
let mut needed_ranges = ranges;
needed_ranges.retain(|r| !self.buffers.has_range(r));
```
##########
parquet/src/file/metadata/parser.rs:
##########
@@ -293,27 +296,27 @@ pub(crate) fn parse_page_index(
fn parse_column_index(
metadata: &ParquetMetaData,
column_index_policy: PageIndexPolicy,
+ mask: &ColumnChunkMask,
page_index_builder: &mut PageIndexBuilder,
- bytes: &Bytes,
- start_offset: u64,
+ bytes: &PushBuffers,
) -> crate::errors::Result<()> {
if column_index_policy == PageIndexPolicy::Skip {
return Ok(());
}
for rg_idx in 0..metadata.num_row_groups() {
+ if !mask.includes_row_group(rg_idx) {
+ continue;
+ }
let rg = metadata.row_group(rg_idx);
for col_idx in 0..rg.num_columns() {
+ if !mask.includes_column(col_idx) {
+ continue;
+ }
Review Comment:
**Perf nit (C15).** For each selected row group, this loop visits every
column and does a `BTreeSet` lookup, also when the mask selects 1 of 10,000
columns. `add_ranges` in `push_decoder.rs` does the same, on every
`try_decode`. A `pub(crate)` helper can iterate the selected set directly
(compiled and tested on this head):
```rust
// reader.rs
impl ColumnChunkMask {
/// Selected column indexes in `0..num_columns`, in order
pub(crate) fn column_indices(
&self,
num_columns: usize,
) -> Box<dyn Iterator<Item = usize> + '_> {
match &self.columns {
None => Box::new(0..num_columns),
Some(set) => {
let end = i32::try_from(num_columns).unwrap_or(i32::MAX);
Box::new(set.range(0..end).map(|&i| i as usize))
}
}
}
}
// here, and in parse_offset_index / add_ranges
for col_idx in mask.column_indices(rg.num_columns()) {
```
##########
parquet/src/file/metadata/parser.rs:
##########
@@ -244,14 +245,16 @@ pub(crate) fn decode_metadata(
/// Required, Optional, Skip).
/// * `offset_index_policy` - The policy for handling offset index parsing
(e.g.,
/// Required, Optional, Skip).
-/// * `bytes` - The byte slice containing the page index data.
+/// * `bytes` - [`PushBuffers`] that should have already been populated with
the bytes containing
+/// the page indexes.
/// * `start_offset` - The offset where `bytes` begin in the file.
Review Comment:
**Nit (C16).** `start_offset` is no longer a parameter.
```suggestion
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
