This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] 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 1514ea4579 fix(parquet): support non-primitive map keys in record
reader (#11229)
1514ea4579 is described below
commit 1514ea4579340f680b460739d5a96fab38b1314c
Author: hsdfat <[email protected]>
AuthorDate: Fri Oct 2 05:01:54 2026 +0000
fix(parquet): support non-primitive map keys in record reader (#11229)
# Which issue does this PR close?
- Closes #7768.
# Rationale for this change
The record API (`RowIter` / `TreeBuilder`) asserted that map keys are
primitive, so reading a file with e.g. struct- or list-keyed maps
panicked. The Parquet spec has no such restriction and the arrow reader
already handles these files.
This picks up the fix from #7769 by @njaremko, which went stale waiting
on a test file in apache/parquet-testing#87. Thanks @njaremko for the
original fix.
# What changes are included in this PR?
- Remove the primitive-key assertion in `TreeBuilder::reader_tree`. The
key reader was already built via `reader_tree` in the same way as the
value reader, so no other reader changes are needed.
- `Map` stores entries as `Vec<(Field, Field)>`, so complex keys need no
`Hash`/`Eq`; `Display` and `to_json_value` already handle non-primitive
keys.
# Are these changes tested?
Yes. `test_compound_map_key` writes a file in memory with `ArrowWriter`
containing a map with struct keys and a map with list keys (including an
empty map and an empty-list key), then reads it back with the record API
and checks the rows. Generating the file in the test avoids depending on
apache/parquet-testing#87 and addresses the review feedback on #7769
(descriptive test name, only the relevant columns). The test panics
without the fix.
# Are there any user-facing changes?
Files with non-primitive map keys can now be read with the record API
instead of panicking. No API changes.
AI usage: this change was prepared with the help of AI assistants
(Claude, OpenAI models). I reviewed the change and the reader's
key/value level handling, and verified the test fails before and passes
after the fix.
Co-authored-by: Ed Seidl <[email protected]>
Co-authored-by: Jeffrey Vo <[email protected]>
---
parquet/src/record/reader.rs | 113 +++++++++++++++++++++++++++++++++++++++++--
1 file changed, 109 insertions(+), 4 deletions(-)
diff --git a/parquet/src/record/reader.rs b/parquet/src/record/reader.rs
index 2e2ddd4db5..f5886e05d1 100644
--- a/parquet/src/record/reader.rs
+++ b/parquet/src/record/reader.rs
@@ -230,10 +230,6 @@ impl TreeBuilder {
path.push(String::from(key_value_type.name()));
let key_type = &key_value_type.get_fields()[0];
- assert!(
- key_type.is_primitive(),
- "Map key type is expected to be a primitive type, but
found {key_type:?}"
- );
let key_reader = self.reader_tree(
key_type.clone(),
path,
@@ -1890,6 +1886,115 @@ mod tests {
}
}
+ #[test]
+ #[cfg(feature = "arrow")]
+ fn test_compound_map_key() {
+ use crate::arrow::ArrowWriter;
+ use arrow_array::builder::{Int32Builder, ListBuilder};
+ use arrow_array::{
+ Array, ArrayRef, Int32Array, MapArray, RecordBatch, StringArray,
StructArray,
+ };
+ use arrow_buffer::OffsetBuffer;
+ use arrow_schema::{DataType, Field as ArrowField, Fields};
+
+ let map_array = |keys: ArrayRef, offsets: Vec<i32>, values: Vec<i32>| {
+ let entries = StructArray::from(vec![
+ (
+ Arc::new(ArrowField::new("key", keys.data_type().clone(),
false)),
+ keys,
+ ),
+ (
+ Arc::new(ArrowField::new("value", DataType::Int32, true)),
+ Arc::new(Int32Array::from(values)) as ArrayRef,
+ ),
+ ]);
+ let entries_field = ArrowField::new("key_value",
entries.data_type().clone(), false);
+ let map = MapArray::try_new(
+ Arc::new(entries_field),
+ OffsetBuffer::new(offsets.into()),
+ entries,
+ None,
+ false,
+ )
+ .unwrap();
+ Arc::new(map) as ArrayRef
+ };
+
+ // struct keys: [{(1, "a") -> 10, (2, "b") -> 20}, {}, {(3, "c") ->
30}]
+ let struct_keys = StructArray::new(
+ Fields::from(vec![
+ ArrowField::new("id", DataType::Int32, false),
+ ArrowField::new("name", DataType::Utf8, false),
+ ]),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2, 3])),
+ Arc::new(StringArray::from(vec!["a", "b", "c"])),
+ ],
+ None,
+ );
+ let struct_map = map_array(Arc::new(struct_keys), vec![0, 2, 2, 3],
vec![10, 20, 30]);
+
+ // list keys: [{[1, 2] -> 10, [3] -> 20}, {}, {[] -> 30}]
+ let mut list_keys = ListBuilder::new(Int32Builder::new());
+ list_keys.append_value([Some(1), Some(2)]);
+ list_keys.append_value([Some(3)]);
+ list_keys.append_value([]);
+ let list_map = map_array(
+ Arc::new(list_keys.finish()),
+ vec![0, 2, 2, 3],
+ vec![10, 20, 30],
+ );
+
+ let batch =
+ RecordBatch::try_from_iter([("struct_map", struct_map),
("list_map", list_map)])
+ .unwrap();
+ let mut buffer = Vec::new();
+ let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(),
None).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
+ let rows: Vec<Row> = reader
+ .get_row_iter(None)
+ .unwrap()
+ .map(|row| row.unwrap())
+ .collect();
+
+ let key = |id: i32, name: &str| {
+ group![
+ ("id".to_string(), Field::Int(id)),
+ ("name".to_string(), Field::Str(name.to_string()))
+ ]
+ };
+ let expected_rows = vec![
+ row![
+ (
+ "struct_map".to_string(),
+ map![(key(1, "a"), Field::Int(10)), (key(2, "b"),
Field::Int(20))]
+ ),
+ (
+ "list_map".to_string(),
+ map![
+ (list![Field::Int(1), Field::Int(2)], Field::Int(10)),
+ (list![Field::Int(3)], Field::Int(20))
+ ]
+ ),
+ ],
+ row![
+ ("struct_map".to_string(), map![]),
+ ("list_map".to_string(), map![]),
+ ],
+ row![
+ (
+ "struct_map".to_string(),
+ map![(key(3, "c"), Field::Int(30))]
+ ),
+ ("list_map".to_string(), map![(list![], Field::Int(30))]),
+ ],
+ ];
+ assert_eq!(rows, expected_rows);
+ }
+
fn test_file_reader_rows(file_name: &str, schema: Option<Type>) ->
Result<Vec<Row>> {
let file = get_test_file(file_name);
let file_reader: Box<dyn FileReader> =
Box::new(SerializedFileReader::new(file)?);
