alamb commented on code in PR #11031:
URL: https://github.com/apache/arrow-rs/pull/11031#discussion_r4097396794
##########
parquet/src/arrow/mod.rs:
##########
@@ -595,6 +595,112 @@ mod test {
assert_eq!(original_metadata, roundtrip_metadata);
}
+ #[test]
+ fn test_metadata_read_write_roundtrip_missing_page_index() {
+ let parquet_bytes = create_parquet_file();
+
+ // read the metadata from the file but skip the page indexes
+ let options =
ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
+ let original_metadata = ParquetMetaDataReader::new()
+ .with_metadata_options(Some(options))
+ .with_page_index_policy(PageIndexPolicy::Skip)
+ .parse_and_finish(&parquet_bytes)
+ .unwrap();
+
+ // metadata_to_bytes_no_page_idx should zero out the page index
locations. if they aren't
+ // then reading metadata_bytes will fail with EOF
+ let metadata_bytes = metadata_to_bytes_no_page_idx(&original_metadata);
+ let options =
ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
+ let roundtrip_metadata = ParquetMetaDataReader::new()
+ .with_metadata_options(Some(options))
+ .with_page_index_policy(PageIndexPolicy::Optional)
+ .parse_and_finish(&metadata_bytes)
+ .expect("page index locations should have been cleared");
+
+ assert!(roundtrip_metadata.page_index().is_none());
+ }
+
+ #[test]
+ fn test_metadata_read_write_roundtrip_offset_index_only() {
+ // `Chunk` statistics: the file has an offset index, but no column
index
+ let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));
+ let batch = RecordBatch::try_from_iter(vec![("id", array)]).unwrap();
+ let props = WriterProperties::builder()
+ .set_statistics_enabled(EnabledStatistics::Chunk)
+ .build();
+ let mut buf = vec![];
+ let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(),
Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let read = |bytes: &Bytes| {
+ let options =
ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
+ ParquetMetaDataReader::new()
+ .with_metadata_options(Some(options))
+ .with_page_index_policy(PageIndexPolicy::Optional)
+ .parse_and_finish(bytes)
+ .unwrap()
+ };
+ let original = read(&Bytes::from(buf));
+ let page_index = original.page_index().unwrap();
+ assert!(!page_index.has_column_indexes() &&
page_index.has_offset_indexes());
+ let roundtrip = read(&metadata_to_bytes(&original));
+ assert_eq!(
+ normalize_locations(original),
+ normalize_locations(roundtrip)
+ );
+ }
+
+ #[test]
+ fn test_metadata_read_write_roundtrip_custom_page_index() {
+ use crate::file::metadata::page_index::PageIndexProvider;
+ use crate::file::page_index::{
+ column_index::ColumnIndexMetaData,
offset_index::OffsetIndexMetaData,
+ };
+
+ /// A custom provider that forwards to another provider
+ #[derive(Debug)]
+ struct Forward(Arc<dyn PageIndexProvider>);
+ impl PageIndexProvider for Forward {
+ fn has_offset_indexes(&self) -> bool {
+ self.0.has_offset_indexes()
+ }
+ fn has_column_indexes(&self) -> bool {
+ self.0.has_column_indexes()
+ }
+ fn column_index(&self, rg: usize, col: usize) ->
Option<&ColumnIndexMetaData> {
+ self.0.column_index(rg, col)
+ }
+ fn offset_index(&self, rg: usize, col: usize) ->
Option<&OffsetIndexMetaData> {
+ self.0.offset_index(rg, col)
+ }
+ fn as_any(&self) -> &dyn std::any::Any {
+ self
+ }
+ }
+
+ let read = |bytes: &Bytes| {
+ let options =
ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
+ ParquetMetaDataReader::new()
+ .with_metadata_options(Some(options))
+ .with_page_index_policy(PageIndexPolicy::Required)
+ .parse_and_finish(bytes)
+ .unwrap()
+ };
+ let original = read(&create_parquet_file());
+ let provider = Forward(original.page_index().unwrap().clone());
+ let custom = original
+ .clone()
+ .into_builder()
+ .set_page_index(Some(Arc::new(provider)))
Review Comment:
nice!
##########
parquet/src/file/metadata/writer.rs:
##########
@@ -54,13 +55,13 @@ 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>>,
Review Comment:
👍
##########
parquet/src/file/metadata/writer.rs:
##########
@@ -443,6 +411,23 @@ impl<'a, W: Write> ParquetMetaDataWriter<'a, W> {
}
}
+ /// Set whether or not to preserve the page index location metadata in the
Thrift
+ /// `ColumnMetaData`.
+ ///
+ /// Because this struct is often used to externalize the footer metadata,
it is
+ /// usually desirable to preserve this location information, even when the
+ /// page indexes are not duplicated (for instance if the provided
`ParquetMetaData`
+ /// returns no `PageIndexProvider`). As such, this defaults to `true`.
+ ///
+ /// Set this to `false` to reset the location metadata if no page indexes
are
Review Comment:
it took me a while to understand this but it makes sense now
The more we deal with PageIndexes the more I feel like it is strange that
BloomFilters are not handled in a unified way (they aren't part of the
ParquetMetadataReader API, etc)
That can be a probelm for another day though
##########
parquet/src/file/metadata/writer.rs:
##########
@@ -270,34 +249,24 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
created_by: Option<String>,
writer_version: i32,
write_path_in_schema: bool,
+ preserve_page_index_locations: bool,
) -> Self {
Self {
buf,
schema_descr,
row_groups,
- column_indexes: None,
- offset_indexes: None,
+ page_index: None,
key_value_metadata: None,
created_by,
object_writer: Default::default(),
writer_version,
write_path_in_schema,
+ preserve_page_index_locations,
}
}
- 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 {
Review Comment:
I always find the `pub` use in this crate confusing -- i think this is
`pub(crate)` effectively -- could we mark it like that explicitly too?
##########
parquet/src/file/metadata/writer.rs:
##########
@@ -443,6 +411,23 @@ impl<'a, W: Write> ParquetMetaDataWriter<'a, W> {
}
}
+ /// Set whether or not to preserve the page index location metadata in the
Thrift
+ /// `ColumnMetaData`.
Review Comment:
```suggestion
/// `ColumnMetaData` (defaults to `true`).
```
--
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]
