Handle child nullability better

Author: AdamGS

encodings/parquet-variant/src/variant_get/mod.rs

@@ -101,6 +101,20 @@ fn variant_get_impl(
             .ok_or_else(|| vortex_err!("variant_get did not return a StructArray"))?,
     )
     .map_err(|e| vortex_err!("failed to create VariantArray from result: {e}"))?;
+    let value_nullable = result_variant
+        .inner()
+        .fields()
+        .iter()
+        .find(|field| field.name() == "value")
+        .map(|field| field.is_nullable())
+        .unwrap_or(false);
+    let typed_value_nullable = result_variant
+        .inner()
+        .fields()
+        .iter()
+        .find(|field| field.name() == "typed_value")
+        .map(|field| field.is_nullable())
+        .unwrap_or(false);
 
     // Ensure the result is always nullable (matching variant_get's return_dtype).
     // Arrow may return a non-nullable result when no nulls are present.
@@ -121,11 +135,11 @@ fn variant_get_impl(
     )?;
     let value = result_variant
         .value_field()
-        .map(|v| ArrayRef::from_arrow(v as &dyn arrow_array::Array, true))
+        .map(|v| ArrayRef::from_arrow(v as &dyn arrow_array::Array, value_nullable))
         .transpose()?;
     let typed_value = result_variant
         .typed_value_field()
-        .map(|tv| ArrayRef::from_arrow(tv.as_ref(), true))
+        .map(|tv| ArrayRef::from_arrow(tv.as_ref(), typed_value_nullable))
         .transpose()?;
 
     let pv = ParquetVariant::try_new(validity, metadata, value, typed_value)?;

encodings/parquet-variant/src/variant_get/tests.rs

@@ -68,6 +68,31 @@ fn vortex_to_arrow_variant(arr: &ArrayRef) -> ArrowVariantArray {
     pv.to_arrow(&mut ctx).unwrap()
 }
 
+fn assert_variant_storage_matches(expected: &ArrowVariantArray, actual: &ArrowVariantArray) {
+    assert_eq!(actual.len(), expected.len(), "length mismatch");
+    assert_eq!(
+        actual.inner().column_names(),
+        expected.inner().column_names(),
+        "column mismatch"
+    );
+    assert_eq!(actual.inner().nulls(), expected.inner().nulls());
+    assert_eq!(
+        actual.inner().fields().len(),
+        expected.inner().fields().len()
+    );
+
+    for (expected, actual) in expected
+        .inner()
+        .fields()
+        .iter()
+        .zip(actual.inner().fields().iter())
+    {
+        assert_eq!(actual.name(), expected.name());
+        assert_eq!(actual.data_type(), expected.data_type());
+        assert_eq!(actual.is_nullable(), expected.is_nullable());
+    }
+}
+
 /// Run variant_get through both Arrow and Vortex on the same input, and assert
 /// the per-row results (value + validity) are identical by comparing at the Arrow level.
 fn assert_matches_arrow(json_rows: &[&str], field: &str) {
@@ -89,11 +114,7 @@ fn assert_matches_arrow(json_rows: &[&str], field: &str) {
     let vortex_result = apply_variant_get!(&vortex_input, field).unwrap();
     let vortex_as_arrow = vortex_to_arrow_variant(&vortex_result);
 
-    assert_eq!(
-        vortex_as_arrow.len(),
-        arrow_result_variant.len(),
-        "length mismatch"
-    );
+    assert_variant_storage_matches(&arrow_result_variant, &vortex_as_arrow);
 
     for i in 0..arrow_result_variant.len() {
         let arrow_is_null = arrow_result_variant.is_null(i);
@@ -139,11 +160,7 @@ fn assert_matches_arrow_with_path(
     let vortex_result = apply_variant_get!(&vortex_input, path).unwrap();
     let vortex_as_arrow = vortex_to_arrow_variant(&vortex_result);
 
-    assert_eq!(
-        vortex_as_arrow.len(),
-        arrow_result_variant.len(),
-        "length mismatch"
-    );
+    assert_variant_storage_matches(&arrow_result_variant, &vortex_as_arrow);
 
     for i in 0..arrow_result_variant.len() {
         let arrow_is_null = arrow_result_variant.is_null(i);
@@ -226,11 +243,7 @@ fn assert_matches_arrow_nullable(json_rows: &[&str], validity: &[bool], field: &
     let vortex_result = apply_variant_get!(&vortex_input, field).unwrap();
     let vortex_as_arrow = vortex_to_arrow_variant(&vortex_result);
 
-    assert_eq!(
-        vortex_as_arrow.len(),
-        arrow_result_variant.len(),
-        "length mismatch"
-    );
+    assert_variant_storage_matches(&arrow_result_variant, &vortex_as_arrow);
 
     for i in 0..arrow_result_variant.len() {
         let arrow_is_null = arrow_result_variant.is_null(i);