This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 5faefa9bf0 fix(parquet): keep virtual columns in the schema reported
with a schema hint (#11047)
5faefa9bf0 is described below
commit 5faefa9bf022471856b7a3e0413a743e53980d8c
Author: Bharadwaj Pendyala <[email protected]>
AuthorDate: Mon Sep 21 04:40:02 2026 -0500
fix(parquet): keep virtual columns in the schema reported with a schema
hint (#11047)
# Which issue does this PR close?
- Closes #11046.
# Rationale for this change
`ArrowReaderMetadata::with_supplied_schema` already treats virtual
columns as extra fields on top of the hint. It passes them to
`parquet_to_arrow_field_levels_with_virtual`, and its own length check
reads `supplied_schema.fields().len() + virtual_columns.len()`. Then it
returns `schema: supplied_schema`, which doesn't have them.
So the two branches of `try_new` disagree. Without a hint you get a
schema with the virtual fields in it; add a hint and they vanish, even
though the reader still decodes them. On the file from the issue that
means `metadata.schema()` reports one field while every batch that
metadata produces has two:
```
metadata.schema().fields() [value]
batch.schema().fields() [value, row_number]
```
@limenilbuz asked for either of two behaviours: include the virtual
columns in the reported schema, or stop erroring when the hint itself
contains them. This does the first. The second is a change to what
`with_schema` accepts, and the length check here already assumes virtual
fields live outside the hint, so the first is the one that makes the
function agree with itself.
# What changes are included in this PR?
The returned schema is now the supplied fields followed by the virtual
fields, keeping the supplied schema's key/value metadata. When no
virtual columns are requested the supplied schema is returned untouched,
so nothing changes for that path.
`parquet_to_arrow_field_levels_with_virtual` appends virtual columns to
the root in the order given and clones them unchanged
(`parquet/src/arrow/schema/mod.rs:223`), so appending them here in the
same order lines the reported schema up with `field_levels`.
# Are these changes tested?
Yes. `test_supplied_schema_keeps_virtual_columns` builds metadata from a
hint plus two virtual fields and checks the field order, that the hint's
schema metadata survives, and that `metadata.schema()` agrees with the
fields of the batch the reader emits. It fails on `f9e02ba` with:
```
left: [Field { name: "value", data_type: Int64 }]
right: [Field { name: "value", .. }, Field { name: "row_number", .. },
Field { name: "row_group_index", .. }]
```
`cargo test -p parquet --lib` is 1381 passed, 0 failed. `cargo fmt --all
-- --check` and `cargo clippy -p parquet --all-targets` are both clean.
# Are there any user-facing changes?
`ArrowReaderMetadata::schema()`, and the builder schema derived from it,
gain the virtual fields when a hint and virtual columns are combined.
That's the fix, but it is a field-count change on a public accessor, so
it's worth calling out. Physical column indices are unaffected and no
crate in the tree combines those two options.
One thing I left alone: with an explicit projection the async reader's
`schema()` drops virtual fields while the batches still carry them. That
reproduces with and without a schema hint, so it's a separate bug from
this one and I didn't touch it here.
---
parquet/src/arrow/arrow_reader/mod.rs | 13 +++++-
.../arrow/arrow_reader/tests/virtual_columns.rs | 46 ++++++++++++++++++++++
2 files changed, 58 insertions(+), 1 deletion(-)
diff --git a/parquet/src/arrow/arrow_reader/mod.rs
b/parquet/src/arrow/arrow_reader/mod.rs
index 5736ff7c01..73bcb2048a 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 9dbdcae6f9..c46a1b89a3 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