diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 5736ff7c0193..73bcb2048ad5 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -1127,9 +1127,20 @@ impl ArrowReaderMetadata { ))); } + // `fields` is the supplied fields followed by the virtual columns, so the reported + // schema has to be built from it or the virtual columns go missing. + let schema = if virtual_columns.is_empty() { + supplied_schema + } else { + Arc::new(Schema::new_with_metadata( + fields, + supplied_schema.metadata().clone(), + )) + }; + Ok(Self { metadata, - schema: supplied_schema, + schema, fields: field_levels.levels.map(Arc::new), }) } diff --git a/parquet/src/arrow/arrow_reader/tests/virtual_columns.rs b/parquet/src/arrow/arrow_reader/tests/virtual_columns.rs index 9dbdcae6f926..c46a1b89a3e0 100644 --- a/parquet/src/arrow/arrow_reader/tests/virtual_columns.rs +++ b/parquet/src/arrow/arrow_reader/tests/virtual_columns.rs @@ -18,6 +18,7 @@ //! Generated row numbers and row-group indices, including ordering and filtering. use super::*; +use std::collections::HashMap; #[test] fn test_read_row_numbers() { @@ -68,6 +69,51 @@ fn test_read_row_numbers() { ); } +#[test] +fn test_supplied_schema_keeps_virtual_columns() { + let file = write_parquet_from_iter(vec![( + "value", + Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef, + )]); + let supplied_fields = Fields::from(vec![Field::new("value", ArrowDataType::Int64, false)]); + let row_number_field = Arc::new( + Field::new("row_number", ArrowDataType::Int64, false).with_extension_type(RowNumber), + ); + let row_group_index_field = Arc::new( + Field::new("row_group_index", ArrowDataType::Int64, false) + .with_extension_type(RowGroupIndex), + ); + let supplied_metadata = HashMap::from([("k".to_string(), "v".to_string())]); + + let options = ArrowReaderOptions::new() + .with_schema(Arc::new(Schema::new_with_metadata( + supplied_fields, + supplied_metadata.clone(), + ))) + .with_virtual_columns(vec![ + row_number_field.clone(), + row_group_index_field.clone(), + ]) + .unwrap(); + let metadata = ArrowReaderMetadata::load(&file, options).unwrap(); + + let expected = Fields::from(vec![ + Arc::new(Field::new("value", ArrowDataType::Int64, false)), + row_number_field, + row_group_index_field, + ]); + assert_eq!(metadata.schema().fields(), &expected); + assert_eq!(metadata.schema().metadata(), &supplied_metadata); + + let batch = ParquetRecordBatchReaderBuilder::new_with_metadata(file, metadata.clone()) + .build() + .unwrap() + .next() + .unwrap() + .unwrap(); + assert_eq!(batch.schema().fields(), metadata.schema().fields()); +} + #[test] fn test_read_only_row_numbers() { let file = write_parquet_from_iter(vec![(