alamb commented on code in PR #10842: URL: https://github.com/apache/arrow-rs/pull/10842#discussion_r3913699561
########## parquet/examples/custom_page_index.rs: ########## @@ -0,0 +1,307 @@ +// 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. + +//! Example of implementing a custom PageIndexProvider +//! +//! This example demonstrates how to create a custom page index provider that: +//! - Only stores page indexes for specified columns (selective storage) +//! - Uses nested HashMaps for efficient storage and lookup +//! - Implements all required PageIndexProvider trait methods +//! +//! This approach can significantly reduce memory usage when working with wide +//! tables where only a few columns need page-level statistics. + +use bytes::Bytes; +use parquet::DecodeResult; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::page_index::PageIndexProvider; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataPushDecoder}; +use parquet::file::page_index::column_index::ColumnIndexMetaData; +use parquet::file::page_index::index_reader::{decode_column_index, decode_offset_index}; +use parquet::file::page_index::offset_index::OffsetIndexMetaData; +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; +use tempfile::TempDir; + +////////////////////////////////////////////// Review Comment: nit: my personal preference is to put helpers at the end so the example starts with the "punchline" and then people can refer to the details if they need. However I am not sure how important that is going forward with coding agents, etc ########## parquet/examples/custom_page_index.rs: ########## @@ -0,0 +1,307 @@ +// 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. + +//! Example of implementing a custom PageIndexProvider +//! +//! This example demonstrates how to create a custom page index provider that: +//! - Only stores page indexes for specified columns (selective storage) +//! - Uses nested HashMaps for efficient storage and lookup +//! - Implements all required PageIndexProvider trait methods +//! +//! This approach can significantly reduce memory usage when working with wide Review Comment: Can we also claim it reduces metadata load time? "This approach can significantly reduce memory usage and metadata load time when working with wide tables and you only access a few columns". Though maybe we should also explain when page level statistics are helpful (evaluating predicates (stats) or fetching specific ranges of rows (after predicates or index application). That might be too nuanced however 🤔 ########## parquet/examples/custom_page_index.rs: ########## @@ -0,0 +1,307 @@ +// 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. + +//! Example of implementing a custom PageIndexProvider +//! +//! This example demonstrates how to create a custom page index provider that: +//! - Only stores page indexes for specified columns (selective storage) +//! - Uses nested HashMaps for efficient storage and lookup +//! - Implements all required PageIndexProvider trait methods +//! +//! This approach can significantly reduce memory usage when working with wide +//! tables where only a few columns need page-level statistics. + +use bytes::Bytes; +use parquet::DecodeResult; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::page_index::PageIndexProvider; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataPushDecoder}; +use parquet::file::page_index::column_index::ColumnIndexMetaData; +use parquet::file::page_index::index_reader::{decode_column_index, decode_offset_index}; +use parquet::file::page_index::offset_index::OffsetIndexMetaData; +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; +use tempfile::TempDir; + +////////////////////////////////////////////// +// helper functions + +fn print_page_index(metadata: &ParquetMetaData, row_group_idx: usize) -> Result<()> { + if let Some(page_index) = metadata.page_index() { + let num_columns = metadata.file_metadata().schema_descr().num_columns(); + + println!("\nIndexes for row group {row_group_idx}:"); + for col_idx in 0..num_columns { + println!( + " Column {col_idx}: has offset idx {}, has column idx {}", + page_index.offset_index(row_group_idx, col_idx).is_some(), + page_index.column_index(row_group_idx, col_idx).is_some() + ); + } + println!(); + } else { + println!("No page index in metadata"); + println!("Note: This example requires a file with page indexes."); + println!("Try using alltypes_tiny_pages.parquet or another file with page indexes."); + return Err(ParquetError::General("no page index".to_string())); + } + Ok(()) +} + +fn create_sample_file(temp_path: &PathBuf) -> Result<()> { + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::ArrowWriter; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; + + println!("Creating sample file: {}", temp_path.display()); + + // Create a sample dataset with multiple columns + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + Field::new("score", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new("amount", DataType::Int32, false), + ])); + + // Create multiple row groups with multiple pages + let file = File::create(temp_path)?; + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_data_page_size_limit(100) // Small pages for demonstration + .set_write_batch_size(10) + .build(); + + let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))?; + + // Write several row groups + for row_group in 0..3 { + for batch_num in 0..5 { + let offset = (row_group * 50) + (batch_num * 10); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from( + (offset..offset + 10).collect::<Vec<i32>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 2).collect::<Vec<i32>>(), + )), + Arc::new(StringArray::from( + (offset..offset + 10) + .map(|x| format!("name{x}")) + .collect::<Vec<String>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 3).collect::<Vec<i32>>(), + )), + Arc::new(StringArray::from( + (offset..offset + 10) + .map(|x| if x % 2 == 0 { "even" } else { "odd" }) + .collect::<Vec<&str>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 4).collect::<Vec<i32>>(), + )), + ], + )?; + writer.write(&batch)?; + } + writer.flush()?; + } + + writer.close()?; + Ok(()) +} + +////////////////////////////////////////////// +// our custom provider + +// A custom PageIndexProvider that only stores indexes for a subset of columns +// +// This provider contains the parsed footer metadata and the entire contents of a +// Parquet file. Indexes are lazily populated as they are requested. +#[derive(Debug, Clone)] +pub struct OnDemandPageIndexProvider { + metadata: ParquetMetaData, + file_bytes: Bytes, + column_indexes: Option<HashMap<usize, HashMap<usize, ColumnIndexMetaData>>>, + offset_indexes: Option<HashMap<usize, HashMap<usize, OffsetIndexMetaData>>>, +} + +impl OnDemandPageIndexProvider { + fn new(metadata: ParquetMetaData, file_bytes: Bytes) -> Self { + Self { + metadata, + file_bytes, + column_indexes: None, + offset_indexes: None, + } + } + + // fetches the bytes for and parses the column index for the given row group and + // column. If already fetched, this does nothing. + fn fetch_column_index(&mut self, row_group_idx: usize, column_idx: usize) -> Result<()> { + let map = self.column_indexes.get_or_insert_with(HashMap::new); + let rg = map.entry(row_group_idx).or_default(); + if let Entry::Vacant(e) = rg.entry(column_idx) { + let column = self.metadata.row_group(row_group_idx).column(column_idx); + let range = column.column_index_range(); + if let Some(range) = range { + let idx_bytes = self + .file_bytes + .slice(range.start as usize..range.end as usize); + let idx = decode_column_index(&idx_bytes, column.column_type())?; + e.insert(idx); + } + } + Ok(()) + } + + // fetches the bytes for and parses the offset index for the given row group and + // column. If already fetched, this does nothing. + fn fetch_offset_index(&mut self, row_group_idx: usize, column_idx: usize) -> Result<()> { + let map = self.offset_indexes.get_or_insert_with(HashMap::new); + let rg = map.entry(row_group_idx).or_default(); + if let Entry::Vacant(e) = rg.entry(column_idx) { + let column = self.metadata.row_group(row_group_idx).column(column_idx); + let range = column.offset_index_range(); Review Comment: ```suggestion // Find the location of the offset index from the file metadata let range = column.offset_index_range(); ``` ########## parquet/src/arrow/arrow_reader/mod.rs: ########## @@ -4905,19 +4906,12 @@ pub(crate) mod tests { ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), ) .unwrap(); - let page_index = builder - .metadata() - .page_index() - .expect("page index should be present"); - let num_columns = builder.metadata().row_group(0).num_columns(); - let offset_indexes = page_index.offset_indexes_for_rowgroup(0); - assert!(offset_indexes.is_some_and(|ois| ois.len() == num_columns)); - let column_indexes = page_index.offset_indexes_for_rowgroup(0); - assert!(column_indexes.is_some_and(|cis| cis.len() == num_columns)); - assert!(page_index.offset_index(0, 0).is_some()); - assert!(page_index.column_index(0, 0).is_some()); - assert!(page_index.page_locations(0, 0).is_some()); - assert_eq!(page_index.num_data_pages(0, 0), Some(325)); + let page_index = builder.metadata().page_index(); + let row_group_page_index = RowGroupPageIndex::new(0, page_index.cloned()); Review Comment: this is so much nicer <img width="245" height="245" alt="Image" src="https://github.com/user-attachments/assets/69f1b6a0-0906-4045-b7c4-90d439f71da0"; /> ########## parquet/src/arrow/async_reader/store.rs: ########## @@ -447,7 +447,7 @@ mod tests { let metadata = reader.get_metadata(Some(&options)).await.unwrap(); // With preload=true, indexes should be loaded since the test file has them - assert!(metadata.page_index().is_some_and(PageIndex::is_complete)); + assert!(metadata.page_index().is_some_and(|idx| idx.is_complete())); Review Comment: Is this change necessary? It isn't wrong, I am just curious ########## parquet/examples/custom_page_index.rs: ########## @@ -0,0 +1,307 @@ +// 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. + +//! Example of implementing a custom PageIndexProvider +//! +//! This example demonstrates how to create a custom page index provider that: +//! - Only stores page indexes for specified columns (selective storage) +//! - Uses nested HashMaps for efficient storage and lookup +//! - Implements all required PageIndexProvider trait methods +//! +//! This approach can significantly reduce memory usage when working with wide +//! tables where only a few columns need page-level statistics. + +use bytes::Bytes; +use parquet::DecodeResult; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::page_index::PageIndexProvider; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataPushDecoder}; +use parquet::file::page_index::column_index::ColumnIndexMetaData; +use parquet::file::page_index::index_reader::{decode_column_index, decode_offset_index}; +use parquet::file::page_index::offset_index::OffsetIndexMetaData; +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; +use tempfile::TempDir; + +////////////////////////////////////////////// +// helper functions + +fn print_page_index(metadata: &ParquetMetaData, row_group_idx: usize) -> Result<()> { + if let Some(page_index) = metadata.page_index() { + let num_columns = metadata.file_metadata().schema_descr().num_columns(); + + println!("\nIndexes for row group {row_group_idx}:"); + for col_idx in 0..num_columns { + println!( + " Column {col_idx}: has offset idx {}, has column idx {}", + page_index.offset_index(row_group_idx, col_idx).is_some(), + page_index.column_index(row_group_idx, col_idx).is_some() + ); + } + println!(); + } else { + println!("No page index in metadata"); + println!("Note: This example requires a file with page indexes."); + println!("Try using alltypes_tiny_pages.parquet or another file with page indexes."); + return Err(ParquetError::General("no page index".to_string())); + } + Ok(()) +} + +fn create_sample_file(temp_path: &PathBuf) -> Result<()> { + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::ArrowWriter; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; + + println!("Creating sample file: {}", temp_path.display()); + + // Create a sample dataset with multiple columns + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + Field::new("score", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new("amount", DataType::Int32, false), + ])); + + // Create multiple row groups with multiple pages + let file = File::create(temp_path)?; + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_data_page_size_limit(100) // Small pages for demonstration + .set_write_batch_size(10) + .build(); + + let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))?; + + // Write several row groups + for row_group in 0..3 { + for batch_num in 0..5 { + let offset = (row_group * 50) + (batch_num * 10); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from( + (offset..offset + 10).collect::<Vec<i32>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 2).collect::<Vec<i32>>(), + )), + Arc::new(StringArray::from( + (offset..offset + 10) + .map(|x| format!("name{x}")) + .collect::<Vec<String>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 3).collect::<Vec<i32>>(), + )), + Arc::new(StringArray::from( + (offset..offset + 10) + .map(|x| if x % 2 == 0 { "even" } else { "odd" }) + .collect::<Vec<&str>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 4).collect::<Vec<i32>>(), + )), + ], + )?; + writer.write(&batch)?; + } + writer.flush()?; + } + + writer.close()?; + Ok(()) +} + +////////////////////////////////////////////// +// our custom provider + +// A custom PageIndexProvider that only stores indexes for a subset of columns +// +// This provider contains the parsed footer metadata and the entire contents of a +// Parquet file. Indexes are lazily populated as they are requested. +#[derive(Debug, Clone)] +pub struct OnDemandPageIndexProvider { + metadata: ParquetMetaData, + file_bytes: Bytes, + column_indexes: Option<HashMap<usize, HashMap<usize, ColumnIndexMetaData>>>, Review Comment: maybe with a comment explaining that the first level is row_group_index and second level is column_index? ########## parquet/src/file/metadata/writer.rs: ########## @@ -453,17 +452,18 @@ impl<'a, W: Write> ParquetMetaDataWriter<'a, W> { self.write_path_in_schema, ); - if let Some(PageIndex { - column_indexes, - offset_indexes, - }) = page_index + // Downcast to PageIndex to access raw index structures for serialization + if let Some(page_index_arc) = self.metadata.page_index.as_ref() + && let Some(page_index) = page_index_arc + .as_any() + .downcast_ref::<crate::file::metadata::PageIndex>() Review Comment: > Downcast to PageIndex to access raw index structures for serialization Claude points out that this means that the writer not write page indexes if the provider is a custom provider. I think we need to either explicitly call that out in docs, or (preferably) actually serialize the page indexes when sourced from a custom provider We could document the limitation in this PR and then fix it in a follow on PR (I bet if we wrote up a ticket someone else would do it) ########## parquet/src/file/metadata/page_index.rs: ########## Review Comment: We can finally have nice things -- a module that has a page index separated -- so nice! ########## parquet/src/file/metadata/mod.rs: ########## @@ -507,19 +229,56 @@ impl ParquetMetaData { #[cfg(not(feature = "encryption"))] let encryption_size = 0usize; + // We can only determine the heap size for PageIndex. Custom providers are + // out of scope. + let page_index_size = if let Some(page_index) = self.page_index.as_ref() { + if let Some(page_index) = page_index.as_any().downcast_ref::<PageIndex>() { + let page_index = Some(Arc::new(page_index.clone())); Review Comment: this `clone` deep copies the page index -- I think it should be something mor elike `std::mem::size_of::<PageIndex>() + page_index.heap_size()` ########## parquet/examples/custom_page_index.rs: ########## @@ -0,0 +1,307 @@ +// 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. + +//! Example of implementing a custom PageIndexProvider +//! +//! This example demonstrates how to create a custom page index provider that: +//! - Only stores page indexes for specified columns (selective storage) +//! - Uses nested HashMaps for efficient storage and lookup +//! - Implements all required PageIndexProvider trait methods +//! +//! This approach can significantly reduce memory usage when working with wide +//! tables where only a few columns need page-level statistics. + +use bytes::Bytes; +use parquet::DecodeResult; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::page_index::PageIndexProvider; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataPushDecoder}; +use parquet::file::page_index::column_index::ColumnIndexMetaData; +use parquet::file::page_index::index_reader::{decode_column_index, decode_offset_index}; +use parquet::file::page_index::offset_index::OffsetIndexMetaData; +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; +use tempfile::TempDir; + +////////////////////////////////////////////// +// helper functions + +fn print_page_index(metadata: &ParquetMetaData, row_group_idx: usize) -> Result<()> { + if let Some(page_index) = metadata.page_index() { + let num_columns = metadata.file_metadata().schema_descr().num_columns(); + + println!("\nIndexes for row group {row_group_idx}:"); + for col_idx in 0..num_columns { + println!( + " Column {col_idx}: has offset idx {}, has column idx {}", + page_index.offset_index(row_group_idx, col_idx).is_some(), + page_index.column_index(row_group_idx, col_idx).is_some() + ); + } + println!(); + } else { + println!("No page index in metadata"); + println!("Note: This example requires a file with page indexes."); + println!("Try using alltypes_tiny_pages.parquet or another file with page indexes."); + return Err(ParquetError::General("no page index".to_string())); + } + Ok(()) +} + +fn create_sample_file(temp_path: &PathBuf) -> Result<()> { + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::ArrowWriter; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; + + println!("Creating sample file: {}", temp_path.display()); + + // Create a sample dataset with multiple columns + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + Field::new("score", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new("amount", DataType::Int32, false), + ])); + + // Create multiple row groups with multiple pages + let file = File::create(temp_path)?; + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_data_page_size_limit(100) // Small pages for demonstration + .set_write_batch_size(10) + .build(); + + let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))?; + + // Write several row groups + for row_group in 0..3 { + for batch_num in 0..5 { + let offset = (row_group * 50) + (batch_num * 10); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from( + (offset..offset + 10).collect::<Vec<i32>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 2).collect::<Vec<i32>>(), + )), + Arc::new(StringArray::from( + (offset..offset + 10) + .map(|x| format!("name{x}")) + .collect::<Vec<String>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 3).collect::<Vec<i32>>(), + )), + Arc::new(StringArray::from( + (offset..offset + 10) + .map(|x| if x % 2 == 0 { "even" } else { "odd" }) + .collect::<Vec<&str>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 4).collect::<Vec<i32>>(), + )), + ], + )?; + writer.write(&batch)?; + } + writer.flush()?; + } + + writer.close()?; + Ok(()) +} + +////////////////////////////////////////////// +// our custom provider + +// A custom PageIndexProvider that only stores indexes for a subset of columns +// +// This provider contains the parsed footer metadata and the entire contents of a +// Parquet file. Indexes are lazily populated as they are requested. +#[derive(Debug, Clone)] +pub struct OnDemandPageIndexProvider { + metadata: ParquetMetaData, + file_bytes: Bytes, + column_indexes: Option<HashMap<usize, HashMap<usize, ColumnIndexMetaData>>>, + offset_indexes: Option<HashMap<usize, HashMap<usize, OffsetIndexMetaData>>>, +} + +impl OnDemandPageIndexProvider { + fn new(metadata: ParquetMetaData, file_bytes: Bytes) -> Self { + Self { + metadata, + file_bytes, + column_indexes: None, + offset_indexes: None, + } + } + + // fetches the bytes for and parses the column index for the given row group and + // column. If already fetched, this does nothing. + fn fetch_column_index(&mut self, row_group_idx: usize, column_idx: usize) -> Result<()> { + let map = self.column_indexes.get_or_insert_with(HashMap::new); + let rg = map.entry(row_group_idx).or_default(); + if let Entry::Vacant(e) = rg.entry(column_idx) { + let column = self.metadata.row_group(row_group_idx).column(column_idx); + let range = column.column_index_range(); Review Comment: ```suggestion // Find the location of the column index from the file metadata let range = column.column_index_range(); ``` ########## parquet/src/file/metadata/mod.rs: ########## @@ -507,19 +229,56 @@ impl ParquetMetaData { #[cfg(not(feature = "encryption"))] let encryption_size = 0usize; + // We can only determine the heap size for PageIndex. Custom providers are Review Comment: I think it is fine to not include heap size in the memory usage calculation in this PR, but we should file a follow on PR to add an API for a custom index provider to report its memory usage (as one of the main points of this PR is to have more efficient caching, for which we need to know how large the memory usage is) ########## parquet/src/file/metadata/page_index.rs: ########## @@ -0,0 +1,672 @@ +// 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. + +//! Page Index structures for efficient page-level skipping + +use crate::file::metadata::memory::HeapSize; +use crate::file::page_index::{ + column_index::ColumnIndexMetaData, + offset_index::{OffsetIndexMetaData, PageLocation}, +}; +use std::sync::Arc; + +/// Trait for accessing Parquet [Page Index] data for efficient page-level skipping +/// +/// The [Page Index] enables query engines to skip irrelevant data pages during scans, +/// significantly improving I/O efficiency. It provides access to two complementary +/// structures: +/// +/// * **[`ColumnIndex`]**: Per-page min/max value boundaries that enable predicate-based +/// page filtering. Allows determining which pages might contain rows matching a query +/// predicate without reading the actual data pages. +/// +/// * **[`OffsetIndex`]**: Physical locations and sizes of data pages, plus the first row +/// index of each page. Used to locate and read only the pages identified as relevant +/// by the ColumnIndex. +/// +/// Together, these indexes enable: +/// - Single-row lookups reading only one data page per column (on sorted columns) +/// - Range scans reading only pages containing values in the query range +/// - Efficient cross-column filtering by skipping corresponding row ranges +/// +/// # Structure +/// +/// Within a Parquet file, both indexes are organized as a two-level structure, with +/// indexes arranged first by row group, and then column. The [`ColumnChunkMetaData`] +/// contains pointers to the indexes for a given column chunk, so they may be +/// populated piecemeal. This trait allows access by row group index and column number +/// ([Self::column_index], [Self::offset_index]). Access by row group is provided by +/// [`RowGroupPageIndex`]. +/// +/// Each entry is `Option<T>` because: +/// - The entire page index might be absent (old files, disabled during write) +/// - Individual columns might lack indexes (unsupported types, statistics disabled) +/// +/// # Example: Checking if Page Index is Available +/// +/// ``` +/// use parquet::file::metadata::ParquetMetaData; +/// # use parquet::errors::Result; +/// +/// fn check_page_index_availability(metadata: &ParquetMetaData) -> Result<()> { +/// if let Some(page_index) = metadata.page_index() { +/// println!("Page index present:"); +/// println!(" Has offset indexes: {}", page_index.has_offset_indexes()); +/// println!(" Has column indexes: {}", page_index.has_column_indexes()); +/// +/// // Check availability for first row group, first column +/// if let Some(col_idx) = page_index.column_index(0, 0) { +/// println!(" Column index found for row group 0, column 0"); +/// println!(" Number of pages: {}", col_idx.num_pages()); +/// } +/// +/// if let Some(offset_idx) = page_index.offset_index(0, 0) { +/// println!(" Offset index found for row group 0, column 0"); +/// println!(" Number of pages: {}", offset_idx.page_locations().len()); +/// } +/// } else { +/// println!("No page index available"); +/// } +/// Ok(()) +/// } +/// ``` +/// +/// # Example: Using Page Index for Predicate Pushdown +/// +/// ``` +/// use parquet::file::metadata::ParquetMetaData; +/// use parquet::file::page_index::column_index::ColumnIndexMetaData; +/// # use parquet::errors::Result; +/// +/// /// Identifies which pages in a column might contain values >= min_value +/// fn find_relevant_pages( +/// metadata: &ParquetMetaData, +/// row_group_idx: usize, +/// column_idx: usize, +/// min_value: i32, +/// ) -> Vec<usize> { +/// let mut relevant_pages = Vec::new(); +/// +/// let Some(page_index) = metadata.page_index() else { +/// // No page index - must read all pages +/// return relevant_pages; +/// }; +/// +/// let Some(column_index) = page_index.column_index(row_group_idx, column_idx) else { +/// // No column index - must read all pages +/// return relevant_pages; +/// }; +/// +/// // Check each page's statistics +/// match column_index { +/// ColumnIndexMetaData::INT32(index) => { +/// for (page_num, max_value) in index.max_values_iter().enumerate() { +/// // Page might contain matching rows if its max >= our min +/// if let Some(max) = max_value { +/// if *max >= min_value { +/// relevant_pages.push(page_num); +/// } +/// } +/// } +/// } +/// _ => { +/// // Wrong column type - read all pages +/// } +/// } +/// +/// relevant_pages +/// } +/// ``` +/// +/// [Page Index]: https://parquet.apache.org/docs/file-format/pageindex/ +/// [`ColumnIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData +/// [`OffsetIndex`]: crate::file::page_index::offset_index::OffsetIndexMetaData +/// [`ColumnChunkMetaData`]: crate::file::metadata::ColumnChunkMetaData +pub trait PageIndexProvider: Send + Sync + std::fmt::Debug { + /// Returns `true` if offset index structures are present Review Comment: It is probably also worth mentioning here that if `has_offset_indexes` should return false only of `offset_index` will *always* return false -- aka if this method reports false, the reader/writer won't even try to load any offset indexes As written it is slightly unclear if it should report false of no indexes are currently loaded but some might be loaded in the future ########## parquet/src/arrow/arrow_reader/mod.rs: ########## @@ -4905,19 +4906,12 @@ pub(crate) mod tests { ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), ) .unwrap(); - let page_index = builder - .metadata() - .page_index() - .expect("page index should be present"); - let num_columns = builder.metadata().row_group(0).num_columns(); - let offset_indexes = page_index.offset_indexes_for_rowgroup(0); - assert!(offset_indexes.is_some_and(|ois| ois.len() == num_columns)); - let column_indexes = page_index.offset_indexes_for_rowgroup(0); - assert!(column_indexes.is_some_and(|cis| cis.len() == num_columns)); - assert!(page_index.offset_index(0, 0).is_some()); - assert!(page_index.column_index(0, 0).is_some()); - assert!(page_index.page_locations(0, 0).is_some()); - assert_eq!(page_index.num_data_pages(0, 0), Some(325)); + let page_index = builder.metadata().page_index(); + let row_group_page_index = RowGroupPageIndex::new(0, page_index.cloned()); Review Comment: We could make it even more beautiful if we made this a method on PageIndex. Something like ```rust let row_group_page_index = page_index.row_group_page_index(0); ``` ########## parquet/examples/custom_page_index.rs: ########## @@ -0,0 +1,307 @@ +// 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. + +//! Example of implementing a custom PageIndexProvider +//! +//! This example demonstrates how to create a custom page index provider that: +//! - Only stores page indexes for specified columns (selective storage) +//! - Uses nested HashMaps for efficient storage and lookup +//! - Implements all required PageIndexProvider trait methods +//! +//! This approach can significantly reduce memory usage when working with wide +//! tables where only a few columns need page-level statistics. + +use bytes::Bytes; +use parquet::DecodeResult; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::page_index::PageIndexProvider; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataPushDecoder}; +use parquet::file::page_index::column_index::ColumnIndexMetaData; +use parquet::file::page_index::index_reader::{decode_column_index, decode_offset_index}; +use parquet::file::page_index::offset_index::OffsetIndexMetaData; +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; +use tempfile::TempDir; + +////////////////////////////////////////////// +// helper functions + +fn print_page_index(metadata: &ParquetMetaData, row_group_idx: usize) -> Result<()> { + if let Some(page_index) = metadata.page_index() { + let num_columns = metadata.file_metadata().schema_descr().num_columns(); + + println!("\nIndexes for row group {row_group_idx}:"); + for col_idx in 0..num_columns { + println!( + " Column {col_idx}: has offset idx {}, has column idx {}", + page_index.offset_index(row_group_idx, col_idx).is_some(), + page_index.column_index(row_group_idx, col_idx).is_some() + ); + } + println!(); + } else { + println!("No page index in metadata"); + println!("Note: This example requires a file with page indexes."); + println!("Try using alltypes_tiny_pages.parquet or another file with page indexes."); + return Err(ParquetError::General("no page index".to_string())); + } + Ok(()) +} + +fn create_sample_file(temp_path: &PathBuf) -> Result<()> { + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::ArrowWriter; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; + + println!("Creating sample file: {}", temp_path.display()); + + // Create a sample dataset with multiple columns + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + Field::new("score", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new("amount", DataType::Int32, false), + ])); + + // Create multiple row groups with multiple pages + let file = File::create(temp_path)?; + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_data_page_size_limit(100) // Small pages for demonstration + .set_write_batch_size(10) + .build(); + + let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))?; + + // Write several row groups + for row_group in 0..3 { + for batch_num in 0..5 { + let offset = (row_group * 50) + (batch_num * 10); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from( + (offset..offset + 10).collect::<Vec<i32>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 2).collect::<Vec<i32>>(), + )), + Arc::new(StringArray::from( + (offset..offset + 10) + .map(|x| format!("name{x}")) + .collect::<Vec<String>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 3).collect::<Vec<i32>>(), + )), + Arc::new(StringArray::from( + (offset..offset + 10) + .map(|x| if x % 2 == 0 { "even" } else { "odd" }) + .collect::<Vec<&str>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 4).collect::<Vec<i32>>(), + )), + ], + )?; + writer.write(&batch)?; + } + writer.flush()?; + } + + writer.close()?; + Ok(()) +} + +////////////////////////////////////////////// +// our custom provider + +// A custom PageIndexProvider that only stores indexes for a subset of columns +// +// This provider contains the parsed footer metadata and the entire contents of a +// Parquet file. Indexes are lazily populated as they are requested. +#[derive(Debug, Clone)] +pub struct OnDemandPageIndexProvider { + metadata: ParquetMetaData, + file_bytes: Bytes, + column_indexes: Option<HashMap<usize, HashMap<usize, ColumnIndexMetaData>>>, + offset_indexes: Option<HashMap<usize, HashMap<usize, OffsetIndexMetaData>>>, +} + +impl OnDemandPageIndexProvider { + fn new(metadata: ParquetMetaData, file_bytes: Bytes) -> Self { + Self { + metadata, + file_bytes, + column_indexes: None, + offset_indexes: None, + } + } + + // fetches the bytes for and parses the column index for the given row group and + // column. If already fetched, this does nothing. + fn fetch_column_index(&mut self, row_group_idx: usize, column_idx: usize) -> Result<()> { + let map = self.column_indexes.get_or_insert_with(HashMap::new); + let rg = map.entry(row_group_idx).or_default(); + if let Entry::Vacant(e) = rg.entry(column_idx) { + let column = self.metadata.row_group(row_group_idx).column(column_idx); + let range = column.column_index_range(); Review Comment: Possibly useful to call out how the file index is fetched -- 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]
