Bug: schema metadata is lost after written to parquet

[niebayes]

I'm using the AsyncArrowWriter for writing record batches to a parquet file. Each record batch has an attached metadata in its schema.
Then I use the ParquetRecordBatchStreamBuilder to read the parquet file and set the skip_arrow_schema in ArrowReaderOptions to false. However, the schema in record batches are lost.

I have written a test for reproducing the bug:

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::sync::Arc;

    use arrow::util::data_gen::create_random_batch;
    use arrow_schema::{DataType, Field, Schema};
    use futures::TryStreamExt;
    use object_store::{LocalFileSystem, ObjectStoreRef, Path};
    use parquet::arrow::AsyncArrowWriter;
    use parquet::arrow::ParquetRecordBatchStreamBuilder;
    use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions};
    use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader};
    use parquet::arrow::async_writer::ParquetObjectWriter;
    use tempfile::tempdir;

    #[tokio::test]
    async fn test_schema_metadata_presence() {
        let temp_dir = tempdir().unwrap();
        let store: ObjectStoreRef =
            Arc::new(LocalFileSystem::new_with_prefix(temp_dir.path()).unwrap());

        let file_schema = Arc::new(
            Schema::new(vec![
                Field::new("row_id", DataType::UInt64, false),
                Field::new("value", DataType::Float32, false),
            ])
            .with_metadata(HashMap::from([(
                "my_key".to_string(),
                "my_value".to_string(),
            )])),
        );
        let batch = create_random_batch(file_schema.clone(), 16, 0.0, 0.0).unwrap();

        // Writes data to a parquet file.
        let path = Path::parse("my.parquet").unwrap();
        let object_writer = ParquetObjectWriter::new(store.clone(), path.clone());
        let mut writer = AsyncArrowWriter::try_new(object_writer, file_schema, None).unwrap();
        writer.write(&batch).await.unwrap();
        writer.finish().await.unwrap();

        // Reads data from the parquet file.
        let mut parquet_reader = ParquetObjectReader::new(store, path);
        let parquet_metadata = parquet_reader.get_metadata(None).await.unwrap();

        let reader_metadata = ArrowReaderMetadata::try_new(
            parquet_metadata,
            ArrowReaderOptions::new().with_skip_arrow_metadata(false),
        )
        .unwrap();

        let stream =
            ParquetRecordBatchStreamBuilder::new_with_metadata(parquet_reader, reader_metadata)
                .build()
                .unwrap();
        let batches = stream.try_collect::<Vec<_>>().await.unwrap();

        for batch in batches {
            assert_eq!(
                batch.schema_ref().metadata().get("my_key").unwrap(),
                "my_value"
            );
        }
    }
}

Arrow version: 56.2.0

[etseidl]

I'll admit to not being familiar with this part of the code, but your metadata survives into the reader_metadata. ParquetRecordBatchStreamBuilder seems to deliberately ignore it, however.

arrow-rs/parquet/src/arrow/async_reader/mod.rs

Lines 517 to 526 in f9dd799

	// Ensure schema of ParquetRecordBatchStream respects projection, and does
	// not store metadata (same as for ParquetRecordBatchReader and emitted RecordBatches)
	let projected_fields = match reader_factory.fields.as_deref().map(|pf| &pf.arrow_type) {
	Some(DataType::Struct(fields)) => {
	fields.filter_leaves(|idx, _| self.projection.leaf_included(idx))
	}
	None => Fields::empty(),
	_ => unreachable!("Must be Struct for root type"),
	};
	let schema = Arc::new(Schema::new(projected_fields));

[niebayes]

@etseidl Thank you. I have checked that the arrow metadata does survive in the ArrowReaderMetadata. However, per-batch arrow metadata still gets lost. Seems they're deliberately ignored, as you said.

[etseidl]

@niebayes I found the discussion in #5135 helpful...this behavior is as designed. Are you able to save the schema metadata and pass it to where it is needed?

[niebayes]

@etseidl I will try.