mbutrovich commented on code in PR #3255:
URL: https://github.com/apache/iceberg-rust/pull/3255#discussion_r4075216628
##########
crates/iceberg/src/arrow/record_batch_transformer.rs:
##########
@@ -192,6 +194,293 @@ pub(crate) enum ColumnSource {
// post-processing step by using the projection mask.
}
+/// Reconciles a source array to the target type by matching nested fields on
+/// field id rather than position, mirroring iceberg-java's nested readers.
+///
+/// The plan is resolved once per file, when the `BatchTransform` is built, so
+/// applying it per batch is just index lookups and array assembly.
+#[derive(Debug)]
+pub(crate) enum PromotePlan {
+ PassThrough,
+ Cast(DataType),
+ Struct {
+ fields: Fields,
+ children: Vec<ChildPlan>,
+ },
+ List {
+ field: FieldRef,
+ element: Box<PromotePlan>,
+ },
+ LargeList {
+ field: FieldRef,
+ element: Box<PromotePlan>,
+ },
+ Map {
+ field: FieldRef,
+ entry_fields: Fields,
+ entries: Vec<ChildPlan>,
+ sorted: bool,
+ },
+}
+
+#[derive(Debug)]
+pub(crate) enum ChildPlan {
+ FromSource {
+ source_index: usize,
+ plan: PromotePlan,
+ },
+ Null {
+ target_type: DataType,
+ },
+}
+
+impl PromotePlan {
+ fn build(
+ source: &DataType,
+ target: &DataType,
+ snapshot_schema: &IcebergSchema,
+ ) -> Result<Self> {
+ if source == target {
+ return Ok(PromotePlan::PassThrough);
+ }
+ match (source, target) {
+ (DataType::Struct(source_fields), DataType::Struct(target_fields))
=> {
+ Ok(PromotePlan::Struct {
+ children: Self::build_struct_children(
+ source_fields,
+ target_fields,
+ snapshot_schema,
+ )?,
+ fields: target_fields.clone(),
+ })
+ }
+ (DataType::List(source_field), DataType::List(target_field)) =>
Ok(PromotePlan::List {
+ element: Box::new(Self::build(
+ source_field.data_type(),
+ target_field.data_type(),
+ snapshot_schema,
+ )?),
+ field: target_field.clone(),
+ }),
+ (DataType::LargeList(source_field),
DataType::LargeList(target_field)) => {
+ Ok(PromotePlan::LargeList {
+ element: Box::new(Self::build(
+ source_field.data_type(),
+ target_field.data_type(),
+ snapshot_schema,
+ )?),
+ field: target_field.clone(),
+ })
+ }
+ (DataType::Map(source_entries, _), DataType::Map(target_entries,
sorted)) => {
+ match (source_entries.data_type(), target_entries.data_type())
{
+ (DataType::Struct(source_fields),
DataType::Struct(target_fields)) => {
+ Ok(PromotePlan::Map {
+ entries: Self::build_struct_children(
+ source_fields,
+ target_fields,
+ snapshot_schema,
+ )?,
+ entry_fields: target_fields.clone(),
+ field: target_entries.clone(),
+ sorted: *sorted,
+ })
+ }
+ _ => Err(Error::new(
+ ErrorKind::Unexpected,
+ format!(
+ "expected struct-typed map entries, got
{source_entries:?} and {target_entries:?}"
+ ),
+ )),
+ }
+ }
+ (_, DataType::Struct(_) | DataType::List(_) |
DataType::LargeList(_))
+ | (_, DataType::Map(_, _)) => Err(Error::new(
+ ErrorKind::DataInvalid,
+ format!("cannot promote {source:?} to {target:?}"),
+ )),
+ _ => Ok(PromotePlan::Cast(target.clone())),
+ }
+ }
+
+ fn build_struct_children(
+ source_fields: &Fields,
+ target_fields: &Fields,
+ snapshot_schema: &IcebergSchema,
+ ) -> Result<Vec<ChildPlan>> {
+ let mut source_by_id = HashMap::with_capacity(source_fields.len());
+ for (idx, field) in source_fields.iter().enumerate() {
+ if let Some(id) = try_get_field_id_from_metadata(field)? {
+ source_by_id.insert(id, idx);
+ }
+ }
+ // Name mapping only assigns top-level ids. Fully id-less children
match by
+ // position when types line up. A same-type reorder cannot be detected
without
+ // ids and stays positional. Any missing id errors instead of nulling
that child.
+ if !source_fields.is_empty() && source_by_id.len() !=
source_fields.len() {
+ if source_by_id.is_empty()
+ && source_fields.len() == target_fields.len()
+ && source_fields
+ .iter()
+ .zip(target_fields.iter())
+ .all(|(s, t)| s.data_type().equals_datatype(t.data_type()))
+ {
Review Comment:
The spec doesn't define positional matching for nested fields. [Column
Projection](https://iceberg.apache.org/spec/#column-projection) resolves a
field id missing from a file through name mapping, and a name mapping carries
`fields` for struct children, map keys and values, and list elements. The
spec-correct fix for id-less nested structs is recursive name mapping, which
#1845 tracks. Until that lands, the positional path here is compatibility
behavior, and it should keep what `main` does today rather than fail reads that
work now.
The `equals_datatype` check came from my suggestion last round, and it fails
a read that works on `main`. Take a file with no field ids where `s` was
written as `struct<x: int>`, read through a table schema where `x` is now
`long`. On `main` that column takes the [positional
cast](https://github.com/apache/iceberg-rust/blob/bb1e4a4861f02377489eff818b75138f414c4cb0/crates/iceberg/src/arrow/record_batch_transformer.rs#L933)
and comes back as `struct<x: long>`. On this head, `process_record_batch`
fails with `DataInvalid => cannot reconcile struct fields by id: source fields
do not all have field ids`.
Could we keep the child-count check and let `build` decide each pair?
Primitive pairs then go through `Cast` as they do on `main`, and a struct or
list paired with an incompatible type still errors in `build`:
```suggestion
// Name mapping only assigns top-level ids (#1845). Until it
recurses, fully
// id-less children match by position when the child counts line up,
as they
// did before, and `build` checks each pair. A same-type reorder
cannot be
// detected without ids. Any missing id errors instead of nulling
that child.
if !source_fields.is_empty() && source_by_id.len() !=
source_fields.len() {
if source_by_id.is_empty() && source_fields.len() ==
target_fields.len() {
```
With this change, the `int` to `long` case returns `struct<x: long>` with
the original values, and the full `cargo test -p iceberg --lib` run still
passes. Please add a test next to `promote_idless_nested_struct_keeps_data`
that drives an id-less `struct<x: int>` through `transform_top_level` against a
`long` target. In the PR description, please update the bullet about id-less
structs whose types don't line up and point the recursive name mapping item
under Out of scope at #1845.
##########
crates/iceberg/src/arrow/record_batch_transformer.rs:
##########
@@ -1144,19 +1424,591 @@ mod test {
use std::collections::HashMap;
use std::sync::Arc;
+ use arrow_array::cast::AsArray;
+ use arrow_array::types::{Int32Type, Int64Type};
use arrow_array::{
- Array, Date32Array, Float32Array, Float64Array, Int32Array,
Int64Array, RecordBatch,
- StringArray,
+ Array, ArrayRef, Date32Array, Float32Array, Float64Array, Int32Array,
Int64Array,
+ LargeListArray, ListArray, MapArray, RecordBatch, StringArray,
StructArray,
};
+ use arrow_buffer::{NullBuffer, OffsetBuffer};
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,
+ PromotePlan, RecordBatchTransformer, RecordBatchTransformerBuilder,
};
- use crate::spec::{Literal, NestedField, PrimitiveType, Schema, Struct,
Type};
+ use crate::spec::{
+ Literal, MapType, NestedField, PrimitiveType, Schema, Struct,
StructType, Type,
+ };
+
+ fn promote(source: &ArrayRef, target: &DataType, schema: &Schema) ->
crate::Result<ArrayRef> {
+ PromotePlan::build(source.data_type(), target, schema)?.apply(source)
+ }
+
+ fn empty_schema() -> Schema {
+ Schema::builder().build().unwrap()
+ }
+
+ fn unevolved_struct_type() -> DataType {
+ DataType::Struct(Fields::from(vec![field_with_id(
+ "x",
+ DataType::Int32,
+ true,
+ 5,
+ )]))
+ }
+
+ fn evolved_struct_type() -> DataType {
+ DataType::Struct(Fields::from(vec![
+ field_with_id("x", DataType::Int32, true, 5),
+ field_with_id("y", DataType::Int32, true, 6),
+ ]))
+ }
+
+ fn unevolved_struct_data(x_values: Vec<i32>) -> Arc<StructArray> {
+ Arc::new(StructArray::new(
+ Fields::from(vec![field_with_id("x", DataType::Int32, true, 5)]),
+ vec![Arc::new(Int32Array::from(x_values)) as ArrayRef],
+ None,
+ ))
+ }
+
+ fn assert_existing_field_kept(s: &StructArray, expected_existing: &[i32]) {
+ assert_eq!(
+ s.column(0).as_primitive::<Int32Type>().values(),
+ expected_existing
+ );
+ }
+
+ fn assert_added_field_null(s: &StructArray) {
+ assert_eq!(s.column(1).null_count(), s.len());
+ }
+
+ fn transform_top_level(column: NestedField, file_column: ArrayRef) ->
crate::Result<ArrayRef> {
+ let name = column.name.clone();
+ let snapshot_schema = Arc::new(
+ Schema::builder()
+ .with_schema_id(1)
+ .with_fields(vec![
+ NestedField::required(1, "id",
Type::Primitive(PrimitiveType::Int)).into(),
+ column.into(),
+ ])
+ .build()
+ .unwrap(),
+ );
+ let mut transformer =
RecordBatchTransformerBuilder::new(snapshot_schema, &[1, 2]).build();
+ let file_schema = Arc::new(ArrowSchema::new(vec![
+ field_with_id("id", DataType::Int32, false, 1),
+ field_with_id(name, file_column.data_type().clone(), true, 2),
+ ]));
+ let batch = RecordBatch::try_new(file_schema, vec![
+ Arc::new(Int32Array::from(vec![1; file_column.len()])) as ArrayRef,
+ file_column,
+ ])
+ .unwrap();
+ Ok(transformer.process_record_batch(batch)?.column(1).clone())
+ }
+
+ #[test]
+ fn promote_struct_fills_added_middle_field_by_id() {
+ let source = Arc::new(StructArray::new(
+ Fields::from(vec![
+ field_with_id("a", DataType::Int32, true, 1),
+ field_with_id("c", DataType::Utf8, true, 3),
+ ]),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef,
+ Arc::new(StringArray::from(vec!["x", "y"])) as ArrayRef,
+ ],
+ None,
+ )) as ArrayRef;
+ let target = DataType::Struct(Fields::from(vec![
+ field_with_id("a", DataType::Int32, true, 1),
+ field_with_id("b", DataType::Int32, true, 2),
+ field_with_id("c", DataType::Utf8, true, 3),
+ ]));
+
+ let out = promote(&source, &target, &empty_schema()).unwrap();
+ let s = out.as_struct();
+ assert_eq!(s.num_columns(), 3);
+ assert_eq!(s.column(0).as_primitive::<Int32Type>().values(), &[1, 2]);
+ assert_eq!(s.column(1).null_count(), 2);
+ let cc = s.column(2).as_string::<i32>();
+ assert_eq!((cc.value(0), cc.value(1)), ("x", "y"));
+ }
+
+ #[test]
+ fn promote_struct_fills_appended_field_by_id() {
+ let source = Arc::new(StructArray::new(
+ Fields::from(vec![
+ field_with_id("a", DataType::Int32, true, 1),
+ field_with_id("b", DataType::Utf8, true, 2),
+ ]),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef,
+ Arc::new(StringArray::from(vec!["x", "y"])) as ArrayRef,
+ ],
+ None,
+ )) as ArrayRef;
+ let target = DataType::Struct(Fields::from(vec![
+ field_with_id("a", DataType::Int32, true, 1),
+ field_with_id("b", DataType::Utf8, true, 2),
+ field_with_id("c", DataType::Int32, true, 3),
+ ]));
+
+ let out = promote(&source, &target, &empty_schema()).unwrap();
+ let s = out.as_struct();
+ assert_eq!(s.num_columns(), 3);
+ assert_eq!(s.column(0).as_primitive::<Int32Type>().values(), &[1, 2]);
+ let bb = s.column(1).as_string::<i32>();
+ assert_eq!((bb.value(0), bb.value(1)), ("x", "y"));
+ assert_eq!(s.column(2).null_count(), 2);
+ }
+
+ #[test]
+ fn promote_struct_missing_field_before_nested_list_struct() {
+ let elem_field = Arc::new(field_with_id("element",
unevolved_struct_type(), true, 4));
+ let list = Arc::new(ListArray::new(
+ elem_field.clone(),
+ OffsetBuffer::new(vec![0, 1, 2].into()),
+ unevolved_struct_data(vec![10, 20]),
+ None,
+ )) as ArrayRef;
+ let source = Arc::new(StructArray::new(
+ Fields::from(vec![
+ field_with_id("s", DataType::Utf8, true, 1),
+ field_with_id("ev", DataType::List(elem_field.clone()), true,
3),
+ ]),
+ vec![
+ Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef,
+ list,
+ ],
+ None,
+ )) as ArrayRef;
+ let target = DataType::Struct(Fields::from(vec![
+ field_with_id("s", DataType::Utf8, true, 1),
+ field_with_id("gap", DataType::Int32, true, 2),
+ field_with_id("ev", DataType::List(elem_field), true, 3),
+ ]));
+
+ let out = promote(&source, &target, &empty_schema()).unwrap();
+ let st = out.as_struct();
+ assert_eq!(st.num_columns(), 3);
+ assert_eq!(st.column(1).null_count(), 2);
+ let ev = st.column(2).as_list::<i32>();
+ assert_eq!(ev.len(), 2);
+ assert_eq!(
+ ev.value(0)
+ .as_struct()
+ .column(0)
+ .as_primitive::<Int32Type>()
+ .value(0),
+ 10
+ );
+ }
+
+ #[test]
+ fn promote_list_element_struct_fills_added_field_by_id() {
+ let source = Arc::new(ListArray::new(
+ Arc::new(field_with_id("element", unevolved_struct_type(), true,
4)),
+ OffsetBuffer::new(vec![0, 1, 2].into()),
+ unevolved_struct_data(vec![10, 20]),
+ None,
+ )) as ArrayRef;
+ let target = DataType::List(Arc::new(field_with_id(
+ "element",
+ evolved_struct_type(),
+ true,
+ 4,
+ )));
+
+ let out = promote(&source, &target, &empty_schema()).unwrap();
+ let lst = out.as_list::<i32>();
+ assert_eq!(lst.len(), 2);
+ let elements = lst.values().as_struct();
+ assert_existing_field_kept(elements, &[10, 20]);
+ assert_added_field_null(elements);
+ }
+
+ #[test]
+ fn promote_map_value_struct_fills_added_field_by_id() {
+ let entries = StructArray::new(
+ Fields::from(vec![
+ field_with_id("key", DataType::Utf8, false, 7),
+ field_with_id("value", unevolved_struct_type(), true, 8),
+ ]),
+ vec![
+ Arc::new(StringArray::from(vec!["k1", "k2"])) as ArrayRef,
+ unevolved_struct_data(vec![100, 200]),
+ ],
+ None,
+ );
+ let source = Arc::new(MapArray::new(
+ Arc::new(Field::new("entries", entries.data_type().clone(),
false)),
+ OffsetBuffer::new(vec![0, 1, 2].into()),
+ entries,
+ None,
+ false,
+ )) as ArrayRef;
+ let target_entries = DataType::Struct(Fields::from(vec![
+ field_with_id("key", DataType::Utf8, false, 7),
+ field_with_id("value", evolved_struct_type(), true, 8),
+ ]));
+ let target = DataType::Map(
+ Arc::new(Field::new("entries", target_entries, false)),
+ false,
+ );
+
+ let out = promote(&source, &target, &empty_schema()).unwrap();
+ let m = out.as_map();
+ assert_eq!(m.len(), 2);
+ let entries = m.entries();
+ let ks = entries.column(0).as_string::<i32>();
+ assert_eq!((ks.value(0), ks.value(1)), ("k1", "k2"));
+ let values = entries.column(1).as_struct();
+ assert_existing_field_kept(values, &[100, 200]);
+ assert_added_field_null(values);
+ }
+
+ #[test]
+ fn promote_large_list_element_struct_fills_added_field_by_id() {
+ let source = Arc::new(LargeListArray::new(
+ Arc::new(field_with_id("element", unevolved_struct_type(), true,
4)),
+ OffsetBuffer::new(vec![0i64, 1, 2].into()),
+ unevolved_struct_data(vec![7, 8]),
+ None,
+ )) as ArrayRef;
+ let target = DataType::LargeList(Arc::new(field_with_id(
+ "element",
+ evolved_struct_type(),
+ true,
+ 4,
+ )));
+
+ let out = promote(&source, &target, &empty_schema()).unwrap();
+ let lst = out.as_list::<i64>();
+ assert_eq!(lst.len(), 2);
+ let elements = lst.values().as_struct();
+ assert_existing_field_kept(elements, &[7, 8]);
+ assert_added_field_null(elements);
+ }
+
+ #[test]
+ fn promote_struct_renames_field_by_id() {
+ let source = Arc::new(StructArray::new(
+ Fields::from(vec![field_with_id("x_old", DataType::Int32, true,
5)]),
+ vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef],
+ None,
+ )) as ArrayRef;
+ let target = DataType::Struct(Fields::from(vec![field_with_id(
+ "x",
+ DataType::Int32,
+ true,
+ 5,
+ )]));
+
+ let out = promote(&source, &target, &empty_schema()).unwrap();
+ let s = out.as_struct();
+ assert_eq!(s.fields()[0].name(), "x");
+ assert_eq!(s.column(0).as_primitive::<Int32Type>().values(), &[1, 2]);
+ }
+
+ #[test]
+ fn promote_struct_dropped_and_readded_same_name_nulls_by_id() {
+ let file = StructArray::new(
+ Fields::from(vec![field_with_id("x", DataType::Int32, true, 5)]),
+ vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef],
+ None,
+ );
+ let out = transform_top_level(
+ NestedField::optional(
+ 2,
+ "s",
+ Type::Struct(StructType::new(vec![
+ NestedField::optional(6, "x",
Type::Primitive(PrimitiveType::Int)).into(),
+ ])),
+ ),
+ Arc::new(file),
+ )
+ .unwrap();
+ assert_eq!(out.as_struct().column(0).null_count(), 2);
+ }
+
+ #[test]
+ fn promote_struct_promotes_child_primitive() {
+ let source = Arc::new(StructArray::new(
+ Fields::from(vec![field_with_id("x", DataType::Int32, true, 5)]),
+ vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef],
+ None,
+ )) as ArrayRef;
+ let target = DataType::Struct(Fields::from(vec![field_with_id(
+ "x",
+ DataType::Int64,
+ true,
+ 5,
+ )]));
+
+ let out = promote(&source, &target, &empty_schema()).unwrap();
+ let s = out.as_struct();
+ assert_eq!(s.column(0).as_primitive::<Int64Type>().values(), &[1, 2]);
+ }
+
+ #[test]
+ fn promote_struct_preserves_null_parent_rows() {
+ let source = Arc::new(StructArray::new(
+ Fields::from(vec![field_with_id("x", DataType::Int32, true, 5)]),
+ vec![Arc::new(Int32Array::from(vec![10, 20])) as ArrayRef],
+ Some(NullBuffer::from(vec![true, false])),
+ )) as ArrayRef;
+
+ let out = promote(&source, &evolved_struct_type(),
&empty_schema()).unwrap();
+ let s = out.as_struct();
+ assert!(!s.is_null(0));
+ assert!(s.is_null(1));
+ assert_eq!(s.column(0).as_primitive::<Int32Type>().value(0), 10);
+ assert_added_field_null(s);
+ }
+
+ #[test]
+ fn promote_struct_without_source_field_ids_errors() {
+ let source = Arc::new(StructArray::new(
+ Fields::from(vec![Field::new("x", DataType::Int32, true)]),
+ vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef],
+ None,
+ )) as ArrayRef;
+
+ let err = promote(&source, &evolved_struct_type(),
&empty_schema()).unwrap_err();
+ assert!(err.to_string().contains("do not all have field ids"));
+ }
+
+ #[test]
+ fn promote_required_nested_field_absent_without_default_errors() {
+ let schema = Schema::builder()
+ .with_fields(vec![
+ NestedField::required(6, "y",
Type::Primitive(PrimitiveType::Int)).into(),
+ ])
+ .build()
+ .unwrap();
+ let source = unevolved_struct_data(vec![1, 2]) as ArrayRef;
+
+ let err = promote(&source, &evolved_struct_type(),
&schema).unwrap_err();
+ assert!(err.to_string().contains("required nested field"));
+ }
+
+ #[test]
+ fn promote_nested_field_with_initial_default_errors() {
+ let schema = Schema::builder()
+ .with_fields(vec![
+ NestedField::optional(6, "y",
Type::Primitive(PrimitiveType::Int))
+ .with_initial_default(Literal::int(42))
+ .into(),
+ ])
+ .build()
+ .unwrap();
+ let source = unevolved_struct_data(vec![1, 2]) as ArrayRef;
+
+ let err = promote(&source, &evolved_struct_type(),
&schema).unwrap_err();
+ assert!(err.to_string().contains("initial-default"));
+ }
+
+ #[test]
+ fn promote_struct_reorders_children_by_id_via_process_record_batch() {
+ let file = StructArray::new(
+ Fields::from(vec![
+ field_with_id("a", DataType::Int32, true, 5),
+ field_with_id("b", DataType::Int32, true, 6),
+ ]),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef,
+ Arc::new(Int32Array::from(vec![100, 200])) as ArrayRef,
+ ],
+ None,
+ );
+ let out = transform_top_level(
+ NestedField::optional(
+ 2,
+ "s",
+ Type::Struct(StructType::new(vec![
+ NestedField::optional(6, "b",
Type::Primitive(PrimitiveType::Int)).into(),
+ NestedField::optional(5, "a",
Type::Primitive(PrimitiveType::Int)).into(),
+ ])),
+ ),
+ Arc::new(file),
+ )
+ .unwrap();
+ let s = out.as_struct();
+ assert_eq!(s.fields()[0].name(), "b");
+ assert_eq!(s.fields()[1].name(), "a");
+ assert_eq!(s.column(0).as_primitive::<Int32Type>().values(), &[
+ 100, 200
+ ]);
+ assert_eq!(s.column(1).as_primitive::<Int32Type>().values(), &[1, 2]);
Review Comment:
Could this build the expected `StructArray` and compare it whole? The bug
this test covers was an array whose child order and names didn't match the
schema, and a whole-array comparison also checks child names, nullability, and
the `PARQUET:field_id` metadata. The same applies to the other new tests that
check names and values one at a time. I ran this version on the head commit and
it passes:
```suggestion
let expected = StructArray::new(
Fields::from(vec![
field_with_id("b", DataType::Int32, true, 6),
field_with_id("a", DataType::Int32, true, 5),
]),
vec![
Arc::new(Int32Array::from(vec![100, 200])) as ArrayRef,
Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef,
],
None,
);
assert_eq!(out.as_struct(), &expected);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
