diff --git a/crates/iceberg/src/arrow/reader/projection.rs b/crates/iceberg/src/arrow/reader/projection.rs index 2fa6bf18f9..b76d60c856 100644 --- a/crates/iceberg/src/arrow/reader/projection.rs +++ b/crates/iceberg/src/arrow/reader/projection.rs @@ -23,7 +23,9 @@ use std::collections::{HashMap, HashSet}; use std::str::FromStr; use std::sync::Arc; -use arrow_schema::{Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef}; +use arrow_schema::{ + DataType, Field, FieldRef, Fields, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, +}; use parquet::arrow::{PARQUET_FIELD_ID_META_KEY, ProjectionMask}; use parquet::schema::types::{SchemaDescriptor, Type as ParquetType}; @@ -32,7 +34,10 @@ use crate::arrow::arrow_schema_to_schema; use crate::error::{Result, invalid_data}; use crate::expr::BoundPredicate; use crate::expr::visitors::bound_predicate_visitor::visit; -use crate::spec::{NameMapping, NestedField, PrimitiveType, Schema, Type}; +use crate::spec::{ + LIST_FIELD_NAME, MAP_KEY_FIELD_NAME, MAP_VALUE_FIELD_NAME, MappedField, NameMapping, + NestedField, PrimitiveType, Schema, Type, +}; use crate::{Error, ErrorKind}; impl ArrowReader { @@ -171,7 +176,7 @@ impl ArrowReader { arrow_schema: &ArrowSchemaRef, type_promotion_is_valid: fn(Option<&PrimitiveType>, Option<&PrimitiveType>) -> bool, ) -> Result { - let mut column_map = HashMap::new(); + let mut column_map = HashMap::with_capacity(leaf_field_ids.len()); let fields = arrow_schema.fields(); // HashSet for O(1) membership checks instead of O(n) slice scans. let leaf_field_id_set: HashSet = leaf_field_ids.iter().copied().collect(); @@ -216,7 +221,7 @@ impl ArrowReader { // We only project existing columns; RecordBatchTransformer adds default/NULL values. // If other leaves are selected, a struct whose requested leaves are all missing gets no // column and reads as NULL (#3378). - let mut indices = vec![]; + let mut indices = Vec::with_capacity(leaf_field_ids.len()); for field_id in leaf_field_ids { if let Some(col_idx) = column_map.get(field_id) { indices.push(*col_idx); @@ -392,7 +397,7 @@ fn parquet_field_id(field: &Field) -> Option { /// Apply name mapping to Arrow schema for Parquet files lacking field IDs. /// -/// Assigns Iceberg field IDs based on column names using the name mapping, +/// Assigns Iceberg field IDs based on names at every depth using the name mapping, /// enabling correct projection on migrated files (e.g., from Hive/Spark via add_files). /// /// Per Iceberg spec Column Projection rule #2: @@ -421,43 +426,20 @@ pub(super) fn apply_name_mapping_to_arrow_schema( "Schema already has field IDs - name mapping should not be applied" ); - let mut field_ids_by_name = HashMap::with_capacity(name_mapping.fields().len()); + let mut mapped_fields_by_name = HashMap::with_capacity(name_mapping.fields().len()); for mapped_field in name_mapping.fields() { for name in mapped_field.names() { // Preserve the existing first-match behavior for duplicate aliases. - field_ids_by_name + mapped_fields_by_name .entry(name.as_str()) - .or_insert(mapped_field.field_id()); + .or_insert(mapped_field); } } - let fields_with_mapped_ids: Vec<_> = arrow_schema - .fields() - .iter() - .map(|field| { - // Look up this column name in name mapping to get the Iceberg field ID. - // Corresponds to Java's ApplyNameMapping visitor which calls - // nameMapping.find(currentPath()) and returns field.withId() if found. - // - // If the field isn't in the mapping, leave it WITHOUT assigning an ID - // (matching Java's behavior of returning the field unchanged). - // Later, during projection, fields without IDs are filtered out. - let mut metadata = field.metadata().clone(); - - if let Some(field_id) = field_ids_by_name - .get(field.name().as_str()) - .copied() - .flatten() - { - // Field found in mapping with a field_id → assign it - metadata.insert(PARQUET_FIELD_ID_META_KEY.to_string(), field_id.to_string()); - } - // If field_id is None, leave the field without an ID (will be filtered by projection) - - Field::new(field.name(), field.data_type().clone(), field.is_nullable()) - .with_metadata(metadata) - }) - .collect(); + let fields_with_mapped_ids = try_map_fields(arrow_schema.fields(), |field| { + let mapped_field = mapped_fields_by_name.get(field.name().as_str()).copied(); + apply_name_mapping_to_field(field, mapped_field) + })?; Ok(Arc::new(ArrowSchema::new_with_metadata( fields_with_mapped_ids, @@ -465,6 +447,92 @@ pub(super) fn apply_name_mapping_to_arrow_schema( ))) } +/// Assigns IDs from `mapped_field` to `field` and its nested fields, as Java's ApplyNameMapping +/// does. List elements match `element`, and map entries match `key` and `value`. A field +/// without an entry, or whose entry has no field ID, keeps no ID, and so do its nested fields. +fn apply_name_mapping_to_field( + field: &FieldRef, + mapped_field: Option<&MappedField>, +) -> Result { + let Some((mapped_field, field_id)) = + mapped_field.zip(mapped_field.and_then(MappedField::field_id)) + else { + return Ok(Arc::clone(field)); + }; + let map_nested = |nested: &FieldRef, name: &str| { + let mapped_nested = find_mapped_field(mapped_field.fields().iter().map(Arc::as_ref), name); + apply_name_mapping_to_field(nested, mapped_nested) + }; + + let data_type = match field.data_type() { + DataType::Struct(fields) => DataType::Struct(try_map_fields(fields, |nested| { + map_nested(nested, nested.name()) + })?), + DataType::List(element) => DataType::List(map_nested(element, LIST_FIELD_NAME)?), + DataType::LargeList(element) => DataType::LargeList(map_nested(element, LIST_FIELD_NAME)?), + DataType::FixedSizeList(element, size) => { + DataType::FixedSizeList(map_nested(element, LIST_FIELD_NAME)?, *size) + } + // The mapping has no entry for Arrow's entries struct. The shape checks are defensive, + // because parquet-rs builds every map as a struct of key and value. + DataType::Map(entries, sorted) => { + let DataType::Struct(entry_fields) = entries.data_type() else { + return Err(invalid_data!( + "Map field `{}` must have struct type", + field.name() + )); + }; + let [key, value] = &entry_fields[..] else { + return Err(invalid_data!( + "Map field `{}` must have exactly 2 fields, found {}", + field.name(), + entry_fields.len() + )); + }; + let entry_fields = Fields::from([ + map_nested(key, MAP_KEY_FIELD_NAME)?, + map_nested(value, MAP_VALUE_FIELD_NAME)?, + ]); + let entries = entries + .as_ref() + .clone() + .with_data_type(DataType::Struct(entry_fields)); + DataType::Map(Arc::new(entries), *sorted) + } + // `Fields::filter_leaves` treats `ListView` and `LargeListView` as leaves, so projection + // cannot select their elements even with IDs. + other => other.clone(), + }; + + let mut field = field.as_ref().clone().with_data_type(data_type); + field + .metadata_mut() + .insert(PARQUET_FIELD_ID_META_KEY.to_string(), field_id.to_string()); + Ok(Arc::new(field)) +} + +/// Returns the first entry in `mapped_fields` that lists `name`. +fn find_mapped_field<'a>( + mapped_fields: impl IntoIterator, + name: &str, +) -> Option<&'a MappedField> { + mapped_fields + .into_iter() + .find(|mapped_field| mapped_field.names().iter().any(|n| n == name)) +} + +/// Maps `fields` into a `Fields` sized up front, which `collect::>()` does not do. +fn try_map_fields( + fields: &Fields, + mut map_field: impl FnMut(&FieldRef) -> Result, +) -> Result { + let mut mapped = Vec::with_capacity(fields.len()); + for field in fields { + mapped.push(map_field(field)?); + } + Ok(Fields::from(mapped)) +} + /// Add position-based fallback field IDs to Arrow schema for Parquet files lacking them. /// Enables projection on migrated files (e.g., from Hive/Spark). /// @@ -1106,6 +1174,225 @@ message schema { assert_eq!(subdept_array.value(0), "Bob"); } + /// Regression test for #3374: a struct in a name-mapped file reads correctly next to a + /// top-level column, and predicates resolve its fields. + #[tokio::test] + async fn test_read_parquet_with_name_mapping_maps_nested_fields() { + use arrow_array::types::Int32Type; + use arrow_array::{Int32Array, StructArray}; + + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::optional(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional( + 2, + "s", + Type::Struct(StructType::new(vec![ + NestedField::optional(3, "a", Type::Primitive(PrimitiveType::Int)) + .into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let name_mapping = Arc::new(NameMapping::new(vec![ + MappedField::new(Some(1), vec!["id".to_string()], vec![]), + MappedField::new(Some(2), vec!["s".to_string()], vec![MappedField::new( + Some(3), + vec!["a".to_string()], + vec![], + )]), + ])); + + // No field IDs at any level. + let s_col = StructArray::from(vec![( + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, + )]); + let to_write = RecordBatch::try_from_iter_with_nullable(vec![ + ( + "id", + Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef, + true, + ), + ("s", Arc::new(s_col) as ArrayRef, true), + ]) + .unwrap(); + let tmp_dir = TempDir::new().unwrap(); + let file_path = format!("{}/1.parquet", tmp_dir.path().to_str().unwrap()); + let file = File::create(&file_path).unwrap(); + let mut writer = ArrowWriter::try_new(file, to_write.schema(), None).unwrap(); + writer.write(&to_write).unwrap(); + writer.close().unwrap(); + + let file_size = std::fs::metadata(&file_path).unwrap().len(); + let tasks = |predicate| { + Box::pin(futures::stream::iter([Ok(FileScanTask::builder() + .with_file_size_in_bytes(file_size) + .with_start(0) + .with_length(0) + .with_data_file_path(file_path.clone()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema.clone()) + .with_project_field_ids(vec![1, 2]) + .with_case_sensitive(false) + .with_name_mapping(Some(name_mapping.clone())) + .with_predicate(predicate) + .build() + .unwrap())])) as FileScanTaskStream + }; + let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build(); + + let result = reader + .clone() + .read(tasks(None)) + .unwrap() + .stream() + .try_collect::>() + .await + .unwrap(); + let s = result[0].column(1).as_struct(); + assert_eq!(s.null_count(), 0); + assert_eq!( + s.column(0).as_primitive::(), + &Int32Array::from(vec![10, 20, 30]) + ); + + // `s.a` resolves, so the row filter rejects the nested predicate as it does for files + // with embedded IDs, rather than matching no rows. + let predicate = Reference::new("s.a") + .greater_than(Datum::int(15)) + .bind(schema.clone(), true) + .unwrap(); + let err = reader + .read(tasks(Some(predicate))) + .unwrap() + .stream() + .try_collect::>() + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } + + /// Every nested Arrow type gets IDs, with list and map entries matched by their spec names. + #[test] + fn test_apply_name_mapping_to_nested_arrow_types() { + use arrow_schema::Fields; + + use super::apply_name_mapping_to_arrow_schema; + use crate::arrow::arrow_schema_to_schema; + + let field = |name: &str, data_type| Field::new(name, data_type, true); + let item = |data_type| Arc::new(field("item", data_type)); + let entries = Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + field("values", DataType::Int32), + ])), + false, + ); + let x = Fields::from(vec![field("x", DataType::Int32)]); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + field("id", DataType::Int32), + // `s.id` must take its ID from the entry for `s`, not from the top-level `id`. + field( + "s", + DataType::Struct(Fields::from(vec![field("id", DataType::Int32)])), + ), + field("list", DataType::List(item(DataType::Int32))), + field("fixed", DataType::FixedSizeList(item(DataType::Int32), 2)), + field("large", DataType::LargeList(item(DataType::Struct(x)))), + field("map", DataType::Map(Arc::new(entries), true)), + ])); + let entry = + |id, name: &str, fields| MappedField::new(Some(id), vec![name.to_string()], fields); + let name_mapping = NameMapping::new(vec![ + entry(1, "id", vec![]), + entry(2, "s", vec![entry(3, "id", vec![])]), + entry(4, "list", vec![entry(5, "element", vec![])]), + entry(6, "fixed", vec![entry(7, "element", vec![])]), + entry(8, "large", vec![entry(9, "element", vec![entry( + 10, + "x", + vec![], + )])]), + entry(11, "map", vec![ + entry(12, "key", vec![]), + entry(13, "value", vec![]), + ]), + ]); + + let mapped = apply_name_mapping_to_arrow_schema(arrow_schema, &name_mapping).unwrap(); + + // The list size and the map ordering carry over. + assert!(matches!( + mapped.field(3).data_type(), + DataType::FixedSizeList(_, 2) + )); + assert!(matches!( + mapped.field(5).data_type(), + DataType::Map(_, true) + )); + // The conversion fails on any field without an ID. + let schema = arrow_schema_to_schema(&mapped).unwrap(); + let names: Vec<_> = (1..=13) + .map(|id| schema.name_by_field_id(id).unwrap_or_default()) + .collect(); + assert_eq!(names, [ + "id", + "s", + "s.id", + "list", + "list.element", + "fixed", + "fixed.element", + "large", + "large.element", + "large.element.x", + "map", + "map.key", + "map.value", + ]); + } + + /// A field without an entry, or whose entry has no field ID, keeps no ID, and so do its + /// nested fields. + #[test] + fn test_apply_name_mapping_skips_unmapped_fields() { + use arrow_schema::Fields; + + use super::{apply_name_mapping_to_arrow_schema, build_field_id_map_from_arrow_schema}; + + let int = |name: &str| Field::new(name, DataType::Int32, true); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "s", + DataType::Struct(Fields::from(vec![int("a"), int("b")])), + true, + ), + Field::new("t", DataType::Struct(Fields::from(vec![int("c")])), true), + ])); + let entry = |id, name: &str, fields| MappedField::new(id, vec![name.to_string()], fields); + let name_mapping = NameMapping::new(vec![ + // No entry for `s.b`. + entry(Some(1), "s", vec![entry(Some(2), "a", vec![])]), + entry(None, "t", vec![entry(Some(3), "c", vec![])]), + ]); + + let mapped = apply_name_mapping_to_arrow_schema(arrow_schema, &name_mapping).unwrap(); + + // Only `s.a`, leaf 0, resolves to an ID. + assert_eq!( + build_field_id_map_from_arrow_schema(&mapped), + HashMap::from([(2, 0)]) + ); + } + /// Test reading Parquet files without field IDs with partial projection. /// Only a subset of columns are requested, verifying position-based fallback /// handles column selection correctly.