laskoviymishka commented on code in PR #2880:
URL: https://github.com/apache/iceberg-rust/pull/2880#discussion_r4469847524
##########
crates/iceberg/src/arrow/reader/projection.rs:
##########
@@ -2482,4 +3245,542 @@ message schema {
.unwrap();
assert_eq!(result, vec![expected]);
}
+
+ // ---------------------------------------------------------------------
+ // End-to-end variant read tests (synthetic Parquet).
+ //
+ // Unless a test says otherwise, the files carry a field id on the variant
group, none on
+ // its leaves, and no Parquet VARIANT annotation. `value` is required, as
in iceberg-java's
+ // unshredded layout.
+ // ---------------------------------------------------------------------
+
+ // Known variant bytes the reader must round-trip unchanged. `value` is
distinct per row,
+ // so a dropped or reordered row is caught, not just the struct shape.
+ const VARIANT_METADATA: [&[u8]; 3] = [b"\x01", b"\x01", b"\x01"];
+ const VARIANT_VALUE: [&[u8]; 3] = [b"\x0a", b"\x0b", b"\x0c"];
+
+ /// The 3-row variant `StructArray` for
[`VARIANT_METADATA`]/[`VARIANT_VALUE`].
+ fn variant_struct_array() -> ArrayRef {
+ use arrow_array::{BinaryArray, StructArray};
+ let metadata = Arc::new(BinaryArray::from(VARIANT_METADATA.to_vec()))
as ArrayRef;
+ let value = Arc::new(BinaryArray::from(VARIANT_VALUE.to_vec())) as
ArrayRef;
+ Arc::new(StructArray::new(
+ variant_arrow_fields(),
+ vec![metadata, value],
+ None,
Review Comment:
Could we add a read where some variant rows are null? Every fixture passes
`None` as the null buffer, and the annotated-file test writes def levels `[1,
1, 1]`. Because the file's `value` is required and the table's is nullable,
these reads go through
[`apply_struct`](https://github.com/apache/iceberg-rust/blob/c2835f9c3e7cf4a03e1cc55da770caad4f2e4055/crates/iceberg/src/arrow/record_batch_transformer.rs#L499-L521),
which re-wraps `source.nulls()` under the non-null `metadata` target. No test
reaches that with a parent null.
##########
crates/iceberg/src/arrow/reader/projection.rs:
##########
@@ -2482,4 +3245,542 @@ message schema {
.unwrap();
assert_eq!(result, vec![expected]);
}
+
+ // ---------------------------------------------------------------------
+ // End-to-end variant read tests (synthetic Parquet).
+ //
+ // Unless a test says otherwise, the files carry a field id on the variant
group, none on
+ // its leaves, and no Parquet VARIANT annotation. `value` is required, as
in iceberg-java's
+ // unshredded layout.
+ // ---------------------------------------------------------------------
+
+ // Known variant bytes the reader must round-trip unchanged. `value` is
distinct per row,
+ // so a dropped or reordered row is caught, not just the struct shape.
+ const VARIANT_METADATA: [&[u8]; 3] = [b"\x01", b"\x01", b"\x01"];
+ const VARIANT_VALUE: [&[u8]; 3] = [b"\x0a", b"\x0b", b"\x0c"];
+
+ /// The 3-row variant `StructArray` for
[`VARIANT_METADATA`]/[`VARIANT_VALUE`].
+ fn variant_struct_array() -> ArrayRef {
+ use arrow_array::{BinaryArray, StructArray};
+ let metadata = Arc::new(BinaryArray::from(VARIANT_METADATA.to_vec()))
as ArrayRef;
+ let value = Arc::new(BinaryArray::from(VARIANT_VALUE.to_vec())) as
ArrayRef;
+ Arc::new(StructArray::new(
+ variant_arrow_fields(),
+ vec![metadata, value],
+ None,
+ )) as ArrayRef
+ }
+
+ /// Asserts `field` is a variant and `column` holds the bytes of
[`variant_struct_array`]
+ /// in the table schema's Arrow type.
+ fn assert_variant(field: &Field, column: &ArrayRef) {
+ use arrow_array::StructArray;
+ let DataType::Struct(table_fields) =
+
crate::arrow::type_to_arrow_type(&Type::Variant(VariantType)).unwrap()
+ else {
+ unreachable!()
+ };
+ let expected = StructArray::new(
+ table_fields,
+ variant_struct_array().as_struct().columns().to_vec(),
+ None,
+ );
+ assert_eq!(field.extension_type_name(), Some("arrow.parquet.variant"));
+ assert_eq!(column.as_struct(), &expected);
+ }
+
+ /// Asserts the top-level variant `v` of `batch` round-tripped.
+ fn assert_top_level_variant(batch: &RecordBatch) {
+ let schema = batch.schema();
+ assert_variant(
+ schema.field_with_name("v").unwrap(),
+ batch.column_by_name("v").unwrap(),
+ );
+ }
+
+ /// Asserts the `id` column holds exactly `[1, 2, 3]`.
+ fn assert_id_values(col: &ArrayRef) {
+ let ids = col.as_primitive::<arrow_array::types::Int32Type>();
+ assert_eq!(ids.values(), &[1, 2, 3]);
+ }
+
+ fn write_variant_parquet(
+ arrow_schema: Arc<ArrowSchema>,
+ columns: Vec<ArrayRef>,
+ dir: &str,
+ ) -> String {
+ let batch = RecordBatch::try_new(arrow_schema.clone(),
columns).unwrap();
+ let path = format!("{dir}/variant.parquet");
+ let props = WriterProperties::builder()
+ .set_compression(Compression::SNAPPY)
+ .build();
+ let mut writer =
+ ArrowWriter::try_new(File::create(&path).unwrap(), arrow_schema,
Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ path
+ }
+
+ async fn try_read_variant(
+ schema: Arc<Schema>,
+ path: &str,
+ project_field_ids: Vec<i32>,
+ name_mapping: Option<Arc<NameMapping>>,
+ ) -> crate::Result<Vec<RecordBatch>> {
+ let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(),
Runtime::current()).build();
+ let task = FileScanTask::builder()
+ .with_file_size_in_bytes(std::fs::metadata(path).unwrap().len())
+ .with_start(0)
+ .with_length(0)
+ .with_data_file_path(path.to_string())
+ .with_data_file_format(DataFileFormat::Parquet)
+ .with_schema(schema)
+ .with_project_field_ids(project_field_ids)
+ .with_name_mapping(name_mapping)
+ .with_case_sensitive(false)
+ .build()?;
+ let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as
FileScanTaskStream;
+ reader.read(tasks)?.stream().try_collect().await
+ }
+
+ /// Reads `path`, checking that each batch's schema matches its columns,
nested field
+ /// metadata included.
+ async fn read_variant(
+ schema: Arc<Schema>,
+ path: &str,
+ project_field_ids: Vec<i32>,
+ ) -> Vec<RecordBatch> {
+ let batches = try_read_variant(schema, path, project_field_ids, None)
+ .await
+ .unwrap();
+ for batch in &batches {
+ RecordBatch::try_new(batch.schema(),
batch.columns().to_vec()).unwrap();
+ }
+ batches
+ }
+
+ /// `id` (int, id 1) + optional `v` (variant, id 2).
+ fn top_level_variant_schema() -> Arc<Schema> {
+ Arc::new(
+ Schema::builder()
+ .with_fields(vec![
+ NestedField::required(1, "id",
Type::Primitive(PrimitiveType::Int)).into(),
+ NestedField::optional(2, "v",
Type::Variant(VariantType)).into(),
+ ])
+ .build()
+ .unwrap(),
+ )
+ }
+
+ /// Writes [`top_level_variant_schema`] rows with `v` stored as
`variant_field`.
+ fn write_top_level_variant_as(dir: &str, id_field: Field, variant_field:
Field) -> String {
+ use arrow_array::{Int32Array, StructArray};
+ let DataType::Struct(storage) = variant_field.data_type() else {
+ unreachable!()
+ };
+ let variant = Arc::new(StructArray::new(
+ storage.clone(),
+ variant_struct_array().as_struct().columns().to_vec(),
+ None,
+ )) as ArrayRef;
+ let arrow_schema = Arc::new(ArrowSchema::new(vec![id_field,
variant_field]));
+ let id = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
+ write_variant_parquet(arrow_schema, vec![id, variant], dir)
+ }
+
+ fn write_top_level_variant(dir: &str) -> Arc<Schema> {
+ write_top_level_variant_as(
+ dir,
+ Field::new("id", DataType::Int32,
false).with_metadata(field_id_meta(1)),
+ variant_arrow_field("v", 2),
+ );
+ top_level_variant_schema()
+ }
+
+ #[tokio::test]
+ async fn test_read_variant_full_scan() {
+ let tmp = TempDir::new().unwrap();
+ let dir = tmp.path().to_str().unwrap();
+ let schema = write_top_level_variant(dir);
+
+ let batches = read_variant(schema, &format!("{dir}/variant.parquet"),
vec![1, 2]).await;
+ assert_eq!(batches.len(), 1);
+ assert_eq!(batches[0].num_rows(), 3);
+ assert_id_values(batches[0].column_by_name("id").unwrap());
+ assert_top_level_variant(&batches[0]);
+ }
+
+ #[tokio::test]
+ async fn test_read_variant_only() {
+ let tmp = TempDir::new().unwrap();
+ let dir = tmp.path().to_str().unwrap();
+ let schema = write_top_level_variant(dir);
+
+ let batches = read_variant(schema, &format!("{dir}/variant.parquet"),
vec![2]).await;
+ assert_eq!(batches[0].num_rows(), 3);
+ assert_top_level_variant(&batches[0]);
+ assert!(batches[0].column_by_name("id").is_none());
+ }
+
+ #[tokio::test]
+ async fn test_read_variant_sibling_only_excludes_variant() {
Review Comment:
Could we add a read of a file written before `v` was added to the table? No
test projects a variant id the file doesn't contain. As I read it, an optional
variant comes back from
[`ColumnSource::Add`](https://github.com/apache/iceberg-rust/blob/c2835f9c3e7cf4a03e1cc55da770caad4f2e4055/crates/iceberg/src/arrow/record_batch_transformer.rs#L1222-L1261)
as an all-null struct with the extension, but `main` rejected every variant
read, so that path has never run for variants.
##########
crates/iceberg/src/arrow/record_batch_transformer.rs:
##########
@@ -233,6 +233,62 @@ pub(crate) enum ChildPlan {
}
impl PromotePlan {
+ /// Plans one field, dispatching on the target field so a variant is seen
by its extension.
+ fn build_field(
+ source: &Field,
+ target: &Field,
+ snapshot_schema: &IcebergSchema,
+ constant_fields: &HashMap<i32, ColumnConstant>,
+ ) -> Result<Self> {
+ if is_variant_field(target) {
+ return Self::build_variant(source.data_type(), target.data_type());
+ }
+ Self::build(
+ source.data_type(),
+ target.data_type(),
+ snapshot_schema,
+ constant_fields,
+ )
+ }
+
+ /// A variant's `metadata` and `value` carry no field ids, so they are
matched by name.
+ /// When reading data files, projection reports unsupported or malformed
storage before this
+ /// runs.
+ fn build_variant(source: &DataType, target: &DataType) -> Result<Self> {
Review Comment:
Could `build_variant` and
[`validate_unshredded_variant`](https://github.com/apache/iceberg-rust/blob/c2835f9c3e7cf4a03e1cc55da770caad4f2e4055/crates/iceberg/src/arrow/reader/projection.rs#L500-L540)
share one check that also covers nullability and names the column? Today
shredded storage is `FeatureUnsupported` there and a generic `DataInvalid`
here. Neither checks nullability, so an optional `metadata` with a null in a
non-null row fails in `StructArray::try_new` as `Unexpected` ("Arrow Schema
Error") instead of `DataInvalid` with the field id.
--
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]
