Better expr

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

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

@@ -7,7 +7,10 @@
 
 use std::sync::Arc;
 
-use parquet_variant::VariantPathElement;
+use arrow_schema::Field;
+use arrow_schema::FieldRef;
+use parquet_variant::VariantPath;
+use parquet_variant::VariantPathElement as ArrowVariantPathElement;
 use parquet_variant_compute::GetOptions;
 use parquet_variant_compute::VariantArray as ArrowVariantArray;
 use vortex_array::ArrayRef;
@@ -19,10 +22,11 @@ use vortex_array::arrays::scalar_fn::ExactScalarFn;
 use vortex_array::arrays::scalar_fn::ScalarFnArrayView;
 use vortex_array::arrow::FromArrowArray;
 use vortex_array::dtype::DType;
-use vortex_array::dtype::FieldName;
 use vortex_array::dtype::Nullability;
 use vortex_array::kernel::ExecuteParentKernel;
 use vortex_array::scalar_fn::fns::variant_get::VariantGet;
+use vortex_array::scalar_fn::fns::variant_get::VariantGetOptions;
+use vortex_array::scalar_fn::fns::variant_get::VariantPathElement as VortexVariantPathElement;
 use vortex_array::validity::Validity;
 use vortex_buffer::BitBuffer;
 use vortex_error::VortexResult;
@@ -47,31 +51,48 @@ impl ExecuteParentKernel<ParquetVariant> for VariantGetExecuteParent {
         _child_idx: usize,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
-        let field_name: &FieldName = parent.options;
-        variant_get_impl(array, field_name, ctx).map(Some)
+        variant_get_impl(array, parent.options, ctx).map(Some)
     }
 }
 
 fn variant_get_impl(
     array: ArrayView<'_, ParquetVariant>,
-    field_name: &FieldName,
+    options: &VariantGetOptions,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
     // Convert to Arrow VariantArray
     let arrow_variant = array.to_arrow(ctx)?;
 
-    // Build path for a single field access
-    let path_element = VariantPathElement::Field {
-        name: field_name.as_ref().into(),
-    };
-    let options = GetOptions::new_with_path(vec![path_element].into());
+    let path = options
+        .path()
+        .iter()
+        .cloned()
+        .map(|element| match element {
+            VortexVariantPathElement::Field(name) => ArrowVariantPathElement::Field {
+                name: name.to_string().into(),
+            },
+            VortexVariantPathElement::Index(index) => ArrowVariantPathElement::Index { index },
+        })
+        .collect::<Vec<_>>();
+    let mut arrow_options = GetOptions::new_with_path(VariantPath::new(path));
+    if let Some(as_dtype) = options.effective_as_dtype() {
+        arrow_options = arrow_options.with_as_type(Some(FieldRef::new(Field::new(
+            "result",
+            as_dtype.to_arrow_dtype()?,
+            as_dtype.is_nullable(),
+        ))));
+    }
 
     // Delegate to the parquet-variant-compute kernel.
     // With as_type = None, the result is itself a VariantArray.
     let inner: Arc<dyn arrow_array::Array> = Arc::new(arrow_variant.into_inner());
-    let arrow_result = parquet_variant_compute::variant_get(&inner, options)
+    let arrow_result = parquet_variant_compute::variant_get(&inner, arrow_options)
         .map_err(|e| vortex_err!("variant_get failed: {e}"))?;
 
+    if options.effective_as_dtype().is_some() {
+        return ArrayRef::from_arrow(arrow_result.as_ref(), true);
+    }
+
     // Convert back to Vortex.
     let result_variant = ArrowVariantArray::try_new(
         arrow_result

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

@@ -8,9 +8,13 @@ use arrow_array::ArrayRef as ArrowArrayRef;
 use arrow_array::StringArray;
 use arrow_array::StructArray;
 use arrow_buffer::NullBuffer;
+use arrow_schema::DataType as ArrowDataType;
+use arrow_schema::Field as ArrowField;
+use arrow_schema::FieldRef;
 use parquet_variant::Variant as PqVariant;
 use parquet_variant::VariantBuilderExt;
 use parquet_variant::VariantPath;
+use parquet_variant::VariantPathElement;
 use parquet_variant_compute::GetOptions;
 use parquet_variant_compute::VariantArray as ArrowVariantArray;
 use parquet_variant_compute::VariantArrayBuilder;
@@ -20,17 +24,36 @@ use vortex_array::ArrayRef;
 use vortex_array::LEGACY_SESSION;
 use vortex_array::VortexSessionExecute;
 use vortex_array::arrays::variant::VariantArrayExt;
+use vortex_array::arrow::FromArrowArray;
+use vortex_array::assert_arrays_eq;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::Nullability;
+use vortex_array::dtype::PType;
 use vortex_array::expr::root;
 use vortex_array::expr::variant_get;
+use vortex_array::expr::variant_get_as;
+use vortex_array::scalar_fn::fns::variant_get::VariantPath as VortexVariantPath;
 use vortex_error::VortexResult;
 
 use crate::ParquetVariant;
 use crate::ParquetVariantArrayExt;
 use crate::ParquetVariantData;
 
 /// Apply variant_get and execute through the full pipeline (including execute_parent).
-fn apply_variant_get(arr: &ArrayRef, field: &str) -> VortexResult<ArrayRef> {
-    let expr = variant_get(field, root());
+fn apply_variant_get(arr: &ArrayRef, path: impl Into<VortexVariantPath>) -> VortexResult<ArrayRef> {
+    let expr = variant_get(path, root());
+    let array = arr.clone().apply(&expr)?;
+    let mut ctx = LEGACY_SESSION.create_execution_ctx();
+    array.execute::<ArrayRef>(&mut ctx)
+}
+
+/// Apply typed variant_get and execute through the full pipeline.
+fn apply_variant_get_as(
+    arr: &ArrayRef,
+    path: impl Into<VortexVariantPath>,
+    as_dtype: DType,
+) -> VortexResult<ArrayRef> {
+    let expr = variant_get_as(path, as_dtype, root());
     let array = arr.clone().apply(&expr)?;
     let mut ctx = LEGACY_SESSION.create_execution_ctx();
     array.execute::<ArrayRef>(&mut ctx)
@@ -47,7 +70,6 @@ fn vortex_to_arrow_variant(arr: &ArrayRef) -> ArrowVariantArray {
 /// 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) {
-    // --- Arrow side ---
     let arrow_strings: ArrowArrayRef = Arc::new(StringArray::from(
         json_rows.iter().map(|s| Some(*s)).collect::<Vec<_>>(),
     ));
@@ -62,12 +84,10 @@ fn assert_matches_arrow(json_rows: &[&str], field: &str) {
         ArrowVariantArray::try_new(arrow_result.as_any().downcast_ref::<StructArray>().unwrap())
             .unwrap();
 
-    // --- Vortex side ---
     let vortex_input = ParquetVariantData::from_arrow_variant(&arrow_variant).unwrap();
     let vortex_result = apply_variant_get(&vortex_input, field).unwrap();
     let vortex_as_arrow = vortex_to_arrow_variant(&vortex_result);
 
-    // --- Compare row-by-row at Arrow Variant level ---
     assert_eq!(
         vortex_as_arrow.len(),
         arrow_result_variant.len(),
@@ -94,10 +114,90 @@ fn assert_matches_arrow(json_rows: &[&str], field: &str) {
     }
 }
 
+/// Run variant_get through both Arrow and Vortex for an explicit nested/index path,
+/// and assert the per-row results match at the Arrow variant level.
+fn assert_matches_arrow_with_path(
+    json_rows: &[&str],
+    path: VortexVariantPath,
+    arrow_path: VariantPath<'static>,
+) {
+    let arrow_strings: ArrowArrayRef = Arc::new(StringArray::from(
+        json_rows.iter().map(|s| Some(*s)).collect::<Vec<_>>(),
+    ));
+    let arrow_variant = json_to_variant(&arrow_strings).unwrap();
+    let arrow_result = parquet_variant_compute::variant_get(
+        &arrow_variant.clone().into(),
+        GetOptions::new_with_path(arrow_path),
+    )
+    .unwrap();
+    let arrow_result_variant =
+        ArrowVariantArray::try_new(arrow_result.as_any().downcast_ref::<StructArray>().unwrap())
+            .unwrap();
+
+    let vortex_input = ParquetVariantData::from_arrow_variant(&arrow_variant).unwrap();
+    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"
+    );
+
+    for i in 0..arrow_result_variant.len() {
+        let arrow_is_null = arrow_result_variant.is_null(i);
+        let vortex_is_null = vortex_as_arrow.is_null(i);
+
+        assert_eq!(
+            vortex_is_null, arrow_is_null,
+            "row {i}: null mismatch (vortex={vortex_is_null}, arrow={arrow_is_null})"
+        );
+
+        if !arrow_is_null {
+            let arrow_value = arrow_result_variant.value(i);
+            let vortex_value = vortex_as_arrow.value(i);
+            assert_eq!(
+                vortex_value, arrow_value,
+                "row {i}: value mismatch\n  vortex: {vortex_value:?}\n  arrow:  {arrow_value:?}"
+            );
+        }
+    }
+}
+
+/// Run typed variant_get through both Arrow and Vortex for an explicit path,
+/// and assert the typed nullable result matches.
+fn assert_typed_matches_arrow_with_path(
+    json_rows: &[&str],
+    path: VortexVariantPath,
+    arrow_path: VariantPath<'static>,
+    as_dtype: DType,
+    arrow_dtype: ArrowDataType,
+) -> VortexResult<()> {
+    let arrow_strings: ArrowArrayRef = Arc::new(StringArray::from(
+        json_rows.iter().map(|s| Some(*s)).collect::<Vec<_>>(),
+    ));
+    let arrow_variant = json_to_variant(&arrow_strings).unwrap();
+    let arrow_result =
+        parquet_variant_compute::variant_get(
+            &arrow_variant.clone().into(),
+            GetOptions::new_with_path(arrow_path).with_as_type(Some(FieldRef::new(
+                ArrowField::new("result", arrow_dtype, true),
+            ))),
+        )
+        .unwrap();
+    let expected = ArrayRef::from_arrow(arrow_result.as_ref(), true)?;
+
+    let vortex_input = ParquetVariantData::from_arrow_variant(&arrow_variant).unwrap();
+    let vortex_result = apply_variant_get_as(&vortex_input, path, as_dtype.clone())?;
+
+    assert_eq!(vortex_result.dtype(), &as_dtype.as_nullable());
+    assert_arrays_eq!(vortex_result, expected);
+    Ok(())
+}
+
 /// Run variant_get through both Arrow and Vortex on nullable input (with NullBuffer),
 /// and assert the results match.
 fn assert_matches_arrow_nullable(json_rows: &[&str], validity: &[bool], field: &str) {
-    // --- Arrow side ---
     let arrow_strings: ArrowArrayRef = Arc::new(StringArray::from(
         json_rows.iter().map(|s| Some(*s)).collect::<Vec<_>>(),
     ));
@@ -121,12 +221,10 @@ fn assert_matches_arrow_nullable(json_rows: &[&str], validity: &[bool], field: &
         ArrowVariantArray::try_new(arrow_result.as_any().downcast_ref::<StructArray>().unwrap())
             .unwrap();
 
-    // --- Vortex side ---
     let vortex_input = ParquetVariantData::from_arrow_variant(&arrow_variant).unwrap();
     let vortex_result = apply_variant_get(&vortex_input, field).unwrap();
     let vortex_as_arrow = vortex_to_arrow_variant(&vortex_result);
 
-    // --- Compare ---
     assert_eq!(
         vortex_as_arrow.len(),
         arrow_result_variant.len(),
@@ -221,6 +319,60 @@ fn test_variant_get_matches_arrow_nested_object_result() {
     );
 }
 
+#[test]
+fn test_variant_get_matches_arrow_nested_path() {
+    let path = VortexVariantPath::from_name("outer").join("inner");
+    let arrow_path = VariantPath::from_iter([
+        VariantPathElement::field("outer".to_string()),
+        VariantPathElement::field("inner".to_string()),
+    ]);
+    assert_matches_arrow_with_path(
+        &[
+            r#"{"outer": {"inner": 42}}"#,
+            r#"{"outer": {"inner": "x"}}"#,
+            r#"{"outer": {"other": true}}"#,
+        ],
+        path,
+        arrow_path,
+    );
+}
+
+#[test]
+fn test_variant_get_matches_arrow_index_path() {
+    let path = VortexVariantPath::from_name("arr").join(1usize);
+    let arrow_path = VariantPath::from_iter([
+        VariantPathElement::field("arr".to_string()),
+        VariantPathElement::index(1),
+    ]);
+    assert_matches_arrow_with_path(
+        &[
+            r#"{"arr": [1, 2, 3]}"#,
+            r#"{"arr": ["a", "b"]}"#,
+            r#"{"arr": [true]}"#,
+        ],
+        path,
+        arrow_path,
+    );
+}
+
+#[test]
+fn test_variant_get_matches_arrow_typed_path() -> VortexResult<()> {
+    assert_typed_matches_arrow_with_path(
+        &[
+            r#"{"outer": {"inner": 42}}"#,
+            r#"{"outer": {"inner": "x"}}"#,
+            r#"{"outer": {"other": true}}"#,
+        ],
+        VortexVariantPath::from_name("outer").join("inner"),
+        VariantPath::from_iter([
+            VariantPathElement::field("outer".to_string()),
+            VariantPathElement::field("inner".to_string()),
+        ]),
+        DType::Primitive(PType::I64, Nullability::NonNullable),
+        ArrowDataType::Int64,
+    )
+}
+
 // ---------------------------------------------------------------------------
 // Original standalone tests
 // ---------------------------------------------------------------------------
@@ -310,6 +462,7 @@ fn test_variant_get_different_field() -> VortexResult<()> {
 // Test data helpers
 // ---------------------------------------------------------------------------
 
+/// Build a small non-null object variant array used by the standalone tests.
 fn make_object_array() -> VortexResult<ArrayRef> {
     let mut builder = VariantArrayBuilder::new(3);
 
@@ -330,6 +483,7 @@ fn make_object_array() -> VortexResult<ArrayRef> {
     ParquetVariantData::from_arrow_variant(&builder.build())
 }
 
+/// Build the same object array shape with an explicit top-level validity bitmap.
 fn make_nullable_object_array() -> VortexResult<ArrayRef> {
     let mut builder = VariantArrayBuilder::new(3);