comphead commented on code in PR #3255:
URL: https://github.com/apache/iceberg-rust/pull/3255#discussion_r4224084527
##########
crates/iceberg/src/arrow/record_batch_transformer.rs:
##########
@@ -193,6 +195,296 @@ 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>,
+ },
+ Map {
+ field: FieldRef,
+ entry_fields: Fields,
+ entries: Vec<ChildPlan>,
+ sorted: bool,
+ },
+}
+
+#[derive(Debug)]
+pub(crate) enum ChildPlan {
+ FromSource {
+ source_index: usize,
+ plan: PromotePlan,
+ },
+ Constant {
+ target_type: DataType,
+ value: PrimitiveLiteral,
+ },
+ Null {
+ target_type: DataType,
+ },
Review Comment:
`Constant` and `Null` both end in `create_column`, which already returns
`new_null_array` for a `None` value (`create_primitive_array_repeated` in
`value.rs`). One variant shaped like `ColumnSource::Add` covers both. It also
removes the `new_null_array` import and an arm in `apply_struct`:
```rust
Add {
target_type: DataType,
value: Option<PrimitiveLiteral>,
},
```
Build it with `value: Some(datum.literal().clone())` for the identity
constant and `value: None` for the null fill. Apply it with
`RecordBatchTransformer::create_column(target_type, value.as_ref(),
source.len())`. Checked locally together with the two comments above: the 47
`arrow::record_batch_transformer` tests pass, and clippy and `cargo fmt
--check` are clean.
##########
crates/iceberg/src/arrow/record_batch_transformer.rs:
##########
@@ -193,6 +195,296 @@ 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>,
+ },
+ Map {
+ field: FieldRef,
+ entry_fields: Fields,
+ entries: Vec<ChildPlan>,
+ sorted: bool,
+ },
+}
+
+#[derive(Debug)]
+pub(crate) enum ChildPlan {
+ FromSource {
+ source_index: usize,
+ plan: PromotePlan,
+ },
+ Constant {
+ target_type: DataType,
+ value: PrimitiveLiteral,
+ },
+ Null {
+ target_type: DataType,
+ },
+}
+
+impl PromotePlan {
+ fn build(
+ source: &DataType,
+ target: &DataType,
+ snapshot_schema: &IcebergSchema,
+ constant_fields: &HashMap<i32, ColumnConstant>,
+ ) -> 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,
+ constant_fields,
+ )?,
+ fields: target_fields.clone(),
+ })
+ }
+ // A file's Arrow schema hint can give either offset width;
`apply` converts it.
+ // Table schemas are always `List`
(`ToArrowSchemaConverter::list`).
+ (
+ DataType::List(source_field) |
DataType::LargeList(source_field),
+ DataType::List(target_field),
+ ) => Ok(PromotePlan::List {
+ element: Box::new(Self::build(
+ source_field.data_type(),
+ target_field.data_type(),
+ snapshot_schema,
+ constant_fields,
+ )?),
+ 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,
+ constant_fields,
+ )?,
+ 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(invalid_data!("cannot promote {source:?} to {target:?}"))
+ }
+ _ => Ok(PromotePlan::Cast(target.clone())),
+ }
+ }
+
+ fn build_struct_children(
+ source_fields: &Fields,
+ target_fields: &Fields,
+ snapshot_schema: &IcebergSchema,
+ constant_fields: &HashMap<i32, ColumnConstant>,
+ ) -> 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 (#1845). Fully id-less
children
+ // match by position when the counts line up, 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() {
+ return source_fields
+ .iter()
+ .zip(target_fields.iter())
+ .enumerate()
+ .map(|(source_index, (source_field, target_field))| {
+ Ok(ChildPlan::FromSource {
+ source_index,
+ plan: Self::build(
+ source_field.data_type(),
+ target_field.data_type(),
+ snapshot_schema,
+ constant_fields,
+ )?,
+ })
+ })
+ .collect();
+ }
+ return Err(invalid_data!(
+ "cannot reconcile struct fields by id: source fields do not
all have field ids"
+ ));
+ }
+
+ target_fields
+ .iter()
+ .map(|target_field| {
+ let field_id = try_get_field_id_from_metadata(target_field)?;
+ match field_id.and_then(|id| source_by_id.get(&id).copied()) {
+ Some(source_index) => Ok(ChildPlan::FromSource {
+ plan: Self::build(
+ source_fields[source_index].data_type(),
+ target_field.data_type(),
+ snapshot_schema,
+ constant_fields,
+ )?,
+ source_index,
+ }),
+ None => {
+ // Identity partition constants apply inside structs
too. A missing
+ // child uses that value before initial-default
(#3261) or null.
+ if let Some(ColumnConstant::Scalar(datum)) =
+ field_id.and_then(|id| constant_fields.get(&id))
+ {
+ return Ok(ChildPlan::Constant {
+ target_type: target_field.data_type().clone(),
+ value: datum.literal().clone(),
+ });
+ }
+ // Missing nested fields are null-filled. Applying an
+ // initial-default here is not implemented yet (#3261).
+ let iceberg_field =
+ field_id.and_then(|id|
snapshot_schema.field_by_id(id));
+ if iceberg_field.is_some_and(|f|
f.initial_default.is_some()) {
+ return Err(Error::new(
+ ErrorKind::FeatureUnsupported,
+ format!(
+ "initial-default of nested field {} is not
supported, see https://github.com/apache/iceberg-rust/issues/3261";,
+ target_field.name()
+ ),
+ ));
+ }
+ if iceberg_field.map_or(!target_field.is_nullable(),
|f| f.required) {
+ return Err(invalid_data!(
+ "required nested field {} is absent from the
data file and has no initial-default",
+ target_field.name()
+ ));
+ }
+ Ok(ChildPlan::Null {
+ target_type: target_field.data_type().clone(),
+ })
+ }
+ }
+ })
+ .collect()
+ }
+
+ fn apply(&self, array: &ArrayRef) -> Result<ArrayRef> {
+ match self {
+ PromotePlan::PassThrough => Ok(Arc::clone(array)),
+ PromotePlan::Cast(target_type) => Ok(cast(array.as_ref(),
target_type)?),
+ PromotePlan::Struct { fields, children } => {
+ let source = array
+ .as_struct_opt()
+ .ok_or_else(|| promote_err(array, "struct"))?;
+ Ok(Arc::new(Self::apply_struct(source, fields, children)?))
+ }
+ PromotePlan::List { field, element } => {
+ let array = match array.data_type() {
+ DataType::LargeList(f) => cast(array.as_ref(),
&DataType::List(Arc::clone(f)))?,
+ _ => Arc::clone(array),
+ };
+ let source = array
+ .as_list_opt::<i32>()
+ .ok_or_else(|| promote_err(&array, "list"))?;
+ Self::apply_list(source, field, element)
Review Comment:
`apply_list` is generic over the offset type, but it has one caller and that
caller always passes a `ListArray`, because a `LargeList` is cast to `List`
just above. Inlining it removes the helper and the `GenericListArray` and
`OffsetSizeTrait` imports (import `ListArray` instead). Delete `apply_list`
below when you apply this.
```suggestion
Ok(Arc::new(ListArray::try_new(
field.clone(),
source.offsets().clone(),
element.apply(source.values())?,
source.nulls().cloned(),
)?))
```
##########
crates/iceberg/src/arrow/record_batch_transformer.rs:
##########
@@ -1147,22 +1434,779 @@ mod test {
use std::collections::HashMap;
use std::sync::Arc;
+ use arrow_array::cast::AsArray;
+ use arrow_array::types::Int32Type;
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, schema_to_arrow_schema};
- use crate::arrow::build_partition_constant;
+ use crate::ErrorKind;
use crate::arrow::record_batch_transformer::{
- RecordBatchTransformer, RecordBatchTransformerBuilder,
+ PromotePlan, RecordBatchTransformer, RecordBatchTransformerBuilder,
};
+ use crate::arrow::{DEFAULT_MAP_FIELD_NAME, build_partition_constant};
use crate::spec::{
- ListType, Literal, MapType, NestedField, PrimitiveType, Schema,
Struct, Type,
+ ListType, 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,
&HashMap::new())?.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 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() {
Review Comment:
Several new tests hit the same `PromotePlan` branch as another new test. The
ones below were removed or merged locally. That takes
`arrow::record_batch_transformer` from 47 to 40 tests and removes about 230
lines. Then 17 single-line mutations were applied to `PromotePlan` (wrong child
index, dropped parent nulls, skipped list or map element apply, no `LargeList`
cast, `Cast` arm as pass-through, `equals_datatype` pass-through, and others).
Every mutation that the original tests catch is still caught by the reduced
set. This is a spot check and not a coverage measurement.
- `promote_struct_fills_appended_field_by_id` differs from this test only in
where the gap sits. One target with a middle and an appended field covers both.
- `promote_struct_missing_field_before_nested_list_struct` is this test with
a pass-through list child.
- `promote_large_list_element_struct_fills_added_field_by_id` is
`promote_list_element_struct_fills_added_field_by_id` plus the `LargeList` to
`List` cast. `promote_large_list_file_column_to_list` covers that cast.
- `promote_struct_renames_field_by_id` repeats
`promote_struct_renames_child_via_process_record_batch`, which also guards the
pass-through decision.
- `promote_struct_promotes_child_primitive` and
`promote_struct_preserves_null_parent_rows` are covered by
`promote_evolved_nested_struct_via_process_record_batch` (an `int` to `long`
child and a null parent row).
- `promote_idless_struct_promotes_child_int_to_long` folds into
`promote_idless_nested_struct_keeps_data` if the table makes the inner `x` a
`long`.
- `promote_evolved_nested_struct_via_process_record_batch` rebuilds the
setup that `transform_top_level` already provides. It can call the helper with
a `StructArray` that carries the null buffer.
##########
crates/iceberg/src/arrow/record_batch_transformer.rs:
##########
@@ -193,6 +195,296 @@ 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>,
+ },
+ Map {
+ field: FieldRef,
+ entry_fields: Fields,
+ entries: Vec<ChildPlan>,
+ sorted: bool,
+ },
+}
+
+#[derive(Debug)]
+pub(crate) enum ChildPlan {
+ FromSource {
+ source_index: usize,
+ plan: PromotePlan,
+ },
+ Constant {
+ target_type: DataType,
+ value: PrimitiveLiteral,
+ },
+ Null {
+ target_type: DataType,
+ },
+}
+
+impl PromotePlan {
+ fn build(
+ source: &DataType,
+ target: &DataType,
+ snapshot_schema: &IcebergSchema,
+ constant_fields: &HashMap<i32, ColumnConstant>,
+ ) -> 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,
+ constant_fields,
+ )?,
+ fields: target_fields.clone(),
+ })
+ }
+ // A file's Arrow schema hint can give either offset width;
`apply` converts it.
+ // Table schemas are always `List`
(`ToArrowSchemaConverter::list`).
+ (
+ DataType::List(source_field) |
DataType::LargeList(source_field),
+ DataType::List(target_field),
+ ) => Ok(PromotePlan::List {
+ element: Box::new(Self::build(
+ source_field.data_type(),
+ target_field.data_type(),
+ snapshot_schema,
+ constant_fields,
+ )?),
+ 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,
+ constant_fields,
+ )?,
+ 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(invalid_data!("cannot promote {source:?} to {target:?}"))
+ }
+ _ => Ok(PromotePlan::Cast(target.clone())),
+ }
+ }
+
+ fn build_struct_children(
+ source_fields: &Fields,
+ target_fields: &Fields,
+ snapshot_schema: &IcebergSchema,
+ constant_fields: &HashMap<i32, ColumnConstant>,
+ ) -> 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 (#1845). Fully id-less
children
+ // match by position when the counts line up, 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() {
Review Comment:
`!source_fields.is_empty()` is redundant. With no source fields
`source_by_id` is empty as well, so `source_by_id.len() != source_fields.len()`
is already false.
```suggestion
if source_by_id.len() != source_fields.len() {
```
--
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]
