c-thiel commented on code in PR #2880:
URL: https://github.com/apache/iceberg-rust/pull/2880#discussion_r4446315031


##########
crates/iceberg/src/arrow/reader/projection.rs:
##########
@@ -599,31 +792,257 @@ message schema {
         assert_eq!(mask, ProjectionMask::leaves(&parquet_schema, vec![0]));
     }
 
+    fn field_id_meta(id: i32) -> HashMap<String, String> {
+        HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), 
id.to_string())])
+    }
+
+    /// Arrow `Struct(metadata: Binary, value: Binary)` — the unshredded 
variant layout.
+    fn variant_arrow_fields() -> Fields {
+        Fields::from(vec![
+            Field::new("metadata", DataType::Binary, false),
+            Field::new("value", DataType::Binary, false),
+        ])
+    }
+
+    /// A variant storage group as it arrives from Parquet: `Struct(metadata, 
value)` with the
+    /// field id on the group and no field ids on the sub-fields. The 
`arrow.parquet.variant`
+    /// extension is intentionally ABSENT — the read path recovers variant 
identity from the
+    /// Iceberg schema and self-attaches it, so this exercises files that lack 
the annotation.
+    fn variant_arrow_field(name: &str, id: i32) -> Field {
+        Field::new(name, DataType::Struct(variant_arrow_fields()), false)
+            .with_metadata(field_id_meta(id))
+    }
+
+    /// A variant is a Parquet group whose leaves carry no field id, so it is 
projected via
+    /// its enclosing group's field id and all of its leaves are read together.
     #[test]
-    fn test_arrow_projection_mask_variant_is_unsupported() {
-        // Reading variant columns is not supported yet: projecting one 
(top-level or
-        // nested) must fail loudly rather than return a partial/incorrect 
batch.
+    fn test_arrow_projection_mask_variant() {
+        // c1 (String, id 1) + v (Variant, id 2).
+        let schema = Arc::new(
+            Schema::builder()
+                .with_schema_id(1)
+                .with_fields(vec![
+                    NestedField::required(1, "c1", 
Type::Primitive(PrimitiveType::String)).into(),
+                    NestedField::required(2, "v", 
Type::Variant(VariantType)).into(),
+                ])
+                .build()
+                .unwrap(),
+        );
+        let arrow_schema = Arc::new(ArrowSchema::new(vec![
+            Field::new("c1", DataType::Utf8, 
false).with_metadata(field_id_meta(1)),
+            variant_arrow_field("v", 2),
+        ]));
+        let message_type = "
+message schema {
+  required binary c1 (STRING) = 1;
+  required group v = 2 {
+    required binary metadata;
+    required binary value;
+  }
+}
+";
+        let parquet_schema =
+            
SchemaDescriptor::new(Arc::new(parse_message_type(message_type).unwrap()));
+
+        // Both fields: all three leaves.
+        let mask = ArrowReader::get_arrow_projection_mask(
+            &[1, 2],
+            &schema,
+            &parquet_schema,
+            &arrow_schema,
+            false,
+        )
+        .expect("projection mask for c1 + v");

Review Comment:
   Thanks, done in c2835f9c3. The tests this PR adds now use `unwrap()` and 
`unwrap_err()`. I left the `expect` calls that were already in the file on main 
as they are, to keep this PR focused on variants.
   



-- 
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]

Reply via email to