From 11c8ef16c7dd0c63d7d8c101ee7fd92492d2f21a Mon Sep 17 00:00:00 2001 From: Manu Zhang Date: Wed, 26 Aug 2026 12:38:57 +0800 Subject: [PATCH 1/3] feat(spec): add unknown primitive type support Co-authored-by: Codex --- crates/catalog/glue/src/schema.rs | 6 + crates/catalog/hms/src/schema.rs | 6 + crates/iceberg/public-api.txt | 1 + crates/iceberg/src/arrow/schema.rs | 11 + crates/iceberg/src/arrow/value.rs | 25 +- crates/iceberg/src/avro/schema.rs | 61 +++-- crates/iceberg/src/spec/datatypes.rs | 98 +++++-- crates/iceberg/src/spec/schema/mod.rs | 329 +++++++++++++++++++++++- crates/iceberg/src/spec/values/datum.rs | 6 + crates/iceberg/src/spec/values/map.rs | 4 + 10 files changed, 507 insertions(+), 40 deletions(-) diff --git a/crates/catalog/glue/src/schema.rs b/crates/catalog/glue/src/schema.rs index 59d51f3dc1..747435611b 100644 --- a/crates/catalog/glue/src/schema.rs +++ b/crates/catalog/glue/src/schema.rs @@ -149,6 +149,12 @@ impl SchemaVisitor for GlueSchemaBuilder { fn primitive(&mut self, p: &PrimitiveType) -> Result { let glue_type = match p { + PrimitiveType::Unknown => { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + format!("Conversion from {p:?} is not supported"), + )); + } PrimitiveType::Boolean => "boolean".to_string(), PrimitiveType::Int => "int".to_string(), PrimitiveType::Long => "bigint".to_string(), diff --git a/crates/catalog/hms/src/schema.rs b/crates/catalog/hms/src/schema.rs index e21a75f976..f04ee11c17 100644 --- a/crates/catalog/hms/src/schema.rs +++ b/crates/catalog/hms/src/schema.rs @@ -100,6 +100,12 @@ impl SchemaVisitor for HiveSchemaBuilder { fn primitive(&mut self, p: &PrimitiveType) -> Result { let hive_type = match p { + PrimitiveType::Unknown => { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + format!("Conversion from {p:?} is not supported"), + )); + } PrimitiveType::Boolean => "boolean".to_string(), PrimitiveType::Int => "int".to_string(), PrimitiveType::Long => "bigint".to_string(), diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 3d9abf2275..18e03cb257 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1655,6 +1655,7 @@ pub iceberg::spec::PrimitiveType::Timestamp pub iceberg::spec::PrimitiveType::TimestampNs pub iceberg::spec::PrimitiveType::Timestamptz pub iceberg::spec::PrimitiveType::TimestamptzNs +pub iceberg::spec::PrimitiveType::Unknown pub iceberg::spec::PrimitiveType::Uuid impl iceberg::spec::PrimitiveType pub fn iceberg::spec::PrimitiveType::compatible(&self, literal: &iceberg::spec::PrimitiveLiteral) -> bool diff --git a/crates/iceberg/src/arrow/schema.rs b/crates/iceberg/src/arrow/schema.rs index dd26824c3a..d1912b159f 100644 --- a/crates/iceberg/src/arrow/schema.rs +++ b/crates/iceberg/src/arrow/schema.rs @@ -191,6 +191,7 @@ fn visit_type(r#type: &DataType, visitor: &mut V) -> Resu | DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + | DataType::Null | DataType::Binary | DataType::LargeBinary | DataType::BinaryView @@ -509,6 +510,7 @@ impl ArrowSchemaVisitor for ArrowSchemaConverter { fn primitive(&mut self, p: &DataType) -> Result { match p { + DataType::Null => Ok(Type::Primitive(PrimitiveType::Unknown)), DataType::Boolean => Ok(Type::Primitive(PrimitiveType::Boolean)), DataType::Int8 | DataType::Int16 | DataType::Int32 => { Ok(Type::Primitive(PrimitiveType::Int)) @@ -697,6 +699,7 @@ impl SchemaVisitor for ToArrowSchemaConverter { fn primitive(&mut self, p: &PrimitiveType) -> Result { match p { + PrimitiveType::Unknown => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Null)), PrimitiveType::Boolean => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Boolean)), PrimitiveType::Int => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Int32)), PrimitiveType::Long => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Int64)), @@ -1209,6 +1212,7 @@ pub(crate) fn primitive_type_to_arrow_type_with_ree(primitive_type: &PrimitiveTy }; match primitive_type { + PrimitiveType::Unknown => make_ree(DataType::Null), PrimitiveType::Boolean => make_ree(DataType::Boolean), PrimitiveType::Int => make_ree(DataType::Int32), PrimitiveType::Long => make_ree(DataType::Int64), @@ -2227,6 +2231,13 @@ mod tests { assert_eq!(iceberg_type, arrow_type_to_type(&arrow_type).unwrap()); } + { + let arrow_type = DataType::Null; + let iceberg_type = Type::Primitive(PrimitiveType::Unknown); + assert_eq!(arrow_type, type_to_arrow_type(&iceberg_type).unwrap()); + assert_eq!(iceberg_type, arrow_type_to_type(&arrow_type).unwrap()); + } + // test struct type { // no metadata will cause error diff --git a/crates/iceberg/src/arrow/value.rs b/crates/iceberg/src/arrow/value.rs index d22e565a4a..0de3b5650a 100644 --- a/crates/iceberg/src/arrow/value.rs +++ b/crates/iceberg/src/arrow/value.rs @@ -206,6 +206,7 @@ impl SchemaWithPartnerVisitor for ArrowArrayToIcebergStructConverter { fn primitive(&mut self, p: &PrimitiveType, partner: &ArrayRef) -> Result>> { match p { + PrimitiveType::Unknown => Ok(vec![None; partner.len()]), PrimitiveType::Boolean => { let array = partner .as_any() @@ -636,6 +637,7 @@ pub(crate) fn create_primitive_array_single_element( prim_lit: &Option, ) -> Result { match (data_type, prim_lit) { + (DataType::Null, None) => Ok(Arc::new(arrow_array::NullArray::new(1))), (DataType::Boolean, Some(PrimitiveLiteral::Boolean(v))) => { Ok(Arc::new(BooleanArray::from(vec![*v]))) } @@ -943,11 +945,10 @@ pub(crate) fn create_primitive_array_repeated( Some(NullBuffer::new_null(num_rows)), )) } - (DataType::Null, _) => Arc::new(arrow_array::NullArray::new(num_rows)), + (DataType::Null, None) => Arc::new(arrow_array::NullArray::new(num_rows)), // --- Catch-all null arm: use arrow-rs new_null_array for any remaining DataType --- (dt, None) => new_null_array(dt, num_rows), - (dt, _) => { return Err(Error::new( ErrorKind::Unexpected, @@ -1866,6 +1867,26 @@ mod test { } } + #[test] + fn test_create_null_array_rejects_non_null_literal() { + let literal = Some(PrimitiveLiteral::Int(1)); + + assert!(create_primitive_array_single_element(&DataType::Null, &literal).is_err()); + assert!(create_primitive_array_repeated(&DataType::Null, &literal, 2).is_err()); + assert_eq!( + create_primitive_array_single_element(&DataType::Null, &None) + .unwrap() + .len(), + 1 + ); + assert_eq!( + create_primitive_array_repeated(&DataType::Null, &None, 2) + .unwrap() + .len(), + 2 + ); + } + #[test] fn test_create_decimal_array_repeated_respects_precision() { // Ensure repeated arrays also respect target precision, not Arrow's default. diff --git a/crates/iceberg/src/avro/schema.rs b/crates/iceberg/src/avro/schema.rs index 528b89b7c1..fc7647fbb7 100644 --- a/crates/iceberg/src/avro/schema.rs +++ b/crates/iceberg/src/avro/schema.rs @@ -74,7 +74,7 @@ impl SchemaVisitor for SchemaToAvroSchema { record.name = Name::from(format!("r{}", field.id).as_str()); } - if !field.required { + if !field.required && !matches!(field_schema, AvroSchema::Null) { field_schema = avro_optional(field_schema)?; } @@ -126,7 +126,7 @@ impl SchemaVisitor for SchemaToAvroSchema { record.name = Name::from(format!("r{}", list.element_field.id).as_str()); } - if !list.element_field.required { + if !list.element_field.required && !matches!(field_schema, AvroSchema::Null) { field_schema = avro_optional(field_schema)?; } @@ -147,7 +147,7 @@ impl SchemaVisitor for SchemaToAvroSchema { ) -> Result { let key_field_schema = key_value.unwrap_left(); let mut value_field_schema = value.unwrap_left(); - if !map.value_field.required { + if !map.value_field.required && !matches!(value_field_schema, AvroSchema::Null) { value_field_schema = avro_optional(value_field_schema)?; } @@ -222,6 +222,7 @@ impl SchemaVisitor for SchemaToAvroSchema { fn primitive(&mut self, p: &PrimitiveType) -> Result { let avro_schema = match p { + PrimitiveType::Unknown => AvroSchema::Null, PrimitiveType::Boolean => AvroSchema::Boolean, PrimitiveType::Int => AvroSchema::Int, PrimitiveType::Long => AvroSchema::Long, @@ -311,6 +312,10 @@ pub(crate) fn avro_decimal_schema(precision: usize, scale: usize) -> Result Result { + if matches!(avro_schema, AvroSchema::Null) { + return Ok(AvroSchema::Null); + } + Ok(AvroSchema::Union(UnionSchema::new(vec![ AvroSchema::Null, avro_schema, @@ -447,10 +452,11 @@ impl AvroSchemaVisitor for AvroSchemaToSchema { let field_id = Self::get_element_id_from_attributes(&avro_field.custom_attributes, FIELD_ID_PROP)?; - let optional = is_avro_optional(&avro_field.schema); + let optional = is_avro_optional(&avro_field.schema) + || matches!(&avro_field.schema, AvroSchema::Null); - let mut field = - NestedField::new(field_id, &avro_field.name, field_type.unwrap(), !optional); + let field_type = field_type.unwrap_or(Type::Primitive(PrimitiveType::Unknown)); + let mut field = NestedField::new(field_id, &avro_field.name, field_type, !optional); if let Some(doc) = &avro_field.doc { field = field.with_doc(doc); @@ -482,7 +488,9 @@ impl AvroSchemaVisitor for AvroSchemaToSchema { } if options.len() == 1 { - Ok(Some(options.remove(0).unwrap())) + Ok(options + .remove(0) + .or(Some(Type::Primitive(PrimitiveType::Unknown)))) } else { Ok(Some(options.remove(1).unwrap())) } @@ -490,10 +498,11 @@ impl AvroSchemaVisitor for AvroSchemaToSchema { fn array(&mut self, array: &ArraySchema, item: Option) -> Result { let element_field_id = Self::get_element_id_from_attributes(&array.attributes, ELEMENT_ID)?; + let item = item.unwrap_or(Type::Primitive(PrimitiveType::Unknown)); let element_field = NestedField::list_element( element_field_id, - item.unwrap(), - !is_avro_optional(&array.items), + item, + !is_avro_optional(&array.items) && !matches!(array.items.as_ref(), AvroSchema::Null), ) .into(); Ok(Some(Type::List(ListType { element_field }))) @@ -504,10 +513,11 @@ impl AvroSchemaVisitor for AvroSchemaToSchema { let key_field = NestedField::map_key_element(key_field_id, Type::Primitive(PrimitiveType::String)); let value_field_id = Self::get_element_id_from_attributes(&map.attributes, VALUE_ID)?; + let value = value.unwrap_or(Type::Primitive(PrimitiveType::Unknown)); let value_field = NestedField::map_value_element( value_field_id, - value.unwrap(), - !is_avro_optional(&map.types), + value, + !is_avro_optional(&map.types) && !matches!(map.types.as_ref(), AvroSchema::Null), ); Ok(Some(Type::Map(MapType { key_field: key_field.into(), @@ -557,12 +567,7 @@ impl AvroSchemaVisitor for AvroSchemaToSchema { "Can't convert avro map schema, missing key schema.", ) })?; - let value = value.ok_or_else(|| { - Error::new( - ErrorKind::DataInvalid, - "Can't convert avro map schema, missing value schema.", - ) - })?; + let value = value.unwrap_or(Type::Primitive(PrimitiveType::Unknown)); let key_id = Self::get_element_id_from_attributes( &array.fields[0].custom_attributes, FIELD_ID_PROP, @@ -575,7 +580,8 @@ impl AvroSchemaVisitor for AvroSchemaToSchema { let value_field = NestedField::map_value_element( value_id, value, - !is_avro_optional(&array.fields[1].schema), + !is_avro_optional(&array.fields[1].schema) + && !matches!(&array.fields[1].schema, AvroSchema::Null), ); Ok(Some(Type::Map(MapType { key_field: key_field.into(), @@ -659,6 +665,25 @@ mod tests { assert_eq!(iceberg_schema, converted_avro_converted_iceberg_schema); } + #[test] + fn test_unknown_type_schema_conversion() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "empty", PrimitiveType::Unknown.into()).into(), + ]) + .build() + .unwrap(); + + let avro_schema = schema_to_avro_schema("table", &schema).unwrap(); + let AvroSchema::Record(record) = &avro_schema else { + panic!("expected avro record schema"); + }; + assert!(matches!(record.fields[0].schema, AvroSchema::Null)); + assert_eq!(record.fields[0].default, Some(Value::Null)); + + assert_eq!(schema, avro_schema_to_schema(&avro_schema).unwrap()); + } + #[test] fn test_manifest_file_v1_schema() { let fields = vec![ diff --git a/crates/iceberg/src/spec/datatypes.rs b/crates/iceberg/src/spec/datatypes.rs index 79c48c1318..bf6d590a53 100644 --- a/crates/iceberg/src/spec/datatypes.rs +++ b/crates/iceberg/src/spec/datatypes.rs @@ -19,7 +19,6 @@ * Data Types */ use std::collections::HashMap; -use std::convert::identity; use std::fmt; use std::ops::Index; use std::sync::{Arc, OnceLock}; @@ -134,8 +133,8 @@ impl Type { /// Minimum [`FormatVersion`] required to support this type, **without** taking /// nested field types into account. /// - /// `TimestampNs` / `TimestamptzNs` / `Variant` require [`FormatVersion::V3`]; every - /// other type is valid from [`FormatVersion::V1`]. Mirrors Java's + /// `Unknown` / `TimestampNs` / `TimestamptzNs` / `Variant` require + /// [`FormatVersion::V3`]; every other type is valid from [`FormatVersion::V1`]. Mirrors Java's /// `Schema.MIN_FORMAT_VERSIONS` (a shallow lookup keyed by type id), so it /// intentionally does not recurse: callers needing the floor for a whole schema /// iterate its flattened fields (see [`Schema::calc_min_compatible_format`]). @@ -143,7 +142,9 @@ impl Type { /// [`Schema::calc_min_compatible_format`]: crate::spec::Schema::calc_min_compatible_format pub(crate) fn min_format_version(&self) -> FormatVersion { match self { - Type::Primitive(PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs) + Type::Primitive( + PrimitiveType::Unknown | PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs, + ) | Type::Variant(_) => FormatVersion::V3, _ => FormatVersion::V1, } @@ -274,6 +275,8 @@ pub enum PrimitiveType { Fixed(u64), /// Arbitrary-length byte array. Binary, + /// Default / null column type used when a more specific type is not known. + Unknown, } impl PrimitiveType { @@ -391,6 +394,7 @@ where S: Serializer { impl fmt::Display for PrimitiveType { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { + PrimitiveType::Unknown => write!(f, "unknown"), PrimitiveType::Boolean => write!(f, "boolean"), PrimitiveType::Int => write!(f, "int"), PrimitiveType::Long => write!(f, "long"), @@ -554,7 +558,7 @@ impl fmt::Display for StructType { } #[derive(Debug, PartialEq, Serialize, Deserialize, Eq, Clone)] -#[serde(from = "SerdeNestedField", into = "SerdeNestedField")] +#[serde(try_from = "SerdeNestedField", into = "SerdeNestedField")] /// A struct is a tuple of typed values. Each field in the tuple is named and has an integer id that is unique in the table schema. /// Each field can be either optional or required, meaning that values can (or cannot) be null. Fields may be any type. /// Fields may have an optional comment or doc string. Fields can have default values. @@ -591,25 +595,72 @@ struct SerdeNestedField { pub write_default: Option, } -impl From for NestedField { - fn from(value: SerdeNestedField) -> Self { - NestedField { +impl TryFrom for NestedField { + type Error = crate::Error; + + fn try_from(value: SerdeNestedField) -> Result { + fn validate_unknown_values(default: &JsonValue, field_type: &Type) -> Result<()> { + if default.is_null() { + return Ok(()); + } + + match (field_type, default) { + (Type::Primitive(PrimitiveType::Unknown), _) => { + ensure_data_valid!(false, "Unknown type only supports null default values",); + } + (Type::Struct(struct_type), JsonValue::Object(object)) => { + for field in struct_type.fields() { + if let Some(value) = object.get(&field.id.to_string()) { + validate_unknown_values(value, &field.field_type)?; + } + } + } + (Type::List(list_type), JsonValue::Array(values)) => { + for value in values { + validate_unknown_values(value, &list_type.element_field.field_type)?; + } + } + (Type::Map(map_type), JsonValue::Object(object)) => { + if let Some(JsonValue::Array(keys)) = object.get("keys") { + for key in keys { + validate_unknown_values(key, &map_type.key_field.field_type)?; + } + } + if let Some(JsonValue::Array(values)) = object.get("values") { + for value in values { + validate_unknown_values(value, &map_type.value_field.field_type)?; + } + } + } + _ => {} + } + + Ok(()) + } + + fn parse_default(default: Option, field_type: &Type) -> Result> { + let Some(default) = default else { + return Ok(None); + }; + validate_unknown_values(&default, field_type)?; + match Literal::try_from_json(default, field_type) { + Ok(default) => Ok(default), + Err(_) => Ok(None), + } + } + + let initial_default = parse_default(value.initial_default, &value.field_type)?; + let write_default = parse_default(value.write_default, &value.field_type)?; + + Ok(NestedField { id: value.id, name: value.name, required: value.required, - initial_default: value.initial_default.and_then(|x| { - Literal::try_from_json(x, &value.field_type) - .ok() - .and_then(identity) - }), - write_default: value.write_default.and_then(|x| { - Literal::try_from_json(x, &value.field_type) - .ok() - .and_then(identity) - }), + initial_default, + write_default, field_type: value.field_type, doc: value.doc, - } + }) } } @@ -936,6 +987,7 @@ mod tests { { "type": "struct", "fields": [ + {"id": 17, "name": "unknown_field", "required": false, "type": "unknown"}, {"id": 1, "name": "bool_field", "required": true, "type": "boolean"}, {"id": 2, "name": "int_field", "required": true, "type": "int"}, {"id": 3, "name": "long_field", "required": true, "type": "long"}, @@ -960,6 +1012,12 @@ mod tests { record, Type::Struct(StructType { fields: vec![ + NestedField::optional( + 17, + "unknown_field", + Type::Primitive(PrimitiveType::Unknown), + ) + .into(), NestedField::required(1, "bool_field", Type::Primitive(PrimitiveType::Boolean)) .into(), NestedField::required(2, "int_field", Type::Primitive(PrimitiveType::Int)) @@ -1341,6 +1399,8 @@ mod tests { for (ty, literal) in pairs { assert!(ty.compatible(&literal)); } + + assert!(!PrimitiveType::Unknown.compatible(&PrimitiveLiteral::Int(1))); } #[test] diff --git a/crates/iceberg/src/spec/schema/mod.rs b/crates/iceberg/src/spec/schema/mod.rs index 652f98b649..ca51f0e6c3 100644 --- a/crates/iceberg/src/spec/schema/mod.rs +++ b/crates/iceberg/src/spec/schema/mod.rs @@ -39,11 +39,11 @@ pub use self::prune_columns::prune_columns; use super::NestedField; use crate::error::Result; use crate::expr::accessor::StructAccessor; -use crate::spec::FormatVersion; use crate::spec::datatypes::{ LIST_FIELD_NAME, ListType, MAP_KEY_FIELD_NAME, MAP_VALUE_FIELD_NAME, MapType, NestedFieldRef, PrimitiveType, StructType, Type, }; +use crate::spec::{FormatVersion, Literal}; use crate::{Error, ErrorKind, ensure_data_valid}; /// Type alias for schema id. @@ -133,6 +133,10 @@ impl SchemaBuilder { /// Builds the schema. pub fn build(self) -> Result { + for field in &self.fields { + Self::validate_unknown_type_field(field)?; + } + let field_id_to_accessor = self.build_accessors(); let r#struct = StructType::new(self.fields); @@ -190,6 +194,77 @@ impl SchemaBuilder { Ok(schema) } + fn validate_unknown_type_field(field: &NestedFieldRef) -> Result<()> { + ensure_data_valid!( + !field + .initial_default + .iter() + .chain(field.write_default.iter()) + .any(|default| { + Self::default_contains_non_null_unknown(default, &field.field_type) + }), + "Field {} cannot have non-null defaults because unknown type requires null defaults", + field.name + ); + + match field.field_type.as_ref() { + Type::Primitive(PrimitiveType::Unknown) => { + ensure_data_valid!( + !field.required, + "Field {} cannot be required because unknown type must be optional", + field.name + ); + } + Type::Struct(struct_type) => { + for nested_field in struct_type.fields() { + Self::validate_unknown_type_field(nested_field)?; + } + } + Type::List(list_type) => { + Self::validate_unknown_type_field(&list_type.element_field)?; + } + Type::Map(map_type) => { + Self::validate_unknown_type_field(&map_type.key_field)?; + Self::validate_unknown_type_field(&map_type.value_field)?; + } + Type::Primitive(_) | Type::Variant(_) => {} + } + + Ok(()) + } + + fn default_contains_non_null_unknown(default: &Literal, field_type: &Type) -> bool { + match (default, field_type) { + (_, Type::Primitive(PrimitiveType::Unknown)) => true, + (Literal::Struct(value), Type::Struct(struct_type)) => value + .iter() + .zip(struct_type.fields()) + .any(|(value, field)| { + value.is_some_and(|value| { + Self::default_contains_non_null_unknown(value, &field.field_type) + }) + }), + (Literal::List(values), Type::List(list_type)) => values.iter().any(|value| { + value.as_ref().is_some_and(|value| { + Self::default_contains_non_null_unknown( + value, + &list_type.element_field.field_type, + ) + }) + }), + (Literal::Map(map), Type::Map(map_type)) => map.iter().any(|(key, value)| { + Self::default_contains_non_null_unknown(key, &map_type.key_field.field_type) + || value.as_ref().is_some_and(|value| { + Self::default_contains_non_null_unknown( + value, + &map_type.value_field.field_type, + ) + }) + }), + _ => false, + } + } + fn build_accessors(&self) -> HashMap> { let mut map = HashMap::new(); @@ -647,6 +722,17 @@ mod tests { ]); assert_eq!(variant.calc_min_compatible_format(), FormatVersion::V3); + // Unknown is a v3-only primitive type. + let unknown = schema_with(vec![ + NestedField::optional(1, "u", Primitive(PrimitiveType::Unknown)).into(), + ]); + assert_eq!(unknown.calc_min_compatible_format(), FormatVersion::V3); + assert!( + unknown + .check_format_compatibility(FormatVersion::V2) + .is_err() + ); + // A v3-only type nested inside a list inside a struct → V3 (flattened fields). let nested = schema_with(vec![ NestedField::required( @@ -1475,4 +1561,245 @@ table { .is_err() ); } + + #[test] + fn test_unknown_type_deserialization_rejects_non_null_default() { + let field_json = serde_json::json!({ + "id": 1, + "name": "empty", + "required": false, + "type": "unknown", + "initial-default": 1 + }); + + let error = serde_json::from_value::(field_json.clone()).unwrap_err(); + assert!( + error + .to_string() + .contains("Unknown type only supports null default values"), + "unexpected error: {error}" + ); + + let schema_json = serde_json::json!({ + "type": "struct", + "schema-id": 1, + "fields": [field_json] + }); + assert!(serde_json::from_value::(schema_json).is_err()); + } + + #[test] + fn test_unknown_type_deserialization_accepts_null_defaults() { + let schema_json = serde_json::json!({ + "type": "struct", + "schema-id": 1, + "fields": [ + { + "id": 1, + "name": "empty", + "required": false, + "type": "unknown", + "initial-default": null, + "write-default": null + } + ] + }); + + serde_json::from_value::(schema_json).unwrap(); + } + + #[test] + fn test_unknown_type_must_be_optional_with_null_defaults() { + assert!( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::optional(1, "empty", Primitive(PrimitiveType::Unknown)).into() + ]) + .build() + .is_ok() + ); + + let required_error = Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required(1, "empty", Primitive(PrimitiveType::Unknown)).into(), + ]) + .build() + .unwrap_err(); + assert!( + required_error + .message() + .contains("unknown type must be optional") + ); + + let default_error = Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::optional(1, "empty", Primitive(PrimitiveType::Unknown)) + .with_initial_default(Literal::int(1)) + .into(), + ]) + .build() + .unwrap_err(); + assert!( + default_error + .message() + .contains("unknown type requires null defaults") + ); + } + + #[test] + fn test_unknown_type_rejects_non_null_container_defaults() { + let cases = [ + ( + "struct", + Struct(StructType::new(vec![ + NestedField::optional(2, "empty", Primitive(PrimitiveType::Unknown)).into(), + ])), + Literal::Struct(crate::spec::Struct::from_iter([Some(Literal::int(1))])), + ), + ( + "list", + List(ListType::new( + NestedField::list_element(2, Primitive(PrimitiveType::Unknown), false).into(), + )), + Literal::List(vec![Some(Literal::int(1))]), + ), + ( + "map", + Map(MapType::optional( + 2, + Primitive(PrimitiveType::String), + 3, + Primitive(PrimitiveType::Unknown), + )), + Literal::Map(MapValue::from([( + Literal::string("key"), + Some(Literal::int(1)), + )])), + ), + ]; + + for (name, field_type, default) in cases { + let error = Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::optional(1, name, field_type) + .with_initial_default(default) + .into(), + ]) + .build() + .unwrap_err(); + assert!( + error + .message() + .contains("unknown type requires null defaults"), + "unexpected error for {name}: {error}" + ); + } + } + + #[test] + fn test_unknown_type_accepts_null_container_defaults() { + let cases = [ + ( + "struct", + Struct(StructType::new(vec![ + NestedField::optional(2, "empty", Primitive(PrimitiveType::Unknown)).into(), + ])), + Literal::Struct(crate::spec::Struct::from_iter([None])), + ), + ( + "list", + List(ListType::new( + NestedField::list_element(2, Primitive(PrimitiveType::Unknown), false).into(), + )), + Literal::List(vec![None]), + ), + ( + "map", + Map(MapType::optional( + 2, + Primitive(PrimitiveType::String), + 3, + Primitive(PrimitiveType::Unknown), + )), + Literal::Map(MapValue::from([(Literal::string("key"), None)])), + ), + ]; + + for (name, field_type, default) in cases { + let schema = Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::optional(1, name, field_type) + .with_write_default(default) + .into(), + ]) + .build() + .unwrap(); + serde_json::to_value(schema).unwrap(); + } + } + + #[test] + fn test_unknown_type_deserialization_rejects_non_null_container_defaults() { + let cases = [ + ( + "struct", + serde_json::json!({ + "type": "struct", + "fields": [{ + "id": 2, + "name": "empty", + "required": false, + "type": "unknown" + }] + }), + serde_json::json!({"2": 1}), + ), + ( + "list", + serde_json::json!({ + "type": "list", + "element-id": 2, + "element-required": false, + "element": "unknown" + }), + serde_json::json!([1]), + ), + ( + "map", + serde_json::json!({ + "type": "map", + "key-id": 2, + "key": "string", + "value-id": 3, + "value-required": false, + "value": "unknown" + }), + serde_json::json!({"keys": ["key"], "values": [1]}), + ), + ]; + + for (name, field_type, default) in cases { + let schema_json = serde_json::json!({ + "type": "struct", + "schema-id": 1, + "fields": [{ + "id": 1, + "name": name, + "required": false, + "type": field_type, + "initial-default": default + }] + }); + + assert!( + serde_json::from_value::(schema_json).is_err(), + "non-null unknown default in {name} should be rejected" + ); + } + } } diff --git a/crates/iceberg/src/spec/values/datum.rs b/crates/iceberg/src/spec/values/datum.rs index f170a09df5..593544b94b 100644 --- a/crates/iceberg/src/spec/values/datum.rs +++ b/crates/iceberg/src/spec/values/datum.rs @@ -368,6 +368,12 @@ impl Datum { /// See [this spec](https://iceberg.apache.org/spec/#binary-single-value-serialization) for reference. pub fn try_from_bytes(bytes: &[u8], data_type: PrimitiveType) -> Result { let literal = match data_type { + PrimitiveType::Unknown => { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "Cannot create datum for unknown type from bytes", + )); + } PrimitiveType::Boolean => { if bytes.len() == 1 && bytes[0] == 0u8 { PrimitiveLiteral::Boolean(false) diff --git a/crates/iceberg/src/spec/values/map.rs b/crates/iceberg/src/spec/values/map.rs index e0f75205f0..81adb8da23 100644 --- a/crates/iceberg/src/spec/values/map.rs +++ b/crates/iceberg/src/spec/values/map.rs @@ -87,6 +87,10 @@ impl Map { self.index.get(key).map(|index| &self.pair[*index].1) } + pub(crate) fn iter(&self) -> impl Iterator)> { + self.pair.iter().map(|(key, value)| (key, value)) + } + /// The order of map is matter, so this method used to compare two maps has same key-value pairs without considering the order. pub fn has_same_content(&self, other: &Map) -> bool { if self.len() != other.len() { From 26cc0c2e2666fd38f143b29885be161b38b10f39 Mon Sep 17 00:00:00 2001 From: Manu Zhang Date: Wed, 26 Aug 2026 12:39:20 +0800 Subject: [PATCH 2/3] feat(parquet): omit unknown fields from writes Co-authored-by: Codex --- .../iceberg/src/arrow/nan_val_cnt_visitor.rs | 93 +++++- crates/iceberg/src/arrow/schema.rs | 301 ++++++++++++++++++ .../src/writer/file_writer/parquet_writer.rs | 281 +++++++++++++++- 3 files changed, 652 insertions(+), 23 deletions(-) diff --git a/crates/iceberg/src/arrow/nan_val_cnt_visitor.rs b/crates/iceberg/src/arrow/nan_val_cnt_visitor.rs index d01f3e9e56..6897f02c06 100644 --- a/crates/iceberg/src/arrow/nan_val_cnt_visitor.rs +++ b/crates/iceberg/src/arrow/nan_val_cnt_visitor.rs @@ -21,15 +21,17 @@ use std::collections::HashMap; use std::collections::hash_map::Entry; use std::sync::Arc; -use arrow_array::{ArrayRef, Float32Array, Float64Array, RecordBatch, StructArray}; +use arrow_array::{ + ArrayRef, Float32Array, Float64Array, ListArray, MapArray, RecordBatch, StructArray, +}; use arrow_schema::DataType; -use crate::Result; -use crate::arrow::{ArrowArrayAccessor, FieldMatchMode}; +use crate::arrow::FieldMatchMode; use crate::spec::{ ListType, MapType, NestedFieldRef, PrimitiveType, Schema, SchemaRef, SchemaWithPartnerVisitor, - StructType, VariantType, visit_struct_with_partner, + StructType, Type, VariantType, }; +use crate::{Error, ErrorKind, Result}; macro_rules! cast_and_update_cnt_map { ($t:ty, $col:ident, $self:ident, $field_id:ident) => { @@ -152,6 +154,78 @@ impl SchemaWithPartnerVisitor for NanValueCountVisitor { } impl NanValueCountVisitor { + fn visit_field(&mut self, field: &NestedFieldRef, array: &ArrayRef) -> Result<()> { + let field_id = field.id; + count_float_nans!(array, self, field_id); + + match field.field_type.as_ref() { + Type::Primitive(_) | Type::Variant(_) => Ok(()), + Type::Struct(struct_type) => self.visit_struct(struct_type, array), + Type::List(list_type) => { + let list_array = array.as_any().downcast_ref::().ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Expected list array for field {}, got {}", + field.id, + array.data_type() + ), + ) + })?; + self.visit_field(&list_type.element_field, list_array.values()) + } + Type::Map(map_type) => { + let map_array = array.as_any().downcast_ref::().ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Expected map array for field {}, got {}", + field.id, + array.data_type() + ), + ) + })?; + self.visit_field(&map_type.key_field, map_array.keys())?; + self.visit_field(&map_type.value_field, map_array.values()) + } + } + } + + fn visit_struct(&mut self, struct_type: &StructType, array: &ArrayRef) -> Result<()> { + let struct_array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!("Expected struct array, got {}", array.data_type()), + ) + })?; + + for field in struct_type.fields() { + if matches!( + field.field_type.as_ref(), + Type::Primitive(PrimitiveType::Unknown) + ) { + continue; + } + + let field_position = struct_array + .fields() + .iter() + .position(|arrow_field| self.match_mode.match_field(arrow_field, field)) + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!("Field id {} not found in struct array", field.id), + ) + })?; + self.visit_field(field, struct_array.column(field_position))?; + } + + Ok(()) + } + /// Creates new instance of NanValueCountVisitor pub fn new() -> Self { Self::new_with_match_mode(FieldMatchMode::Id) @@ -167,17 +241,8 @@ impl NanValueCountVisitor { /// Compute nan value counts in given schema and record batch pub fn compute(&mut self, schema: SchemaRef, batch: RecordBatch) -> Result<()> { - let arrow_arr_partner_accessor = ArrowArrayAccessor::new_with_match_mode(self.match_mode); - let struct_arr = Arc::new(StructArray::from(batch)) as ArrayRef; - visit_struct_with_partner( - schema.as_struct(), - &struct_arr, - self, - &arrow_arr_partner_accessor, - )?; - - Ok(()) + self.visit_struct(schema.as_struct(), &struct_arr) } } diff --git a/crates/iceberg/src/arrow/schema.rs b/crates/iceberg/src/arrow/schema.rs index d1912b159f..004ce6296d 100644 --- a/crates/iceberg/src/arrow/schema.rs +++ b/crates/iceberg/src/arrow/schema.rs @@ -793,6 +793,164 @@ pub fn schema_to_arrow_schema(schema: &Schema) -> Result { } } +fn parquet_arrow_field(field: &NestedFieldRef, arrow_field: &FieldRef) -> Result> { + let data_type = match field.field_type.as_ref() { + Type::Primitive(PrimitiveType::Unknown) => return Ok(None), + Type::Struct(struct_type) => { + let DataType::Struct(arrow_fields) = arrow_field.data_type() else { + return Err(Error::new( + ErrorKind::Unexpected, + format!( + "Expected Arrow struct for Iceberg field {}, got {}", + field.id, + arrow_field.data_type() + ), + )); + }; + if struct_type.fields().len() != arrow_fields.len() { + return Err(Error::new( + ErrorKind::Unexpected, + format!( + "Arrow and Iceberg struct field counts differ for field {}", + field.id + ), + )); + } + + let fields = struct_type + .fields() + .iter() + .zip(arrow_fields.iter()) + .filter_map(|(field, arrow_field)| { + parquet_arrow_field(field, arrow_field).transpose() + }) + .collect::>>()?; + if fields.is_empty() { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + format!( + "Cannot write struct field {} with no Parquet physical fields", + field.id + ), + )); + } + DataType::Struct(fields.into()) + } + Type::List(list_type) => { + let DataType::List(element_field) = arrow_field.data_type() else { + return Err(Error::new( + ErrorKind::Unexpected, + format!( + "Expected Arrow list for Iceberg field {}, got {}", + field.id, + arrow_field.data_type() + ), + )); + }; + let element_field = parquet_arrow_field(&list_type.element_field, element_field)? + .ok_or_else(|| { + Error::new( + ErrorKind::FeatureUnsupported, + format!( + "Cannot write list element {} with no Parquet physical fields", + list_type.element_field.id + ), + ) + })?; + DataType::List(element_field) + } + Type::Map(map_type) => { + let DataType::Map(entries_field, ordered) = arrow_field.data_type() else { + return Err(Error::new( + ErrorKind::Unexpected, + format!( + "Expected Arrow map for Iceberg field {}, got {}", + field.id, + arrow_field.data_type() + ), + )); + }; + let DataType::Struct(entry_fields) = entries_field.data_type() else { + return Err(Error::new( + ErrorKind::Unexpected, + format!( + "Expected Arrow map entries struct for Iceberg field {}", + field.id + ), + )); + }; + if entry_fields.len() != 2 { + return Err(Error::new( + ErrorKind::Unexpected, + format!( + "Expected two Arrow map entry fields for Iceberg field {}", + field.id + ), + )); + } + + let key_field = parquet_arrow_field(&map_type.key_field, &entry_fields[0])? + .ok_or_else(|| { + Error::new( + ErrorKind::FeatureUnsupported, + format!( + "Cannot write map key {} with no Parquet physical fields", + map_type.key_field.id + ), + ) + })?; + let value_field = parquet_arrow_field(&map_type.value_field, &entry_fields[1])? + .ok_or_else(|| { + Error::new( + ErrorKind::FeatureUnsupported, + format!( + "Cannot write map value {} with no Parquet physical fields", + map_type.value_field.id + ), + ) + })?; + let entries_field = Arc::new( + entries_field + .as_ref() + .clone() + .with_data_type(DataType::Struct(vec![key_field, value_field].into())), + ); + DataType::Map(entries_field, *ordered) + } + Type::Primitive(_) | Type::Variant(_) => arrow_field.data_type().clone(), + }; + + Ok(Some(Arc::new( + arrow_field.as_ref().clone().with_data_type(data_type), + ))) +} + +/// Convert an Iceberg schema to the Arrow schema used for Parquet writes. +/// +/// Unknown fields are omitted because Iceberg has no Parquet physical mapping for them. Structs, +/// list elements, and map keys/values left with no physical fields cannot be omitted without +/// losing container semantics, so those schemas are rejected. +pub(crate) fn schema_to_arrow_schema_for_parquet(schema: &Schema) -> Result { + let arrow_schema = schema_to_arrow_schema(schema)?; + let fields = schema + .as_struct() + .fields() + .iter() + .zip(arrow_schema.fields().iter()) + .filter_map(|(field, arrow_field)| parquet_arrow_field(field, arrow_field).transpose()) + .collect::>>()?; + if fields.is_empty() { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "Cannot write a schema with no Parquet physical fields", + )); + } + Ok(ArrowSchema::new_with_metadata( + fields, + arrow_schema.metadata().clone(), + )) +} + /// Convert iceberg type to an arrow type. pub fn type_to_arrow_type(ty: &Type) -> Result { let mut converter = ToArrowSchemaConverter; @@ -2221,6 +2379,149 @@ mod tests { ); } + #[test] + fn test_parquet_arrow_schema_omits_unknown_struct_fields() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "unknown", PrimitiveType::Unknown.into()).into(), + NestedField::optional( + 2, + "struct", + Type::Struct(StructType::new(vec![ + NestedField::optional(3, "unknown", PrimitiveType::Unknown.into()).into(), + NestedField::optional(4, "known", PrimitiveType::Int.into()).into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(); + + let arrow_schema = schema_to_arrow_schema_for_parquet(&schema).unwrap(); + assert_eq!(arrow_schema.fields().len(), 1); + assert_eq!(arrow_schema.field(0).name(), "struct"); + let DataType::Struct(fields) = arrow_schema.field(0).data_type() else { + panic!("expected struct field"); + }; + assert_eq!(fields.len(), 1); + assert_eq!(fields[0].name(), "known"); + } + + #[test] + fn test_parquet_arrow_schema_rejects_empty_struct() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "known", PrimitiveType::Int.into()).into(), + NestedField::optional( + 2, + "empty_struct", + Type::Struct(StructType::new(vec![ + NestedField::optional(3, "unknown", PrimitiveType::Unknown.into()).into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(); + + assert!( + schema_to_arrow_schema_for_parquet(&schema) + .unwrap_err() + .message() + .contains("struct field 2") + ); + } + + #[test] + fn test_parquet_arrow_schema_rejects_all_unknown_fields() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "unknown", PrimitiveType::Unknown.into()).into(), + ]) + .build() + .unwrap(); + + assert!( + schema_to_arrow_schema_for_parquet(&schema) + .unwrap_err() + .message() + .contains("no Parquet physical fields") + ); + } + + #[test] + fn test_parquet_arrow_schema_rejects_unknown_container_values() { + let list_schema = Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "list", + Type::List(ListType::new( + NestedField::optional(2, "element", PrimitiveType::Unknown.into()).into(), + )), + ) + .into(), + ]) + .build() + .unwrap(); + assert!( + schema_to_arrow_schema_for_parquet(&list_schema) + .unwrap_err() + .message() + .contains("list element") + ); + + let list_of_empty_struct_schema = Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "list", + Type::List(ListType::new( + NestedField::optional( + 2, + "element", + Type::Struct(StructType::new(vec![ + NestedField::optional(3, "unknown", PrimitiveType::Unknown.into()) + .into(), + ])), + ) + .into(), + )), + ) + .into(), + ]) + .build() + .unwrap(); + assert!( + schema_to_arrow_schema_for_parquet(&list_of_empty_struct_schema) + .unwrap_err() + .message() + .contains("struct field 2") + ); + + let map_schema = Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "map", + Type::Map(MapType::new( + NestedField::map_key_element(2, PrimitiveType::String.into()).into(), + NestedField::map_value_element(3, PrimitiveType::Unknown.into(), false) + .into(), + )), + ) + .into(), + ]) + .build() + .unwrap(); + assert!( + schema_to_arrow_schema_for_parquet(&map_schema) + .unwrap_err() + .message() + .contains("map value") + ); + } + #[test] fn test_type_conversion() { // test primitive type diff --git a/crates/iceberg/src/writer/file_writer/parquet_writer.rs b/crates/iceberg/src/writer/file_writer/parquet_writer.rs index 95654c89a7..a6274f2761 100644 --- a/crates/iceberg/src/writer/file_writer/parquet_writer.rs +++ b/crates/iceberg/src/writer/file_writer/parquet_writer.rs @@ -20,13 +20,16 @@ use std::collections::HashMap; use std::sync::Arc; -use arrow_schema::SchemaRef as ArrowSchemaRef; +use arrow_array::{ + Array, ArrayRef, ListArray, MapArray, RecordBatch, RecordBatchOptions, StructArray, +}; +use arrow_schema::{DataType, FieldRef, Fields, SchemaRef as ArrowSchemaRef}; use bytes::Bytes; use futures::future::BoxFuture; use itertools::Itertools; -use parquet::arrow::AsyncArrowWriter; use parquet::arrow::async_reader::AsyncFileReader; use parquet::arrow::async_writer::AsyncFileWriter as ArrowAsyncFileWriter; +use parquet::arrow::{AsyncArrowWriter, PARQUET_FIELD_ID_META_KEY}; use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}; use parquet::encryption::encrypt::FileEncryptionProperties; use parquet::file::metadata::ParquetMetaData; @@ -37,6 +40,7 @@ use super::{FileWriter, FileWriterBuilder}; use crate::arrow::{ ArrowFileReader, DEFAULT_MAP_FIELD_NAME, FieldMatchMode, NanValueCountVisitor, get_parquet_stat_max_as_datum, get_parquet_stat_min_as_datum, + schema_to_arrow_schema_for_parquet, }; use crate::compression::CompressionCodec; use crate::encryption::{EncryptionManager, StandardKeyMetadata}; @@ -167,6 +171,8 @@ impl FileWriterBuilder for ParquetWriterBuilder { resolve_writer_properties(self.props.clone(), key_metadata.as_ref())?; Ok(ParquetWriter { schema: self.schema.clone(), + arrow_schema: Arc::new(schema_to_arrow_schema_for_parquet(&self.schema)?), + match_mode: self.match_mode, inner_writer: None, writer_properties, current_row_num: 0, @@ -177,6 +183,172 @@ impl FileWriterBuilder for ParquetWriterBuilder { } } +fn field_id(field: &FieldRef) -> Option<&str> { + field + .metadata() + .get(PARQUET_FIELD_ID_META_KEY) + .map(String::as_str) +} + +fn find_field_index( + fields: &Fields, + target: &FieldRef, + match_mode: FieldMatchMode, +) -> Option { + match match_mode { + FieldMatchMode::Id => field_id(target).and_then(|target_id| { + fields + .iter() + .position(|field| field_id(field) == Some(target_id)) + }), + FieldMatchMode::Name => fields + .iter() + .position(|field| field.name() == target.name()), + } +} + +fn project_array_for_parquet( + array: &ArrayRef, + target: &FieldRef, + match_mode: FieldMatchMode, +) -> Result { + if array.data_type() == target.data_type() { + return Ok(array.clone()); + } + + match target.data_type() { + DataType::Struct(target_fields) => { + let source = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Expected struct array for Parquet field {}, got {}", + target.name(), + array.data_type() + ), + ) + })?; + let columns = target_fields + .iter() + .map(|target_field| { + let index = find_field_index(source.fields(), target_field, match_mode) + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Field {} is missing from struct array for Parquet write", + target_field.name() + ), + ) + })?; + project_array_for_parquet(source.column(index), target_field, match_mode) + }) + .collect::>>()?; + Ok(Arc::new(StructArray::try_new_with_length( + target_fields.clone(), + columns, + source.nulls().cloned(), + source.len(), + )?)) + } + DataType::List(target_element) => { + let source = array.as_any().downcast_ref::().ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Expected list array for Parquet field {}, got {}", + target.name(), + array.data_type() + ), + ) + })?; + let values = project_array_for_parquet(source.values(), target_element, match_mode)?; + Ok(Arc::new(ListArray::try_new( + target_element.clone(), + source.offsets().clone(), + values, + source.nulls().cloned(), + )?)) + } + DataType::Map(target_entries, ordered) => { + let source = array.as_any().downcast_ref::().ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Expected map array for Parquet field {}, got {}", + target.name(), + array.data_type() + ), + ) + })?; + let source_entries: ArrayRef = Arc::new(source.entries().clone()); + let entries = project_array_for_parquet(&source_entries, target_entries, match_mode)?; + let entries = entries + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "Projected Parquet map entries are not a struct array", + ) + })? + .clone(); + Ok(Arc::new(MapArray::try_new( + target_entries.clone(), + source.offsets().clone(), + entries, + source.nulls().cloned(), + *ordered, + )?)) + } + _ => Err(Error::new( + ErrorKind::DataInvalid, + format!( + "Cannot project Arrow type {} to {} for Parquet field {}", + array.data_type(), + target.data_type(), + target.name() + ), + )), + } +} + +fn project_batch_for_parquet( + batch: &RecordBatch, + target_schema: ArrowSchemaRef, + match_mode: FieldMatchMode, +) -> Result { + if batch.schema_ref() == &target_schema { + return Ok(batch.clone()); + } + + let source_schema = batch.schema(); + let columns = target_schema + .fields() + .iter() + .map(|target_field| { + let index = find_field_index(source_schema.fields(), target_field, match_mode) + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Field {} is missing from record batch for Parquet write", + target_field.name() + ), + ) + })?; + project_array_for_parquet(batch.column(index), target_field, match_mode) + }) + .collect::>>()?; + let options = RecordBatchOptions::default() + .with_match_field_names(false) + .with_row_count(Some(batch.num_rows())); + RecordBatch::try_new_with_options(target_schema, columns, &options).map_err(Into::into) +} + /// A mapping from Parquet column path names to internal field id struct IndexByParquetPathName { name_to_id: HashMap, @@ -305,6 +477,8 @@ impl SchemaVisitor for IndexByParquetPathName { /// `ParquetWriter`` is used to write arrow data into parquet file on storage. pub struct ParquetWriter { schema: SchemaRef, + arrow_schema: ArrowSchemaRef, + match_mode: FieldMatchMode, output_file: OutputFile, inner_writer: Option>, writer_properties: WriterProperties, @@ -610,7 +784,7 @@ fn resolve_writer_properties( } impl FileWriter for ParquetWriter { - async fn write(&mut self, batch: &arrow_array::RecordBatch) -> Result<()> { + async fn write(&mut self, batch: &RecordBatch) -> Result<()> { // Skip empty batch if batch.num_rows() == 0 { return Ok(()); @@ -618,20 +792,19 @@ impl FileWriter for ParquetWriter { self.current_row_num += batch.num_rows(); - let batch_c = batch.clone(); + let batch = project_batch_for_parquet(batch, self.arrow_schema.clone(), self.match_mode)?; self.nan_value_count_visitor - .compute(self.schema.clone(), batch_c)?; + .compute(self.schema.clone(), batch.clone())?; // Lazy initialize the writer let writer = if let Some(writer) = &mut self.inner_writer { writer } else { - let arrow_schema: ArrowSchemaRef = Arc::new(self.schema.as_ref().try_into()?); let inner_writer = self.output_file.writer().await?; let async_writer = AsyncFileWriter::new(inner_writer); let writer = AsyncArrowWriter::try_new( async_writer, - arrow_schema.clone(), + self.arrow_schema.clone(), Some(self.writer_properties.clone()), ) .map_err(|err| { @@ -642,7 +815,7 @@ impl FileWriter for ParquetWriter { self.inner_writer.as_mut().unwrap() }; - writer.write(batch).await.map_err(|err| { + writer.write(&batch).await.map_err(|err| { Error::new( ErrorKind::Unexpected, "Failed to write using parquet writer.", @@ -754,7 +927,7 @@ mod tests { use arrow_array::types::{Float32Type, Int64Type}; use arrow_array::{ Array, ArrayRef, BooleanArray, Decimal128Array, Float32Array, Float64Array, Int32Array, - Int64Array, ListArray, MapArray, RecordBatch, StructArray, + Int64Array, ListArray, MapArray, NullArray, RecordBatch, StructArray, }; use arrow_schema::{DataType, Field, Fields, SchemaRef as ArrowSchemaRef}; use arrow_select::concat::concat_batches; @@ -841,6 +1014,96 @@ mod tests { .unwrap() } + #[test] + fn test_project_batch_for_parquet_omits_unknown_fields() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "unknown", PrimitiveType::Unknown.into()).into(), + NestedField::optional( + 2, + "struct", + Type::Struct(StructType::new(vec![ + NestedField::optional(3, "unknown", PrimitiveType::Unknown.into()).into(), + NestedField::optional(4, "known", PrimitiveType::Int.into()).into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(); + let source_schema = Arc::new(schema_to_arrow_schema(&schema).unwrap()); + let DataType::Struct(struct_fields) = source_schema.field(1).data_type() else { + panic!("expected struct field"); + }; + let struct_array = Arc::new(StructArray::new( + struct_fields.clone(), + vec![ + Arc::new(NullArray::new(2)), + Arc::new(Int32Array::from(vec![Some(1), Some(2)])), + ], + None, + )); + let batch = RecordBatch::try_new(source_schema, vec![ + Arc::new(NullArray::new(2)), + struct_array, + ]) + .unwrap(); + let target_schema = Arc::new(schema_to_arrow_schema_for_parquet(&schema).unwrap()); + + let projected = + project_batch_for_parquet(&batch, target_schema, FieldMatchMode::Id).unwrap(); + + assert_eq!(projected.num_rows(), 2); + assert_eq!(projected.num_columns(), 1); + let projected_struct = projected + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(projected_struct.num_columns(), 1); + assert_eq!(projected_struct.fields()[0].name(), "known"); + + let mut nan_visitor = NanValueCountVisitor::new(); + nan_visitor.compute(Arc::new(schema), projected).unwrap(); + assert!(nan_visitor.nan_value_counts.is_empty()); + } + + #[test] + fn test_project_batch_for_parquet_honors_name_match_mode() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "known", PrimitiveType::Int.into()).into(), + ]) + .build() + .unwrap(); + let source_schema = Arc::new(arrow_schema::Schema::new(vec![ + Field::new("known", DataType::Int32, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "2".to_string(), + )])), + Field::new("wrong", DataType::Int32, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + let batch = RecordBatch::try_new(source_schema, vec![ + Arc::new(Int32Array::from(vec![10])), + Arc::new(Int32Array::from(vec![20])), + ]) + .unwrap(); + let target_schema = Arc::new(schema_to_arrow_schema_for_parquet(&schema).unwrap()); + + let projected = + project_batch_for_parquet(&batch, target_schema, FieldMatchMode::Name).unwrap(); + + let values = projected + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), 10); + } + fn nested_schema_for_test() -> Schema { // Int, Struct(Int,Int), String, List(Int), Struct(Struct(Int)), Map(String, List(Int)) Schema::builder() From 049573e9021f75604ae5dcfbe74012637b8e275f Mon Sep 17 00:00:00 2001 From: Manu Zhang Date: Wed, 26 Aug 2026 12:41:13 +0800 Subject: [PATCH 3/3] feat(parquet): reconstruct unknown fields on reads Co-authored-by: Codex --- crates/iceberg/src/arrow/reader/pipeline.rs | 6 +- crates/iceberg/src/arrow/reader/projection.rs | 605 +++++++++++++-- .../src/arrow/record_batch_transformer.rs | 729 +++++++++++++++++- crates/iceberg/src/arrow/value.rs | 140 +++- 4 files changed, 1384 insertions(+), 96 deletions(-) diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index b954055a28..4f39f3a243 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -180,7 +180,7 @@ impl FileScanTaskReader { } else { // Branch 3: No name mapping - use position-based fallback IDs // Corresponds to Java's ParquetSchemaUtil.addFallbackIds() - add_fallback_field_ids_to_arrow_schema(arrow_metadata.schema()) + add_fallback_field_ids_to_arrow_schema(arrow_metadata.schema(), &task.schema)? }; let options = ArrowReaderOptions::new().with_schema(arrow_schema); @@ -345,7 +345,8 @@ impl FileScanTaskReader { // that come back from the file, such as type promotion, default column insertion, // column re-ordering, partition constants, and virtual field addition (like _file) let mut record_batch_transformer_builder = - RecordBatchTransformerBuilder::new(task.schema_ref(), task.project_field_ids()); + RecordBatchTransformerBuilder::new(task.schema_ref(), task.project_field_ids()) + .with_position_fallback(use_position_fallback); // Add the _file metadata column if it's in the projected fields if task.project_field_ids().contains(&RESERVED_FIELD_ID_FILE) { @@ -518,6 +519,7 @@ impl FileScanTaskReader { record_batch_stream_builder.parquet_schema(), record_batch_stream_builder.schema(), &predicate, + &task.schema, use_position_fallback, )?; diff --git a/crates/iceberg/src/arrow/reader/projection.rs b/crates/iceberg/src/arrow/reader/projection.rs index 9fdc9fed65..26e9071e63 100644 --- a/crates/iceberg/src/arrow/reader/projection.rs +++ b/crates/iceberg/src/arrow/reader/projection.rs @@ -40,6 +40,7 @@ impl ArrowReader { parquet_schema: &SchemaDescriptor, arrow_schema: &ArrowSchemaRef, predicate: &BoundPredicate, + iceberg_schema: &Schema, use_position_fallback: bool, ) -> Result<(HashSet, HashMap)> { // Collects all Iceberg field IDs referenced in the filter predicate @@ -50,15 +51,26 @@ impl ArrowReader { let iceberg_field_ids = collector.field_ids(); - let field_id_map = match build_field_id_map(parquet_schema)? { + let mut field_id_map = match build_field_id_map(parquet_schema)? { Some(map) => map, // No embedded field IDs and no name mapping: position-based fallback - None if use_position_fallback => build_fallback_field_id_map(parquet_schema), + None if use_position_fallback => build_field_id_map_from_arrow_schema(arrow_schema), // No embedded field IDs, but a name mapping assigned them to the Arrow // schema: resolve columns through the mapped Arrow field-id metadata None => build_field_id_map_from_arrow_schema(arrow_schema), }; + // Unknown fields are always logically null and must not be read from data files, even if + // a non-conforming or older file contains a physical column with the same field ID. + field_id_map.retain(|field_id, _| { + !iceberg_schema.field_by_id(*field_id).is_some_and(|field| { + matches!( + field.field_type.as_ref(), + Type::Primitive(PrimitiveType::Unknown) + ) + }) + }); + Ok((iceberg_field_ids, field_id_map)) } @@ -66,6 +78,7 @@ impl ArrowReader { /// Nested types (struct/list/map) are flattened in Parquet's columnar format. fn include_leaf_field_id(field: &NestedField, field_ids: &mut Vec) { match field.field_type.as_ref() { + Type::Primitive(PrimitiveType::Unknown) => {} Type::Primitive(_) => { field_ids.push(field.id); } @@ -104,6 +117,7 @@ impl ArrowReader { (Some(lhs), Some(rhs)) if lhs == rhs => true, (Some(PrimitiveType::Int), Some(PrimitiveType::Long)) => true, (Some(PrimitiveType::Float), Some(PrimitiveType::Double)) => true, + (Some(PrimitiveType::Unknown), Some(_)) => true, ( Some(PrimitiveType::Decimal { precision: file_precision, @@ -139,7 +153,7 @@ impl ArrowReader { if use_fallback { // Position-based projection necessary because file lacks embedded field IDs - Self::get_arrow_projection_mask_fallback(field_ids, parquet_schema) + Self::get_arrow_projection_mask_fallback(field_ids, parquet_schema, arrow_schema) } else { // Field-ID-based projection using embedded field IDs from Parquet metadata @@ -229,36 +243,38 @@ impl ArrowReader { } if indices.is_empty() { - // Edge case: All requested columns are new (don't exist in file). - // Project all columns so RecordBatchTransformer has a batch to transform. - Ok(ProjectionMask::all()) + // All requested columns are new or have no physical representation. An empty + // projection still preserves row counts for RecordBatchTransformer. + Ok(ProjectionMask::none(parquet_schema.num_columns())) } else { Ok(ProjectionMask::leaves(parquet_schema, indices)) } } /// Fallback projection for Parquet files without field IDs. - /// Uses position-based matching: field ID N → column position N-1. + /// Uses the fallback IDs assigned to physical top-level Arrow fields. /// Projects entire top-level columns (including nested content) for iceberg-java compatibility. fn get_arrow_projection_mask_fallback( field_ids: &[i32], parquet_schema: &SchemaDescriptor, + arrow_schema: &ArrowSchemaRef, ) -> Result { - // Position-based: field_id N → column N-1 (field IDs are 1-indexed) - let parquet_root_fields = parquet_schema.root_schema().get_fields(); + let field_id_set: HashSet = field_ids.iter().copied().collect(); let mut root_indices = vec![]; - for field_id in field_ids.iter() { - let parquet_pos = (*field_id - 1) as usize; - - if parquet_pos < parquet_root_fields.len() { - root_indices.push(parquet_pos); + for (root_index, field) in arrow_schema.fields().iter().enumerate() { + if field + .metadata() + .get(PARQUET_FIELD_ID_META_KEY) + .and_then(|field_id| i32::from_str(field_id).ok()) + .is_some_and(|field_id| field_id_set.contains(&field_id)) + { + root_indices.push(root_index); } - // RecordBatchTransformer adds missing columns with NULL values } if root_indices.is_empty() { - Ok(ProjectionMask::all()) + Ok(ProjectionMask::none(parquet_schema.num_columns())) } else { Ok(ProjectionMask::roots(parquet_schema, root_indices)) } @@ -312,44 +328,6 @@ pub(super) fn build_field_id_map( Ok(Some(column_map)) } -/// Build a fallback field ID map for Parquet files without embedded field IDs. -/// -/// Returns the number of primitive (leaf) columns in a Parquet type, recursing into groups. -fn leaf_count(ty: &parquet::schema::types::Type) -> usize { - if ty.is_primitive() { - 1 - } else { - ty.get_fields().iter().map(|f| leaf_count(f)).sum() - } -} - -/// Builds a mapping from fallback field IDs to leaf column indices for Parquet files -/// without embedded field IDs. Returns entries only for primitive top-level fields. -/// -/// Must use top-level field positions (not leaf column positions) to stay consistent -/// with `add_fallback_field_ids_to_arrow_schema`, which assigns ordinal IDs to -/// top-level Arrow fields. Using leaf positions instead would produce wrong indices -/// when nested types (struct/list/map) expand into multiple leaf columns. -/// -/// Mirrors iceberg-java's ParquetSchemaUtil.addFallbackIds() which iterates -/// fileSchema.getFields() assigning ordinal IDs to top-level fields. -pub(super) fn build_fallback_field_id_map( - parquet_schema: &SchemaDescriptor, -) -> HashMap { - let mut column_map = HashMap::new(); - let mut leaf_idx = 0; - - for (top_pos, field) in parquet_schema.root_schema().get_fields().iter().enumerate() { - let field_id = (top_pos + 1) as i32; - if field.is_primitive() { - column_map.insert(field_id, leaf_idx); - } - leaf_idx += leaf_count(field); - } - - column_map -} - /// Builds a mapping from field IDs to leaf column indices using the field-id metadata /// carried by the Arrow schema. /// @@ -448,9 +426,17 @@ pub(super) fn apply_name_mapping_to_arrow_schema( /// Why at schema level (not per-batch): Efficiency - avoids repeated schema modification. /// Why only top-level: Nested projection uses leaf column indices, not parent struct IDs. /// Why 1-indexed: Compatibility with iceberg-java's ParquetSchemaUtil.addFallbackIds(). +/// Unknown fields have no Parquet column, so their actual field IDs are skipped in the positional +/// fallback sequence. The remaining IDs stay position-based; inferring identity from current field +/// names could expose stale values when a deleted name is reused by a new field. +/// If a fallback ID is absent from the current top-level schema, the file is ambiguous: that gap +/// may be a dropped physical field or an Unknown field omitted by an older schema. Without a name +/// mapping or embedded field IDs, assigning later columns would risk returning them under the wrong +/// Iceberg IDs, so reject the fallback instead. pub(super) fn add_fallback_field_ids_to_arrow_schema( arrow_schema: &ArrowSchemaRef, -) -> Arc { + iceberg_schema: &Schema, +) -> Result> { debug_assert!( arrow_schema .fields() @@ -460,24 +446,69 @@ pub(super) fn add_fallback_field_ids_to_arrow_schema( "Schema already has field IDs" ); + let omitted_field_ids: HashSet = iceberg_schema + .as_struct() + .fields() + .iter() + .filter(|field| !type_has_parquet_physical_field(&field.field_type)) + .map(|field| field.id) + .collect(); + let top_level_field_ids: HashSet = iceberg_schema + .as_struct() + .fields() + .iter() + .map(|field| field.id) + .collect(); + let mut fallback_field_ids = (1_i32..).filter(|field_id| !omitted_field_ids.contains(field_id)); let fields_with_fallback_ids: Vec<_> = arrow_schema .fields() .iter() .enumerate() - .map(|(pos, field)| { + .map(|(position, field)| { let mut metadata = field.metadata().clone(); - let field_id = (pos + 1) as i32; // 1-indexed for Java compatibility + let field_id = fallback_field_ids.next().unwrap(); + if !top_level_field_ids.contains(&field_id) { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "Cannot safely apply positional fallback: Parquet field {} at position {} maps to missing top-level Iceberg field ID {}; provide a name mapping", + field.name(), + position + 1, + field_id + ), + )); + } metadata.insert(PARQUET_FIELD_ID_META_KEY.to_string(), field_id.to_string()); - Field::new(field.name(), field.data_type().clone(), field.is_nullable()) - .with_metadata(metadata) + Ok( + Field::new(field.name(), field.data_type().clone(), field.is_nullable()) + .with_metadata(metadata), + ) }) - .collect(); + .collect::>>()?; - Arc::new(ArrowSchema::new_with_metadata( + Ok(Arc::new(ArrowSchema::new_with_metadata( fields_with_fallback_ids, arrow_schema.metadata().clone(), - )) + ))) +} + +fn type_has_parquet_physical_field(field_type: &Type) -> bool { + match field_type { + Type::Primitive(PrimitiveType::Unknown) => false, + Type::Primitive(_) | Type::Variant(_) => true, + Type::Struct(struct_type) => struct_type + .fields() + .iter() + .any(|field| type_has_parquet_physical_field(&field.field_type)), + Type::List(list_type) => { + type_has_parquet_physical_field(&list_type.element_field.field_type) + } + Type::Map(map_type) => { + type_has_parquet_physical_field(&map_type.key_field.field_type) + && type_has_parquet_physical_field(&map_type.value_field.field_type) + } + } } #[cfg(test)] @@ -487,7 +518,7 @@ mod tests { use std::sync::Arc; use arrow_array::cast::AsArray; - use arrow_array::{Array, ArrayRef, RecordBatch, StringArray}; + use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit}; use futures::TryStreamExt; use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY, ProjectionMask}; @@ -497,6 +528,7 @@ mod tests { use parquet::schema::types::SchemaDescriptor; use tempfile::TempDir; + use super::add_fallback_field_ids_to_arrow_schema; use crate::arrow::{ArrowReader, ArrowReaderBuilder}; use crate::expr::{Bind, Reference}; use crate::io::FileIO; @@ -507,6 +539,453 @@ mod tests { }; use crate::{ErrorKind, Runtime}; + #[test] + fn test_fallback_field_ids_skip_actual_unknown_field_ids() { + let iceberg_schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "unknown", PrimitiveType::Unknown.into()).into(), + NestedField::optional(2, "known_before", PrimitiveType::Int.into()).into(), + NestedField::optional(3, "known_after", PrimitiveType::Int.into()).into(), + ]) + .build() + .unwrap(); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("known_before", DataType::Int32, true), + Field::new("known_after", DataType::Int32, true), + ])); + + let schema_with_ids = + add_fallback_field_ids_to_arrow_schema(&arrow_schema, &iceberg_schema).unwrap(); + + assert_eq!( + schema_with_ids.fields()[0] + .metadata() + .get(PARQUET_FIELD_ID_META_KEY) + .map(String::as_str), + Some("2") + ); + assert_eq!( + schema_with_ids.fields()[1] + .metadata() + .get(PARQUET_FIELD_ID_META_KEY) + .map(String::as_str), + Some("3") + ); + } + + #[test] + fn test_fallback_field_ids_ignore_unknown_reordering() { + let iceberg_schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "a", PrimitiveType::Int.into()).into(), + NestedField::optional(3, "unknown", PrimitiveType::Unknown.into()).into(), + NestedField::optional(2, "b", PrimitiveType::Int.into()).into(), + ]) + .build() + .unwrap(); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + + let schema_with_ids = + add_fallback_field_ids_to_arrow_schema(&arrow_schema, &iceberg_schema).unwrap(); + + assert_eq!( + schema_with_ids.fields()[0] + .metadata() + .get(PARQUET_FIELD_ID_META_KEY) + .map(String::as_str), + Some("1") + ); + assert_eq!( + schema_with_ids.fields()[1] + .metadata() + .get(PARQUET_FIELD_ID_META_KEY) + .map(String::as_str), + Some("2") + ); + } + + #[test] + fn test_fallback_field_ids_reject_dropped_unknown_gap() { + // The file was written while top-level field ID 3 was Unknown and therefore omitted from + // Parquet. After that field is dropped, the current schema can no longer distinguish the + // historical gap from a dropped physical field, so assigning ID 3 to the third column + // would silently hide field ID 4. + let iceberg_schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "a", PrimitiveType::Int.into()).into(), + NestedField::optional(2, "b", PrimitiveType::Int.into()).into(), + NestedField::optional(4, "c", PrimitiveType::Int.into()).into(), + ]) + .build() + .unwrap(); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ])); + + let error = + add_fallback_field_ids_to_arrow_schema(&arrow_schema, &iceberg_schema).unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::DataInvalid); + assert!( + error + .message() + .contains("missing top-level Iceberg field ID 3"), + "unexpected error: {error}" + ); + } + + #[test] + fn test_unknown_only_projection_reads_no_physical_columns() { + let iceberg_schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "unknown", PrimitiveType::Unknown.into()).into(), + NestedField::optional(2, "known", PrimitiveType::Int.into()).into(), + ]) + .build() + .unwrap(); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("known", DataType::Int32, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "2".to_string(), + )])), + ])); + let parquet_schema = SchemaDescriptor::new(Arc::new( + parse_message_type("message schema { optional int32 known = 2; }").unwrap(), + )); + + for use_fallback in [false, true] { + let mask = ArrowReader::get_arrow_projection_mask( + &[1], + &iceberg_schema, + &parquet_schema, + &arrow_schema, + use_fallback, + ) + .unwrap(); + assert_eq!(mask, ProjectionMask::none(parquet_schema.num_columns())); + } + } + + #[tokio::test] + async fn test_fallback_projection_after_omitted_unknown_field() { + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "unknown", PrimitiveType::Unknown.into()).into(), + NestedField::optional(2, "known", PrimitiveType::Int.into()).into(), + ]) + .build() + .unwrap(), + ); + let file_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "known", + DataType::Int32, + true, + )])); + let file_batch = + RecordBatch::try_new(file_schema.clone(), vec![Arc::new(Int32Array::from(vec![ + 10, 20, + ]))]) + .unwrap(); + let tmp_dir = TempDir::new().unwrap(); + let file_path = tmp_dir.path().join("unknown-fallback.parquet"); + let file = File::create(&file_path).unwrap(); + let mut writer = ArrowWriter::try_new(file, file_schema, None).unwrap(); + writer.write(&file_batch).unwrap(); + writer.close().unwrap(); + + let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build(); + let tasks = Box::pin(futures::stream::iter(vec![Ok(FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(file_path.to_string_lossy().into_owned()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema) + .with_project_field_ids(vec![1, 2]) + .with_case_sensitive(false) + .build())])) as FileScanTaskStream; + + let batches = reader + .read(tasks) + .unwrap() + .stream() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].column(0).data_type(), &DataType::Null); + assert_eq!(batches[0].column(0).logical_null_count(), 2); + let known = batches[0] + .column(1) + .as_primitive::(); + assert_eq!(known.values(), &[10, 20]); + } + + #[tokio::test] + async fn test_fallback_projection_preserves_ids_after_unknown_reordering() { + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "a", PrimitiveType::Int.into()).into(), + NestedField::optional(3, "unknown", PrimitiveType::Unknown.into()).into(), + NestedField::optional(2, "b", PrimitiveType::Int.into()).into(), + ]) + .build() + .unwrap(), + ); + let file_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + let file_batch = RecordBatch::try_new(file_schema.clone(), vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![10, 20])), + ]) + .unwrap(); + let tmp_dir = TempDir::new().unwrap(); + let file_path = tmp_dir.path().join("unknown-reordered-fallback.parquet"); + let file = File::create(&file_path).unwrap(); + let mut writer = ArrowWriter::try_new(file, file_schema, None).unwrap(); + writer.write(&file_batch).unwrap(); + writer.close().unwrap(); + + let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build(); + let tasks = Box::pin(futures::stream::iter(vec![Ok(FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(file_path.to_string_lossy().into_owned()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema) + .with_project_field_ids(vec![1, 3, 2]) + .with_case_sensitive(false) + .build())])) as FileScanTaskStream; + + let batches = reader + .read(tasks) + .unwrap() + .stream() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(batches.len(), 1); + let a = batches[0] + .column(0) + .as_primitive::(); + assert_eq!(a.values(), &[1, 2]); + assert_eq!(batches[0].column(1).data_type(), &DataType::Null); + assert_eq!(batches[0].column(1).logical_null_count(), 2); + let b = batches[0] + .column(2) + .as_primitive::(); + assert_eq!(b.values(), &[10, 20]); + } + + #[tokio::test] + async fn test_fallback_projection_does_not_reuse_deleted_field_name() { + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional(2, "reused", PrimitiveType::Int.into()).into(), + ]) + .build() + .unwrap(), + ); + let file_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "reused", + DataType::Int32, + true, + )])); + let file_batch = + RecordBatch::try_new(file_schema.clone(), vec![Arc::new(Int32Array::from(vec![ + 10, 20, + ]))]) + .unwrap(); + let tmp_dir = TempDir::new().unwrap(); + let file_path = tmp_dir.path().join("reused-name-fallback.parquet"); + let file = File::create(&file_path).unwrap(); + let mut writer = ArrowWriter::try_new(file, file_schema, None).unwrap(); + writer.write(&file_batch).unwrap(); + writer.close().unwrap(); + + let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build(); + let tasks = Box::pin(futures::stream::iter(vec![Ok(FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(file_path.to_string_lossy().into_owned()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema) + .with_project_field_ids(vec![2]) + .with_case_sensitive(false) + .build())])) as FileScanTaskStream; + + let error = reader + .read(tasks) + .unwrap() + .stream() + .try_collect::>() + .await + .unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::DataInvalid); + assert!( + error + .message() + .contains("Cannot safely apply positional fallback"), + "unexpected error: {error}" + ); + } + + #[tokio::test] + async fn test_position_fallback_nested_reused_name_uses_default() { + use arrow_array::StructArray; + use arrow_schema::Fields; + + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "nested", + Type::Struct(StructType::new(vec![ + NestedField::optional(2, "known", PrimitiveType::Int.into()).into(), + NestedField::optional(4, "unknown", PrimitiveType::Unknown.into()) + .into(), + NestedField::optional(3, "reused", PrimitiveType::Int.into()) + .with_initial_default(crate::spec::Literal::int(99)) + .into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let source_fields = Fields::from(vec![Field::new("reused", DataType::Int32, true)]); + let file_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "nested", + DataType::Struct(source_fields.clone()), + true, + )])); + let file_batch = + RecordBatch::try_new(file_schema.clone(), vec![Arc::new(StructArray::new( + source_fields, + vec![Arc::new(Int32Array::from(vec![10, 20]))], + None, + ))]) + .unwrap(); + let tmp_dir = TempDir::new().unwrap(); + let file_path = tmp_dir.path().join("nested-reused-name-fallback.parquet"); + let file = File::create(&file_path).unwrap(); + let mut writer = ArrowWriter::try_new(file, file_schema, None).unwrap(); + writer.write(&file_batch).unwrap(); + writer.close().unwrap(); + + let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build(); + let tasks = Box::pin(futures::stream::iter(vec![Ok(FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(file_path.to_string_lossy().into_owned()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema) + .with_project_field_ids(vec![1]) + .with_case_sensitive(false) + .build())])) as FileScanTaskStream; + + let batches = reader + .read(tasks) + .unwrap() + .stream() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(batches.len(), 1); + let nested = batches[0].column(0).as_struct(); + let reused = nested + .column(2) + .as_primitive::(); + assert_eq!(reused.values(), &[99, 99]); + } + + #[tokio::test] + async fn test_unknown_filter_ignores_physical_column() { + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "unknown", PrimitiveType::Unknown.into()).into(), + NestedField::optional(2, "known", PrimitiveType::Int.into()).into(), + ]) + .build() + .unwrap(), + ); + let file_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("unknown", DataType::Int32, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + Field::new("known", DataType::Int32, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "2".to_string(), + )])), + ])); + let file_batch = RecordBatch::try_new(file_schema.clone(), vec![ + Arc::new(Int32Array::from(vec![7, 8])), + Arc::new(Int32Array::from(vec![10, 20])), + ]) + .unwrap(); + let tmp_dir = TempDir::new().unwrap(); + let file_path = tmp_dir.path().join("physical-unknown.parquet"); + let file = File::create(&file_path).unwrap(); + let mut writer = ArrowWriter::try_new(file, file_schema, None).unwrap(); + writer.write(&file_batch).unwrap(); + writer.close().unwrap(); + + let predicate = Reference::new("unknown").is_null(); + let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) + .with_row_group_filtering_enabled(true) + .with_row_selection_enabled(true) + .build(); + let tasks = Box::pin(futures::stream::iter(vec![Ok(FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(file_path.to_string_lossy().into_owned()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema.clone()) + .with_project_field_ids(vec![1, 2]) + .with_case_sensitive(false) + .with_predicate(Some(predicate.bind(schema, true).unwrap())) + .build())])) as FileScanTaskStream; + + let batches = reader + .read(tasks) + .unwrap() + .stream() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 2); + assert_eq!(batches[0].column(0).data_type(), &DataType::Null); + assert_eq!(batches[0].column(0).logical_null_count(), 2); + let known = batches[0] + .column(1) + .as_primitive::(); + assert_eq!(known.values(), &[10, 20]); + } + #[test] fn test_arrow_projection_mask() { let schema = Arc::new( diff --git a/crates/iceberg/src/arrow/record_batch_transformer.rs b/crates/iceberg/src/arrow/record_batch_transformer.rs index b06c656e54..d13de70754 100644 --- a/crates/iceberg/src/arrow/record_batch_transformer.rs +++ b/crates/iceberg/src/arrow/record_batch_transformer.rs @@ -20,8 +20,8 @@ use std::sync::Arc; use arrow_arith::boolean::is_not_null; use arrow_array::{ - Array as ArrowArray, ArrayRef, Int32Array, Int64Array, RecordBatch, RecordBatchOptions, - RunArray, StructArray, + Array as ArrowArray, ArrayRef, Int32Array, Int64Array, ListArray, MapArray, RecordBatch, + RecordBatchOptions, RunArray, StructArray, new_null_array, }; use arrow_cast::cast; use arrow_schema::{ @@ -31,7 +31,10 @@ use arrow_schema::{ use arrow_select::zip::zip; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; -use crate::arrow::value::{create_primitive_array_repeated, create_primitive_array_single_element}; +use crate::arrow::value::{ + create_literal_array_repeated, create_primitive_array_repeated, + create_primitive_array_single_element, +}; use crate::arrow::{ datum_to_arrow_type_with_ree, primitive_type_to_arrow_type_with_ree, schema_to_arrow_schema, type_to_arrow_type, @@ -161,7 +164,7 @@ pub(crate) enum ColumnSource { // a preceding operation. Add { target_type: DataType, - value: Option, + value: Option, }, // A struct column where each child is a constant primitive value. @@ -237,6 +240,7 @@ pub(crate) struct RecordBatchTransformerBuilder { projected_iceberg_field_ids: Vec, constant_fields: HashMap, virtual_fields: HashSet, + use_position_fallback: bool, } /// How a metadata (or identity-partition) column's values are supplied. @@ -311,9 +315,16 @@ impl RecordBatchTransformerBuilder { projected_iceberg_field_ids: projected_iceberg_field_ids.to_vec(), constant_fields: HashMap::new(), virtual_fields: HashSet::new(), + use_position_fallback: false, } } + /// Use positional, rather than name-first, matching for nested fields without IDs. + pub(crate) fn with_position_fallback(mut self, use_position_fallback: bool) -> Self { + self.use_position_fallback = use_position_fallback; + self + } + /// Add a scalar constant value for a specific field ID. /// This is used for virtual/metadata fields like _file that have constant values per batch. pub(crate) fn with_constant(mut self, field_id: i32, datum: Datum) -> Self { @@ -404,6 +415,7 @@ impl RecordBatchTransformerBuilder { projected_iceberg_field_ids: self.projected_iceberg_field_ids, constant_fields: self.constant_fields, virtual_fields: self.virtual_fields, + use_position_fallback: self.use_position_fallback, batch_transform: None, } } @@ -451,12 +463,211 @@ pub(crate) struct RecordBatchTransformer { // Iceberg projection rules (name mapping / initial-default / null) virtual_fields: HashSet, + // True for the no-ID/no-name-mapping reader branch. Nested fields in that branch must follow + // positional fallback semantics instead of inferring identity from current names. + use_position_fallback: bool, + // BatchTransform gets lazily constructed based on the schema of // the first RecordBatch we receive from the file batch_transform: Option, } impl RecordBatchTransformer { + fn data_type_contains_null(data_type: &DataType) -> bool { + match data_type { + DataType::Null => true, + DataType::Struct(fields) => fields + .iter() + .any(|field| Self::data_type_contains_null(field.data_type())), + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) => Self::data_type_contains_null(field.data_type()), + DataType::Map(entries, _) => Self::data_type_contains_null(entries.data_type()), + _ => false, + } + } + + fn transform_array(&self, array: &ArrayRef, target_type: &DataType) -> Result { + if array.data_type().equals_datatype(target_type) { + return Ok(array.clone()); + } + + match target_type { + DataType::Struct(target_fields) => { + let source = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Expected struct array while transforming {} to {target_type}", + array.data_type() + ), + ) + })?; + let source_has_field_ids = source + .fields() + .iter() + .any(|field| field.metadata().contains_key(PARQUET_FIELD_ID_META_KEY)); + let columns = target_fields + .iter() + .enumerate() + .map(|(target_index, target_field)| { + let target_field_id = + target_field.metadata().get(PARQUET_FIELD_ID_META_KEY); + let source_index = target_field_id + .and_then(|target_field_id| { + source.fields().iter().position(|source_field| { + source_field.metadata().get(PARQUET_FIELD_ID_META_KEY) + == Some(target_field_id) + }) + }) + .or_else(|| { + if source_has_field_ids { + return None; + } + + let position = + (!matches!(target_field.data_type(), DataType::Null)).then( + || { + target_fields + .iter() + .take(target_index) + .filter(|field| { + !matches!(field.data_type(), DataType::Null) + }) + .count() + }, + ); + if self.use_position_fallback { + position + } else { + source + .fields() + .iter() + .position(|source_field| { + source_field.name() == target_field.name() + }) + .or(position) + } + }) + .filter(|source_index| *source_index < source.num_columns()); + + match source_index { + Some(source_index) => self.transform_array( + source.column(source_index), + target_field.data_type(), + ), + None => self.create_missing_nested_column(target_field, source.len()), + } + }) + .collect::>>()?; + + Ok(Arc::new(StructArray::try_new_with_length( + target_fields.clone(), + columns, + source.nulls().cloned(), + source.len(), + )?)) + } + DataType::List(target_element) => { + let source = array.as_any().downcast_ref::().ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Expected list array while transforming {} to {target_type}", + array.data_type() + ), + ) + })?; + let values = self.transform_array(source.values(), target_element.data_type())?; + Ok(Arc::new(ListArray::try_new( + target_element.clone(), + source.offsets().clone(), + values, + source.nulls().cloned(), + )?)) + } + DataType::Map(target_entries, ordered) => { + let source = array.as_any().downcast_ref::().ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Expected map array while transforming {} to {target_type}", + array.data_type() + ), + ) + })?; + let source_entries: ArrayRef = Arc::new(source.entries().clone()); + let entries = self.transform_array(&source_entries, target_entries.data_type())?; + let entries = entries + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "Transformed map entries are not a struct array", + ) + })? + .clone(); + Ok(Arc::new(MapArray::try_new( + target_entries.clone(), + source.offsets().clone(), + entries, + source.nulls().cloned(), + *ordered, + )?)) + } + _ => Ok(cast(array.as_ref(), target_type)?), + } + } + + fn create_missing_nested_column( + &self, + target_field: &FieldRef, + num_rows: usize, + ) -> Result { + let Some(field_id) = target_field.metadata().get(PARQUET_FIELD_ID_META_KEY) else { + return if matches!(target_field.data_type(), DataType::Null) { + Ok(new_null_array(target_field.data_type(), num_rows)) + } else { + Err(Error::new( + ErrorKind::DataInvalid, + format!( + "Field {} is missing while transforming nested struct", + target_field.name() + ), + )) + }; + }; + let field_id = field_id.parse::().map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("field id not parseable as an i32: {e}"), + ) + })?; + let iceberg_field = self.snapshot_schema.field_by_id(field_id).ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!("Field {field_id} not found in snapshot schema"), + ) + })?; + + if iceberg_field.initial_default.is_none() && iceberg_field.required { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("Missing required field: {}", iceberg_field.name), + )); + } + + create_literal_array_repeated( + target_field.data_type(), + iceberg_field.initial_default.as_ref(), + num_rows, + ) + } + pub(crate) fn process_record_batch( &mut self, record_batch: RecordBatch, @@ -472,7 +683,11 @@ impl RecordBatchTransformer { .with_row_count(Some(record_batch.num_rows())); RecordBatch::try_new_with_options( Arc::clone(target_schema), - self.transform_columns(record_batch.columns(), operations)?, + self.transform_columns( + record_batch.columns(), + operations, + record_batch.num_rows(), + )?, &options, )? } @@ -748,7 +963,7 @@ impl RecordBatchTransformer { }; return Ok(ColumnSource::Add { - value: Some(datum.literal().clone()), + value: Some(Literal::Primitive(datum.literal().clone())), target_type: arrow_type, }); } @@ -864,16 +1079,8 @@ impl RecordBatchTransformer { )); } - let default_value = iceberg_field.initial_default.as_ref().and_then(|lit| { - if let Literal::Primitive(prim) = lit { - Some(prim.clone()) - } else { - None - } - }); - ColumnSource::Add { - value: default_value, + value: iceberg_field.initial_default.clone(), target_type: target_type.clone(), } }; @@ -910,25 +1117,28 @@ impl RecordBatchTransformer { &self, columns: &[Arc], operations: &[ColumnSource], + num_rows: usize, ) -> Result>> { - if columns.is_empty() { - return Ok(columns.to_vec()); - } - let num_rows = columns[0].len(); - operations .iter() .map(|op| { Ok(match op { ColumnSource::PassThrough { source_index } => columns[*source_index].clone(), + ColumnSource::Promote { + target_type, + source_index, + } if Self::data_type_contains_null(target_type) => { + self.transform_array(&columns[*source_index], target_type)? + } + ColumnSource::Promote { target_type, source_index, } => cast(&*columns[*source_index], target_type)?, ColumnSource::Add { target_type, value } => { - Self::create_column(target_type, value, num_rows)? + Self::create_column(target_type, value.as_ref(), num_rows)? } ColumnSource::AddStructConstant { @@ -981,11 +1191,22 @@ impl RecordBatchTransformer { fn create_column( target_type: &DataType, - prim_lit: &Option, + literal: Option<&Literal>, num_rows: usize, ) -> Result { // Check if this is a RunEndEncoded type (for constant fields) if let DataType::RunEndEncoded(_, values_field) = target_type { + let prim_lit = match literal { + Some(Literal::Primitive(value)) => Some(value.clone()), + None => None, + Some(value) => { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("Run-end encoded constant must be primitive, got {value:?}"), + )); + } + }; + // Helper to create a Run-End Encoded array let create_ree_array = |values_array: ArrayRef| -> Result { let run_ends = if num_rows == 0 { @@ -1006,13 +1227,13 @@ impl RecordBatchTransformer { // Create the values array using the helper function let values_array = - create_primitive_array_single_element(values_field.data_type(), prim_lit)?; + create_primitive_array_single_element(values_field.data_type(), &prim_lit)?; // Wrap in Run-End Encoding create_ree_array(values_array) } else { // Non-REE type (simple arrays for non-constant fields) - create_primitive_array_repeated(target_type, prim_lit, num_rows) + create_literal_array_repeated(target_type, literal, num_rows) } } @@ -1106,18 +1327,21 @@ mod test { use std::sync::Arc; use arrow_array::{ - Array, Date32Array, Float32Array, Float64Array, Int32Array, Int64Array, RecordBatch, - StringArray, + Array, Date32Array, Float32Array, Float64Array, Int32Array, Int64Array, ListArray, + MapArray, RecordBatch, RecordBatchOptions, StringArray, StructArray, }; use arrow_cast::cast; - use arrow_schema::{DataType, Field, Schema as ArrowSchema}; + use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema}; use super::field_with_id; use crate::arrow::build_partition_constant; use crate::arrow::record_batch_transformer::{ RecordBatchTransformer, RecordBatchTransformerBuilder, }; - use crate::spec::{Literal, NestedField, PrimitiveType, Schema, Struct, Type}; + use crate::spec::{ + ListType, Literal, Map as MapValue, MapType, NestedField, PrimitiveType, Schema, Struct, + Type, + }; /// Helper to extract string values from either StringArray or RunEndEncoded /// Returns empty string for null values @@ -1219,6 +1443,453 @@ mod test { assert_eq!(result, expected); } + #[test] + fn processor_materializes_unknown_from_zero_column_batch() { + let snapshot_schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "unknown", PrimitiveType::Unknown.into()).into(), + ]) + .build() + .unwrap(), + ); + let options = RecordBatchOptions::default().with_row_count(Some(2)); + let source_batch = + RecordBatch::try_new_with_options(Arc::new(ArrowSchema::empty()), vec![], &options) + .unwrap(); + let mut transformer = RecordBatchTransformerBuilder::new(snapshot_schema, &[1]).build(); + + let result = transformer.process_record_batch(source_batch).unwrap(); + + assert_eq!(result.num_rows(), 2); + assert_eq!(result.num_columns(), 1); + assert_eq!(result.column(0).data_type(), &DataType::Null); + assert_eq!(result.column(0).logical_null_count(), 2); + } + + #[test] + fn processor_rebuilds_unknown_children_omitted_from_parquet_struct() { + let snapshot_schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "nested", + Type::Struct(crate::spec::StructType::new(vec![ + NestedField::optional(2, "known", PrimitiveType::Int.into()).into(), + NestedField::optional(3, "unknown", PrimitiveType::Unknown.into()) + .into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let source_fields = Fields::from(vec![field_with_id("known", DataType::Int32, true, 2)]); + let source_struct = Arc::new(StructArray::new( + source_fields.clone(), + vec![Arc::new(Int32Array::from(vec![Some(1), Some(2)]))], + None, + )); + let source_schema = Arc::new(ArrowSchema::new(vec![field_with_id( + "nested", + DataType::Struct(source_fields), + true, + 1, + )])); + let source_batch = RecordBatch::try_new(source_schema, vec![source_struct]).unwrap(); + let mut transformer = RecordBatchTransformerBuilder::new(snapshot_schema, &[1]).build(); + + let result = transformer.process_record_batch(source_batch).unwrap(); + + let nested = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(nested.num_columns(), 2); + assert_eq!(nested.fields()[1].data_type(), &DataType::Null); + assert_eq!(nested.column(1).logical_null_count(), 2); + } + + #[test] + fn processor_materializes_missing_nested_children_during_unknown_reconstruction() { + let snapshot_schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "nested", + Type::Struct(crate::spec::StructType::new(vec![ + NestedField::optional(2, "known", PrimitiveType::Int.into()).into(), + NestedField::optional(3, "unknown", PrimitiveType::Unknown.into()) + .into(), + NestedField::optional( + 4, + "added_optional", + PrimitiveType::String.into(), + ) + .into(), + NestedField::required(5, "added_default", PrimitiveType::Long.into()) + .with_initial_default(Literal::long(7)) + .into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let source_fields = Fields::from(vec![field_with_id("known", DataType::Int32, true, 2)]); + let source_struct = Arc::new(StructArray::new( + source_fields.clone(), + vec![Arc::new(Int32Array::from(vec![Some(1), Some(2)]))], + None, + )); + let source_schema = Arc::new(ArrowSchema::new(vec![field_with_id( + "nested", + DataType::Struct(source_fields), + true, + 1, + )])); + let source_batch = RecordBatch::try_new(source_schema, vec![source_struct]).unwrap(); + let mut transformer = RecordBatchTransformerBuilder::new(snapshot_schema, &[1]).build(); + + let result = transformer.process_record_batch(source_batch).unwrap(); + + let nested = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(nested.num_columns(), 4); + assert_eq!(nested.column(1).data_type(), &DataType::Null); + assert_eq!(nested.column(1).logical_null_count(), 2); + assert_eq!(nested.column(2).logical_null_count(), 2); + let added_default = nested + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(added_default.values(), &[7, 7]); + } + + #[test] + fn processor_position_fallback_does_not_prefer_reused_nested_name() { + let snapshot_schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "nested", + Type::Struct(crate::spec::StructType::new(vec![ + NestedField::optional(2, "known", PrimitiveType::Int.into()).into(), + NestedField::optional(4, "unknown", PrimitiveType::Unknown.into()) + .into(), + NestedField::optional(3, "reused", PrimitiveType::Int.into()) + .with_initial_default(Literal::int(99)) + .into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let source_fields = Fields::from(vec![Field::new("reused", DataType::Int32, true)]); + let source_struct = Arc::new(StructArray::new( + source_fields.clone(), + vec![Arc::new(Int32Array::from(vec![10, 20]))], + None, + )); + let source_schema = Arc::new(ArrowSchema::new(vec![field_with_id( + "nested", + DataType::Struct(source_fields), + true, + 1, + )])); + let source_batch = RecordBatch::try_new(source_schema, vec![source_struct]).unwrap(); + let mut transformer = RecordBatchTransformerBuilder::new(snapshot_schema, &[1]) + .with_position_fallback(true) + .build(); + + let result = transformer.process_record_batch(source_batch).unwrap(); + + let nested = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let reused = nested + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(reused.values(), &[99, 99]); + } + + #[test] + fn processor_materializes_nested_container_defaults() { + let struct_default = Type::Struct(crate::spec::StructType::new(vec![ + NestedField::optional(5, "value", PrimitiveType::Int.into()).into(), + ])); + let list_default = Type::List(ListType::new( + NestedField::list_element(7, PrimitiveType::Int.into(), false).into(), + )); + let map_default = Type::Map(MapType::optional( + 9, + PrimitiveType::String.into(), + 10, + PrimitiveType::Int.into(), + )); + let snapshot_schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "nested", + Type::Struct(crate::spec::StructType::new(vec![ + NestedField::optional(2, "known", PrimitiveType::Int.into()).into(), + NestedField::optional(3, "unknown", PrimitiveType::Unknown.into()) + .into(), + NestedField::required(4, "added_struct", struct_default) + .with_initial_default(Literal::Struct(Struct::from_iter([Some( + Literal::int(8), + )]))) + .into(), + NestedField::required(6, "added_list", list_default) + .with_initial_default(Literal::List(vec![ + Some(Literal::int(1)), + Some(Literal::int(2)), + ])) + .into(), + NestedField::required(8, "added_map", map_default) + .with_initial_default(Literal::Map(MapValue::from_iter([( + Literal::string("key"), + Some(Literal::int(9)), + )]))) + .into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let source_fields = Fields::from(vec![field_with_id("known", DataType::Int32, true, 2)]); + let source_struct = Arc::new(StructArray::new( + source_fields.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + None, + )); + let source_schema = Arc::new(ArrowSchema::new(vec![field_with_id( + "nested", + DataType::Struct(source_fields), + true, + 1, + )])); + let source_batch = RecordBatch::try_new(source_schema, vec![source_struct]).unwrap(); + let mut transformer = RecordBatchTransformerBuilder::new(snapshot_schema, &[1]).build(); + + let result = transformer.process_record_batch(source_batch).unwrap(); + + let nested = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let added_struct = nested + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let struct_values = added_struct + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(struct_values.values(), &[8, 8]); + + let added_list = nested + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(added_list.value_offsets(), &[0, 2, 4]); + let list_values = added_list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(list_values.values(), &[1, 2, 1, 2]); + + let added_map = nested + .column(4) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(added_map.value_offsets(), &[0, 1, 2]); + let map_keys = added_map + .keys() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(map_keys.iter().collect::>(), vec![Some("key"); 2]); + let map_values = added_map + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(map_values.values(), &[9, 9]); + } + + #[test] + fn processor_materializes_top_level_container_defaults() { + let snapshot_schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "known", PrimitiveType::Int.into()).into(), + NestedField::required( + 2, + "added_struct", + Type::Struct(crate::spec::StructType::new(vec![ + NestedField::required(3, "value", PrimitiveType::Int.into()).into(), + ])), + ) + .with_initial_default(Literal::Struct(Struct::from_iter([Some(Literal::int( + 8, + ))]))) + .into(), + NestedField::required( + 4, + "added_list", + Type::List(ListType::new( + NestedField::list_element(5, PrimitiveType::Int.into(), true).into(), + )), + ) + .with_initial_default(Literal::List(vec![ + Some(Literal::int(1)), + Some(Literal::int(2)), + ])) + .into(), + NestedField::required( + 6, + "added_map", + Type::Map(MapType::required( + 7, + PrimitiveType::String.into(), + 8, + PrimitiveType::Int.into(), + )), + ) + .with_initial_default(Literal::Map(MapValue::from_iter([( + Literal::string("key"), + Some(Literal::int(9)), + )]))) + .into(), + ]) + .build() + .unwrap(), + ); + let source_schema = Arc::new(ArrowSchema::new(vec![field_with_id( + "known", + DataType::Int32, + false, + 1, + )])); + let source_batch = + RecordBatch::try_new(source_schema, vec![Arc::new(Int32Array::from(vec![1, 2]))]) + .unwrap(); + let mut transformer = + RecordBatchTransformerBuilder::new(snapshot_schema, &[1, 2, 4, 6]).build(); + + let result = transformer.process_record_batch(source_batch).unwrap(); + + let added_struct = result + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let struct_values = added_struct + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(struct_values.values(), &[8, 8]); + + let added_list = result + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(added_list.value_offsets(), &[0, 2, 4]); + let list_values = added_list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(list_values.values(), &[1, 2, 1, 2]); + + let added_map = result + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(added_map.value_offsets(), &[0, 1, 2]); + let map_keys = added_map + .keys() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(map_keys.iter().collect::>(), vec![Some("key"); 2]); + let map_values = added_map + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(map_values.values(), &[9, 9]); + } + + #[test] + fn processor_does_not_reuse_nested_name_when_field_ids_differ() { + let snapshot_schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "nested", + Type::Struct(crate::spec::StructType::new(vec![ + NestedField::required(3, "reused", PrimitiveType::Long.into()).into(), + NestedField::optional(4, "unknown", PrimitiveType::Unknown.into()) + .into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let source_fields = Fields::from(vec![field_with_id("reused", DataType::Int32, true, 2)]); + let source_struct = Arc::new(StructArray::new( + source_fields.clone(), + vec![Arc::new(Int32Array::from(vec![Some(10), Some(20)]))], + None, + )); + let source_schema = Arc::new(ArrowSchema::new(vec![field_with_id( + "nested", + DataType::Struct(source_fields), + true, + 1, + )])); + let source_batch = RecordBatch::try_new(source_schema, vec![source_struct]).unwrap(); + let mut transformer = RecordBatchTransformerBuilder::new(snapshot_schema, &[1]).build(); + + let error = transformer.process_record_batch(source_batch).unwrap_err(); + + assert!(error.message().contains("Missing required field: reused")); + } + #[test] fn schema_evolution_adds_date_column_with_nulls() { // Reproduces TestSelect.readAndWriteWithBranchAfterSchemaChange from iceberg-spark. @@ -1399,7 +2070,7 @@ mod test { let struct_column = result .column(2) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); assert!(struct_column.is_null(0)); assert!(struct_column.is_null(1)); diff --git a/crates/iceberg/src/arrow/value.rs b/crates/iceberg/src/arrow/value.rs index 0de3b5650a..43d4eab2ad 100644 --- a/crates/iceberg/src/arrow/value.rs +++ b/crates/iceberg/src/arrow/value.rs @@ -21,10 +21,13 @@ use arrow_array::{ Array, ArrayRef, BinaryArray, BooleanArray, Date32Array, Decimal128Array, FixedSizeBinaryArray, FixedSizeListArray, Float32Array, Float64Array, Int32Array, Int64Array, LargeBinaryArray, LargeListArray, LargeStringArray, ListArray, MapArray, StringArray, StructArray, - Time64MicrosecondArray, TimestampMicrosecondArray, TimestampNanosecondArray, new_null_array, + Time64MicrosecondArray, TimestampMicrosecondArray, TimestampNanosecondArray, UInt32Array, + new_empty_array, new_null_array, }; -use arrow_buffer::NullBuffer; +use arrow_buffer::{NullBuffer, OffsetBuffer}; use arrow_schema::{DataType, FieldRef, TimeUnit}; +use arrow_select::concat::concat; +use arrow_select::take::take; use uuid::Uuid; use super::get_field_id_from_metadata; @@ -958,6 +961,139 @@ pub(crate) fn create_primitive_array_repeated( }) } +/// Create an array by repeating an Iceberg literal, including nested container literals. +pub(crate) fn create_literal_array_repeated( + data_type: &DataType, + literal: Option<&Literal>, + num_rows: usize, +) -> Result { + let Some(literal) = literal else { + return create_primitive_array_repeated(data_type, &None, num_rows); + }; + + if let Literal::Primitive(primitive) = literal { + return create_primitive_array_repeated(data_type, &Some(primitive.clone()), num_rows); + } + + let single = create_literal_array_single(data_type, literal)?; + if num_rows == 1 { + return Ok(single); + } + let indices = UInt32Array::from(vec![0; num_rows]); + Ok(take(single.as_ref(), &indices, None)?) +} + +fn create_literal_array_single(data_type: &DataType, literal: &Literal) -> Result { + match (data_type, literal) { + (DataType::Struct(fields), Literal::Struct(value)) => { + if fields.len() != value.fields().len() { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "Struct default has {} values but Arrow type has {} fields", + value.fields().len(), + fields.len() + ), + )); + } + + let columns = fields + .iter() + .zip(value.iter()) + .map(|(field, value)| match value { + Some(value) => create_literal_array_single(field.data_type(), value), + None => Ok(new_null_array(field.data_type(), 1)), + }) + .collect::>>()?; + Ok(Arc::new(StructArray::try_new( + fields.clone(), + columns, + None, + )?)) + } + (DataType::List(element), Literal::List(values)) => { + let values = + create_literal_values(element.data_type(), values.iter().map(Option::as_ref))?; + Ok(Arc::new(ListArray::try_new( + element.clone(), + OffsetBuffer::from_lengths([values.len()]), + values, + None, + )?)) + } + (DataType::LargeList(element), Literal::List(values)) => { + let values = + create_literal_values(element.data_type(), values.iter().map(Option::as_ref))?; + Ok(Arc::new(LargeListArray::try_new( + element.clone(), + OffsetBuffer::from_lengths([values.len()]), + values, + None, + )?)) + } + (DataType::Map(entries, ordered), Literal::Map(value)) => { + let DataType::Struct(entry_fields) = entries.data_type() else { + return Err(Error::new( + ErrorKind::Unexpected, + "Map entries field must be a struct", + )); + }; + if entry_fields.len() != 2 { + return Err(Error::new( + ErrorKind::Unexpected, + "Map entries struct must contain key and value fields", + )); + } + + let keys = create_literal_values( + entry_fields[0].data_type(), + value.iter().map(|(key, _)| Some(key)), + )?; + let values = create_literal_values( + entry_fields[1].data_type(), + value.iter().map(|(_, value)| value.as_ref()), + )?; + let entries_array = + StructArray::try_new(entry_fields.clone(), vec![keys, values], None)?; + Ok(Arc::new(MapArray::try_new( + entries.clone(), + OffsetBuffer::from_lengths([value.len()]), + entries_array, + None, + *ordered, + )?)) + } + (_, Literal::Primitive(primitive)) => { + create_primitive_array_single_element(data_type, &Some(primitive.clone())) + } + _ => Err(Error::new( + ErrorKind::DataInvalid, + format!("Default literal {literal:?} does not match Arrow type {data_type}"), + )), + } +} + +fn create_literal_values<'a>( + data_type: &DataType, + values: impl IntoIterator>, +) -> Result { + let arrays = values + .into_iter() + .map(|value| match value { + Some(value) => create_literal_array_single(data_type, value), + None => Ok(new_null_array(data_type, 1)), + }) + .collect::>>()?; + if arrays.is_empty() { + return Ok(new_empty_array(data_type)); + } + let arrays = arrays + .iter() + .map(|array| array.as_ref()) + .collect::>(); + Ok(concat(&arrays)?) +} + #[cfg(test)] mod test { use std::collections::HashMap;