c-thiel commented on code in PR #2880:
URL: https://github.com/apache/iceberg-rust/pull/2880#discussion_r4491838348
##########
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:
Added test_read_variant_missing_from_file in 524f5fddc. The file only has
id. Reading it with id and v, and with v alone, returns v as an all-null struct
with the extension, as you 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]
