[Variant] Improve variant_get performance on a perfect shredding

parquet-variant-compute/benches/variant_kernels.rs

@@ -15,10 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{Array, ArrayRef, StringArray};
+use arrow::array::{Array, ArrayRef, BinaryViewArray, StringArray, StructArray};
 use arrow::util::test_util::seedable_rng;
+use arrow_schema::{DataType, Field, FieldRef, Fields};
 use criterion::{Criterion, criterion_group, criterion_main};
-use parquet_variant::{Variant, VariantBuilder};
+use parquet_variant::{EMPTY_VARIANT_METADATA_BYTES, Variant, VariantBuilder};
 use parquet_variant_compute::{
     GetOptions, VariantArray, VariantArrayBuilder, json_to_variant, variant_get,
 };
@@ -98,9 +99,26 @@ pub fn variant_get_bench(c: &mut Criterion) {
     });
 }
 
+pub fn variant_get_shredded_utf8_bench(c: &mut Criterion) {
+    let variant_array = create_shredded_utf8_variant_array(8192);
+    let input = ArrayRef::from(variant_array);
+
+    let field: FieldRef = Arc::new(Field::new("typed_value", DataType::Utf8, true));
+    let options = GetOptions {
+        path: vec![].into(),
+        as_type: Some(field),
+        cast_options: Default::default(),
+    };
+
+    c.bench_function("variant_get_shredded_utf8", |b| {
+        b.iter(|| variant_get(&input.clone(), options.clone()))
+    });
+}
+
 criterion_group!(
     benches,
     variant_get_bench,
+    variant_get_shredded_utf8_bench,
     benchmark_batch_json_string_to_variant
 );
 criterion_main!(benches);
@@ -121,6 +139,35 @@ fn create_primitive_variant_array(size: usize) -> VariantArray {
     variant_builder.build()
 }
 
+/// Creates a `VariantArray` where the values are already shredded as UTF8.
+fn create_shredded_utf8_variant_array(size: usize) -> VariantArray {
+    let metadata =
+        BinaryViewArray::from_iter_values(std::iter::repeat_n(EMPTY_VARIANT_METADATA_BYTES, size));
+    let typed_value = StringArray::from_iter_values((0..size).map(|i| format!("value_{i}")));
+
+    let metadata_ref: ArrayRef = Arc::new(metadata);
+    let typed_value_ref: ArrayRef = Arc::new(typed_value);
+
+    let fields = Fields::from(vec![
+        Arc::new(Field::new(
+            "metadata",
+            metadata_ref.data_type().clone(),
+            false,
+        )),
+        Arc::new(Field::new(
+            "typed_value",
+            typed_value_ref.data_type().clone(),
+            true,
+        )),
+    ]);
+
+    let struct_array = StructArray::new(fields, vec![metadata_ref, typed_value_ref], None);
+    let struct_array_ref: ArrayRef = Arc::new(struct_array);
+
+    VariantArray::try_new(struct_array_ref.as_ref())
+        .expect("created struct should be a valid shredded variant")
+}
+
 /// Return an iterator off JSON strings, each representing a person
 /// with random first name, last name, and age.
 ///

parquet-variant-compute/src/variant_get.rs

@@ -208,6 +208,11 @@ fn shredded_get_path(
         return Ok(ArrayRef::from(target));
     };
 
+    // Try to return the typed value directly when we have a perfect shredding match.
+    if let Some(shredded) = try_perfect_shredding(&target, as_field) {
+        return Ok(shredded);
+    }
+
     // Structs are special. Recurse into each field separately, hoping to follow the shredding even
     // further, and build up the final struct from those individually shredded results.
     if let DataType::Struct(fields) = as_field.data_type() {
@@ -236,6 +241,28 @@ fn shredded_get_path(
     shred_basic_variant(target, VariantPath::default(), Some(as_field))
 }
 
+fn try_perfect_shredding(variant_array: &VariantArray, as_field: &Field) -> Option<ArrayRef> {
+    // Try to return the typed value directly when we have a perfect shredding match.
+    if matches!(as_field.data_type(), DataType::Struct(_)) {
+        return None;
+    }
+    if let Some(typed_value) = variant_array.typed_value_field()
+        && typed_value.data_type() == as_field.data_type()
+        && variant_array
+            .value_field()
+            .is_none_or(|v| v.null_count() == v.len())
+    {
+        // Here we need to gate against the case where the `typed_value` is null but data is in the `value` column.
+        // 1. If the `value` column is null, or
+        // 2. If every row in the `value` column is null
+
+        // This is a perfect shredding, where the value is entirely shredded out,
+        // so we can just return the typed value.
+        return Some(typed_value.clone());
+    }
+    None
+}
+
 /// 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
@@ -3799,4 +3826,18 @@ mod test {
             assert!(result.is_null(i));
         }
     }
+
+    #[test]
+    fn test_perfect_shredding_returns_same_arc_ptr() {
+        let variant_array = perfectly_shredded_int32_variant_array();
+
+        let variant_array_ref = VariantArray::try_new(&variant_array).unwrap();
+        let typed_value_arc = variant_array_ref.typed_value_field().unwrap().clone();
+
+        let field = Field::new("result", DataType::Int32, true);
+        let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
+        let result = variant_get(&variant_array, options).unwrap();
+
+        assert!(Arc::ptr_eq(&typed_value_arc, &result));
+    }
 }