ewoodbury commented on code in PR #3255:
URL: https://github.com/apache/iceberg-rust/pull/3255#discussion_r4226904213
##########
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:
Thanks - agree with this cleanup, applied in ee881aa0f, so the module is
down from 47 to 40 tests.
The appended field is folded into `promote_struct_fills_added_fields_by_id`
(middle + appended gap). Inner `x` in `promote_idless_nested_struct_keeps_data`
is now long, so it still guards the id-less int to long case from the last
round.
`promote_evolved_nested_struct_via_process_record_batch` now goes through
`transform_top_level` as well. I spot-checked the reduced set with a few
mutations (old equals_datatype guard, dropped parent nulls, wrong child index,
no LargeList cast, ignored partition constant), and each is still caught.
--
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]
