Skip to content

Allow for serialization of Parquet Page Index via custom PageIndexProvider #11031

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
etseidl wants to merge 2 commits into
base: main
Choose a base branch
from
+132 −83
Open

Allow for serialization of Parquet Page Index via custom PageIndexProvider #11031

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1c77984
allow for serialization of custom page index providers
etseidl Sep 8, 2026
b2d3a1a
clippy
etseidl Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 4 additions & 17 deletions parquet/src/file/metadata/page_index.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Vec<Option<ColumnIndexMetaData>>>> {
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<Vec<Option<OffsetIndexMetaData>>>> {
self.offset_indexes.as_ref()
}
}

impl PageIndexProvider for PageIndex {
Expand Down Expand Up @@ -536,7 +520,10 @@ impl PageIndexBuilder {
///
/// # 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>>>> {
pub(crate) fn empty_index<T>(
num_row_groups: usize,
num_columns: usize,
) -> Option<Vec<Vec<Option<T>>>> {
Some(
(0..num_row_groups)
.map(|_| {
Expand Down
121 changes: 58 additions & 63 deletions parquet/src/file/metadata/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -54,8 +55,7 @@ pub(crate) struct ThriftMetadataWriter<'a, W: Write> {
buf: &'a mut TrackedWrite<W>,
schema_descr: &'a SchemaDescPtr,
row_groups: Vec<RowGroupMetaData>,
column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>,
offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>,
page_index: Option<Arc<dyn PageIndexProvider>>,
key_value_metadata: Option<Vec<KeyValue>>,
created_by: Option<String>,
object_writer: MetadataObjectWriter,
Expand All @@ -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<Option<OffsetIndexMetaData>>],
) -> Result<()> {
page_index: &Arc<dyn PageIndexProvider>,
) -> Result<Option<Vec<Vec<Option<OffsetIndexMetaData>>>>> {
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,
Expand All @@ -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`
Expand All @@ -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<Option<ColumnIndexMetaData>>],
) -> Result<()> {
page_index: &Arc<dyn PageIndexProvider>,
) -> Result<Option<Vec<Vec<Option<ColumnIndexMetaData>>>>> {
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(
Expand All @@ -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<ParquetColumnIndex>`
fn finalize_column_indexes(&mut self) -> Result<Option<Vec<Vec<Option<ColumnIndexMetaData>>>>> {
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<Vec<Vec<Option<ColumnIndexMetaData>>>>`
fn finalize_column_indexes(
&mut self,
page_index: Option<&Arc<dyn PageIndexProvider>>,
) -> Result<Option<Vec<Vec<Option<ColumnIndexMetaData>>>>> {
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()
Expand All @@ -155,13 +177,19 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
}

/// Serialize the offset indexes and transform to `Option<ParquetOffsetIndex>`
fn finalize_offset_indexes(&mut self) -> Result<Option<Vec<Vec<Option<OffsetIndexMetaData>>>>> {
let offset_indexes = std::mem::take(&mut self.offset_indexes);
fn finalize_offset_indexes(
&mut self,
page_index: Option<&Arc<dyn PageIndexProvider>>,
) -> Result<Option<Vec<Vec<Option<OffsetIndexMetaData>>>>> {
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
Expand All @@ -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.as_ref())?;
let offset_indexes = self.finalize_offset_indexes(page_index.as_ref())?;

// We only include ColumnOrder for leaf nodes.
let column_orders = self
Expand Down Expand Up @@ -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(),
Expand All @@ -285,19 +313,8 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
}
}

pub fn with_column_indexes(
mut self,
column_indexes: Vec<Vec<Option<ColumnIndexMetaData>>>,
) -> Self {
self.column_indexes = Some(column_indexes);
self
}

pub fn with_offset_indexes(
mut self,
offset_indexes: Vec<Vec<Option<OffsetIndexMetaData>>>,
) -> Self {
self.offset_indexes = Some(offset_indexes);
pub fn with_page_index(mut self, page_index: Arc<dyn PageIndexProvider>) -> Self {
self.page_index = Some(page_index);
self
}

Expand Down Expand Up @@ -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.
///
/// <div class="warning">
///
/// **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
/// <https://github.com/apache/arrow-rs/issues/11030> for more details.
///
/// </div>
///
/// # Output Format
///
/// The format of the metadata is as follows:
Expand Down Expand Up @@ -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
// <https://github.com/apache/arrow-rs/issues/11030>
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>()
&& (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 {
Expand Down
73 changes: 70 additions & 3 deletions parquet/src/file/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -377,10 +378,27 @@ impl<W: Write + Send> SerializedFileWriter<W> {
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()
}

Expand Down Expand Up @@ -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::<Int32Type>();
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::<Int32Type>();
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 = "
Expand Down
Loading