From d604008950a00e0c4797621640acd3b546be4039 Mon Sep 17 00:00:00 2001 From: Bharadwaj Pendyala Date: Thu, 10 Sep 2026 08:47:35 -0500 Subject: [PATCH 1/3] fix(parquet): keep virtual columns in the schema reported with a schema hint with_supplied_schema passes virtual_columns to parquet_to_arrow_field_levels_with_virtual and counts them in its own length check, but returned the bare supplied schema, so ArrowReaderMetadata::schema() dropped them. The no-hint branch of try_new keeps them, and the reader decodes them either way. Append the virtual fields to the supplied fields, preserving the supplied schema's metadata. Closes #11046 --- parquet/src/arrow/arrow_reader/mod.rs | 19 +++++++- .../arrow_reader/tests/virtual_columns.rs | 45 +++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 5736ff7c0193..8c0e20849401 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -1127,9 +1127,26 @@ impl ArrowReaderMetadata { ))); } + // Virtual columns are appended to `field_levels` and counted in the length check + // above, so the reported schema has to carry them too. + let schema = if virtual_columns.is_empty() { + supplied_schema + } else { + let fields = supplied_schema + .fields() + .iter() + .cloned() + .chain(virtual_columns.iter().cloned()) + .collect::>(); + 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..6987e2d0e07d 100644 --- a/parquet/src/arrow/arrow_reader/tests/virtual_columns.rs +++ b/parquet/src/arrow/arrow_reader/tests/virtual_columns.rs @@ -68,6 +68,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![( From 2b5300a12e935ff04388d0f94ab933da51bba1cb Mon Sep 17 00:00:00 2001 From: Bharadwaj Pendyala Date: Fri, 18 Sep 2026 21:54:22 -0500 Subject: [PATCH 2/3] refactor(parquet): build the hinted schema from field_levels --- parquet/src/arrow/arrow_reader/mod.rs | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 8c0e20849401..73bcb2048ad5 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -1127,17 +1127,11 @@ impl ArrowReaderMetadata { ))); } - // Virtual columns are appended to `field_levels` and counted in the length check - // above, so the reported schema has to carry them too. + // `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 { - let fields = supplied_schema - .fields() - .iter() - .cloned() - .chain(virtual_columns.iter().cloned()) - .collect::>(); Arc::new(Schema::new_with_metadata( fields, supplied_schema.metadata().clone(), From 7e315dec5d0f7a09fcfc8b5d4efd2c828f328e59 Mon Sep 17 00:00:00 2001 From: Bharadwaj Pendyala Date: Sat, 19 Sep 2026 13:19:01 -0500 Subject: [PATCH 3/3] test(parquet): import HashMap in the virtual column tests tests/mod.rs no longer brings HashMap into scope after #11103, so the supplied-metadata test failed to compile against main. --- parquet/src/arrow/arrow_reader/tests/virtual_columns.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/parquet/src/arrow/arrow_reader/tests/virtual_columns.rs b/parquet/src/arrow/arrow_reader/tests/virtual_columns.rs index 6987e2d0e07d..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() {