From 1c7798468aff9b57a0171cec262e7518334ba9c9 Mon Sep 17 00:00:00 2001 From: seidl Date: Tue, 8 Sep 2026 15:00:56 -0700 Subject: [PATCH 1/2] allow for serialization of custom page index providers --- parquet/src/file/metadata/page_index.rs | 21 +--- parquet/src/file/metadata/writer.rs | 121 ++++++++++++------------ parquet/src/file/writer.rs | 73 +++++++++++++- 3 files changed, 132 insertions(+), 83 deletions(-) diff --git a/parquet/src/file/metadata/page_index.rs b/parquet/src/file/metadata/page_index.rs index 9ded57d0a28e..83b595daf141 100644 --- a/parquet/src/file/metadata/page_index.rs +++ b/parquet/src/file/metadata/page_index.rs @@ -460,22 +460,6 @@ impl PageIndex { pub fn into_builder(self) -> PageIndexBuilder { self.into() } - - /// Returns a reference to the raw column indexes structure - /// - /// This method provides access to the underlying column index data for serialization - /// and other low-level operations. - pub(crate) fn column_indexes_raw(&self) -> Option<&Vec>>> { - self.column_indexes.as_ref() - } - - /// Returns a reference to the raw offset indexes structure - /// - /// This method provides access to the underlying offset index data for serialization - /// and other low-level operations. - pub(crate) fn offset_indexes_raw(&self) -> Option<&Vec>>> { - self.offset_indexes.as_ref() - } } impl PageIndexProvider for PageIndex { @@ -536,7 +520,10 @@ impl PageIndexBuilder { /// /// # Type Parameters /// * `T` - The type of index this is to be, either `ColumnIndexMetaData` or `OffsetIndexMetaData` - fn empty_index(num_row_groups: usize, num_columns: usize) -> Option>>> { + pub(crate) fn empty_index( + num_row_groups: usize, + num_columns: usize, + ) -> Option>>> { Some( (0..num_row_groups) .map(|_| { diff --git a/parquet/src/file/metadata/writer.rs b/parquet/src/file/metadata/writer.rs index a1708ccf2c3a..d6de7e45c8c5 100644 --- a/parquet/src/file/metadata/writer.rs +++ b/parquet/src/file/metadata/writer.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::file::metadata::page_index::{PageIndexBuilder, PageIndexProvider}; use crate::file::metadata::thrift::FileMeta; use crate::file::metadata::{ColumnChunkMetaData, PageIndex, RowGroupMetaData}; use crate::schema::types::{SchemaDescPtr, SchemaDescriptor}; @@ -54,8 +55,7 @@ pub(crate) struct ThriftMetadataWriter<'a, W: Write> { buf: &'a mut TrackedWrite, schema_descr: &'a SchemaDescPtr, row_groups: Vec, - column_indexes: Option>>>, - offset_indexes: Option>>>, + page_index: Option>, key_value_metadata: Option>, created_by: Option, object_writer: MetadataObjectWriter, @@ -71,14 +71,21 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { /// of the serialized offset indexes. fn write_offset_indexes( &mut self, - offset_indexes: &[Vec>], - ) -> Result<()> { + page_index: &Arc, + ) -> Result>>>> { + let mut offset_indexes = + PageIndexBuilder::empty_index(self.row_groups.len(), self.schema_descr.num_columns()); + let offidx_vec = offset_indexes.as_mut().unwrap(); + + // we've already checked before calling that the offset indexes are populated + assert!(page_index.has_offset_indexes()); + // iter row group // iter each column // write offset index to the file for (row_group_idx, row_group) in self.row_groups.iter_mut().enumerate() { for (column_idx, column_metadata) in row_group.columns.iter_mut().enumerate() { - if let Some(offset_index) = &offset_indexes[row_group_idx][column_idx] { + if let Some(offset_index) = page_index.offset_index(row_group_idx, column_idx) { let start_offset = self.buf.bytes_written(); self.object_writer.write_offset_index( offset_index, @@ -91,10 +98,11 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { // set offset and index for offset index column_metadata.offset_index_offset = Some(start_offset as i64); column_metadata.offset_index_length = Some((end_offset - start_offset) as i32); + offidx_vec[row_group_idx][column_idx] = Some(offset_index.clone()); } } } - Ok(()) + Ok(offset_indexes) } /// Serialize all the column indexes to the `self.buf` @@ -104,14 +112,21 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { /// of the serialized column indexes. fn write_column_indexes( &mut self, - column_indexes: &[Vec>], - ) -> Result<()> { + page_index: &Arc, + ) -> Result>>>> { + let mut column_indexes = + PageIndexBuilder::empty_index(self.row_groups.len(), self.schema_descr.num_columns()); + let colidx_vec = column_indexes.as_mut().unwrap(); + + // we've already checked before calling that the column indexes are populated + assert!(page_index.has_column_indexes()); + // iter row group // iter each column // write column index to the file for (row_group_idx, row_group) in self.row_groups.iter_mut().enumerate() { for (column_idx, column_metadata) in row_group.columns.iter_mut().enumerate() { - if let Some(column_index) = &column_indexes[row_group_idx][column_idx] { + if let Some(column_index) = page_index.column_index(row_group_idx, column_idx) { let start_offset = self.buf.bytes_written(); // only update column_metadata if the write succeeds if self.object_writer.write_column_index( @@ -127,21 +142,28 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { column_metadata.column_index_length = Some((end_offset - start_offset) as i32); } + colidx_vec[row_group_idx][column_idx] = Some(column_index.clone()); } } } - Ok(()) + Ok(column_indexes) } - /// Serialize the column indexes and transform to `Option` - fn finalize_column_indexes(&mut self) -> Result>>>> { - let column_indexes = std::mem::take(&mut self.column_indexes); - - // Write column indexes to file - if let Some(column_indexes) = column_indexes.as_ref() { - self.write_column_indexes(column_indexes)?; + /// Serialize the column indexes and transform to `Option>>>` + fn finalize_column_indexes( + &mut self, + page_index: &Option>, + ) -> Result>>>> { + if page_index + .as_ref() + .is_none_or(|pi| !pi.has_column_indexes()) + { + return Ok(None); } + // Write column indexes to file while assembling a column index + let column_indexes = self.write_column_indexes(page_index.as_ref().unwrap())?; + // check to see if the index is `None` for every row group and column chunk let all_none = column_indexes .as_ref() @@ -155,13 +177,19 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { } /// Serialize the offset indexes and transform to `Option` - fn finalize_offset_indexes(&mut self) -> Result>>>> { - let offset_indexes = std::mem::take(&mut self.offset_indexes); + fn finalize_offset_indexes( + &mut self, + page_index: &Option>, + ) -> Result>>>> { + if page_index + .as_ref() + .is_none_or(|pi| !pi.has_offset_indexes()) + { + return Ok(None); + } // Write offset indexes to file - if let Some(offset_indexes) = offset_indexes.as_ref() { - self.write_offset_indexes(offset_indexes)?; - } + let offset_indexes = self.write_offset_indexes(page_index.as_ref().unwrap())?; // check to see if the index is `None` for every row group and column chunk let all_none = offset_indexes @@ -180,8 +208,9 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { let num_rows = self.row_groups.iter().map(|x| x.num_rows).sum(); // serialize page indexes and transform to the proper form for use in ParquetMetaData - let column_indexes = self.finalize_column_indexes()?; - let offset_indexes = self.finalize_offset_indexes()?; + let page_index = self.page_index.take(); + let column_indexes = self.finalize_column_indexes(&page_index)?; + let offset_indexes = self.finalize_offset_indexes(&page_index)?; // We only include ColumnOrder for leaf nodes. let column_orders = self @@ -275,8 +304,7 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { buf, schema_descr, row_groups, - column_indexes: None, - offset_indexes: None, + page_index: None, key_value_metadata: None, created_by, object_writer: Default::default(), @@ -285,19 +313,8 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { } } - pub fn with_column_indexes( - mut self, - column_indexes: Vec>>, - ) -> Self { - self.column_indexes = Some(column_indexes); - self - } - - pub fn with_offset_indexes( - mut self, - offset_indexes: Vec>>, - ) -> Self { - self.offset_indexes = Some(offset_indexes); + pub fn with_page_index(mut self, page_index: Arc) -> Self { + self.page_index = Some(page_index); self } @@ -333,17 +350,6 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { /// metadata writer. Then set the corresponding `bloom_filter_offset` and /// `bloom_filter_length` on [`ColumnChunkMetaData`] passed to this writer. /// -///
-/// -/// **NOTE:** -/// The serialization of custom [`PageIndexProvider`]s is not currently supported. -/// The only supported page index structure is [`PageIndex`]. If the metadata -/// contains any other [`PageIndexProvider`] implementation, the [`ColumnIndex`] -/// and [`OffsetIndex`] structures are silently omitted from the output. See -/// for more details. -/// -///
-/// /// # Output Format /// /// The format of the metadata is as follows: @@ -464,21 +470,10 @@ impl<'a, W: Write> ParquetMetaDataWriter<'a, W> { self.write_path_in_schema, ); - // Downcast to PageIndex to access raw index structures for serialization. - // Page indexes from custom PageIndexProviders are not written. See - // if let Some(page_index_arc) = self.metadata.page_index.as_ref() - && let Some(page_index) = page_index_arc - .as_any() - .downcast_ref::() + && (page_index_arc.has_column_indexes() || page_index_arc.has_offset_indexes()) { - if let Some(column_indexes) = page_index.column_indexes_raw() { - encoder = encoder.with_column_indexes(column_indexes.clone()); - } - - if let Some(offset_indexes) = page_index.offset_indexes_raw() { - encoder = encoder.with_offset_indexes(offset_indexes.clone()); - } + encoder = encoder.with_page_index(page_index_arc.clone()); } if let Some(key_value_metadata) = key_value_metadata { diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs index 6c0f3e8dd283..18b7e6e9999e 100644 --- a/parquet/src/file/writer.rs +++ b/parquet/src/file/writer.rs @@ -18,6 +18,7 @@ //! [`SerializedFileWriter`]: Low level Parquet writer API use crate::bloom_filter::Sbbf; +use crate::file::metadata::page_index::PageIndex; use crate::file::metadata::thrift::PageHeader; use crate::file::page_index::column_index::ColumnIndexMetaData; use crate::file::page_index::offset_index::OffsetIndexMetaData; @@ -377,10 +378,27 @@ impl SerializedFileWriter { encoder = encoder.with_key_value_metadata(key_value_metadata) } - encoder = encoder.with_column_indexes(column_indexes); - if !self.props.offset_index_disabled() { - encoder = encoder.with_offset_indexes(offset_indexes); + // check for empty column index + let column_indexes = if column_indexes.is_empty() + || column_indexes + .iter() + .all(|cis| cis.iter().all(|ci| ci.is_none())) + { + None + } else { + Some(column_indexes) + }; + // offset index will always be created unless explicitly disabled + let offset_indexes = if self.props.offset_index_disabled() { + None + } else { + Some(offset_indexes) + }; + if column_indexes.is_some() || offset_indexes.is_some() { + let page_index = PageIndex::new(column_indexes, offset_indexes); + encoder = encoder.with_page_index(Arc::new(page_index)); } + encoder.finish() } @@ -2201,6 +2219,55 @@ mod tests { assert!(b_idx.is_none(), "{b_idx:?}"); } + #[test] + fn test_offset_index_disabled() { + let message_type = " + message test_schema { + REQUIRED INT32 a; + REQUIRED INT32 b; + } + "; + // write file with indexes disabled (including offset indexes) + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::None) + .set_offset_index_disabled(true) + .build(); + let mut file = Vec::with_capacity(1024); + let mut file_writer = + SerializedFileWriter::new(&mut file, schema, Arc::new(props)).unwrap(); + + let mut row_group_writer = file_writer.next_row_group().unwrap(); + let mut a_writer = row_group_writer.next_column().unwrap().unwrap(); + let col_writer = a_writer.typed::(); + col_writer.write_batch(&[1, 2, 3], None, None).unwrap(); + a_writer.close().unwrap(); + + let mut b_writer = row_group_writer.next_column().unwrap().unwrap(); + let col_writer = b_writer.typed::(); + col_writer.write_batch(&[4, 5, 6], None, None).unwrap(); + b_writer.close().unwrap(); + row_group_writer.close().unwrap(); + + let metadata = file_writer.finish().unwrap(); + assert_eq!(metadata.num_row_groups(), 1); + let row_group = metadata.row_group(0); + assert_eq!(row_group.num_columns(), 2); + // no page indexes should exist + assert!(row_group.column(0).offset_index_offset().is_none()); + assert!(row_group.column(0).column_index_offset().is_none()); + assert!(row_group.column(1).offset_index_offset().is_none()); + assert!(row_group.column(1).column_index_offset().is_none()); + + drop(file_writer); + + // read file and request page index...should be `None` + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(Bytes::from(file), options).unwrap(); + + assert!(reader.metadata().page_index().is_none()); + } + #[test] fn test_byte_array_size_statistics() { let message_type = " From b2d3a1a40b749a0fada964a49b52d9395913c73d Mon Sep 17 00:00:00 2001 From: seidl Date: Tue, 8 Sep 2026 15:11:03 -0700 Subject: [PATCH 2/2] clippy --- parquet/src/file/metadata/writer.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/parquet/src/file/metadata/writer.rs b/parquet/src/file/metadata/writer.rs index d6de7e45c8c5..63cdf540d146 100644 --- a/parquet/src/file/metadata/writer.rs +++ b/parquet/src/file/metadata/writer.rs @@ -152,7 +152,7 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { /// Serialize the column indexes and transform to `Option>>>` fn finalize_column_indexes( &mut self, - page_index: &Option>, + page_index: Option<&Arc>, ) -> Result>>>> { if page_index .as_ref() @@ -179,7 +179,7 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { /// Serialize the offset indexes and transform to `Option` fn finalize_offset_indexes( &mut self, - page_index: &Option>, + page_index: Option<&Arc>, ) -> Result>>>> { if page_index .as_ref() @@ -209,8 +209,8 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { // serialize page indexes and transform to the proper form for use in ParquetMetaData let page_index = self.page_index.take(); - let column_indexes = self.finalize_column_indexes(&page_index)?; - let offset_indexes = self.finalize_offset_indexes(&page_index)?; + let column_indexes = self.finalize_column_indexes(page_index.as_ref())?; + let offset_indexes = self.finalize_offset_indexes(page_index.as_ref())?; // We only include ColumnOrder for leaf nodes. let column_orders = self