Add variant_get compute kernel

Samyak2 · Samyak2 · commit 184e1bff24d9 · 2025-07-14T00:24:41.000+05:30
In parquet-variant: - Add a new function `Variant::get_path`: this traverses the path to create a new Variant (does not cast any of it). - Add a new module `parquet_variant::path`: adds structs/enums to define a path to access a variant value deeply. In parquet-variant-compute: - Add a new compute kernel `variant_get`: does the path traversal over a `VariantArray`. In the future, this would also cast the values to a specified type. - Includes some basic unit tests. Not comprehensive. - Includes a simple micro-benchmark for reference. Current limitations: - It can only return another VariantArray. Casts are not implemented yet. - Only top-level object/list access is supported. It panics on finding a nested object/list. Needs apache#7914 to fix this. - Perf is a TODO.
diff --git a/parquet-variant-compute/Cargo.toml b/parquet-variant-compute/Cargo.toml
@@ -42,3 +42,9 @@ name = "parquet_variant_compute"
 bench = false
 
 [dev-dependencies]
+criterion = { version = "0.6", default-features = false }
+rand = { version = "0.9.1" }
+
+[[bench]]
+name = "variant_get"
+harness = false
diff --git a/parquet-variant-compute/benches/variant_get.rs b/parquet-variant-compute/benches/variant_get.rs
@@ -0,0 +1,43 @@
+use std::sync::Arc;
+
+use arrow::array::ArrayRef;
+use criterion::{criterion_group, criterion_main, Criterion};
+use parquet_variant::{Variant, VariantBuilder};
+use parquet_variant_compute::{
+    variant_get::{variant_get, GetOptions},
+    VariantArray, VariantArrayBuilder,
+};
+use rand::{rngs::StdRng, Rng, SeedableRng};
+
+fn create_primitive_variant(size: usize) -> VariantArray {
+    let mut rng = StdRng::seed_from_u64(42);
+
+    let mut variant_builder = VariantArrayBuilder::new(1);
+
+    for _ in 0..size {
+        let mut builder = VariantBuilder::new();
+        builder.append_value(rng.random::<i64>());
+        let (metadata, value) = builder.finish();
+        variant_builder.append_variant(Variant::try_new(&metadata, &value).unwrap());
+    }
+
+    variant_builder.build()
+}
+
+pub fn variant_get_bench(c: &mut Criterion) {
+    let variant_array = create_primitive_variant(8192);
+    let input: ArrayRef = Arc::new(variant_array);
+
+    let options = GetOptions {
+        path: vec![].into(),
+        as_type: None,
+        cast_options: Default::default(),
+    };
+
+    c.bench_function("variant_get_primitive", |b| {
+        b.iter(|| variant_get(&input.clone(), options.clone()))
+    });
+}
+
+criterion_group!(benches, variant_get_bench);
+criterion_main!(benches);
diff --git a/parquet-variant-compute/src/lib.rs b/parquet-variant-compute/src/lib.rs
@@ -19,6 +19,7 @@ mod from_json;
 mod to_json;
 mod variant_array;
 mod variant_array_builder;
+pub mod variant_get;
 
 pub use variant_array::VariantArray;
 pub use variant_array_builder::VariantArrayBuilder;
diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs
@@ -0,0 +1,158 @@
+use std::sync::Arc;
+
+use arrow::{
+    array::{Array, ArrayRef},
+    compute::CastOptions,
+    error::Result,
+};
+use arrow_schema::{ArrowError, Field};
+use parquet_variant::path::VariantPath;
+
+use crate::{VariantArray, VariantArrayBuilder};
+
+/// Returns an array with the specified path extracted from the variant values.
+///
+/// The return array type depends on the `as_type` field of the options parameter
+/// 1. `as_type: None`: a VariantArray is returned. The values in this new VariantArray will point
+///    to the specified path.
+/// 2. `as_type: Some(<specific field>)`: an array of the specified type is returned.
+pub fn variant_get(input: &ArrayRef, options: GetOptions) -> Result<ArrayRef> {
+    let variant_array: &VariantArray = input.as_any().downcast_ref().ok_or_else(|| {
+        ArrowError::InvalidArgumentError(
+            "expected a VariantArray as the input for variant_get".to_owned(),
+        )
+    })?;
+
+    if let Some(as_type) = options.as_type {
+        return Err(ArrowError::NotYetImplemented(format!(
+            "getting a {} from a VariantArray is not implemented yet",
+            as_type
+        )));
+    }
+
+    let mut builder = VariantArrayBuilder::new(variant_array.len());
+    for i in 0..variant_array.len() {
+        let new_variant = variant_array.value(i);
+        // TODO: perf?
+        let new_variant = new_variant.get_path(&options.path);
+        if let Some(new_variant) = new_variant {
+            // TODO: we're decoding the value and doing a copy into a variant value again. This
+            // copy can be much smarter.
+            builder.append_variant(new_variant);
+        } else {
+            builder.append_null();
+        }
+    }
+
+    Ok(Arc::new(builder.build()))
+}
+
+/// Controls the action of the variant_get kernel.
+#[derive(Debug, Clone)]
+pub struct GetOptions<'a> {
+    /// What path to extract
+    pub path: VariantPath,
+    /// if `as_type` is None, the returned array will itself be a VariantArray.
+    ///
+    /// if `as_type` is `Some(type)` the field is returned as the specified type if possible. To specify returning
+    /// a Variant, pass a Field with variant type in the metadata.
+    pub as_type: Option<Field>,
+    /// Controls the casting behavior (e.g. error vs substituting null on cast error).
+    pub cast_options: CastOptions<'a>,
+}
+
+impl<'a> GetOptions<'a> {
+    /// Construct options to get the specified path as a variant.
+    pub fn new_with_path(path: VariantPath) -> Self {
+        Self {
+            path,
+            as_type: None,
+            cast_options: Default::default(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::sync::Arc;
+
+    use arrow::array::{Array, ArrayRef};
+    use parquet_variant::{path::VariantPathElement, VariantBuilder};
+
+    use crate::{VariantArray, VariantArrayBuilder};
+
+    use super::{variant_get, GetOptions};
+
+    #[test]
+    fn get_primitive_variant() {
+        let mut builder = VariantBuilder::new();
+        builder.add_field_name("some_field");
+        let mut object = builder.new_object();
+        object.insert("some_field", 1234i64);
+        object.finish().unwrap();
+        let (metadata, value) = builder.finish();
+
+        let mut builder = VariantArrayBuilder::new(1);
+        builder.append_variant_buffers(&metadata, &value);
+
+        let variant_array = builder.build();
+
+        let input = Arc::new(variant_array) as ArrayRef;
+
+        let result = variant_get(
+            &input,
+            GetOptions::new_with_path(
+                vec![VariantPathElement::field("some_field".to_owned())].into(),
+            ),
+        )
+        .unwrap();
+
+        let result: &VariantArray = result.as_any().downcast_ref().unwrap();
+        assert!(result.nulls().is_none());
+        let result = result.value(0);
+        assert_eq!(result.as_int64().unwrap(), 1234);
+    }
+
+    #[test]
+    #[should_panic(
+        expected = "Nested values are handled specially by ObjectBuilder and ListBuilder"
+    )]
+    fn get_complex_variant() {
+        let mut builder = VariantBuilder::new();
+        builder.add_field_name("top_level_field");
+        builder.add_field_name("inner_field");
+
+        let mut object = builder.new_object();
+        let mut inner_object = object.new_object("top_level_field");
+        inner_object.insert("inner_field", 1234i64);
+        inner_object.finish().unwrap();
+        object.finish().unwrap();
+        let (metadata, value) = builder.finish();
+
+        let mut builder = VariantArrayBuilder::new(1);
+        builder.append_variant_buffers(&metadata, &value);
+
+        let variant_array = builder.build();
+
+        let input = Arc::new(variant_array) as ArrayRef;
+
+        let result = variant_get(
+            &input,
+            GetOptions::new_with_path(
+                vec![VariantPathElement::field("top_level_field".to_owned())].into(),
+            ),
+        )
+        .unwrap();
+
+        // uncomment once implemented
+        todo!("{:?}", result);
+        // let result: &VariantArray = result.as_any().downcast_ref().unwrap();
+        // assert!(result.nulls().is_none());
+        // let result = result.value(0);
+        // let result = result.as_object().unwrap();
+        // let binding = result.get("top_level_field").unwrap();
+        // let result = binding.as_object().unwrap();
+        // let result = result.get("inner_field").unwrap();
+        // assert_eq!(result.as_int64().unwrap(), 1234);
+    }
+}
diff --git a/parquet-variant/src/lib.rs b/parquet-variant/src/lib.rs
@@ -31,6 +31,7 @@ mod builder;
 mod decoder;
 mod utils;
 mod variant;
+pub mod path;
 
 pub use builder::*;
 pub use variant::*;
diff --git a/parquet-variant/src/path.rs b/parquet-variant/src/path.rs
@@ -0,0 +1,48 @@
+use std::ops::Deref;
+
+/// Represents a qualified path to a potential subfield or index of a variant value.
+#[derive(Debug, Clone)]
+pub struct VariantPath(Vec<VariantPathElement>);
+
+impl VariantPath {
+    pub fn new(path: Vec<VariantPathElement>) -> Self {
+        Self(path)
+    }
+
+    pub fn path(&self) -> &Vec<VariantPathElement> {
+        &self.0
+    }
+}
+
+impl From<Vec<VariantPathElement>> for VariantPath {
+    fn from(value: Vec<VariantPathElement>) -> Self {
+        Self::new(value)
+    }
+}
+
+impl Deref for VariantPath {
+    type Target = Vec<VariantPathElement>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+/// Element of a path
+#[derive(Debug, Clone)]
+pub enum VariantPathElement {
+    /// Access field with name `name`
+    Field { name: String },
+    /// Access the list element at `index`
+    Index { index: usize },
+}
+
+impl VariantPathElement {
+    pub fn field(name: String) -> VariantPathElement {
+        VariantPathElement::Field { name }
+    }
+
+    pub fn index(index: usize) -> VariantPathElement {
+        VariantPathElement::Index { index }
+    }
+}
diff --git a/parquet-variant/src/variant.rs b/parquet-variant/src/variant.rs
@@ -22,6 +22,7 @@ pub use self::object::VariantObject;
 use crate::decoder::{
     self, get_basic_type, get_primitive_type, VariantBasicType, VariantPrimitiveType,
 };
+use crate::path::{VariantPath, VariantPathElement};
 use crate::utils::{first_byte_from_slice, slice_from_slice};
 use std::ops::Deref;
 
@@ -1063,6 +1064,34 @@ impl<'m, 'v> Variant<'m, 'v> {
             _ => None,
         }
     }
+
+    /// Return a new Variant with the path followed.
+    ///
+    /// If the path is not found, `None` is returned.
+    pub fn get_path(&self, path: &VariantPath) -> Option<Variant> {
+        let mut output = self.clone();
+        for element in path.deref() {
+            match element {
+                VariantPathElement::Field { name } => {
+                    let Variant::Object(variant_object) = output else {
+                        return None;
+                    };
+                    let field = variant_object.get(name);
+                    let field = field?;
+                    output = field;
+                }
+                VariantPathElement::Index { index } => {
+                    let Variant::List(variant_list) = output else {
+                        return None;
+                    };
+                    let index = variant_list.get(*index);
+                    let index = index?;
+                    output = index;
+                }
+            }
+        }
+        Some(output)
+    }
 }
 
 impl From<()> for Variant<'_, '_> {
