diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index 38c577564dea..76666b26fb46 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -20,11 +20,11 @@ use arrow::{ UInt64Array, make_array, }, buffer::NullBuffer, - compute::{CastOptions, take}, + compute::{CastOptions, can_cast_types, cast_with_options, take}, datatypes::Field, error::Result, }; -use arrow_schema::{ArrowError, DataType, FieldRef}; +use arrow_schema::{ArrowError, DataType, FieldRef, TimeUnit}; use parquet_variant::{VariantPath, VariantPathElement}; use crate::ShreddingState; @@ -219,11 +219,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 } }; } @@ -280,7 +281,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); } @@ -323,44 +324,167 @@ fn shredded_get_path( shred_basic_variant(target, VariantPath::default(), Some(as_field)) } -fn try_perfect_shredding(variant_array: &VariantArray, as_field: &Field) -> Option { +fn try_perfect_shredding( + variant_array: &VariantArray, + as_field: &Field, + cast_options: &CastOptions, +) -> Result> { // 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) + }; - return Some(target_array); + if target_array.data_type() == as_field.data_type() { + return Ok(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. + cast_with_options(target_array.as_ref(), as_field.data_type(), cast_options).map(Some) +} + +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(..) + ) + } + Float16 | Float32 | Float64 => { + is_non_decimal_numeric_or_bool(to_type) || to_type.is_decimal() + } + Decimal32(..) | Decimal64(..) | Decimal128(..) | Decimal256(..) => { + to_type.is_decimal() + || (is_decimal_with_non_negative_scale(from_type) + && is_non_decimal_numeric(to_type)) + } + 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) + || is_decimal_with_non_negative_scale(to_type) + } + Binary | LargeBinary | BinaryView => matches!(to_type, Binary | LargeBinary | BinaryView), + FixedSizeBinary(16) => matches!(to_type, FixedSizeBinary(16)), + List(_) | LargeList(_) | ListView(_) | LargeListView(_) => match (from_type, to_type) { + ( + List(from_field) + | LargeList(from_field) + | ListView(from_field) + | LargeListView(from_field), + List(to_field) | LargeList(to_field) | ListView(to_field) | LargeListView(to_field), + ) => can_use_perfect_shredding_arrow_cast(from_field.data_type(), to_field.data_type()), + _ => false, + }, + _ => false, + } +} + +fn is_decimal_with_non_negative_scale(data_type: &DataType) -> bool { + use DataType::*; + + matches!( + data_type, + Decimal32(_, scale) | Decimal64(_, scale) | Decimal128(_, scale) | Decimal256(_, scale) + if *scale >= 0 + ) +} + +fn is_non_decimal_numeric(data_type: &DataType) -> bool { + use DataType::*; + + matches!( + data_type, + Int8 | Int16 + | Int32 + | Int64 + | UInt8 + | UInt16 + | UInt32 + | UInt64 + | Float16 + | Float32 + | Float64 + ) +} + +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. @@ -432,7 +556,9 @@ mod test { use std::str::FromStr; use std::sync::Arc; - use super::{GetOptions, variant_get}; + use super::{ + GetOptions, can_use_perfect_shredding_arrow_cast, try_perfect_shredding, variant_get, + }; use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder}; use crate::{ ShreddedSchemaBuilder, VariantArray, VariantArrayBuilder, cast_to_variant, json_to_variant, @@ -448,7 +574,7 @@ mod test { Time64NanosecondArray, }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; - use arrow::compute::CastOptions; + use arrow::compute::{CastOptions, can_cast_types}; use arrow::datatypes::DataType::{Int16, Int32, Int64}; use arrow::datatypes::i256; use arrow::util::display::FormatOptions; @@ -4287,6 +4413,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 @@ -4512,6 +4659,162 @@ mod test { } } + #[test] + fn test_perfect_shredding_list_cast_gate_uses_variant_element_semantics() { + let int64_item = Arc::new(Field::new("item", DataType::Int64, true)); + let utf8_item = Arc::new(Field::new("item", DataType::Utf8, true)); + assert!(!can_use_perfect_shredding_arrow_cast( + &DataType::List(int64_item), + &DataType::List(utf8_item), + )); + + let int32_item = Arc::new(Field::new("item", DataType::Int32, true)); + let int64_item = Arc::new(Field::new("item", DataType::Int64, true)); + assert!(can_use_perfect_shredding_arrow_cast( + &DataType::List(int32_item), + &DataType::List(int64_item), + )); + } + + #[test] + fn test_perfect_shredding_cast_gate_allows_supported_variant_casts() { + use DataType::*; + + let supported_casts = vec![ + (Null, Utf8), + (Boolean, Int64), + (Boolean, Float64), + (Int32, UInt64), + (Int64, Decimal128(38, 2)), + (Float64, Decimal128(38, 2)), + (Date32, Date64), + (Time64(TimeUnit::Microsecond), Time64(TimeUnit::Nanosecond)), + ( + Timestamp(TimeUnit::Microsecond, None), + Timestamp(TimeUnit::Nanosecond, None), + ), + ( + Timestamp(TimeUnit::Microsecond, Some(Arc::from("+00:00"))), + Timestamp(TimeUnit::Nanosecond, Some(Arc::from("+00:00"))), + ), + (Utf8, LargeUtf8), + (Utf8View, Decimal128(38, 2)), + (Binary, LargeBinary), + (FixedSizeBinary(16), FixedSizeBinary(16)), + ( + List(Arc::new(Field::new("item", Int32, true))), + LargeList(Arc::new(Field::new("item", Int64, true))), + ), + ]; + + for (from_type, to_type) in supported_casts { + assert!( + can_cast_types(&from_type, &to_type), + "Arrow should support cast from {from_type} to {to_type}" + ); + assert!( + can_use_perfect_shredding_arrow_cast(&from_type, &to_type), + "perfect shredding should use Arrow cast from {from_type} to {to_type}" + ); + } + } + + #[test] + fn test_perfect_shredding_cast_gate_rejects_arrow_casts_without_variant_semantics() { + use DataType::*; + + let unsupported_arrow_casts = vec![ + (Boolean, Utf8), + (Int64, Utf8), + (UInt64, Decimal128(38, 2)), + (Utf8, Int64), + (Utf8, Date32), + (Utf8, Timestamp(TimeUnit::Microsecond, None)), + (Date32, Int64), + (Time64(TimeUnit::Microsecond), Int64), + ( + Timestamp(TimeUnit::Microsecond, Some(Arc::from("+00:00"))), + Timestamp(TimeUnit::Microsecond, None), + ), + (Timestamp(TimeUnit::Microsecond, None), Int64), + ( + Timestamp(TimeUnit::Second, None), + Timestamp(TimeUnit::Microsecond, None), + ), + (Binary, Utf8), + (FixedSizeBinary(16), Binary), + (Dictionary(Box::new(Int32), Box::new(Utf8)), Utf8), + (List(Arc::new(Field::new("item", Int64, true))), Utf8), + ( + List(Arc::new(Field::new("item", Int64, true))), + List(Arc::new(Field::new("item", Utf8, true))), + ), + (Decimal32(9, 2), Utf8), + (Decimal32(9, -1), Int64), + (Utf8, Decimal32(9, -1)), + ]; + + for (from_type, to_type) in unsupported_arrow_casts { + assert!( + can_cast_types(&from_type, &to_type), + "Arrow should support cast from {from_type} to {to_type}" + ); + assert!( + !can_use_perfect_shredding_arrow_cast(&from_type, &to_type), + "perfect shredding should fall back for {from_type} to {to_type}" + ); + } + } + + #[test] + fn test_perfect_shredding_cast_gate_matches_decimal_semantics() { + use DataType::*; + + let decimal_targets = [ + Decimal32(9, 2), + Decimal64(18, 2), + Decimal128(38, 2), + Decimal256(76, 2), + ]; + for from_type in [Float16, Float32, Float64, Utf8, LargeUtf8, Utf8View] { + for to_type in &decimal_targets { + assert!( + can_use_perfect_shredding_arrow_cast(&from_type, to_type), + "{from_type} should use Arrow cast to {to_type}" + ); + } + } + + for from_type in [ + Decimal32(9, 2), + Decimal64(18, 2), + Decimal128(38, 2), + Decimal256(76, 2), + ] { + for to_type in [ + Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64, Float16, Float32, Float64, + ] { + assert!( + can_use_perfect_shredding_arrow_cast(&from_type, &to_type), + "{from_type} should use Arrow cast to {to_type}" + ); + } + } + + assert!(!can_use_perfect_shredding_arrow_cast( + &Utf8, + &Decimal32(9, -1), + )); + assert!(!can_use_perfect_shredding_arrow_cast( + &Decimal32(9, -1), + &Int64, + )); + assert!(!can_use_perfect_shredding_arrow_cast( + &Decimal32(9, 2), + &Utf8, + )); + } + #[test] fn test_variant_get_nested_list() { use arrow::datatypes::Int64Type;