nice cast

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

Author: AdamGS

parquet-variant-compute/src/variant_get.rs

@@ -17,11 +17,11 @@
 use arrow::{
     array::{self, Array, ArrayRef, StructArray, make_array},
     buffer::NullBuffer,
-    compute::CastOptions,
+    compute::{CastOptions, can_cast_types, cast_with_options},
     datatypes::Field,
     error::Result,
 };
-use arrow_schema::{ArrowError, DataType, FieldRef};
+use arrow_schema::{ArrowError, DataType, FieldRef, TimeUnit};
 use parquet_variant::{VariantPath, VariantPathElement};
 
 use crate::VariantArray;
@@ -140,11 +140,12 @@ fn shredded_get_path(
                     let value = target.try_value(i)?;
                     builder.append_value(value)?;
                 } else {
-                    let _ = match target.try_value(i) {
-                        Ok(v) => builder.append_value(v)?,
+                    match target.try_value(i) {
+                        Ok(v) => {
+                            builder.append_value(v)?;
+                        }
                         Err(_) => {
                             builder.append_null()?;
-                            false // add this to make match arms have the same return type
                         }
                     };
                 }
@@ -201,7 +202,7 @@ fn shredded_get_path(
     };
 
     // Try to return the typed value directly when we have a perfect shredding match.
-    if let Some(shredded) = try_perfect_shredding(&target, as_field) {
+    if let Some(shredded) = try_perfect_shredding(&target, as_field, cast_options)? {
         return Ok(shredded);
     }
 
@@ -244,44 +245,131 @@ 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> {
+fn try_perfect_shredding(
+    variant_array: &VariantArray,
+    as_field: &Field,
+    cast_options: &CastOptions,
+) -> Result<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;
+        return Ok(None);
     }
-    let typed_value = variant_array.typed_value_field()?;
+    let Some(typed_value) = variant_array.typed_value_field() else {
+        return Ok(None);
+    };
 
-    if typed_value.data_type() == as_field.data_type()
-        && variant_array
-            .value_field()
-            .is_none_or(|v| v.null_count() == v.len())
+    if variant_array
+        .value_field()
+        .is_some_and(|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 after merging the accumulated nulls.
-        let parent_nulls = variant_array.nulls();
-
-        // If we have no nulls OR the shredded array is `Null`, which doesn't support external nulls.
-        let target_array = if parent_nulls.is_none() || typed_value.data_type().is_null() {
-            typed_value.clone()
-        } else {
-            let merged_nulls = NullBuffer::union(parent_nulls, typed_value.nulls());
-            let data = typed_value
-                .to_data()
-                .into_builder()
-                .nulls(merged_nulls)
-                .build()
-                .ok()?;
-            make_array(data)
-        };
+        // We can only use the perfect-shredding fast path when the value is entirely shredded out.
+        return Ok(None);
+    }
+
+    // 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 reuse the
+    // typed value after merging the accumulated nulls from the traversed parent fields.
+    let parent_nulls = variant_array.nulls();
+
+    // If we have no nulls OR the shredded array is `Null`, which doesn't support external nulls.
+    let target_array = if parent_nulls.is_none() || typed_value.data_type().is_null() {
+        typed_value.clone()
+    } else {
+        let merged_nulls = NullBuffer::union(parent_nulls, typed_value.nulls());
+        let data = typed_value
+            .to_data()
+            .into_builder()
+            .nulls(merged_nulls)
+            .build()?;
+        make_array(data)
+    };
+
+    if target_array.data_type() == as_field.data_type() {
+        return Ok(Some(target_array));
+    }
 
-        return Some(target_array);
+    if !can_use_perfect_shredding_arrow_cast(target_array.data_type(), as_field.data_type()) {
+        return Ok(None);
     }
 
-    None
+    // Use Arrow's vectorized cast when it cleanly matches the shredded representation. If not,
+    // fall back to row-wise extraction to preserve the existing variant-specific semantics.
+    Ok(cast_with_options(target_array.as_ref(), as_field.data_type(), cast_options).ok())
+}
+
+fn can_use_perfect_shredding_arrow_cast(from_type: &DataType, to_type: &DataType) -> bool {
+    use DataType::*;
+
+    if !can_cast_types(from_type, to_type) {
+        return false;
+    }
+
+    match from_type {
+        Null => true,
+        Boolean => is_non_decimal_numeric_or_bool(to_type),
+        Int8 | Int16 | Int32 | Int64 => {
+            is_non_decimal_numeric_or_bool(to_type)
+                || matches!(
+                    to_type,
+                    Decimal32(..) | Decimal64(..) | Decimal128(..) | Decimal256(..)
+                )
+        }
+        Float32 | Float64 => is_non_decimal_numeric_or_bool(to_type),
+        Decimal32(..) | Decimal64(..) | Decimal128(..) | Decimal256(..) => {
+            matches!(
+                to_type,
+                Decimal32(..) | Decimal64(..) | Decimal128(..) | Decimal256(..)
+            )
+        }
+        Date32 => matches!(to_type, Date32 | Date64),
+        Time64(TimeUnit::Microsecond) => matches!(
+            to_type,
+            Time64(TimeUnit::Microsecond) | Time64(TimeUnit::Nanosecond)
+        ),
+        Timestamp(TimeUnit::Microsecond, from_tz) => matches!(
+            to_type,
+            Timestamp(TimeUnit::Microsecond | TimeUnit::Nanosecond, to_tz)
+                if from_tz.is_some() == to_tz.is_some()
+        ),
+        Timestamp(TimeUnit::Nanosecond, from_tz) => matches!(
+            to_type,
+            Timestamp(TimeUnit::Nanosecond, to_tz) if from_tz.is_some() == to_tz.is_some()
+        ),
+        Utf8 | LargeUtf8 | Utf8View => matches!(to_type, Utf8 | LargeUtf8 | Utf8View),
+        Binary | LargeBinary | BinaryView => matches!(to_type, Binary | LargeBinary | BinaryView),
+        FixedSizeBinary(16) => matches!(to_type, FixedSizeBinary(16)),
+        List(_) | LargeList(_) | ListView(_) | LargeListView(_) => {
+            matches!(
+                to_type,
+                List(_) | LargeList(_) | ListView(_) | LargeListView(_)
+            )
+        }
+        _ => false,
+    }
+}
+
+fn is_non_decimal_numeric_or_bool(data_type: &DataType) -> bool {
+    use DataType::*;
+
+    matches!(
+        data_type,
+        Boolean
+            | Int8
+            | Int16
+            | Int32
+            | Int64
+            | UInt8
+            | UInt16
+            | UInt32
+            | UInt64
+            | Float16
+            | Float32
+            | Float64
+    )
 }
 
 /// Returns an array with the specified path extracted from the variant values.
@@ -353,7 +441,7 @@ mod test {
     use std::str::FromStr;
     use std::sync::Arc;
 
-    use super::{GetOptions, variant_get};
+    use super::{GetOptions, try_perfect_shredding, variant_get};
     use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder};
     use crate::{
         VariantArray, VariantArrayBuilder, cast_to_variant, json_to_variant, shred_variant,
@@ -3974,6 +4062,27 @@ mod test {
         assert!(Arc::ptr_eq(&typed_value_arc, &result));
     }
 
+    #[test]
+    fn test_try_perfect_shredding_casts_typed_value_and_preserves_parent_nulls() {
+        let typed_value: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)]));
+        let metadata = Arc::new(BinaryViewArray::from_iter_values(std::iter::repeat_n(
+            EMPTY_VARIANT_METADATA_BYTES,
+            typed_value.len(),
+        )));
+        let parent_nulls = Some(NullBuffer::from(vec![true, false, true]));
+        let variant_array =
+            VariantArray::from_parts(metadata, None, Some(typed_value), parent_nulls);
+        let requested_field = Field::new("result", DataType::Int64, true);
+
+        let result =
+            try_perfect_shredding(&variant_array, &requested_field, &CastOptions::default())
+                .unwrap()
+                .unwrap();
+        let expected: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]));
+
+        assert_eq!(&result, &expected);
+    }
+
     #[test]
     fn test_perfect_shredding_three_typed_value_columns() {
         // Column 1: perfectly shredded primitive with all nulls