alamb commented on code in PR #11475:
URL: https://github.com/apache/arrow-rs/pull/11475#discussion_r4551381106
##########
parquet/tests/arrow_reader/page_index.rs:
##########
@@ -87,6 +87,92 @@ fn test_parse_with_page_index_mask() {
);
}
+#[test]
+fn test_sparse_page_index_missing_ranges_preserve_coordinates() {
+ let file = create_test_file();
+ let original = ParquetMetaDataReader::new()
+ .with_page_index_policy(PageIndexPolicy::Required)
+ .parse_and_finish(&file)
+ .unwrap();
+ let metadata = ParquetMetaDataReader::new()
+ .with_page_index_policy(PageIndexPolicy::Skip)
+ .parse_and_finish(&file)
+ .unwrap();
+ let mut builder = metadata.into_builder();
+ let row_groups = builder
Review Comment:
I found this test hard to follow as the behavior to be validated (the row /
columns) is interleaved with the mechnics of setting up the test
I would personally find it easier to read if there was some sort of human
readable string that showed what the contents of the input row groups and the
output row groups
##########
parquet/src/file/metadata/column_chunk.rs:
##########
@@ -0,0 +1,380 @@
+// 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.
+
+//! Sparse storage keyed by row-group and column coordinates.
+
+use crate::errors::{ParquetError, Result};
+use crate::file::metadata::memory::HeapSize;
+use std::collections::HashMap;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub(crate) struct ColumnChunkKey(u64);
+
+impl ColumnChunkKey {
+ pub(crate) fn try_new(row_group_idx: usize, column_idx: usize) ->
Result<Self> {
+ let row_group_idx = u32::try_from(row_group_idx).map_err(|_| {
+ ParquetError::General(format!(
+ "row group index exceeds packed coordinate limit:
{row_group_idx}"
+ ))
+ })?;
+ let column_idx = u32::try_from(column_idx).map_err(|_| {
+ ParquetError::General(format!(
+ "column index exceeds packed coordinate limit: {column_idx}"
+ ))
+ })?;
+ Ok(Self(
+ (u64::from(row_group_idx) << u32::BITS) | u64::from(column_idx),
+ ))
+ }
+
+ pub(crate) fn coordinates(self) -> (usize, usize) {
+ ((self.0 >> u32::BITS) as usize, (self.0 as u32) as usize)
+ }
+}
+
+impl HeapSize for ColumnChunkKey {
+ fn heap_size(&self) -> usize {
+ 0
+ }
+}
+
+fn validate_dimensions(num_row_groups: usize, num_columns: usize) ->
Result<()> {
+ // Parquet metadata collection lengths are represented by signed 32-bit
Thrift values.
+ i32::try_from(num_row_groups).map_err(|_| {
+ ParquetError::General(format!(
+ "row group count exceeds i32::MAX: {num_row_groups}"
+ ))
+ })?;
+ i32::try_from(num_columns).map_err(|_| {
+ ParquetError::General(format!("column count exceeds i32::MAX:
{num_columns}"))
+ })?;
+ Ok(())
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct ColumnChunkMap<T> {
+ num_row_groups: usize,
+ num_columns: usize,
+ entries: HashMap<ColumnChunkKey, T>,
+}
+
+impl<T> ColumnChunkMap<T> {
+ pub(crate) fn try_new(num_row_groups: usize, num_columns: usize) ->
Result<Self> {
+ validate_dimensions(num_row_groups, num_columns)?;
+ Ok(Self {
+ num_row_groups,
+ num_columns,
+ entries: HashMap::new(),
+ })
+ }
+
+ pub(crate) fn num_row_groups(&self) -> usize {
+ self.num_row_groups
+ }
+
+ pub(crate) fn num_columns(&self) -> usize {
+ self.num_columns
+ }
+
+ pub(crate) fn is_empty(&self) -> bool {
+ self.entries.is_empty()
+ }
+
+ fn key(&self, row_group_idx: usize, column_idx: usize) ->
Result<ColumnChunkKey> {
+ if row_group_idx >= self.num_row_groups || column_idx >=
self.num_columns {
+ return Err(ParquetError::General(format!(
+ "column chunk coordinate out of bounds: row group
{row_group_idx}, column \
+ {column_idx}; dimensions are {} by {}",
+ self.num_row_groups, self.num_columns
+ )));
+ }
+ ColumnChunkKey::try_new(row_group_idx, column_idx)
+ }
+
+ /// Inserts or replaces an entry.
+ pub(crate) fn try_insert(
+ &mut self,
+ row_group_idx: usize,
+ column_idx: usize,
+ value: T,
+ ) -> Result<Option<T>> {
+ let key = self.key(row_group_idx, column_idx)?;
+ if let Some(existing) = self.entries.get_mut(&key) {
+ return Ok(Some(std::mem::replace(existing, value)));
+ }
+
+ self.entries.try_reserve(1).map_err(|err| {
+ ParquetError::General(format!(
+ "failed to reserve storage for column chunk at row group
{row_group_idx}, \
+ column {column_idx}: {err}"
+ ))
+ })?;
+ Ok(self.entries.insert(key, value))
+ }
+
+ pub(crate) fn try_freeze(self) -> Result<ColumnChunkEntries<T>> {
+ validate_dimensions(self.num_row_groups, self.num_columns)?;
+ let num_entries = self.entries.len();
+ let mut entries = Vec::new();
+ entries.try_reserve_exact(num_entries).map_err(|err| {
Review Comment:
I don't understand the rationale for checking for an allocation failure --
we don't do that in the rest of the crate, do we? I think it is fine, I just
found it kind of strange
##########
parquet/tests/arrow_reader/page_index.rs:
##########
@@ -382,9 +468,9 @@ fn test_page_index_sizes() {
let metadata = reader.finish().unwrap();
assert!(metadata.page_index().is_some());
#[cfg(not(feature = "encryption"))]
- assert_eq!(metadata.memory_size(), 13929);
+ assert_eq!(metadata.memory_size(), 13993);
Review Comment:
this is slightly larger because there is some extra overhead per entry,
right?
##########
parquet/tests/arrow_reader/page_index.rs:
##########
@@ -412,9 +498,9 @@ fn test_page_index_sizes() {
let metadata = reader.finish().unwrap();
assert!(metadata.page_index().is_some());
#[cfg(not(feature = "encryption"))]
- assert_eq!(metadata.memory_size(), 12088);
+ assert_eq!(metadata.memory_size(), 9656);
Review Comment:
that certainly looks a lot nicer 🎉
##########
parquet/src/file/metadata/column_chunk.rs:
##########
@@ -0,0 +1,380 @@
+// 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.
+
+//! Sparse storage keyed by row-group and column coordinates.
+
+use crate::errors::{ParquetError, Result};
+use crate::file::metadata::memory::HeapSize;
+use std::collections::HashMap;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub(crate) struct ColumnChunkKey(u64);
+
+impl ColumnChunkKey {
+ pub(crate) fn try_new(row_group_idx: usize, column_idx: usize) ->
Result<Self> {
+ let row_group_idx = u32::try_from(row_group_idx).map_err(|_| {
+ ParquetError::General(format!(
+ "row group index exceeds packed coordinate limit:
{row_group_idx}"
+ ))
+ })?;
+ let column_idx = u32::try_from(column_idx).map_err(|_| {
+ ParquetError::General(format!(
+ "column index exceeds packed coordinate limit: {column_idx}"
+ ))
+ })?;
+ Ok(Self(
+ (u64::from(row_group_idx) << u32::BITS) | u64::from(column_idx),
+ ))
+ }
+
+ pub(crate) fn coordinates(self) -> (usize, usize) {
+ ((self.0 >> u32::BITS) as usize, (self.0 as u32) as usize)
+ }
+}
+
+impl HeapSize for ColumnChunkKey {
+ fn heap_size(&self) -> usize {
+ 0
+ }
+}
+
+fn validate_dimensions(num_row_groups: usize, num_columns: usize) ->
Result<()> {
+ // Parquet metadata collection lengths are represented by signed 32-bit
Thrift values.
+ i32::try_from(num_row_groups).map_err(|_| {
+ ParquetError::General(format!(
+ "row group count exceeds i32::MAX: {num_row_groups}"
+ ))
+ })?;
+ i32::try_from(num_columns).map_err(|_| {
+ ParquetError::General(format!("column count exceeds i32::MAX:
{num_columns}"))
+ })?;
+ Ok(())
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct ColumnChunkMap<T> {
+ num_row_groups: usize,
+ num_columns: usize,
+ entries: HashMap<ColumnChunkKey, T>,
+}
+
+impl<T> ColumnChunkMap<T> {
+ pub(crate) fn try_new(num_row_groups: usize, num_columns: usize) ->
Result<Self> {
+ validate_dimensions(num_row_groups, num_columns)?;
+ Ok(Self {
+ num_row_groups,
+ num_columns,
+ entries: HashMap::new(),
+ })
+ }
+
+ pub(crate) fn num_row_groups(&self) -> usize {
+ self.num_row_groups
+ }
+
+ pub(crate) fn num_columns(&self) -> usize {
+ self.num_columns
+ }
+
+ pub(crate) fn is_empty(&self) -> bool {
+ self.entries.is_empty()
+ }
+
+ fn key(&self, row_group_idx: usize, column_idx: usize) ->
Result<ColumnChunkKey> {
+ if row_group_idx >= self.num_row_groups || column_idx >=
self.num_columns {
+ return Err(ParquetError::General(format!(
+ "column chunk coordinate out of bounds: row group
{row_group_idx}, column \
+ {column_idx}; dimensions are {} by {}",
+ self.num_row_groups, self.num_columns
+ )));
+ }
+ ColumnChunkKey::try_new(row_group_idx, column_idx)
+ }
+
+ /// Inserts or replaces an entry.
+ pub(crate) fn try_insert(
+ &mut self,
+ row_group_idx: usize,
+ column_idx: usize,
+ value: T,
+ ) -> Result<Option<T>> {
+ let key = self.key(row_group_idx, column_idx)?;
+ if let Some(existing) = self.entries.get_mut(&key) {
+ return Ok(Some(std::mem::replace(existing, value)));
+ }
+
+ self.entries.try_reserve(1).map_err(|err| {
+ ParquetError::General(format!(
+ "failed to reserve storage for column chunk at row group
{row_group_idx}, \
+ column {column_idx}: {err}"
+ ))
+ })?;
+ Ok(self.entries.insert(key, value))
+ }
+
+ pub(crate) fn try_freeze(self) -> Result<ColumnChunkEntries<T>> {
+ validate_dimensions(self.num_row_groups, self.num_columns)?;
+ let num_entries = self.entries.len();
+ let mut entries = Vec::new();
+ entries.try_reserve_exact(num_entries).map_err(|err| {
+ ParquetError::General(format!(
+ "failed to allocate storage for {num_entries} column chunk
entries: {err}"
+ ))
+ })?;
+ for entry in self.entries {
+ entries.push(entry);
+ }
+ entries.sort_unstable_by_key(|(key, _)| *key);
+ ColumnChunkEntries::try_from_sorted_entries(self.num_row_groups,
self.num_columns, entries)
+ }
+}
+
+impl<T: HeapSize> HeapSize for ColumnChunkMap<T> {
+ fn heap_size(&self) -> usize {
+ self.entries.heap_size()
+ }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct ColumnChunkEntries<T> {
Review Comment:
This seems like it is very similar to the Map above, except that it has a
single contiguous allocation -- what is the ratinale for the difference? Is it
to reduce the in memory / heap size? If so perhaps we can document that
rationale
##########
parquet/src/file/metadata/page_index.rs:
##########
@@ -506,77 +562,138 @@ impl HeapSize for PageIndex {
/// - Populating column indexes for predicate columns (for page filtering)
/// - Populating offset indexes for projected columns (for direct I/O)
/// - Automatic conversion of empty structures to `None` to save memory
-#[derive(Default)]
+#[derive(Default, Debug)]
pub struct PageIndexBuilder {
- column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>,
- offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>,
+ column_indexes: Option<ColumnChunkMap<ColumnIndexMetaData>>,
+ offset_indexes: Option<ColumnChunkMap<OffsetIndexMetaData>>,
}
impl PageIndexBuilder {
- /// Creates an empty index structure with space for the specified number
of row groups and columns
+ /// Creates a new [`PageIndexBuilder`] for the specified number of row
groups and columns.
///
- /// Returns `Some` containing a nested vector structure where all entries
are initialized to `None`.
- /// The outer vector has one entry per row group, and each inner vector
has one entry per column.
+ /// The dimensions are used to validate inserted coordinates. Storage is
allocated only for
+ /// entries populated using
+ /// [`put_column_index`](Self::put_column_index) and
[`put_offset_index`](Self::put_offset_index).
///
- /// # Type Parameters
- /// * `T` - The type of index this is to be, either `ColumnIndexMetaData`
or `OffsetIndexMetaData`
- fn empty_index<T>(num_row_groups: usize, num_columns: usize) ->
Option<Vec<Vec<Option<T>>>> {
- Some(
- (0..num_row_groups)
- .map(|_| {
- let mut idx = Vec::with_capacity(num_columns);
- idx.resize_with(num_columns, || None);
- idx
- })
- .collect(),
- )
+ /// # Panics
+ ///
+ /// Panics if either dimension exceeds `i32::MAX`, the maximum collection
size representable
+ /// by Thrift.
+ ///
+ /// Use [`try_new`](Self::try_new) for a fallible alternative.
Review Comment:
allocating more than i32::MAX entries does sound like a bad idea to me and
not a case that is important
--
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]
