adriangb commented on code in PR #11159:
URL: https://github.com/apache/arrow-rs/pull/11159#discussion_r4078133601
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -363,6 +365,183 @@ impl RowGroupPageIndex {
}
}
+/// A memory-efficient sparse set representation storing sorted deduplicated
indexes
+///
+/// Stores which positions are set in a sparse vector. For example:
+/// `[None, None, Some(...), Some(...), None, Some(...)]` becomes `[2, 3, 5]`
+///
+/// Position checking uses binary search for O(log n) lookup.
Review Comment:
**Nit (D10).** `position` uses a linear scan for 32 indexes or fewer
(L427-433).
```suggestion
/// Position checking uses a linear scan for up to 32 indexes, and binary
search above that.
```
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -363,6 +365,183 @@ impl RowGroupPageIndex {
}
}
+/// A memory-efficient sparse set representation storing sorted deduplicated
indexes
+///
+/// Stores which positions are set in a sparse vector. For example:
+/// `[None, None, Some(...), Some(...), None, Some(...)]` becomes `[2, 3, 5]`
+///
+/// Position checking uses binary search for O(log n) lookup.
+///
+/// Implementation note: we can downsize to `i32` here because thrift encodes
vector
+/// sizes with an `i32`.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Keep {
+ /// Sorted, deduplicated indexes of set positions
+ /// None means all positions in the span are set
+ kept: Option<Arc<[i32]>>,
+ /// Total span of positions (0..span)
+ span: i32,
+}
+
+impl Keep {
+ // an empty set or a set whose size equals "span" will consider all
positions in the
+ // span are set.
+ pub(crate) fn new(set: &BTreeSet<i32>, span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ // this should preserve the BTreeSet ordering
+ let kept = set
+ .iter()
+ .filter(|&&idx| idx < span as i32 && idx >= 0)
+ .copied()
+ .collect::<Vec<_>>();
+ let kept = if kept.is_empty() || kept.len() == span {
+ None
+ } else {
+ Some(Arc::from(kept))
+ };
Review Comment:
**Footgun (D2) + pre-allocation (D7).** A mask whose indexes are all out of
range keeps *every* position. The reader uses the opposite meaning
(`includes_column` is `false`):
```text
ColumnChunkMask::columns([100]) on 4 columns:
mask.includes_column(0) = false
new_with_mask(..) then put_column_index(_, 0, 0) -> stored
ColumnChunkMask::columns([1]):
new_with_mask(..) then put_column_index(_, 0, 0) -> dropped
```
Through the reader (1000 columns × 100 row groups, CI `Optional`, OI
`Required`), `with_column_index_mask(columns([5000]))` peaks at 30.4 MB, the
same as `all()`. With `columns([0])` the peak is 7.2 MB. The parser loads
nothing into the extra 23.2 MB grid.
Cause: L400 treats "no index left after the filter" like "empty set" (= all).
Fix: keep "empty set = all" (as in `ColumnChunkMask`), but let "no index in
range" keep nothing. The same change sizes the `Vec` up front: `set.len()` is
an upper bound, and `range` replaces the filter. Tested: `cargo test -p parquet
--features arrow,async --lib file::metadata` passes (also
`test_empty_keep_selects_all`), and the peak above drops to 7.2 MB.
```suggestion
// An empty set, or a set that covers the whole span, keeps all
positions.
// Indexes outside `0..span` are ignored: if none is in range, no
position is kept.
pub(crate) fn new(set: &BTreeSet<i32>, span: usize) -> Self {
assert!(
i32::try_from(span).is_ok(),
"Keep cannot have a span that exceeds the storage of an i32: got
{span}"
);
// `range` yields the in-range indexes in sorted order; `set.len()`
is an upper bound
let mut kept = Vec::with_capacity(set.len());
kept.extend(set.range(0..span as i32));
let kept = (!set.is_empty() && kept.len() != span).then(||
Arc::from(kept));
```
`Arc::from` still copies once. If `ColumnChunkMask` stored a sorted
`Arc<[i32]>` instead of `Arc<BTreeSet<i32>>`, `Keep` could share it: 0
allocations, and one slice for both grids.
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -363,6 +365,183 @@ impl RowGroupPageIndex {
}
}
+/// A memory-efficient sparse set representation storing sorted deduplicated
indexes
+///
+/// Stores which positions are set in a sparse vector. For example:
+/// `[None, None, Some(...), Some(...), None, Some(...)]` becomes `[2, 3, 5]`
+///
+/// Position checking uses binary search for O(log n) lookup.
+///
+/// Implementation note: we can downsize to `i32` here because thrift encodes
vector
+/// sizes with an `i32`.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Keep {
+ /// Sorted, deduplicated indexes of set positions
+ /// None means all positions in the span are set
+ kept: Option<Arc<[i32]>>,
+ /// Total span of positions (0..span)
+ span: i32,
+}
+
+impl Keep {
+ // an empty set or a set whose size equals "span" will consider all
positions in the
+ // span are set.
+ pub(crate) fn new(set: &BTreeSet<i32>, span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ // this should preserve the BTreeSet ordering
+ let kept = set
+ .iter()
+ .filter(|&&idx| idx < span as i32 && idx >= 0)
+ .copied()
+ .collect::<Vec<_>>();
+ let kept = if kept.is_empty() || kept.len() == span {
+ None
+ } else {
+ Some(Arc::from(kept))
+ };
+
+ Self {
+ kept,
+ span: span as i32,
+ }
+ }
+
+ // shortened version for a full keep set
+ pub(crate) fn new_full(span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ Self {
+ kept: None,
+ span: span as i32,
+ }
+ }
+
+ /// Retrieve a position if set
+ fn position(&self, idx: usize) -> Option<usize> {
+ let needle = i32::try_from(idx).ok()?;
+ // below CUTOFF elements, use linear search
+ const CUTOFF: usize = 32;
+ match self.kept.as_ref() {
+ None => (idx < self.span as usize).then_some(idx),
+ Some(k) if k.len() > CUTOFF => k.binary_search(&needle).ok(),
+ Some(k) => k.iter().position(|&i| i == needle),
+ }
+ }
+
+ /// Returns the number of set positions
+ fn len(&self) -> usize {
+ match &self.kept {
+ None => self.span as usize,
+ Some(indexes) => indexes.len(),
+ }
+ }
+}
+
+impl HeapSize for Arc<[i32]> {
+ fn heap_size(&self) -> usize {
+ // Arc stores weak and strong counts on the heap alongside an instance
of T
+ // T = [i32], so that should be the size of a pointer + the size of
the allocation
+ 2 * std::mem::size_of::<usize>()
+ + std::mem::size_of::<*mut i32>()
+ + std::mem::size_of_val(self.as_ref())
Review Comment:
**Nit (D8).** The fat pointer is stored inline in `Keep`, not in the heap
block. The block is also padded to the alignment of `usize`. Measured with a
counting allocator (`Arc::<[i32]>::from`, 64-bit):
| len | allocated | this PR | suggestion |
|---|---|---|---|
| 1 | 24 | 28 | 24 |
| 2 | 24 | 32 | 24 |
| 3 | 32 | 36 | 32 |
```suggestion
// The heap block holds the strong and weak counts, then the slice,
padded to the
// alignment of `usize`. The (fat) pointer itself is stored inline,
not on the heap.
(2 * std::mem::size_of::<usize>() +
std::mem::size_of_val(self.as_ref()))
.next_multiple_of(std::mem::align_of::<usize>())
```
`test_page_index_sizes` then expects 10664, 9680, 8342 (was 10668, 9692,
8362). With `encryption`: 11088, 10104, 8766 (was 11092, 10116, 8786). I ran
both.
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -363,6 +365,183 @@ impl RowGroupPageIndex {
}
}
+/// A memory-efficient sparse set representation storing sorted deduplicated
indexes
+///
+/// Stores which positions are set in a sparse vector. For example:
+/// `[None, None, Some(...), Some(...), None, Some(...)]` becomes `[2, 3, 5]`
+///
+/// Position checking uses binary search for O(log n) lookup.
+///
+/// Implementation note: we can downsize to `i32` here because thrift encodes
vector
+/// sizes with an `i32`.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Keep {
+ /// Sorted, deduplicated indexes of set positions
+ /// None means all positions in the span are set
+ kept: Option<Arc<[i32]>>,
+ /// Total span of positions (0..span)
+ span: i32,
+}
+
+impl Keep {
+ // an empty set or a set whose size equals "span" will consider all
positions in the
+ // span are set.
+ pub(crate) fn new(set: &BTreeSet<i32>, span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ // this should preserve the BTreeSet ordering
+ let kept = set
+ .iter()
+ .filter(|&&idx| idx < span as i32 && idx >= 0)
+ .copied()
+ .collect::<Vec<_>>();
+ let kept = if kept.is_empty() || kept.len() == span {
+ None
+ } else {
+ Some(Arc::from(kept))
+ };
+
+ Self {
+ kept,
+ span: span as i32,
+ }
+ }
+
+ // shortened version for a full keep set
+ pub(crate) fn new_full(span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ Self {
+ kept: None,
+ span: span as i32,
+ }
+ }
+
+ /// Retrieve a position if set
+ fn position(&self, idx: usize) -> Option<usize> {
+ let needle = i32::try_from(idx).ok()?;
+ // below CUTOFF elements, use linear search
+ const CUTOFF: usize = 32;
+ match self.kept.as_ref() {
+ None => (idx < self.span as usize).then_some(idx),
+ Some(k) if k.len() > CUTOFF => k.binary_search(&needle).ok(),
+ Some(k) => k.iter().position(|&i| i == needle),
+ }
+ }
+
+ /// Returns the number of set positions
+ fn len(&self) -> usize {
+ match &self.kept {
+ None => self.span as usize,
+ Some(indexes) => indexes.len(),
+ }
+ }
+}
+
+impl HeapSize for Arc<[i32]> {
+ fn heap_size(&self) -> usize {
+ // Arc stores weak and strong counts on the heap alongside an instance
of T
+ // T = [i32], so that should be the size of a pointer + the size of
the allocation
+ 2 * std::mem::size_of::<usize>()
+ + std::mem::size_of::<*mut i32>()
+ + std::mem::size_of_val(self.as_ref())
+ }
+}
+
+impl HeapSize for Keep {
+ fn heap_size(&self) -> usize {
+ self.kept.heap_size()
+ }
+}
+
+/// A memory-efficient 2D sparse grid using Keep structures for rows and
columns
+///
+/// Maps (row_group_idx, column_idx) to values efficiently for sparse access
patterns.
+/// This is particularly useful when only a few columns are accessed from wide
schemas.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Grid<T> {
+ /// Set of row group indexes that have any values
+ rows: Keep,
+ /// Set of column indexes that have any values
+ cols: Keep,
Review Comment:
**Nit (D10).** These fields hold the positions that have storage (from the
mask). A kept position can have no value.
```suggestion
/// Row group indexes that have storage (from the mask); a cell can
still be `None`
rows: Keep,
/// Column indexes that have storage (from the mask); a cell can still
be `None`
cols: Keep,
```
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -363,6 +365,183 @@ impl RowGroupPageIndex {
}
}
+/// A memory-efficient sparse set representation storing sorted deduplicated
indexes
+///
+/// Stores which positions are set in a sparse vector. For example:
+/// `[None, None, Some(...), Some(...), None, Some(...)]` becomes `[2, 3, 5]`
+///
+/// Position checking uses binary search for O(log n) lookup.
+///
+/// Implementation note: we can downsize to `i32` here because thrift encodes
vector
+/// sizes with an `i32`.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Keep {
+ /// Sorted, deduplicated indexes of set positions
+ /// None means all positions in the span are set
+ kept: Option<Arc<[i32]>>,
+ /// Total span of positions (0..span)
+ span: i32,
+}
+
+impl Keep {
+ // an empty set or a set whose size equals "span" will consider all
positions in the
+ // span are set.
+ pub(crate) fn new(set: &BTreeSet<i32>, span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ // this should preserve the BTreeSet ordering
+ let kept = set
+ .iter()
+ .filter(|&&idx| idx < span as i32 && idx >= 0)
+ .copied()
+ .collect::<Vec<_>>();
+ let kept = if kept.is_empty() || kept.len() == span {
+ None
+ } else {
+ Some(Arc::from(kept))
+ };
+
+ Self {
+ kept,
+ span: span as i32,
+ }
+ }
+
+ // shortened version for a full keep set
+ pub(crate) fn new_full(span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ Self {
+ kept: None,
+ span: span as i32,
+ }
+ }
+
+ /// Retrieve a position if set
+ fn position(&self, idx: usize) -> Option<usize> {
+ let needle = i32::try_from(idx).ok()?;
+ // below CUTOFF elements, use linear search
+ const CUTOFF: usize = 32;
+ match self.kept.as_ref() {
+ None => (idx < self.span as usize).then_some(idx),
+ Some(k) if k.len() > CUTOFF => k.binary_search(&needle).ok(),
+ Some(k) => k.iter().position(|&i| i == needle),
+ }
+ }
+
+ /// Returns the number of set positions
+ fn len(&self) -> usize {
+ match &self.kept {
+ None => self.span as usize,
+ Some(indexes) => indexes.len(),
+ }
+ }
+}
+
+impl HeapSize for Arc<[i32]> {
+ fn heap_size(&self) -> usize {
+ // Arc stores weak and strong counts on the heap alongside an instance
of T
+ // T = [i32], so that should be the size of a pointer + the size of
the allocation
+ 2 * std::mem::size_of::<usize>()
+ + std::mem::size_of::<*mut i32>()
+ + std::mem::size_of_val(self.as_ref())
+ }
+}
+
+impl HeapSize for Keep {
+ fn heap_size(&self) -> usize {
+ self.kept.heap_size()
+ }
+}
+
+/// A memory-efficient 2D sparse grid using Keep structures for rows and
columns
+///
+/// Maps (row_group_idx, column_idx) to values efficiently for sparse access
patterns.
+/// This is particularly useful when only a few columns are accessed from wide
schemas.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Grid<T> {
+ /// Set of row group indexes that have any values
+ rows: Keep,
+ /// Set of column indexes that have any values
+ cols: Keep,
+ /// Flattened cells stored in row-major order
+ /// cells[row_offset * cols.len() + col_offset] = value at (row, col)
+ /// where row_offset = position of row in rows.kept
+ /// and col_offset = position of col in cols.kept
+ cells: Vec<Option<T>>,
+}
+
+impl<T: Clone> Grid<T> {
+ /// Creates a new empty Grid with the specified dimensions
+ fn new(rows: Keep, cols: Keep) -> Self {
+ let size = rows.len() * cols.len();
+ let cells = vec![None; size];
+ Self { rows, cols, cells }
+ }
+
+ pub(crate) fn new_dense(num_row_groups: usize, num_columns: usize) -> Self
{
+ let rows = Keep::new_full(num_row_groups);
+ let cols = Keep::new_full(num_columns);
+ Self::new(rows, cols)
+ }
+
+ pub(crate) fn from_vec(index: Vec<Vec<Option<T>>>) -> Self {
+ let num_row_groups = index.len();
+ let num_columns = if index.is_empty() { 0 } else { index[0].len() };
+ let mut result = Self::new_dense(num_row_groups, num_columns);
+ for (rg_idx, row_group) in index.into_iter().enumerate() {
+ for (col_idx, idx) in row_group.into_iter().enumerate() {
+ if let Some(idx) = idx {
+ result.insert(rg_idx, col_idx, idx);
+ }
+ }
+ }
+
+ result
+ }
Review Comment:
**Pre-allocation + footgun (D6).** The size is known (n × m). This code
fills n × m `None` values, then calls `insert` for each cell (2
`Keep::position` lookups each). It also takes the width from row 0, and it
drops cells of longer rows without an error (observed on this head):
```text
Grid::from_vec(vec![vec![None], vec![None, Some(1)]]).get(1, 1) -> None
Grid::from_vec(vec![vec![], vec![Some(1)]]).is_empty() -> true
```
The only caller (`SerializedFileWriter::write_metadata`) passes rectangular
input today, so the data loss is latent. This version moves each row into one
exact allocation and rejects ragged input. `--lib file::metadata` and `--test
arrow_writer` pass.
```suggestion
pub(crate) fn from_vec(index: Vec<Vec<Option<T>>>) -> Self {
let num_row_groups = index.len();
let num_columns = index.first().map_or(0, Vec::len);
let mut cells = Vec::with_capacity(num_row_groups * num_columns);
for row_group in index {
assert_eq!(row_group.len(), num_columns, "ragged page index");
cells.extend(row_group);
}
Self {
rows: Keep::new_full(num_row_groups),
cols: Keep::new_full(num_columns),
cells,
}
}
```
Keep `assert_eq!`, not `debug_assert_eq!`: in a release build, ragged input
would move cells to the wrong position.
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -554,9 +712,34 @@ impl PageIndexBuilder {
/// All index entries are initialized to `None` and can be populated using
/// [`put_column_index`](Self::put_column_index) and
[`put_offset_index`](Self::put_offset_index).
Review Comment:
**Docs (D9).** `Keep::new_full` asserts that the span fits in an `i32`. So
this method now panics for a dimension above `i32::MAX`, even when the other
dimension is 0:
```text
PageIndexBuilder::new(0, i32::MAX as usize + 1)
panicked at parquet/src/file/metadata/page_index.rs:414:9:
Keep cannot have a span that exceeds the storage of an i32: got 2147483648
```
On the [#11157](https://github.com/apache/arrow-rs/pull/11157) head it
returns an empty builder. Real files cannot have that many row groups or
columns, so a `# Panics` section is enough. The same applies to `new_with_mask`
and `allocate_*`.
```suggestion
/// [`put_column_index`](Self::put_column_index) and
[`put_offset_index`](Self::put_offset_index).
///
/// # Panics
///
/// Panics if `num_row_groups` or `num_columns` is greater than
`i32::MAX`.
```
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -554,9 +712,34 @@ impl PageIndexBuilder {
/// All index entries are initialized to `None` and can be populated using
/// [`put_column_index`](Self::put_column_index) and
[`put_offset_index`](Self::put_offset_index).
pub fn new(num_row_groups: usize, num_columns: usize) -> Self {
+ let keep_cols = Keep::new_full(num_columns);
+ let keep_rows = Keep::new_full(num_row_groups);
Self {
- column_indexes: Self::empty_index(num_row_groups, num_columns),
- offset_indexes: Self::empty_index(num_row_groups, num_columns),
+ column_indexes: Some(Grid::new(keep_rows.clone(),
keep_cols.clone())),
+ offset_indexes: Some(Grid::new(keep_rows, keep_cols)),
+ }
+ }
+
+ /// Creates a new [`PageIndexBuilder`] where storage is defined by the
policy
+ ///
+ /// For sparse indexes, this can save a great deal of memory
+ pub fn new_with_mask(
Review Comment:
**API (D3).** This is new public API, but the PR description says "No, only
internal storage structures are changed". After a release, the signature and
the "drop outside the mask" behavior (D4) are hard to change. The parser is the
only caller. I suggest `pub(crate)` until there is a public use case. The doc
also says "policy", but the storage comes from the masks.
```suggestion
/// Creates a new [`PageIndexBuilder`] where storage is defined by the
masks
///
/// For sparse indexes, this can save a great deal of memory. The
`put_*` methods
/// ignore positions outside the masks.
pub(crate) fn new_with_mask(
```
The D1 patch (review body) also changes this signature to
`Option<&ColumnChunkMask>`. With `pub(crate)`, that change costs nothing.
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -363,6 +365,183 @@ impl RowGroupPageIndex {
}
}
+/// A memory-efficient sparse set representation storing sorted deduplicated
indexes
+///
+/// Stores which positions are set in a sparse vector. For example:
+/// `[None, None, Some(...), Some(...), None, Some(...)]` becomes `[2, 3, 5]`
+///
+/// Position checking uses binary search for O(log n) lookup.
+///
+/// Implementation note: we can downsize to `i32` here because thrift encodes
vector
+/// sizes with an `i32`.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Keep {
+ /// Sorted, deduplicated indexes of set positions
+ /// None means all positions in the span are set
+ kept: Option<Arc<[i32]>>,
+ /// Total span of positions (0..span)
+ span: i32,
+}
+
+impl Keep {
+ // an empty set or a set whose size equals "span" will consider all
positions in the
+ // span are set.
+ pub(crate) fn new(set: &BTreeSet<i32>, span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ // this should preserve the BTreeSet ordering
+ let kept = set
+ .iter()
+ .filter(|&&idx| idx < span as i32 && idx >= 0)
+ .copied()
+ .collect::<Vec<_>>();
+ let kept = if kept.is_empty() || kept.len() == span {
+ None
+ } else {
+ Some(Arc::from(kept))
+ };
+
+ Self {
+ kept,
+ span: span as i32,
+ }
+ }
+
+ // shortened version for a full keep set
+ pub(crate) fn new_full(span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ Self {
+ kept: None,
+ span: span as i32,
+ }
+ }
+
+ /// Retrieve a position if set
+ fn position(&self, idx: usize) -> Option<usize> {
+ let needle = i32::try_from(idx).ok()?;
+ // below CUTOFF elements, use linear search
+ const CUTOFF: usize = 32;
+ match self.kept.as_ref() {
+ None => (idx < self.span as usize).then_some(idx),
+ Some(k) if k.len() > CUTOFF => k.binary_search(&needle).ok(),
+ Some(k) => k.iter().position(|&i| i == needle),
+ }
+ }
+
+ /// Returns the number of set positions
+ fn len(&self) -> usize {
+ match &self.kept {
+ None => self.span as usize,
+ Some(indexes) => indexes.len(),
+ }
+ }
+}
+
+impl HeapSize for Arc<[i32]> {
+ fn heap_size(&self) -> usize {
+ // Arc stores weak and strong counts on the heap alongside an instance
of T
+ // T = [i32], so that should be the size of a pointer + the size of
the allocation
+ 2 * std::mem::size_of::<usize>()
+ + std::mem::size_of::<*mut i32>()
+ + std::mem::size_of_val(self.as_ref())
+ }
+}
+
+impl HeapSize for Keep {
+ fn heap_size(&self) -> usize {
+ self.kept.heap_size()
+ }
+}
+
+/// A memory-efficient 2D sparse grid using Keep structures for rows and
columns
+///
+/// Maps (row_group_idx, column_idx) to values efficiently for sparse access
patterns.
+/// This is particularly useful when only a few columns are accessed from wide
schemas.
+#[derive(Debug, Clone, PartialEq)]
Review Comment:
**API (D5).** The derived `PartialEq` compares the storage shape.
`PageIndex` and `ParquetMetaData` use it, so indexes with equal content can
compare unequal. With the dense layout of 60.0.0, this did not happen.
```text
a = PageIndexBuilder::new(1, 2) + put CI/OI
at (0, 0)
b = PageIndexBuilder::new_with_mask(1, 2, columns([0]), columns([0])) + put
CI/OI at (0, 0)
same answer for every (rg, col): true
a == b: false
```
The same happens for `ParquetMetaData` from a full read and a
`with_column_index_mask(columns([0]))` read of a file where column 1 has no
column index.
Fix: compare the content (tested; the check above then gives `a == b: true`,
and the lib tests pass). Remove `PartialEq` from this derive and add:
```rust
/// Two grids are equal if they return the same value for every position,
/// whatever their storage shape
impl<T: PartialEq> PartialEq for Grid<T> {
fn eq(&self, other: &Self) -> bool {
if self.rows == other.rows && self.cols == other.cols {
return self.cells == other.cells;
}
let rows = self.rows.span.max(other.rows.span) as usize;
let cols = self.cols.span.max(other.cols.span) as usize;
(0..rows).all(|r| (0..cols).all(|c| self.get(r, c) == other.get(r,
c)))
}
}
```
The fast path keeps today's cost when the shapes are equal.
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -363,6 +365,183 @@ impl RowGroupPageIndex {
}
}
+/// A memory-efficient sparse set representation storing sorted deduplicated
indexes
+///
+/// Stores which positions are set in a sparse vector. For example:
+/// `[None, None, Some(...), Some(...), None, Some(...)]` becomes `[2, 3, 5]`
+///
+/// Position checking uses binary search for O(log n) lookup.
+///
+/// Implementation note: we can downsize to `i32` here because thrift encodes
vector
+/// sizes with an `i32`.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Keep {
+ /// Sorted, deduplicated indexes of set positions
+ /// None means all positions in the span are set
+ kept: Option<Arc<[i32]>>,
+ /// Total span of positions (0..span)
+ span: i32,
+}
+
+impl Keep {
+ // an empty set or a set whose size equals "span" will consider all
positions in the
+ // span are set.
+ pub(crate) fn new(set: &BTreeSet<i32>, span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ // this should preserve the BTreeSet ordering
+ let kept = set
+ .iter()
+ .filter(|&&idx| idx < span as i32 && idx >= 0)
+ .copied()
+ .collect::<Vec<_>>();
+ let kept = if kept.is_empty() || kept.len() == span {
+ None
+ } else {
+ Some(Arc::from(kept))
+ };
+
+ Self {
+ kept,
+ span: span as i32,
+ }
+ }
+
+ // shortened version for a full keep set
+ pub(crate) fn new_full(span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ Self {
+ kept: None,
+ span: span as i32,
+ }
+ }
+
+ /// Retrieve a position if set
+ fn position(&self, idx: usize) -> Option<usize> {
+ let needle = i32::try_from(idx).ok()?;
+ // below CUTOFF elements, use linear search
+ const CUTOFF: usize = 32;
+ match self.kept.as_ref() {
+ None => (idx < self.span as usize).then_some(idx),
+ Some(k) if k.len() > CUTOFF => k.binary_search(&needle).ok(),
+ Some(k) => k.iter().position(|&i| i == needle),
+ }
+ }
+
+ /// Returns the number of set positions
+ fn len(&self) -> usize {
+ match &self.kept {
+ None => self.span as usize,
+ Some(indexes) => indexes.len(),
+ }
+ }
+}
+
+impl HeapSize for Arc<[i32]> {
+ fn heap_size(&self) -> usize {
+ // Arc stores weak and strong counts on the heap alongside an instance
of T
+ // T = [i32], so that should be the size of a pointer + the size of
the allocation
+ 2 * std::mem::size_of::<usize>()
+ + std::mem::size_of::<*mut i32>()
+ + std::mem::size_of_val(self.as_ref())
+ }
+}
+
+impl HeapSize for Keep {
+ fn heap_size(&self) -> usize {
+ self.kept.heap_size()
+ }
+}
+
+/// A memory-efficient 2D sparse grid using Keep structures for rows and
columns
+///
+/// Maps (row_group_idx, column_idx) to values efficiently for sparse access
patterns.
+/// This is particularly useful when only a few columns are accessed from wide
schemas.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Grid<T> {
+ /// Set of row group indexes that have any values
+ rows: Keep,
+ /// Set of column indexes that have any values
+ cols: Keep,
+ /// Flattened cells stored in row-major order
+ /// cells[row_offset * cols.len() + col_offset] = value at (row, col)
+ /// where row_offset = position of row in rows.kept
+ /// and col_offset = position of col in cols.kept
+ cells: Vec<Option<T>>,
+}
+
+impl<T: Clone> Grid<T> {
+ /// Creates a new empty Grid with the specified dimensions
+ fn new(rows: Keep, cols: Keep) -> Self {
+ let size = rows.len() * cols.len();
+ let cells = vec![None; size];
+ Self { rows, cols, cells }
+ }
Review Comment:
**Nit (D10).** `Clone` is necessary only for `vec![None; size]`.
`resize_with` keeps the allocation exact and removes the bound. Then also
remove `T: Clone` from `storage_for_selection` (L693) and `is_empty_index`
(L816). Tested with `cargo clippy -p parquet --features arrow,async --lib
--tests -- -D warnings`.
```suggestion
impl<T> Grid<T> {
/// Creates a new empty Grid with the specified dimensions
fn new(rows: Keep, cols: Keep) -> Self {
let size = rows.len() * cols.len();
let mut cells = Vec::with_capacity(size);
cells.resize_with(size, || None);
Self { rows, cols, cells }
}
```
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -363,6 +365,183 @@ impl RowGroupPageIndex {
}
}
+/// A memory-efficient sparse set representation storing sorted deduplicated
indexes
+///
+/// Stores which positions are set in a sparse vector. For example:
+/// `[None, None, Some(...), Some(...), None, Some(...)]` becomes `[2, 3, 5]`
+///
+/// Position checking uses binary search for O(log n) lookup.
+///
+/// Implementation note: we can downsize to `i32` here because thrift encodes
vector
+/// sizes with an `i32`.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Keep {
+ /// Sorted, deduplicated indexes of set positions
+ /// None means all positions in the span are set
+ kept: Option<Arc<[i32]>>,
+ /// Total span of positions (0..span)
+ span: i32,
+}
+
+impl Keep {
+ // an empty set or a set whose size equals "span" will consider all
positions in the
+ // span are set.
+ pub(crate) fn new(set: &BTreeSet<i32>, span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ // this should preserve the BTreeSet ordering
+ let kept = set
+ .iter()
+ .filter(|&&idx| idx < span as i32 && idx >= 0)
+ .copied()
+ .collect::<Vec<_>>();
+ let kept = if kept.is_empty() || kept.len() == span {
+ None
+ } else {
+ Some(Arc::from(kept))
+ };
+
+ Self {
+ kept,
+ span: span as i32,
+ }
+ }
+
+ // shortened version for a full keep set
+ pub(crate) fn new_full(span: usize) -> Self {
+ assert!(
+ i32::try_from(span).is_ok(),
+ "Keep cannot have a span that exceeds the storage of an i32: got
{span}"
+ );
+ Self {
+ kept: None,
+ span: span as i32,
+ }
+ }
+
+ /// Retrieve a position if set
+ fn position(&self, idx: usize) -> Option<usize> {
+ let needle = i32::try_from(idx).ok()?;
+ // below CUTOFF elements, use linear search
+ const CUTOFF: usize = 32;
+ match self.kept.as_ref() {
+ None => (idx < self.span as usize).then_some(idx),
+ Some(k) if k.len() > CUTOFF => k.binary_search(&needle).ok(),
+ Some(k) => k.iter().position(|&i| i == needle),
+ }
+ }
+
+ /// Returns the number of set positions
+ fn len(&self) -> usize {
+ match &self.kept {
+ None => self.span as usize,
+ Some(indexes) => indexes.len(),
+ }
+ }
+}
+
+impl HeapSize for Arc<[i32]> {
+ fn heap_size(&self) -> usize {
+ // Arc stores weak and strong counts on the heap alongside an instance
of T
+ // T = [i32], so that should be the size of a pointer + the size of
the allocation
+ 2 * std::mem::size_of::<usize>()
+ + std::mem::size_of::<*mut i32>()
+ + std::mem::size_of_val(self.as_ref())
+ }
+}
+
+impl HeapSize for Keep {
+ fn heap_size(&self) -> usize {
+ self.kept.heap_size()
+ }
+}
+
+/// A memory-efficient 2D sparse grid using Keep structures for rows and
columns
+///
+/// Maps (row_group_idx, column_idx) to values efficiently for sparse access
patterns.
+/// This is particularly useful when only a few columns are accessed from wide
schemas.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct Grid<T> {
+ /// Set of row group indexes that have any values
+ rows: Keep,
+ /// Set of column indexes that have any values
+ cols: Keep,
+ /// Flattened cells stored in row-major order
+ /// cells[row_offset * cols.len() + col_offset] = value at (row, col)
+ /// where row_offset = position of row in rows.kept
+ /// and col_offset = position of col in cols.kept
+ cells: Vec<Option<T>>,
+}
+
+impl<T: Clone> Grid<T> {
+ /// Creates a new empty Grid with the specified dimensions
+ fn new(rows: Keep, cols: Keep) -> Self {
+ let size = rows.len() * cols.len();
+ let cells = vec![None; size];
+ Self { rows, cols, cells }
+ }
+
+ pub(crate) fn new_dense(num_row_groups: usize, num_columns: usize) -> Self
{
+ let rows = Keep::new_full(num_row_groups);
+ let cols = Keep::new_full(num_columns);
+ Self::new(rows, cols)
+ }
+
+ pub(crate) fn from_vec(index: Vec<Vec<Option<T>>>) -> Self {
+ let num_row_groups = index.len();
+ let num_columns = if index.is_empty() { 0 } else { index[0].len() };
+ let mut result = Self::new_dense(num_row_groups, num_columns);
+ for (rg_idx, row_group) in index.into_iter().enumerate() {
+ for (col_idx, idx) in row_group.into_iter().enumerate() {
+ if let Some(idx) = idx {
+ result.insert(rg_idx, col_idx, idx);
+ }
+ }
+ }
+
+ result
+ }
+
+ /// Gets a value at the specified row and column
+ pub(crate) fn get(&self, row: usize, col: usize) -> Option<&T> {
+ // Find the offset of this row in the kept rows
+ let row_offset = self.rows.position(row)?;
+ let col_offset = self.cols.position(col)?;
+
+ let index = row_offset * self.cols.len() + col_offset;
+ self.cells.get(index)?.as_ref()
+ }
+
+ /// Sets a value at the specified row and column
+ pub(crate) fn insert(&mut self, row: usize, col: usize, value: T) {
+ let row_offset = self.rows.position(row);
+ let col_offset = self.cols.position(col);
+ if let Some(row_offset) = row_offset
+ && let Some(col_offset) = col_offset
+ {
+ // update the existing cell
+ let index = row_offset * self.cols.len() + col_offset;
+
+ if index < self.cells.len() {
+ self.cells[index] = Some(value);
+ }
+ }
+ }
+
+ /// Returns true if the grid has no values
+ pub(crate) fn is_empty(&self) -> bool {
+ self.cells.iter().all(|cell| cell.is_none())
+ }
+}
+
+impl<T: HeapSize> HeapSize for Grid<T> {
+ fn heap_size(&self) -> usize {
+ self.rows.heap_size() + self.cols.heap_size() + self.cells.heap_size()
+ }
+}
+
/// Struct to encapsulate the Parquet [Page Index]
///
/// This struct provides a dense representation of the Page Index. It is
Review Comment:
**Nit (D10).** The storage is no longer dense.
```suggestion
/// This struct provides a sparse representation of the Page Index: it has
storage only for
/// the column chunks selected when it was built (all chunks by default). It
is
```
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -605,11 +792,8 @@ impl PageIndexBuilder {
row_group_idx: usize,
column_idx: usize,
) {
- if let Some(ref mut indexes) = self.column_indexes
- && let Some(row_group) = indexes.get_mut(row_group_idx)
- && let Some(column_slot) = row_group.get_mut(column_idx)
- {
- *column_slot = Some(column_index);
+ if let Some(ref mut indexes) = self.column_indexes {
+ indexes.insert(row_group_idx, column_idx, column_index);
}
Review Comment:
**Footgun (D4).** `put_*` drops a value outside the mask without a signal.
This also happens after `into_builder()` on a masked index, because the builder
keeps the fixed shape:
```text
new_with_mask(2, 3, columns([0]), columns([0])) + put_offset_index(oi, 0, 1)
-> offset_index(0, 1) = None
masked_index.into_builder() + put_column_index(ci, 1, 2)
-> column_index(1, 2) = None
```
The doc at L787-788 lists only "not allocated" and "out of bounds". Also,
`allocate_*` (L765-783) replaces a masked grid with a dense one and drops its
cells. The doc does not say this.
Options: document the three cases, or return `bool` so that a caller can see
the drop. The `bool` version (tested, clippy clean) also removes the `index <
self.cells.len()` check at L527, which cargo-mutants shows is dead code:
```rust
// Grid
/// Returns `false`, and drops `value`, if the grid has no storage for
the position
pub(crate) fn insert(&mut self, row: usize, col: usize, value: T) ->
bool {
let (Some(row_offset), Some(col_offset)) =
(self.rows.position(row), self.cols.position(col))
else {
return false;
};
self.cells[row_offset * self.cols.len() + col_offset] = Some(value);
true
}
// PageIndexBuilder::put_column_index (put_offset_index is the same)
) -> bool {
self.column_indexes
.as_mut()
.is_some_and(|indexes| indexes.insert(row_group_idx, column_idx,
column_index))
}
```
`()` → `bool` on these 60.0.0 methods is a small API change.
--
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]
