From c5c282a09a8abca4ec9dfdd384794812732c37ec Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Wed, 22 Apr 2026 14:29:43 -0400 Subject: [PATCH 1/5] remove `BorrowedShreddingState` --- parquet-variant-compute/src/lib.rs | 2 +- .../src/unshred_variant.rs | 231 ++++++++++-------- parquet-variant-compute/src/variant_array.rs | 73 +----- parquet-variant-compute/src/variant_get.rs | 16 +- 4 files changed, 144 insertions(+), 178 deletions(-) diff --git a/parquet-variant-compute/src/lib.rs b/parquet-variant-compute/src/lib.rs index b05d0e023653..066fe15b7c8b 100644 --- a/parquet-variant-compute/src/lib.rs +++ b/parquet-variant-compute/src/lib.rs @@ -51,7 +51,7 @@ mod variant_array_builder; mod variant_get; mod variant_to_arrow; -pub use variant_array::{BorrowedShreddingState, ShreddingState, VariantArray, VariantType}; +pub use variant_array::{ShreddingState, VariantArray, VariantType}; pub use variant_array_builder::{VariantArrayBuilder, VariantValueArrayBuilder}; pub use cast_to_variant::{cast_to_variant, cast_to_variant_with_options}; diff --git a/parquet-variant-compute/src/unshred_variant.rs b/parquet-variant-compute/src/unshred_variant.rs index 7bf31e2256df..07932463698d 100644 --- a/parquet-variant-compute/src/unshred_variant.rs +++ b/parquet-variant-compute/src/unshred_variant.rs @@ -18,7 +18,7 @@ //! Module for unshredding VariantArray by folding typed_value columns back into the value column. use crate::variant_array::binary_array_value; -use crate::{BorrowedShreddingState, VariantArray, VariantValueArrayBuilder}; +use crate::{ShreddingState, VariantArray, VariantValueArrayBuilder}; use arrow::array::{ Array, ArrayRef, AsArray as _, BinaryArray, BinaryViewArray, BooleanArray, FixedSizeBinaryArray, FixedSizeListArray, GenericListArray, GenericListViewArray, @@ -67,8 +67,8 @@ pub fn unshred_variant(array: &VariantArray) -> Result { // NOTE: None/None at top-level is technically invalid, but the shredding spec requires us to // emit `Variant::Null` when a required value is missing. let nulls = array.nulls(); - let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.shredding_state().borrow())? - .unwrap_or_else(|| UnshredVariantRowBuilder::null(nulls)); + let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.shredding_state())? + .unwrap_or_else(|| UnshredVariantRowBuilder::null(nulls.cloned())); let metadata = array.metadata_field(); let mut value_builder = VariantValueArrayBuilder::new(array.len()); @@ -97,41 +97,41 @@ pub fn unshred_variant(array: &VariantArray) -> Result { } /// Row builder for converting shredded VariantArray rows back to unshredded form -enum UnshredVariantRowBuilder<'a> { - PrimitiveInt8(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), - PrimitiveInt16(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), - PrimitiveInt32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), - PrimitiveInt64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), - PrimitiveFloat32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), - PrimitiveFloat64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), - Decimal32(DecimalUnshredRowBuilder<'a, Decimal32Type, VariantDecimal4>), - Decimal64(DecimalUnshredRowBuilder<'a, Decimal64Type, VariantDecimal8>), - Decimal128(DecimalUnshredRowBuilder<'a, Decimal128Type, VariantDecimal16>), - PrimitiveDate32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), - PrimitiveTime64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), - TimestampMicrosecond(TimestampUnshredRowBuilder<'a, TimestampMicrosecondType>), - TimestampNanosecond(TimestampUnshredRowBuilder<'a, TimestampNanosecondType>), - PrimitiveBoolean(UnshredPrimitiveRowBuilder<'a, BooleanArray>), - PrimitiveString(UnshredPrimitiveRowBuilder<'a, StringArray>), - PrimitiveStringView(UnshredPrimitiveRowBuilder<'a, StringViewArray>), - PrimitiveLargeString(UnshredPrimitiveRowBuilder<'a, LargeStringArray>), - PrimitiveBinary(UnshredPrimitiveRowBuilder<'a, BinaryArray>), - PrimitiveBinaryView(UnshredPrimitiveRowBuilder<'a, BinaryViewArray>), - PrimitiveLargeBinary(UnshredPrimitiveRowBuilder<'a, LargeBinaryArray>), - PrimitiveUuid(UnshredPrimitiveRowBuilder<'a, FixedSizeBinaryArray>), - List(ListUnshredVariantBuilder<'a, GenericListArray>), - LargeList(ListUnshredVariantBuilder<'a, GenericListArray>), - ListView(ListUnshredVariantBuilder<'a, GenericListViewArray>), - LargeListView(ListUnshredVariantBuilder<'a, GenericListViewArray>), - FixedSizeList(ListUnshredVariantBuilder<'a, FixedSizeListArray>), - Struct(StructUnshredVariantBuilder<'a>), - ValueOnly(ValueOnlyUnshredVariantBuilder<'a>), - Null(NullUnshredVariantBuilder<'a>), +enum UnshredVariantRowBuilder { + PrimitiveInt8(UnshredPrimitiveRowBuilder>), + PrimitiveInt16(UnshredPrimitiveRowBuilder>), + PrimitiveInt32(UnshredPrimitiveRowBuilder>), + PrimitiveInt64(UnshredPrimitiveRowBuilder>), + PrimitiveFloat32(UnshredPrimitiveRowBuilder>), + PrimitiveFloat64(UnshredPrimitiveRowBuilder>), + Decimal32(DecimalUnshredRowBuilder), + Decimal64(DecimalUnshredRowBuilder), + Decimal128(DecimalUnshredRowBuilder), + PrimitiveDate32(UnshredPrimitiveRowBuilder>), + PrimitiveTime64(UnshredPrimitiveRowBuilder>), + TimestampMicrosecond(TimestampUnshredRowBuilder), + TimestampNanosecond(TimestampUnshredRowBuilder), + PrimitiveBoolean(UnshredPrimitiveRowBuilder), + PrimitiveString(UnshredPrimitiveRowBuilder), + PrimitiveStringView(UnshredPrimitiveRowBuilder), + PrimitiveLargeString(UnshredPrimitiveRowBuilder), + PrimitiveBinary(UnshredPrimitiveRowBuilder), + PrimitiveBinaryView(UnshredPrimitiveRowBuilder), + PrimitiveLargeBinary(UnshredPrimitiveRowBuilder), + PrimitiveUuid(UnshredPrimitiveRowBuilder), + List(ListUnshredVariantBuilder>), + LargeList(ListUnshredVariantBuilder>), + ListView(ListUnshredVariantBuilder>), + LargeListView(ListUnshredVariantBuilder>), + FixedSizeList(ListUnshredVariantBuilder), + Struct(StructUnshredVariantBuilder), + ValueOnly(ValueOnlyUnshredVariantBuilder), + Null(NullUnshredVariantBuilder), } -impl<'a> UnshredVariantRowBuilder<'a> { +impl UnshredVariantRowBuilder { /// Creates an all-null row builder. - fn null(nulls: Option<&'a NullBuffer>) -> Self { + fn null(nulls: Option) -> Self { Self::Null(NullUnshredVariantBuilder::new(nulls)) } @@ -175,12 +175,11 @@ impl<'a> UnshredVariantRowBuilder<'a> { } } - /// Creates a new UnshredVariantRowBuilder from shredding state - /// Returns None for None/None case - caller decides how to handle based on context - fn try_new_opt(shredding_state: BorrowedShreddingState<'a>) -> Result> { - let value = shredding_state.value_field(); - let typed_value = shredding_state.typed_value_field(); - let Some(typed_value) = typed_value else { + /// Creates a new UnshredVariantRowBuilder from the given shredding state. + /// Returns None for the None/None case - caller decides how to handle based on context. + fn try_new_opt(shredding_state: &ShreddingState) -> Result> { + let value = shredding_state.value_field().cloned(); + let Some(typed_value) = shredding_state.typed_value_field().cloned() else { // Copy the value across directly, if present. Else caller decides what to do. return Ok(value.map(|v| Self::ValueOnly(ValueOnlyUnshredVariantBuilder::new(v)))); }; @@ -190,7 +189,7 @@ impl<'a> UnshredVariantRowBuilder<'a> { ($enum_variant:ident, $cast_fn:ident) => { Self::$enum_variant(UnshredPrimitiveRowBuilder::new( value, - typed_value.$cast_fn(), + typed_value.$cast_fn().clone(), )) }; } @@ -203,13 +202,25 @@ impl<'a> UnshredVariantRowBuilder<'a> { DataType::Float32 => primitive_builder!(PrimitiveFloat32, as_primitive), DataType::Float64 => primitive_builder!(PrimitiveFloat64, as_primitive), DataType::Decimal32(p, s) if VariantDecimal4::is_valid_precision_and_scale(p, s) => { - Self::Decimal32(DecimalUnshredRowBuilder::new(value, typed_value, *s as _)) + Self::Decimal32(DecimalUnshredRowBuilder::new( + value, + typed_value.as_primitive().clone(), + *s as _, + )) } DataType::Decimal64(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => { - Self::Decimal64(DecimalUnshredRowBuilder::new(value, typed_value, *s as _)) + Self::Decimal64(DecimalUnshredRowBuilder::new( + value, + typed_value.as_primitive().clone(), + *s as _, + )) } DataType::Decimal128(p, s) if VariantDecimal16::is_valid_precision_and_scale(p, s) => { - Self::Decimal128(DecimalUnshredRowBuilder::new(value, typed_value, *s as _)) + Self::Decimal128(DecimalUnshredRowBuilder::new( + value, + typed_value.as_primitive().clone(), + *s as _, + )) } DataType::Decimal32(_, _) | DataType::Decimal64(_, _) @@ -229,12 +240,20 @@ impl<'a> UnshredVariantRowBuilder<'a> { "Time64({time_unit}) is not a valid variant shredding type", ))); } - DataType::Timestamp(TimeUnit::Microsecond, timezone) => Self::TimestampMicrosecond( - TimestampUnshredRowBuilder::new(value, typed_value, timezone.is_some()), - ), - DataType::Timestamp(TimeUnit::Nanosecond, timezone) => Self::TimestampNanosecond( - TimestampUnshredRowBuilder::new(value, typed_value, timezone.is_some()), - ), + DataType::Timestamp(TimeUnit::Microsecond, timezone) => { + Self::TimestampMicrosecond(TimestampUnshredRowBuilder::new( + value, + typed_value.as_primitive().clone(), + timezone.is_some(), + )) + } + DataType::Timestamp(TimeUnit::Nanosecond, timezone) => { + Self::TimestampNanosecond(TimestampUnshredRowBuilder::new( + value, + typed_value.as_primitive().clone(), + timezone.is_some(), + )) + } DataType::Timestamp(time_unit, _) => { return Err(ArrowError::InvalidArgumentError(format!( "Timestamp({time_unit}) is not a valid variant shredding type", @@ -257,27 +276,30 @@ impl<'a> UnshredVariantRowBuilder<'a> { } DataType::Struct(_) => Self::Struct(StructUnshredVariantBuilder::try_new( value, - typed_value.as_struct(), + typed_value.as_struct().clone(), )?), DataType::List(_) => Self::List(ListUnshredVariantBuilder::try_new( value, - typed_value.as_list(), + typed_value.as_list::().clone(), )?), DataType::LargeList(_) => Self::LargeList(ListUnshredVariantBuilder::try_new( value, - typed_value.as_list(), + typed_value.as_list::().clone(), )?), DataType::ListView(_) => Self::ListView(ListUnshredVariantBuilder::try_new( value, - typed_value.as_list_view(), + typed_value.as_list_view::().clone(), )?), DataType::LargeListView(_) => Self::LargeListView(ListUnshredVariantBuilder::try_new( value, - typed_value.as_list_view(), + typed_value.as_list_view::().clone(), )?), - DataType::FixedSizeList(_, _) => Self::FixedSizeList( - ListUnshredVariantBuilder::try_new(value, typed_value.as_fixed_size_list())?, - ), + DataType::FixedSizeList(_, _) => { + Self::FixedSizeList(ListUnshredVariantBuilder::try_new( + value, + typed_value.as_fixed_size_list().clone(), + )?) + } _ => { return Err(ArrowError::NotYetImplemented(format!( "Unshredding not yet supported for type: {}", @@ -290,12 +312,12 @@ impl<'a> UnshredVariantRowBuilder<'a> { } /// Builder for arrays with neither typed_value nor value (all NULL/Variant::Null) -struct NullUnshredVariantBuilder<'a> { - nulls: Option<&'a NullBuffer>, +struct NullUnshredVariantBuilder { + nulls: Option, } -impl<'a> NullUnshredVariantBuilder<'a> { - fn new(nulls: Option<&'a NullBuffer>) -> Self { +impl NullUnshredVariantBuilder { + fn new(nulls: Option) -> Self { Self { nulls } } @@ -305,7 +327,11 @@ impl<'a> NullUnshredVariantBuilder<'a> { _metadata: &VariantMetadata, index: usize, ) -> Result<()> { - if self.nulls.is_some_and(|nulls| nulls.is_null(index)) { + if self + .nulls + .as_ref() + .is_some_and(|nulls| nulls.is_null(index)) + { builder.append_null(); } else { builder.append_value(Variant::Null); @@ -315,12 +341,12 @@ impl<'a> NullUnshredVariantBuilder<'a> { } /// Builder for arrays that only have value column (already unshredded) -struct ValueOnlyUnshredVariantBuilder<'a> { - value: &'a ArrayRef, +struct ValueOnlyUnshredVariantBuilder { + value: ArrayRef, } -impl<'a> ValueOnlyUnshredVariantBuilder<'a> { - fn new(value: &'a ArrayRef) -> Self { +impl ValueOnlyUnshredVariantBuilder { + fn new(value: ArrayRef) -> Self { Self { value } } @@ -394,13 +420,13 @@ macro_rules! handle_unshredded_case { } /// Generic unshred builder that works with any Array implementing AppendToVariantBuilder -struct UnshredPrimitiveRowBuilder<'a, T> { - value: Option<&'a ArrayRef>, - typed_value: &'a T, +struct UnshredPrimitiveRowBuilder { + value: Option, + typed_value: T, } -impl<'a, T: AppendToVariantBuilder> UnshredPrimitiveRowBuilder<'a, T> { - fn new(value: Option<&'a ArrayRef>, typed_value: &'a T) -> Self { +impl UnshredPrimitiveRowBuilder { + fn new(value: Option, typed_value: T) -> Self { Self { value, typed_value } } @@ -497,17 +523,17 @@ impl TimestampType for TimestampNanosecondType { } /// Generic builder for timestamp types that handles timezone-aware conversion -struct TimestampUnshredRowBuilder<'a, T: TimestampType> { - value: Option<&'a ArrayRef>, - typed_value: &'a PrimitiveArray, +struct TimestampUnshredRowBuilder { + value: Option, + typed_value: PrimitiveArray, has_timezone: bool, } -impl<'a, T: TimestampType> TimestampUnshredRowBuilder<'a, T> { - fn new(value: Option<&'a ArrayRef>, typed_value: &'a dyn Array, has_timezone: bool) -> Self { +impl TimestampUnshredRowBuilder { + fn new(value: Option, typed_value: PrimitiveArray, has_timezone: bool) -> Self { Self { value, - typed_value: typed_value.as_primitive(), + typed_value, has_timezone, } } @@ -533,24 +559,24 @@ impl<'a, T: TimestampType> TimestampUnshredRowBuilder<'a, T> { } /// Generic builder for decimal unshredding -struct DecimalUnshredRowBuilder<'a, A: DecimalType, V> +struct DecimalUnshredRowBuilder where V: VariantDecimalType, { - value: Option<&'a ArrayRef>, - typed_value: &'a PrimitiveArray, + value: Option, + typed_value: PrimitiveArray, scale: i8, _phantom: PhantomData, } -impl<'a, A: DecimalType, V> DecimalUnshredRowBuilder<'a, A, V> +impl DecimalUnshredRowBuilder where V: VariantDecimalType, { - fn new(value: Option<&'a ArrayRef>, typed_value: &'a dyn Array, scale: i8) -> Self { + fn new(value: Option, typed_value: PrimitiveArray, scale: i8) -> Self { Self { value, - typed_value: typed_value.as_primitive(), + typed_value, scale, _phantom: PhantomData, } @@ -572,14 +598,14 @@ where } /// Builder for unshredding struct/object types with nested fields -struct StructUnshredVariantBuilder<'a> { - value: Option<&'a ArrayRef>, - typed_value: &'a arrow::array::StructArray, - field_unshredders: IndexMap<&'a str, Option>>, +struct StructUnshredVariantBuilder { + value: Option, + typed_value: StructArray, + field_unshredders: IndexMap>, } -impl<'a> StructUnshredVariantBuilder<'a> { - fn try_new(value: Option<&'a ArrayRef>, typed_value: &'a StructArray) -> Result { +impl StructUnshredVariantBuilder { + fn try_new(value: Option, typed_value: StructArray) -> Result { // Create unshredders for each field in constructor let mut field_unshredders = IndexMap::new(); for (field, field_array) in typed_value.fields().iter().zip(typed_value.columns()) { @@ -590,8 +616,9 @@ impl<'a> StructUnshredVariantBuilder<'a> { field_array.data_type() ))); }; - let field_unshredder = UnshredVariantRowBuilder::try_new_opt(field_array.try_into()?)?; - field_unshredders.insert(field.name().as_ref(), field_unshredder); + let field_state = ShreddingState::try_from(field_array)?; + let field_unshredder = UnshredVariantRowBuilder::try_new_opt(&field_state)?; + field_unshredders.insert(field.name().clone(), field_unshredder); } Ok(Self { @@ -615,7 +642,8 @@ impl<'a> StructUnshredVariantBuilder<'a> { // Process typed fields (skip empty builders that indicate missing fields) for (field_name, field_unshredder_opt) in &mut self.field_unshredders { if let Some(field_unshredder) = field_unshredder_opt { - let mut field_builder = ObjectFieldBuilder::new(field_name, &mut object_builder); + let mut field_builder = + ObjectFieldBuilder::new(field_name.as_str(), &mut object_builder); field_unshredder.append_row(&mut field_builder, metadata, index)?; } } @@ -645,14 +673,14 @@ impl<'a> StructUnshredVariantBuilder<'a> { } /// Builder for unshredding list/array types with recursive element processing -struct ListUnshredVariantBuilder<'a, L: ListLikeArray> { - value: Option<&'a ArrayRef>, - typed_value: &'a L, - element_unshredder: Box>, +struct ListUnshredVariantBuilder { + value: Option, + typed_value: L, + element_unshredder: Box, } -impl<'a, L: ListLikeArray> ListUnshredVariantBuilder<'a, L> { - fn try_new(value: Option<&'a ArrayRef>, typed_value: &'a L) -> Result { +impl ListUnshredVariantBuilder { + fn try_new(value: Option, typed_value: L) -> Result { // Create a recursive unshredder for the list elements // The element type comes from the values array of the list let element_values = typed_value.values(); @@ -670,7 +698,8 @@ impl<'a, L: ListLikeArray> ListUnshredVariantBuilder<'a, L> { // // NOTE: A None/None array element is technically invalid, but the shredding spec // requires us to emit `Variant::Null` when a required value is missing. - let element_unshredder = UnshredVariantRowBuilder::try_new_opt(element_values.try_into()?)? + let element_state = ShreddingState::try_from(element_values)?; + let element_unshredder = UnshredVariantRowBuilder::try_new_opt(&element_state)? .unwrap_or_else(|| UnshredVariantRowBuilder::null(None)); Ok(Self { diff --git a/parquet-variant-compute/src/variant_array.rs b/parquet-variant-compute/src/variant_array.rs index 5f52939e712c..a82ed11c5565 100644 --- a/parquet-variant-compute/src/variant_array.rs +++ b/parquet-variant-compute/src/variant_array.rs @@ -843,14 +843,6 @@ impl ShreddingState { self.typed_value.as_ref() } - /// Returns a borrowed version of this shredding state - pub fn borrow(&self) -> BorrowedShreddingState<'_> { - BorrowedShreddingState { - value: self.value_field(), - typed_value: self.typed_value_field(), - } - } - /// Slice all the underlying arrays pub fn slice(&self, offset: usize, length: usize) -> Self { Self { @@ -860,74 +852,19 @@ impl ShreddingState { } } -/// Similar to [`ShreddingState`] except it holds borrowed references of the target arrays. Useful -/// for avoiding clone operations when the caller does not need a self-standing shredding state. -#[derive(Clone, Debug)] -pub struct BorrowedShreddingState<'a> { - value: Option<&'a ArrayRef>, - typed_value: Option<&'a ArrayRef>, -} - -impl<'a> BorrowedShreddingState<'a> { - /// Create a new `BorrowedShreddingState` from the given `value` and `typed_value` fields - /// - /// Note you can create a `BorrowedShreddingState` from a &[`StructArray`] using - /// `BorrowedShreddingState::try_from(&struct_array)`, for example: - /// - /// ```no_run - /// # use arrow::array::StructArray; - /// # use parquet_variant_compute::BorrowedShreddingState; - /// # fn get_struct_array() -> StructArray { - /// # unimplemented!() - /// # } - /// let struct_array: StructArray = get_struct_array(); - /// let shredding_state = BorrowedShreddingState::try_from(&struct_array).unwrap(); - /// ``` - pub fn new(value: Option<&'a ArrayRef>, typed_value: Option<&'a ArrayRef>) -> Self { - Self { value, typed_value } - } - - /// Return a reference to the value field, if present - pub fn value_field(&self) -> Option<&'a ArrayRef> { - self.value - } - - /// Return a reference to the typed_value field, if present - pub fn typed_value_field(&self) -> Option<&'a ArrayRef> { - self.typed_value - } -} - -impl<'a> TryFrom<&'a StructArray> for BorrowedShreddingState<'a> { +impl TryFrom<&StructArray> for ShreddingState { type Error = ArrowError; - fn try_from(inner_struct: &'a StructArray) -> Result { + fn try_from(inner_struct: &StructArray) -> Result { // The `value` column need not exist, but if it does it must be a binary type. let value = if let Some(value_col) = inner_struct.column_by_name("value") { validate_binary_array(value_col.as_ref(), "value")?; - Some(value_col) + Some(value_col.clone()) } else { None }; - let typed_value = inner_struct.column_by_name("typed_value"); - Ok(BorrowedShreddingState::new(value, typed_value)) - } -} - -impl TryFrom<&StructArray> for ShreddingState { - type Error = ArrowError; - - fn try_from(inner_struct: &StructArray) -> Result { - Ok(BorrowedShreddingState::try_from(inner_struct)?.into()) - } -} - -impl From> for ShreddingState { - fn from(state: BorrowedShreddingState<'_>) -> Self { - ShreddingState { - value: state.value_field().cloned(), - typed_value: state.typed_value_field().cloned(), - } + let typed_value = inner_struct.column_by_name("typed_value").cloned(); + Ok(ShreddingState::new(value, typed_value)) } } diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index e95c774f2926..172f41a1767e 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -25,15 +25,15 @@ use arrow_schema::{ArrowError, DataType, FieldRef}; use parquet_variant::{VariantPath, VariantPathElement}; use crate::VariantArray; -use crate::variant_array::BorrowedShreddingState; +use crate::variant_array::ShreddingState; use crate::variant_to_arrow::make_variant_to_arrow_row_builder; use arrow::array::AsArray; use std::sync::Arc; -pub(crate) enum ShreddedPathStep<'a> { +pub(crate) enum ShreddedPathStep { /// Path step succeeded, return the new shredding state - Success(BorrowedShreddingState<'a>), + Success(ShreddingState), /// The path element is not present in the `typed_value` column and there is no `value` column, /// so we know it does not exist. It, and all paths under it, are all-NULL. Missing, @@ -49,11 +49,11 @@ pub(crate) enum ShreddedPathStep<'a> { /// a missing-path step (`Missing` or `NotShredded` depending on whether `value` exists). /// /// TODO: Support `VariantPathElement::Index`? It wouldn't be easy, and maybe not even possible. -pub(crate) fn follow_shredded_path_element<'a>( - shredding_state: &BorrowedShreddingState<'a>, +pub(crate) fn follow_shredded_path_element( + shredding_state: &ShreddingState, path_element: &VariantPathElement<'_>, _cast_options: &CastOptions, -) -> Result> { +) -> Result { // If the requested path element is not present in `typed_value`, and `value` is missing, then // we know it does not exist; it, and all paths under it, are all-NULL. let missing_path_step = || match shredding_state.value_field() { @@ -90,7 +90,7 @@ pub(crate) fn follow_shredded_path_element<'a>( )) })?; - let state = BorrowedShreddingState::try_from(struct_array)?; + let state = ShreddingState::try_from(struct_array)?; Ok(ShreddedPathStep::Success(state)) } VariantPathElement::Index { .. } => { @@ -154,7 +154,7 @@ fn shredded_get_path( // Peel away the prefix of path elements that traverses the shredded parts of this variant // column. Shredding will traverse the rest of the path on a per-row basis. - let mut shredding_state = input.shredding_state().borrow(); + let mut shredding_state = input.shredding_state().clone(); let mut accumulated_nulls = input.inner().nulls().cloned(); let mut path_index = 0; for path_element in path { From 14662dd9ca01db0308411d3ca800b15b352ac87c Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Wed, 22 Apr 2026 15:21:12 -0400 Subject: [PATCH 2/5] Remove dead code --- .../src/unshred_variant.rs | 29 +++++-------------- 1 file changed, 7 insertions(+), 22 deletions(-) diff --git a/parquet-variant-compute/src/unshred_variant.rs b/parquet-variant-compute/src/unshred_variant.rs index 07932463698d..c035fbf65c9e 100644 --- a/parquet-variant-compute/src/unshred_variant.rs +++ b/parquet-variant-compute/src/unshred_variant.rs @@ -25,7 +25,6 @@ use arrow::array::{ LargeBinaryArray, LargeStringArray, ListLikeArray, PrimitiveArray, StringArray, StringViewArray, StructArray, }; -use arrow::buffer::NullBuffer; use arrow::datatypes::{ ArrowPrimitiveType, DataType, Date32Type, Decimal32Type, Decimal64Type, Decimal128Type, DecimalType, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, @@ -68,7 +67,7 @@ pub fn unshred_variant(array: &VariantArray) -> Result { // emit `Variant::Null` when a required value is missing. let nulls = array.nulls(); let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.shredding_state())? - .unwrap_or_else(|| UnshredVariantRowBuilder::null(nulls.cloned())); + .unwrap_or_else(UnshredVariantRowBuilder::null); let metadata = array.metadata_field(); let mut value_builder = VariantValueArrayBuilder::new(array.len()); @@ -131,8 +130,8 @@ enum UnshredVariantRowBuilder { impl UnshredVariantRowBuilder { /// Creates an all-null row builder. - fn null(nulls: Option) -> Self { - Self::Null(NullUnshredVariantBuilder::new(nulls)) + fn null() -> Self { + Self::Null(NullUnshredVariantBuilder) } /// Appends a single row at the given value index to the supplied builder. @@ -312,30 +311,16 @@ impl UnshredVariantRowBuilder { } /// Builder for arrays with neither typed_value nor value (all NULL/Variant::Null) -struct NullUnshredVariantBuilder { - nulls: Option, -} +struct NullUnshredVariantBuilder; impl NullUnshredVariantBuilder { - fn new(nulls: Option) -> Self { - Self { nulls } - } - fn append_row( &mut self, builder: &mut impl VariantBuilderExt, _metadata: &VariantMetadata, - index: usize, + _index: usize, ) -> Result<()> { - if self - .nulls - .as_ref() - .is_some_and(|nulls| nulls.is_null(index)) - { - builder.append_null(); - } else { - builder.append_value(Variant::Null); - } + builder.append_value(Variant::Null); Ok(()) } } @@ -700,7 +685,7 @@ impl ListUnshredVariantBuilder { // requires us to emit `Variant::Null` when a required value is missing. let element_state = ShreddingState::try_from(element_values)?; let element_unshredder = UnshredVariantRowBuilder::try_new_opt(&element_state)? - .unwrap_or_else(|| UnshredVariantRowBuilder::null(None)); + .unwrap_or_else(UnshredVariantRowBuilder::null); Ok(Self { value, From 5b7d1e9411f8e3ac6b074845b0259fae1f837107 Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Thu, 23 Apr 2026 19:17:36 -0400 Subject: [PATCH 3/5] return lifetimes --- .../src/unshred_variant.rs | 220 ++++++++---------- parquet-variant-compute/src/variant_array.rs | 2 +- 2 files changed, 101 insertions(+), 121 deletions(-) diff --git a/parquet-variant-compute/src/unshred_variant.rs b/parquet-variant-compute/src/unshred_variant.rs index c035fbf65c9e..f4bfa73bded1 100644 --- a/parquet-variant-compute/src/unshred_variant.rs +++ b/parquet-variant-compute/src/unshred_variant.rs @@ -17,8 +17,8 @@ //! Module for unshredding VariantArray by folding typed_value columns back into the value column. -use crate::variant_array::binary_array_value; -use crate::{ShreddingState, VariantArray, VariantValueArrayBuilder}; +use crate::variant_array::{binary_array_value, validate_binary_array}; +use crate::{VariantArray, VariantValueArrayBuilder}; use arrow::array::{ Array, ArrayRef, AsArray as _, BinaryArray, BinaryViewArray, BooleanArray, FixedSizeBinaryArray, FixedSizeListArray, GenericListArray, GenericListViewArray, @@ -66,7 +66,7 @@ pub fn unshred_variant(array: &VariantArray) -> Result { // NOTE: None/None at top-level is technically invalid, but the shredding spec requires us to // emit `Variant::Null` when a required value is missing. let nulls = array.nulls(); - let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.shredding_state())? + let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.inner())? .unwrap_or_else(UnshredVariantRowBuilder::null); let metadata = array.metadata_field(); @@ -96,39 +96,39 @@ pub fn unshred_variant(array: &VariantArray) -> Result { } /// Row builder for converting shredded VariantArray rows back to unshredded form -enum UnshredVariantRowBuilder { - PrimitiveInt8(UnshredPrimitiveRowBuilder>), - PrimitiveInt16(UnshredPrimitiveRowBuilder>), - PrimitiveInt32(UnshredPrimitiveRowBuilder>), - PrimitiveInt64(UnshredPrimitiveRowBuilder>), - PrimitiveFloat32(UnshredPrimitiveRowBuilder>), - PrimitiveFloat64(UnshredPrimitiveRowBuilder>), - Decimal32(DecimalUnshredRowBuilder), - Decimal64(DecimalUnshredRowBuilder), - Decimal128(DecimalUnshredRowBuilder), - PrimitiveDate32(UnshredPrimitiveRowBuilder>), - PrimitiveTime64(UnshredPrimitiveRowBuilder>), - TimestampMicrosecond(TimestampUnshredRowBuilder), - TimestampNanosecond(TimestampUnshredRowBuilder), - PrimitiveBoolean(UnshredPrimitiveRowBuilder), - PrimitiveString(UnshredPrimitiveRowBuilder), - PrimitiveStringView(UnshredPrimitiveRowBuilder), - PrimitiveLargeString(UnshredPrimitiveRowBuilder), - PrimitiveBinary(UnshredPrimitiveRowBuilder), - PrimitiveBinaryView(UnshredPrimitiveRowBuilder), - PrimitiveLargeBinary(UnshredPrimitiveRowBuilder), - PrimitiveUuid(UnshredPrimitiveRowBuilder), - List(ListUnshredVariantBuilder>), - LargeList(ListUnshredVariantBuilder>), - ListView(ListUnshredVariantBuilder>), - LargeListView(ListUnshredVariantBuilder>), - FixedSizeList(ListUnshredVariantBuilder), - Struct(StructUnshredVariantBuilder), - ValueOnly(ValueOnlyUnshredVariantBuilder), +enum UnshredVariantRowBuilder<'a> { + PrimitiveInt8(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + PrimitiveInt16(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + PrimitiveInt32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + PrimitiveInt64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + PrimitiveFloat32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + PrimitiveFloat64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + Decimal32(DecimalUnshredRowBuilder<'a, Decimal32Type, VariantDecimal4>), + Decimal64(DecimalUnshredRowBuilder<'a, Decimal64Type, VariantDecimal8>), + Decimal128(DecimalUnshredRowBuilder<'a, Decimal128Type, VariantDecimal16>), + PrimitiveDate32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + PrimitiveTime64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + TimestampMicrosecond(TimestampUnshredRowBuilder<'a, TimestampMicrosecondType>), + TimestampNanosecond(TimestampUnshredRowBuilder<'a, TimestampNanosecondType>), + PrimitiveBoolean(UnshredPrimitiveRowBuilder<'a, BooleanArray>), + PrimitiveString(UnshredPrimitiveRowBuilder<'a, StringArray>), + PrimitiveStringView(UnshredPrimitiveRowBuilder<'a, StringViewArray>), + PrimitiveLargeString(UnshredPrimitiveRowBuilder<'a, LargeStringArray>), + PrimitiveBinary(UnshredPrimitiveRowBuilder<'a, BinaryArray>), + PrimitiveBinaryView(UnshredPrimitiveRowBuilder<'a, BinaryViewArray>), + PrimitiveLargeBinary(UnshredPrimitiveRowBuilder<'a, LargeBinaryArray>), + PrimitiveUuid(UnshredPrimitiveRowBuilder<'a, FixedSizeBinaryArray>), + List(ListUnshredVariantBuilder<'a, GenericListArray>), + LargeList(ListUnshredVariantBuilder<'a, GenericListArray>), + ListView(ListUnshredVariantBuilder<'a, GenericListViewArray>), + LargeListView(ListUnshredVariantBuilder<'a, GenericListViewArray>), + FixedSizeList(ListUnshredVariantBuilder<'a, FixedSizeListArray>), + Struct(StructUnshredVariantBuilder<'a>), + ValueOnly(ValueOnlyUnshredVariantBuilder<'a>), Null(NullUnshredVariantBuilder), } -impl UnshredVariantRowBuilder { +impl<'a> UnshredVariantRowBuilder<'a> { /// Creates an all-null row builder. fn null() -> Self { Self::Null(NullUnshredVariantBuilder) @@ -174,11 +174,17 @@ impl UnshredVariantRowBuilder { } } - /// Creates a new UnshredVariantRowBuilder from the given shredding state. - /// Returns None for the None/None case - caller decides how to handle based on context. - fn try_new_opt(shredding_state: &ShreddingState) -> Result> { - let value = shredding_state.value_field().cloned(); - let Some(typed_value) = shredding_state.typed_value_field().cloned() else { + /// Creates a new UnshredVariantRowBuilder from the `(value, typed_value)` pair of a shredded + /// variant struct. Returns None for the None/None case - caller decides how to handle based on + /// context. + fn try_new_opt(inner_struct: &'a StructArray) -> Result> { + let value = if let Some(value_col) = inner_struct.column_by_name("value") { + validate_binary_array(value_col.as_ref(), "value")?; + Some(value_col) + } else { + None + }; + let Some(typed_value) = inner_struct.column_by_name("typed_value") else { // Copy the value across directly, if present. Else caller decides what to do. return Ok(value.map(|v| Self::ValueOnly(ValueOnlyUnshredVariantBuilder::new(v)))); }; @@ -188,7 +194,7 @@ impl UnshredVariantRowBuilder { ($enum_variant:ident, $cast_fn:ident) => { Self::$enum_variant(UnshredPrimitiveRowBuilder::new( value, - typed_value.$cast_fn().clone(), + typed_value.$cast_fn(), )) }; } @@ -201,25 +207,13 @@ impl UnshredVariantRowBuilder { DataType::Float32 => primitive_builder!(PrimitiveFloat32, as_primitive), DataType::Float64 => primitive_builder!(PrimitiveFloat64, as_primitive), DataType::Decimal32(p, s) if VariantDecimal4::is_valid_precision_and_scale(p, s) => { - Self::Decimal32(DecimalUnshredRowBuilder::new( - value, - typed_value.as_primitive().clone(), - *s as _, - )) + Self::Decimal32(DecimalUnshredRowBuilder::new(value, typed_value, *s as _)) } DataType::Decimal64(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => { - Self::Decimal64(DecimalUnshredRowBuilder::new( - value, - typed_value.as_primitive().clone(), - *s as _, - )) + Self::Decimal64(DecimalUnshredRowBuilder::new(value, typed_value, *s as _)) } DataType::Decimal128(p, s) if VariantDecimal16::is_valid_precision_and_scale(p, s) => { - Self::Decimal128(DecimalUnshredRowBuilder::new( - value, - typed_value.as_primitive().clone(), - *s as _, - )) + Self::Decimal128(DecimalUnshredRowBuilder::new(value, typed_value, *s as _)) } DataType::Decimal32(_, _) | DataType::Decimal64(_, _) @@ -239,20 +233,12 @@ impl UnshredVariantRowBuilder { "Time64({time_unit}) is not a valid variant shredding type", ))); } - DataType::Timestamp(TimeUnit::Microsecond, timezone) => { - Self::TimestampMicrosecond(TimestampUnshredRowBuilder::new( - value, - typed_value.as_primitive().clone(), - timezone.is_some(), - )) - } - DataType::Timestamp(TimeUnit::Nanosecond, timezone) => { - Self::TimestampNanosecond(TimestampUnshredRowBuilder::new( - value, - typed_value.as_primitive().clone(), - timezone.is_some(), - )) - } + DataType::Timestamp(TimeUnit::Microsecond, timezone) => Self::TimestampMicrosecond( + TimestampUnshredRowBuilder::new(value, typed_value, timezone.is_some()), + ), + DataType::Timestamp(TimeUnit::Nanosecond, timezone) => Self::TimestampNanosecond( + TimestampUnshredRowBuilder::new(value, typed_value, timezone.is_some()), + ), DataType::Timestamp(time_unit, _) => { return Err(ArrowError::InvalidArgumentError(format!( "Timestamp({time_unit}) is not a valid variant shredding type", @@ -275,30 +261,27 @@ impl UnshredVariantRowBuilder { } DataType::Struct(_) => Self::Struct(StructUnshredVariantBuilder::try_new( value, - typed_value.as_struct().clone(), + typed_value.as_struct(), )?), DataType::List(_) => Self::List(ListUnshredVariantBuilder::try_new( value, - typed_value.as_list::().clone(), + typed_value.as_list(), )?), DataType::LargeList(_) => Self::LargeList(ListUnshredVariantBuilder::try_new( value, - typed_value.as_list::().clone(), + typed_value.as_list(), )?), DataType::ListView(_) => Self::ListView(ListUnshredVariantBuilder::try_new( value, - typed_value.as_list_view::().clone(), + typed_value.as_list_view(), )?), DataType::LargeListView(_) => Self::LargeListView(ListUnshredVariantBuilder::try_new( value, - typed_value.as_list_view::().clone(), + typed_value.as_list_view(), )?), - DataType::FixedSizeList(_, _) => { - Self::FixedSizeList(ListUnshredVariantBuilder::try_new( - value, - typed_value.as_fixed_size_list().clone(), - )?) - } + DataType::FixedSizeList(_, _) => Self::FixedSizeList( + ListUnshredVariantBuilder::try_new(value, typed_value.as_fixed_size_list())?, + ), _ => { return Err(ArrowError::NotYetImplemented(format!( "Unshredding not yet supported for type: {}", @@ -310,7 +293,7 @@ impl UnshredVariantRowBuilder { } } -/// Builder for arrays with neither typed_value nor value (all NULL/Variant::Null) +/// Builder for arrays with neither typed_value nor value (all Variant::Null) struct NullUnshredVariantBuilder; impl NullUnshredVariantBuilder { @@ -326,12 +309,12 @@ impl NullUnshredVariantBuilder { } /// Builder for arrays that only have value column (already unshredded) -struct ValueOnlyUnshredVariantBuilder { - value: ArrayRef, +struct ValueOnlyUnshredVariantBuilder<'a> { + value: &'a ArrayRef, } -impl ValueOnlyUnshredVariantBuilder { - fn new(value: ArrayRef) -> Self { +impl<'a> ValueOnlyUnshredVariantBuilder<'a> { + fn new(value: &'a ArrayRef) -> Self { Self { value } } @@ -405,13 +388,13 @@ macro_rules! handle_unshredded_case { } /// Generic unshred builder that works with any Array implementing AppendToVariantBuilder -struct UnshredPrimitiveRowBuilder { - value: Option, - typed_value: T, +struct UnshredPrimitiveRowBuilder<'a, T> { + value: Option<&'a ArrayRef>, + typed_value: &'a T, } -impl UnshredPrimitiveRowBuilder { - fn new(value: Option, typed_value: T) -> Self { +impl<'a, T: AppendToVariantBuilder> UnshredPrimitiveRowBuilder<'a, T> { + fn new(value: Option<&'a ArrayRef>, typed_value: &'a T) -> Self { Self { value, typed_value } } @@ -508,17 +491,17 @@ impl TimestampType for TimestampNanosecondType { } /// Generic builder for timestamp types that handles timezone-aware conversion -struct TimestampUnshredRowBuilder { - value: Option, - typed_value: PrimitiveArray, +struct TimestampUnshredRowBuilder<'a, T: TimestampType> { + value: Option<&'a ArrayRef>, + typed_value: &'a PrimitiveArray, has_timezone: bool, } -impl TimestampUnshredRowBuilder { - fn new(value: Option, typed_value: PrimitiveArray, has_timezone: bool) -> Self { +impl<'a, T: TimestampType> TimestampUnshredRowBuilder<'a, T> { + fn new(value: Option<&'a ArrayRef>, typed_value: &'a dyn Array, has_timezone: bool) -> Self { Self { value, - typed_value, + typed_value: typed_value.as_primitive(), has_timezone, } } @@ -544,24 +527,24 @@ impl TimestampUnshredRowBuilder { } /// Generic builder for decimal unshredding -struct DecimalUnshredRowBuilder +struct DecimalUnshredRowBuilder<'a, A: DecimalType, V> where V: VariantDecimalType, { - value: Option, - typed_value: PrimitiveArray, + value: Option<&'a ArrayRef>, + typed_value: &'a PrimitiveArray, scale: i8, _phantom: PhantomData, } -impl DecimalUnshredRowBuilder +impl<'a, A: DecimalType, V> DecimalUnshredRowBuilder<'a, A, V> where V: VariantDecimalType, { - fn new(value: Option, typed_value: PrimitiveArray, scale: i8) -> Self { + fn new(value: Option<&'a ArrayRef>, typed_value: &'a dyn Array, scale: i8) -> Self { Self { value, - typed_value, + typed_value: typed_value.as_primitive(), scale, _phantom: PhantomData, } @@ -583,14 +566,14 @@ where } /// Builder for unshredding struct/object types with nested fields -struct StructUnshredVariantBuilder { - value: Option, - typed_value: StructArray, - field_unshredders: IndexMap>, +struct StructUnshredVariantBuilder<'a> { + value: Option<&'a ArrayRef>, + typed_value: &'a arrow::array::StructArray, + field_unshredders: IndexMap<&'a str, Option>>, } -impl StructUnshredVariantBuilder { - fn try_new(value: Option, typed_value: StructArray) -> Result { +impl<'a> StructUnshredVariantBuilder<'a> { + fn try_new(value: Option<&'a ArrayRef>, typed_value: &'a StructArray) -> Result { // Create unshredders for each field in constructor let mut field_unshredders = IndexMap::new(); for (field, field_array) in typed_value.fields().iter().zip(typed_value.columns()) { @@ -601,9 +584,8 @@ impl StructUnshredVariantBuilder { field_array.data_type() ))); }; - let field_state = ShreddingState::try_from(field_array)?; - let field_unshredder = UnshredVariantRowBuilder::try_new_opt(&field_state)?; - field_unshredders.insert(field.name().clone(), field_unshredder); + let field_unshredder = UnshredVariantRowBuilder::try_new_opt(field_array)?; + field_unshredders.insert(field.name().as_ref(), field_unshredder); } Ok(Self { @@ -627,8 +609,7 @@ impl StructUnshredVariantBuilder { // Process typed fields (skip empty builders that indicate missing fields) for (field_name, field_unshredder_opt) in &mut self.field_unshredders { if let Some(field_unshredder) = field_unshredder_opt { - let mut field_builder = - ObjectFieldBuilder::new(field_name.as_str(), &mut object_builder); + let mut field_builder = ObjectFieldBuilder::new(field_name, &mut object_builder); field_unshredder.append_row(&mut field_builder, metadata, index)?; } } @@ -658,14 +639,14 @@ impl StructUnshredVariantBuilder { } /// Builder for unshredding list/array types with recursive element processing -struct ListUnshredVariantBuilder { - value: Option, - typed_value: L, - element_unshredder: Box, +struct ListUnshredVariantBuilder<'a, L: ListLikeArray> { + value: Option<&'a ArrayRef>, + typed_value: &'a L, + element_unshredder: Box>, } -impl ListUnshredVariantBuilder { - fn try_new(value: Option, typed_value: L) -> Result { +impl<'a, L: ListLikeArray> ListUnshredVariantBuilder<'a, L> { + fn try_new(value: Option<&'a ArrayRef>, typed_value: &'a L) -> Result { // Create a recursive unshredder for the list elements // The element type comes from the values array of the list let element_values = typed_value.values(); @@ -683,8 +664,7 @@ impl ListUnshredVariantBuilder { // // NOTE: A None/None array element is technically invalid, but the shredding spec // requires us to emit `Variant::Null` when a required value is missing. - let element_state = ShreddingState::try_from(element_values)?; - let element_unshredder = UnshredVariantRowBuilder::try_new_opt(&element_state)? + let element_unshredder = UnshredVariantRowBuilder::try_new_opt(element_values)? .unwrap_or_else(UnshredVariantRowBuilder::null); Ok(Self { diff --git a/parquet-variant-compute/src/variant_array.rs b/parquet-variant-compute/src/variant_array.rs index a82ed11c5565..a6ee281002cc 100644 --- a/parquet-variant-compute/src/variant_array.rs +++ b/parquet-variant-compute/src/variant_array.rs @@ -64,7 +64,7 @@ pub(crate) fn variant_from_arrays_at<'m, 'v>( } /// Validates that an array has a binary-like data type. -fn validate_binary_array(array: &dyn Array, field_name: &str) -> Result<()> { +pub(crate) fn validate_binary_array(array: &dyn Array, field_name: &str) -> Result<()> { match array.data_type() { DataType::Binary | DataType::LargeBinary | DataType::BinaryView => Ok(()), _ => Err(ArrowError::InvalidArgumentError(format!( From 8aadf4010c08cde4e4336af37ffe6e9a201a5efe Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Fri, 15 May 2026 15:25:29 -0400 Subject: [PATCH 4/5] Cow the shredding state --- .../src/unshred_variant.rs | 28 ++-- parquet-variant-compute/src/variant_array.rs | 151 +++++++++++------ parquet-variant-compute/src/variant_get.rs | 157 +++++++++++------- 3 files changed, 215 insertions(+), 121 deletions(-) diff --git a/parquet-variant-compute/src/unshred_variant.rs b/parquet-variant-compute/src/unshred_variant.rs index f4bfa73bded1..797a93f0e86a 100644 --- a/parquet-variant-compute/src/unshred_variant.rs +++ b/parquet-variant-compute/src/unshred_variant.rs @@ -17,7 +17,7 @@ //! Module for unshredding VariantArray by folding typed_value columns back into the value column. -use crate::variant_array::{binary_array_value, validate_binary_array}; +use crate::variant_array::{ShreddingState, binary_array_value}; use crate::{VariantArray, VariantValueArrayBuilder}; use arrow::array::{ Array, ArrayRef, AsArray as _, BinaryArray, BinaryViewArray, BooleanArray, @@ -66,7 +66,7 @@ pub fn unshred_variant(array: &VariantArray) -> Result { // NOTE: None/None at top-level is technically invalid, but the shredding spec requires us to // emit `Variant::Null` when a required value is missing. let nulls = array.nulls(); - let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.inner())? + let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.shredding_state())? .unwrap_or_else(UnshredVariantRowBuilder::null); let metadata = array.metadata_field(); @@ -174,17 +174,15 @@ impl<'a> UnshredVariantRowBuilder<'a> { } } - /// Creates a new UnshredVariantRowBuilder from the `(value, typed_value)` pair of a shredded - /// variant struct. Returns None for the None/None case - caller decides how to handle based on - /// context. - fn try_new_opt(inner_struct: &'a StructArray) -> Result> { - let value = if let Some(value_col) = inner_struct.column_by_name("value") { - validate_binary_array(value_col.as_ref(), "value")?; - Some(value_col) - } else { - None - }; - let Some(typed_value) = inner_struct.column_by_name("typed_value") else { + /// Creates a new UnshredVariantRowBuilder from a shredded `(value, typed_value)` pair. + /// Returns None for the None/None case - caller decides how to handle based on context. + /// + /// The `ShreddingState` is consumed so its `'a` lifetime is preserved through to the borrowed + /// fields stored in the returned builder. The state must be `Cow::Borrowed`; the unshred path + /// only ever sees borrowed inputs (rooted at the top-level `VariantArray`'s inner struct). + fn try_new_opt(shredding_state: ShreddingState<'a>) -> Result> { + let (value, typed_value) = shredding_state.into_borrowed_fields(); + let Some(typed_value) = typed_value else { // Copy the value across directly, if present. Else caller decides what to do. return Ok(value.map(|v| Self::ValueOnly(ValueOnlyUnshredVariantBuilder::new(v)))); }; @@ -584,7 +582,7 @@ impl<'a> StructUnshredVariantBuilder<'a> { field_array.data_type() ))); }; - let field_unshredder = UnshredVariantRowBuilder::try_new_opt(field_array)?; + let field_unshredder = UnshredVariantRowBuilder::try_new_opt(field_array.try_into()?)?; field_unshredders.insert(field.name().as_ref(), field_unshredder); } @@ -664,7 +662,7 @@ impl<'a, L: ListLikeArray> ListUnshredVariantBuilder<'a, L> { // // NOTE: A None/None array element is technically invalid, but the shredding spec // requires us to emit `Variant::Null` when a required value is missing. - let element_unshredder = UnshredVariantRowBuilder::try_new_opt(element_values)? + let element_unshredder = UnshredVariantRowBuilder::try_new_opt(element_values.try_into()?)? .unwrap_or_else(UnshredVariantRowBuilder::null); Ok(Self { diff --git a/parquet-variant-compute/src/variant_array.rs b/parquet-variant-compute/src/variant_array.rs index a6ee281002cc..7e5b525580a3 100644 --- a/parquet-variant-compute/src/variant_array.rs +++ b/parquet-variant-compute/src/variant_array.rs @@ -257,8 +257,10 @@ pub struct VariantArray { /// The metadata column of this variant (Binary, LargeBinary, or BinaryView) metadata: ArrayRef, - /// how is this variant array shredded? - shredding_state: ShreddingState, + /// Cached shredding state derived from `inner`. Stored as `'static` (i.e. owned `ArrayRef` + /// clones) so per-row accessors don't redo name lookups. Borrowed views are produced on + /// demand by `shredding_state()`. + shredding_state: ShreddingState<'static>, } impl VariantArray { @@ -308,11 +310,14 @@ impl VariantArray { }; validate_binary_array(metadata_col.as_ref(), "metadata")?; + // Cache the (value, typed_value) lookup. try_from validates the value column type. + let shredding_state = ShreddingState::try_from(inner)?.into_static(); + // Note these clones are cheap, they just bump the ref count Ok(Self { inner: inner.clone(), metadata: metadata_col.clone(), - shredding_state: ShreddingState::try_from(inner)?, + shredding_state, }) } @@ -323,11 +328,11 @@ impl VariantArray { nulls: Option, ) -> Self { let mut builder = StructArrayBuilder::new().with_field("metadata", metadata.clone(), false); - if let Some(value) = value.clone() { - builder = builder.with_field("value", value, true); + if let Some(ref value) = value { + builder = builder.with_field("value", value.clone(), true); } - if let Some(typed_value) = typed_value.clone() { - builder = builder.with_field("typed_value", typed_value, true); + if let Some(ref typed_value) = typed_value { + builder = builder.with_field("typed_value", typed_value.clone(), true); } if let Some(nulls) = nulls { builder = builder.with_nulls(nulls); @@ -350,9 +355,13 @@ impl VariantArray { self.inner } - /// Return the shredding state of this `VariantArray` - pub fn shredding_state(&self) -> &ShreddingState { - &self.shredding_state + /// Return a borrowed view of this `VariantArray`'s shredding state, backed by the cached + /// column lookup so accessors are O(1). + pub fn shredding_state(&self) -> ShreddingState<'_> { + ShreddingState { + value: self.shredding_state.value_field().map(Cow::Borrowed), + typed_value: self.shredding_state.typed_value_field().map(Cow::Borrowed), + } } /// Return the [`Variant`] instance stored at the given row @@ -645,7 +654,9 @@ impl<'a> ExactSizeIterator for VariantArrayIter<'a> {} pub struct ShreddedVariantFieldArray { /// Reference to the underlying StructArray inner: StructArray, - shredding_state: ShreddingState, + + /// Cached shredding state derived from `inner` (see `VariantArray::shredding_state`). + shredding_state: ShreddingState<'static>, } #[allow(unused)] @@ -676,16 +687,23 @@ impl ShreddedVariantFieldArray { )); }; + // Cache the (value, typed_value) lookup. try_from validates the value column type. + let shredding_state = ShreddingState::try_from(inner_struct)?.into_static(); + // Note this clone is cheap, it just bumps the ref count Ok(Self { inner: inner_struct.clone(), - shredding_state: ShreddingState::try_from(inner_struct)?, + shredding_state, }) } - /// Return the shredding state of this `VariantArray` - pub fn shredding_state(&self) -> &ShreddingState { - &self.shredding_state + /// Return a borrowed view of this `ShreddedVariantFieldArray`'s shredding state, backed by + /// the cached column lookup. + pub fn shredding_state(&self) -> ShreddingState<'_> { + ShreddingState { + value: self.shredding_state.value_field().map(Cow::Borrowed), + typed_value: self.shredding_state.typed_value_field().map(Cow::Borrowed), + } } /// Return a reference to the value field of the `StructArray` @@ -709,11 +727,11 @@ impl ShreddedVariantFieldArray { nulls: Option, ) -> Self { let mut builder = StructArrayBuilder::new(); - if let Some(value) = value.clone() { - builder = builder.with_field("value", value, true); + if let Some(ref value) = value { + builder = builder.with_field("value", value.clone(), true); } - if let Some(typed_value) = typed_value.clone() { - builder = builder.with_field("typed_value", typed_value, true); + if let Some(ref typed_value) = typed_value { + builder = builder.with_field("typed_value", typed_value.clone(), true); } if let Some(nulls) = nulls { builder = builder.with_nulls(nulls); @@ -809,62 +827,97 @@ impl From for StructArray { /// /// [Parquet Variant Shredding Spec]: https://github.com/apache/parquet-format/blob/master/VariantShredding.md#value-shredding #[derive(Debug, Clone)] -pub struct ShreddingState { - value: Option, - typed_value: Option, +pub struct ShreddingState<'a> { + value: Option>, + typed_value: Option>, } -impl ShreddingState { - /// Create a new `ShreddingState` from the given `value` and `typed_value` fields - /// - /// Note you can create a `ShreddingState` from a &[`StructArray`] using - /// `ShreddingState::try_from(&struct_array)`, for example: +impl ShreddingState<'static> { + /// Create a new owned `ShreddingState` from the given `value` and `typed_value` fields /// - /// ```no_run - /// # use arrow::array::StructArray; - /// # use parquet_variant_compute::ShreddingState; - /// # fn get_struct_array() -> StructArray { - /// # unimplemented!() - /// # } - /// let struct_array: StructArray = get_struct_array(); - /// let shredding_state = ShreddingState::try_from(&struct_array).unwrap(); - /// ``` + /// Use this when constructing a `ShreddingState` from freshly produced arrays (e.g., the + /// result of a `take` kernel call). For borrowing an existing `StructArray`'s columns, use + /// `ShreddingState::try_from(&struct_array)`. pub fn new(value: Option, typed_value: Option) -> Self { - Self { value, typed_value } + Self { + value: value.map(Cow::Owned), + typed_value: typed_value.map(Cow::Owned), + } } +} +impl<'a> ShreddingState<'a> { /// Return a reference to the value field, if present pub fn value_field(&self) -> Option<&ArrayRef> { - self.value.as_ref() + self.value.as_deref() } /// Return a reference to the typed_value field, if present pub fn typed_value_field(&self) -> Option<&ArrayRef> { - self.typed_value.as_ref() + self.typed_value.as_deref() } - /// Slice all the underlying arrays - pub fn slice(&self, offset: usize, length: usize) -> Self { - Self { - value: self.value.as_ref().map(|v| v.slice(offset, length)), - typed_value: self.typed_value.as_ref().map(|tv| tv.slice(offset, length)), + /// Consume self and return the `(value, typed_value)` Cows so callers can extract long-lived + /// borrows from `Cow::Borrowed` variants. + pub(crate) fn into_fields(self) -> (Option>, Option>) { + (self.value, self.typed_value) + } + + /// Convert into a `'static`-lifetime state by cloning any borrowed Arcs. Cheap (just Arc + /// refcount bumps for each `Cow::Borrowed`; no-op for `Cow::Owned`). + pub fn into_static(self) -> ShreddingState<'static> { + ShreddingState { + value: self.value.map(|cow| Cow::Owned(cow.into_owned())), + typed_value: self.typed_value.map(|cow| Cow::Owned(cow.into_owned())), + } + } + + /// Consume self and return the `(value, typed_value)` borrowed references, asserting that + /// the state was constructed from borrowed data. Panics on `Cow::Owned`. + /// + /// Use this in code paths (like unshred) that only ever receive borrowed state and need + /// `&'a ArrayRef` references to store in lifetime-parameterized builders. + pub(crate) fn into_borrowed_fields(self) -> (Option<&'a ArrayRef>, Option<&'a ArrayRef>) { + fn expect_borrowed<'a>(opt: Option>) -> Option<&'a ArrayRef> { + opt.map(|cow| match cow { + Cow::Borrowed(r) => r, + Cow::Owned(_) => unreachable!("expected Cow::Borrowed ShreddingState"), + }) + } + ( + expect_borrowed(self.value), + expect_borrowed(self.typed_value), + ) + } + + /// Slice all the underlying arrays. Produces owned data. + pub fn slice(&self, offset: usize, length: usize) -> ShreddingState<'static> { + ShreddingState { + value: self + .value_field() + .map(|v| Cow::Owned(v.slice(offset, length))), + typed_value: self + .typed_value_field() + .map(|tv| Cow::Owned(tv.slice(offset, length))), } } } -impl TryFrom<&StructArray> for ShreddingState { +impl<'a> TryFrom<&'a StructArray> for ShreddingState<'a> { type Error = ArrowError; - fn try_from(inner_struct: &StructArray) -> Result { + fn try_from(inner_struct: &'a StructArray) -> Result { // The `value` column need not exist, but if it does it must be a binary type. let value = if let Some(value_col) = inner_struct.column_by_name("value") { validate_binary_array(value_col.as_ref(), "value")?; - Some(value_col.clone()) + Some(Cow::Borrowed(value_col)) } else { None }; - let typed_value = inner_struct.column_by_name("typed_value").cloned(); - Ok(ShreddingState::new(value, typed_value)) + let typed_value = inner_struct + .column_by_name("typed_value") + .map(Cow::Borrowed); + Ok(ShreddingState { value, typed_value }) } } diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index 172f41a1767e..3c216e8ec1b5 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -25,22 +25,52 @@ use arrow_schema::{ArrowError, DataType, FieldRef}; use parquet_variant::{VariantPath, VariantPathElement}; use crate::VariantArray; -use crate::variant_array::ShreddingState; +use crate::variant_array::{ShreddingState, validate_binary_array}; use crate::variant_to_arrow::make_variant_to_arrow_row_builder; use arrow::array::AsArray; +use std::borrow::Cow; use std::sync::Arc; -pub(crate) enum ShreddedPathStep { - /// Path step succeeded, return the new shredding state - Success(ShreddingState), +pub(crate) enum ShreddedPathStep<'a> { + /// Path step succeeded. + Success { + /// The shredding state for the next step. + next_state: ShreddingState<'a>, + /// The `typed_value` column the step descended through. Returned so the caller can union + /// its nulls into the accumulated null buffer without needing to re-borrow from the + /// (now-consumed) input state. + consumed_typed_value: Cow<'a, ArrayRef>, + }, /// The path element is not present in the `typed_value` column and there is no `value` column, /// so we know it does not exist. It, and all paths under it, are all-NULL. Missing, - /// The path element is not present in the `typed_value` column and must be retrieved from the `value` - /// column instead. The caller should be prepared to handle any value, including the requested - /// type, an arbitrary "wrong" type, or `Variant::Null`. - NotShredded, + /// The path element is not present in the `typed_value` column and must be retrieved from the + /// `value` column instead. The caller should be prepared to handle any value, including the + /// requested type, an arbitrary "wrong" type, or `Variant::Null`. The consumed `value` column + /// is returned so the caller can use it without re-borrowing from the (now-consumed) input. + NotShredded { + consumed_value: Option>, + }, +} + +/// Walk into `typed_value`'s nested `StructArray` and look up `name`. Returns +/// `Ok(Some(nested_struct))` on success, `Ok(None)` for a missing path (non-struct or missing +/// field), and `Err` on an invalid nested struct. +fn walk_field_step<'b>(typed_value: &'b ArrayRef, name: &str) -> Result> { + let Some(struct_array) = typed_value.as_struct_opt() else { + return Ok(None); + }; + let Some(field) = struct_array.column_by_name(name) else { + return Ok(None); + }; + let nested_struct = field.as_struct_opt().ok_or_else(|| { + ArrowError::InvalidArgumentError(format!( + "Expected Struct array while following path, got {}", + field.data_type(), + )) + })?; + Ok(Some(nested_struct)) } /// Given a shredded variant field -- a `(value?, typed_value?)` pair -- try to take one path step @@ -48,51 +78,62 @@ pub(crate) enum ShreddedPathStep { /// `typed_value` is not a struct, or if the requested field name does not exist, traversal returns /// a missing-path step (`Missing` or `NotShredded` depending on whether `value` exists). /// +/// The state is consumed by value so its lifetime `'a` (rooted at the original input) is preserved +/// onto the returned step. For `Cow::Borrowed` inputs this lets the chain walk all the way down +/// without per-step Arc bumps; for `Cow::Owned` inputs (e.g., `take`-kernel results) the function +/// clones the new state's columns since the owned typed_value goes out of scope on return. +/// /// TODO: Support `VariantPathElement::Index`? It wouldn't be easy, and maybe not even possible. -pub(crate) fn follow_shredded_path_element( - shredding_state: &ShreddingState, +pub(crate) fn follow_shredded_path_element<'a>( + shredding_state: ShreddingState<'a>, path_element: &VariantPathElement<'_>, _cast_options: &CastOptions, -) -> Result { - // If the requested path element is not present in `typed_value`, and `value` is missing, then - // we know it does not exist; it, and all paths under it, are all-NULL. - let missing_path_step = || match shredding_state.value_field() { - Some(_) => ShreddedPathStep::NotShredded, +) -> Result> { + let (value, typed_value) = shredding_state.into_fields(); + let missing_path_step = |value: Option>| match value { + Some(_) => ShreddedPathStep::NotShredded { + consumed_value: value, + }, None => ShreddedPathStep::Missing, }; - let Some(typed_value) = shredding_state.typed_value_field() else { - return Ok(missing_path_step()); + let Some(typed_value) = typed_value else { + return Ok(missing_path_step(value)); }; match path_element { - VariantPathElement::Field { name } => { - // Try to step into the requested field name of a struct. - // First, try to downcast to StructArray - let Some(struct_array) = typed_value.as_any().downcast_ref::() else { - // Object field path step follows JSONPath semantics and returns missing path step (NotShredded/Missing) on non-struct path - return Ok(missing_path_step()); - }; - - // Now try to find the column - missing column in a present struct is just missing data - let Some(field) = struct_array.column_by_name(name) else { - // Missing column in a present struct is just missing, not wrong - return Ok - return Ok(missing_path_step()); - }; - - let struct_array = field.as_struct_opt().ok_or_else(|| { - // TODO: Should we blow up? Or just end the traversal and let the normal - // variant pathing code sort out the mess that it must anyway be - // prepared to handle? - ArrowError::InvalidArgumentError(format!( - "Expected Struct array while following path, got {}", - field.data_type(), - )) - })?; - - let state = ShreddingState::try_from(struct_array)?; - Ok(ShreddedPathStep::Success(state)) - } + VariantPathElement::Field { name } => match typed_value { + Cow::Borrowed(typed_value_ref) => { + let Some(nested) = walk_field_step(typed_value_ref, name)? else { + return Ok(missing_path_step(value)); + }; + // nested: &'a StructArray, so try_from produces ShreddingState<'a> with + // Cow::Borrowed views into `nested`. No Arc bumps on this path. + let next_state = ShreddingState::try_from(nested)?; + Ok(ShreddedPathStep::Success { + next_state, + consumed_typed_value: Cow::Borrowed(typed_value_ref), + }) + } + Cow::Owned(typed_value_arc) => { + let Some(nested) = walk_field_step(&typed_value_arc, name)? else { + return Ok(missing_path_step(value)); + }; + // `typed_value_arc` is the owned input and is about to move into the returned + // `consumed_typed_value`. We cannot also borrow into it (self-referential), so we + // clone `nested`'s columns to produce a `ShreddingState<'static>` (coerces to 'a). + let value_col = nested.column_by_name("value"); + if let Some(v) = value_col { + validate_binary_array(v.as_ref(), "value")?; + } + let typed_value_col = nested.column_by_name("typed_value"); + let next_state = ShreddingState::new(value_col.cloned(), typed_value_col.cloned()); + Ok(ShreddedPathStep::Success { + next_state, + consumed_typed_value: Cow::Owned(typed_value_arc), + }) + } + }, VariantPathElement::Index { .. } => { // TODO: Support array indexing. Among other things, it will require slicing not // only the array we have here, but also the corresponding metadata and null masks. @@ -154,18 +195,19 @@ fn shredded_get_path( // Peel away the prefix of path elements that traverses the shredded parts of this variant // column. Shredding will traverse the rest of the path on a per-row basis. - let mut shredding_state = input.shredding_state().clone(); + let mut shredding_state = input.shredding_state(); let mut accumulated_nulls = input.inner().nulls().cloned(); let mut path_index = 0; for path_element in path { - match follow_shredded_path_element(&shredding_state, path_element, cast_options)? { - ShreddedPathStep::Success(state) => { - // Union nulls from the typed_value we just accessed - if let Some(typed_value) = shredding_state.typed_value_field() { - accumulated_nulls = - NullBuffer::union(accumulated_nulls.as_ref(), typed_value.nulls()); - } - shredding_state = state; + match follow_shredded_path_element(shredding_state, path_element, cast_options)? { + ShreddedPathStep::Success { + next_state, + consumed_typed_value, + } => { + // Union nulls from the typed_value we just descended through. + accumulated_nulls = + NullBuffer::union(accumulated_nulls.as_ref(), consumed_typed_value.nulls()); + shredding_state = next_state; path_index += 1; continue; } @@ -177,9 +219,9 @@ fn shredded_get_path( }; return Ok(arr); } - ShreddedPathStep::NotShredded => { + ShreddedPathStep::NotShredded { consumed_value } => { let target = make_target_variant( - shredding_state.value_field().cloned(), + consumed_value.map(Cow::into_owned), None, accumulated_nulls, ); @@ -189,9 +231,10 @@ fn shredded_get_path( } // Path exhausted! Create a new `VariantArray` for the location we landed on. + let (final_value, final_typed_value) = shredding_state.into_fields(); let target = make_target_variant( - shredding_state.value_field().cloned(), - shredding_state.typed_value_field().cloned(), + final_value.map(Cow::into_owned), + final_typed_value.map(Cow::into_owned), accumulated_nulls, ); From cde17dea8871c537fe58cf04a668f7ca98fbce0e Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Fri, 15 May 2026 21:41:34 -0400 Subject: [PATCH 5/5] Revert "Cow the shredding state" This reverts commit 8aadf4010c08cde4e4336af37ffe6e9a201a5efe. --- .../src/unshred_variant.rs | 28 ++-- parquet-variant-compute/src/variant_array.rs | 151 ++++++----------- parquet-variant-compute/src/variant_get.rs | 157 +++++++----------- 3 files changed, 121 insertions(+), 215 deletions(-) diff --git a/parquet-variant-compute/src/unshred_variant.rs b/parquet-variant-compute/src/unshred_variant.rs index 797a93f0e86a..f4bfa73bded1 100644 --- a/parquet-variant-compute/src/unshred_variant.rs +++ b/parquet-variant-compute/src/unshred_variant.rs @@ -17,7 +17,7 @@ //! Module for unshredding VariantArray by folding typed_value columns back into the value column. -use crate::variant_array::{ShreddingState, binary_array_value}; +use crate::variant_array::{binary_array_value, validate_binary_array}; use crate::{VariantArray, VariantValueArrayBuilder}; use arrow::array::{ Array, ArrayRef, AsArray as _, BinaryArray, BinaryViewArray, BooleanArray, @@ -66,7 +66,7 @@ pub fn unshred_variant(array: &VariantArray) -> Result { // NOTE: None/None at top-level is technically invalid, but the shredding spec requires us to // emit `Variant::Null` when a required value is missing. let nulls = array.nulls(); - let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.shredding_state())? + let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.inner())? .unwrap_or_else(UnshredVariantRowBuilder::null); let metadata = array.metadata_field(); @@ -174,15 +174,17 @@ impl<'a> UnshredVariantRowBuilder<'a> { } } - /// Creates a new UnshredVariantRowBuilder from a shredded `(value, typed_value)` pair. - /// Returns None for the None/None case - caller decides how to handle based on context. - /// - /// The `ShreddingState` is consumed so its `'a` lifetime is preserved through to the borrowed - /// fields stored in the returned builder. The state must be `Cow::Borrowed`; the unshred path - /// only ever sees borrowed inputs (rooted at the top-level `VariantArray`'s inner struct). - fn try_new_opt(shredding_state: ShreddingState<'a>) -> Result> { - let (value, typed_value) = shredding_state.into_borrowed_fields(); - let Some(typed_value) = typed_value else { + /// Creates a new UnshredVariantRowBuilder from the `(value, typed_value)` pair of a shredded + /// variant struct. Returns None for the None/None case - caller decides how to handle based on + /// context. + fn try_new_opt(inner_struct: &'a StructArray) -> Result> { + let value = if let Some(value_col) = inner_struct.column_by_name("value") { + validate_binary_array(value_col.as_ref(), "value")?; + Some(value_col) + } else { + None + }; + let Some(typed_value) = inner_struct.column_by_name("typed_value") else { // Copy the value across directly, if present. Else caller decides what to do. return Ok(value.map(|v| Self::ValueOnly(ValueOnlyUnshredVariantBuilder::new(v)))); }; @@ -582,7 +584,7 @@ impl<'a> StructUnshredVariantBuilder<'a> { field_array.data_type() ))); }; - let field_unshredder = UnshredVariantRowBuilder::try_new_opt(field_array.try_into()?)?; + let field_unshredder = UnshredVariantRowBuilder::try_new_opt(field_array)?; field_unshredders.insert(field.name().as_ref(), field_unshredder); } @@ -662,7 +664,7 @@ impl<'a, L: ListLikeArray> ListUnshredVariantBuilder<'a, L> { // // NOTE: A None/None array element is technically invalid, but the shredding spec // requires us to emit `Variant::Null` when a required value is missing. - let element_unshredder = UnshredVariantRowBuilder::try_new_opt(element_values.try_into()?)? + let element_unshredder = UnshredVariantRowBuilder::try_new_opt(element_values)? .unwrap_or_else(UnshredVariantRowBuilder::null); Ok(Self { diff --git a/parquet-variant-compute/src/variant_array.rs b/parquet-variant-compute/src/variant_array.rs index 7e5b525580a3..a6ee281002cc 100644 --- a/parquet-variant-compute/src/variant_array.rs +++ b/parquet-variant-compute/src/variant_array.rs @@ -257,10 +257,8 @@ pub struct VariantArray { /// The metadata column of this variant (Binary, LargeBinary, or BinaryView) metadata: ArrayRef, - /// Cached shredding state derived from `inner`. Stored as `'static` (i.e. owned `ArrayRef` - /// clones) so per-row accessors don't redo name lookups. Borrowed views are produced on - /// demand by `shredding_state()`. - shredding_state: ShreddingState<'static>, + /// how is this variant array shredded? + shredding_state: ShreddingState, } impl VariantArray { @@ -310,14 +308,11 @@ impl VariantArray { }; validate_binary_array(metadata_col.as_ref(), "metadata")?; - // Cache the (value, typed_value) lookup. try_from validates the value column type. - let shredding_state = ShreddingState::try_from(inner)?.into_static(); - // Note these clones are cheap, they just bump the ref count Ok(Self { inner: inner.clone(), metadata: metadata_col.clone(), - shredding_state, + shredding_state: ShreddingState::try_from(inner)?, }) } @@ -328,11 +323,11 @@ impl VariantArray { nulls: Option, ) -> Self { let mut builder = StructArrayBuilder::new().with_field("metadata", metadata.clone(), false); - if let Some(ref value) = value { - builder = builder.with_field("value", value.clone(), true); + if let Some(value) = value.clone() { + builder = builder.with_field("value", value, true); } - if let Some(ref typed_value) = typed_value { - builder = builder.with_field("typed_value", typed_value.clone(), true); + if let Some(typed_value) = typed_value.clone() { + builder = builder.with_field("typed_value", typed_value, true); } if let Some(nulls) = nulls { builder = builder.with_nulls(nulls); @@ -355,13 +350,9 @@ impl VariantArray { self.inner } - /// Return a borrowed view of this `VariantArray`'s shredding state, backed by the cached - /// column lookup so accessors are O(1). - pub fn shredding_state(&self) -> ShreddingState<'_> { - ShreddingState { - value: self.shredding_state.value_field().map(Cow::Borrowed), - typed_value: self.shredding_state.typed_value_field().map(Cow::Borrowed), - } + /// Return the shredding state of this `VariantArray` + pub fn shredding_state(&self) -> &ShreddingState { + &self.shredding_state } /// Return the [`Variant`] instance stored at the given row @@ -654,9 +645,7 @@ impl<'a> ExactSizeIterator for VariantArrayIter<'a> {} pub struct ShreddedVariantFieldArray { /// Reference to the underlying StructArray inner: StructArray, - - /// Cached shredding state derived from `inner` (see `VariantArray::shredding_state`). - shredding_state: ShreddingState<'static>, + shredding_state: ShreddingState, } #[allow(unused)] @@ -687,23 +676,16 @@ impl ShreddedVariantFieldArray { )); }; - // Cache the (value, typed_value) lookup. try_from validates the value column type. - let shredding_state = ShreddingState::try_from(inner_struct)?.into_static(); - // Note this clone is cheap, it just bumps the ref count Ok(Self { inner: inner_struct.clone(), - shredding_state, + shredding_state: ShreddingState::try_from(inner_struct)?, }) } - /// Return a borrowed view of this `ShreddedVariantFieldArray`'s shredding state, backed by - /// the cached column lookup. - pub fn shredding_state(&self) -> ShreddingState<'_> { - ShreddingState { - value: self.shredding_state.value_field().map(Cow::Borrowed), - typed_value: self.shredding_state.typed_value_field().map(Cow::Borrowed), - } + /// Return the shredding state of this `VariantArray` + pub fn shredding_state(&self) -> &ShreddingState { + &self.shredding_state } /// Return a reference to the value field of the `StructArray` @@ -727,11 +709,11 @@ impl ShreddedVariantFieldArray { nulls: Option, ) -> Self { let mut builder = StructArrayBuilder::new(); - if let Some(ref value) = value { - builder = builder.with_field("value", value.clone(), true); + if let Some(value) = value.clone() { + builder = builder.with_field("value", value, true); } - if let Some(ref typed_value) = typed_value { - builder = builder.with_field("typed_value", typed_value.clone(), true); + if let Some(typed_value) = typed_value.clone() { + builder = builder.with_field("typed_value", typed_value, true); } if let Some(nulls) = nulls { builder = builder.with_nulls(nulls); @@ -827,97 +809,62 @@ impl From for StructArray { /// /// [Parquet Variant Shredding Spec]: https://github.com/apache/parquet-format/blob/master/VariantShredding.md#value-shredding #[derive(Debug, Clone)] -pub struct ShreddingState<'a> { - value: Option>, - typed_value: Option>, +pub struct ShreddingState { + value: Option, + typed_value: Option, } -impl ShreddingState<'static> { - /// Create a new owned `ShreddingState` from the given `value` and `typed_value` fields +impl ShreddingState { + /// Create a new `ShreddingState` from the given `value` and `typed_value` fields + /// + /// Note you can create a `ShreddingState` from a &[`StructArray`] using + /// `ShreddingState::try_from(&struct_array)`, for example: /// - /// Use this when constructing a `ShreddingState` from freshly produced arrays (e.g., the - /// result of a `take` kernel call). For borrowing an existing `StructArray`'s columns, use - /// `ShreddingState::try_from(&struct_array)`. + /// ```no_run + /// # use arrow::array::StructArray; + /// # use parquet_variant_compute::ShreddingState; + /// # fn get_struct_array() -> StructArray { + /// # unimplemented!() + /// # } + /// let struct_array: StructArray = get_struct_array(); + /// let shredding_state = ShreddingState::try_from(&struct_array).unwrap(); + /// ``` pub fn new(value: Option, typed_value: Option) -> Self { - Self { - value: value.map(Cow::Owned), - typed_value: typed_value.map(Cow::Owned), - } + Self { value, typed_value } } -} -impl<'a> ShreddingState<'a> { /// Return a reference to the value field, if present pub fn value_field(&self) -> Option<&ArrayRef> { - self.value.as_deref() + self.value.as_ref() } /// Return a reference to the typed_value field, if present pub fn typed_value_field(&self) -> Option<&ArrayRef> { - self.typed_value.as_deref() + self.typed_value.as_ref() } - /// Consume self and return the `(value, typed_value)` Cows so callers can extract long-lived - /// borrows from `Cow::Borrowed` variants. - pub(crate) fn into_fields(self) -> (Option>, Option>) { - (self.value, self.typed_value) - } - - /// Convert into a `'static`-lifetime state by cloning any borrowed Arcs. Cheap (just Arc - /// refcount bumps for each `Cow::Borrowed`; no-op for `Cow::Owned`). - pub fn into_static(self) -> ShreddingState<'static> { - ShreddingState { - value: self.value.map(|cow| Cow::Owned(cow.into_owned())), - typed_value: self.typed_value.map(|cow| Cow::Owned(cow.into_owned())), - } - } - - /// Consume self and return the `(value, typed_value)` borrowed references, asserting that - /// the state was constructed from borrowed data. Panics on `Cow::Owned`. - /// - /// Use this in code paths (like unshred) that only ever receive borrowed state and need - /// `&'a ArrayRef` references to store in lifetime-parameterized builders. - pub(crate) fn into_borrowed_fields(self) -> (Option<&'a ArrayRef>, Option<&'a ArrayRef>) { - fn expect_borrowed<'a>(opt: Option>) -> Option<&'a ArrayRef> { - opt.map(|cow| match cow { - Cow::Borrowed(r) => r, - Cow::Owned(_) => unreachable!("expected Cow::Borrowed ShreddingState"), - }) - } - ( - expect_borrowed(self.value), - expect_borrowed(self.typed_value), - ) - } - - /// Slice all the underlying arrays. Produces owned data. - pub fn slice(&self, offset: usize, length: usize) -> ShreddingState<'static> { - ShreddingState { - value: self - .value_field() - .map(|v| Cow::Owned(v.slice(offset, length))), - typed_value: self - .typed_value_field() - .map(|tv| Cow::Owned(tv.slice(offset, length))), + /// Slice all the underlying arrays + pub fn slice(&self, offset: usize, length: usize) -> Self { + Self { + value: self.value.as_ref().map(|v| v.slice(offset, length)), + typed_value: self.typed_value.as_ref().map(|tv| tv.slice(offset, length)), } } } -impl<'a> TryFrom<&'a StructArray> for ShreddingState<'a> { +impl TryFrom<&StructArray> for ShreddingState { type Error = ArrowError; - fn try_from(inner_struct: &'a StructArray) -> Result { + fn try_from(inner_struct: &StructArray) -> Result { // The `value` column need not exist, but if it does it must be a binary type. let value = if let Some(value_col) = inner_struct.column_by_name("value") { validate_binary_array(value_col.as_ref(), "value")?; - Some(Cow::Borrowed(value_col)) + Some(value_col.clone()) } else { None }; - let typed_value = inner_struct - .column_by_name("typed_value") - .map(Cow::Borrowed); - Ok(ShreddingState { value, typed_value }) + let typed_value = inner_struct.column_by_name("typed_value").cloned(); + Ok(ShreddingState::new(value, typed_value)) } } diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index 3c216e8ec1b5..172f41a1767e 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -25,52 +25,22 @@ use arrow_schema::{ArrowError, DataType, FieldRef}; use parquet_variant::{VariantPath, VariantPathElement}; use crate::VariantArray; -use crate::variant_array::{ShreddingState, validate_binary_array}; +use crate::variant_array::ShreddingState; use crate::variant_to_arrow::make_variant_to_arrow_row_builder; use arrow::array::AsArray; -use std::borrow::Cow; use std::sync::Arc; -pub(crate) enum ShreddedPathStep<'a> { - /// Path step succeeded. - Success { - /// The shredding state for the next step. - next_state: ShreddingState<'a>, - /// The `typed_value` column the step descended through. Returned so the caller can union - /// its nulls into the accumulated null buffer without needing to re-borrow from the - /// (now-consumed) input state. - consumed_typed_value: Cow<'a, ArrayRef>, - }, +pub(crate) enum ShreddedPathStep { + /// Path step succeeded, return the new shredding state + Success(ShreddingState), /// The path element is not present in the `typed_value` column and there is no `value` column, /// so we know it does not exist. It, and all paths under it, are all-NULL. Missing, - /// The path element is not present in the `typed_value` column and must be retrieved from the - /// `value` column instead. The caller should be prepared to handle any value, including the - /// requested type, an arbitrary "wrong" type, or `Variant::Null`. The consumed `value` column - /// is returned so the caller can use it without re-borrowing from the (now-consumed) input. - NotShredded { - consumed_value: Option>, - }, -} - -/// Walk into `typed_value`'s nested `StructArray` and look up `name`. Returns -/// `Ok(Some(nested_struct))` on success, `Ok(None)` for a missing path (non-struct or missing -/// field), and `Err` on an invalid nested struct. -fn walk_field_step<'b>(typed_value: &'b ArrayRef, name: &str) -> Result> { - let Some(struct_array) = typed_value.as_struct_opt() else { - return Ok(None); - }; - let Some(field) = struct_array.column_by_name(name) else { - return Ok(None); - }; - let nested_struct = field.as_struct_opt().ok_or_else(|| { - ArrowError::InvalidArgumentError(format!( - "Expected Struct array while following path, got {}", - field.data_type(), - )) - })?; - Ok(Some(nested_struct)) + /// The path element is not present in the `typed_value` column and must be retrieved from the `value` + /// column instead. The caller should be prepared to handle any value, including the requested + /// type, an arbitrary "wrong" type, or `Variant::Null`. + NotShredded, } /// Given a shredded variant field -- a `(value?, typed_value?)` pair -- try to take one path step @@ -78,62 +48,51 @@ fn walk_field_step<'b>(typed_value: &'b ArrayRef, name: &str) -> Result( - shredding_state: ShreddingState<'a>, +pub(crate) fn follow_shredded_path_element( + shredding_state: &ShreddingState, path_element: &VariantPathElement<'_>, _cast_options: &CastOptions, -) -> Result> { - let (value, typed_value) = shredding_state.into_fields(); - let missing_path_step = |value: Option>| match value { - Some(_) => ShreddedPathStep::NotShredded { - consumed_value: value, - }, +) -> Result { + // If the requested path element is not present in `typed_value`, and `value` is missing, then + // we know it does not exist; it, and all paths under it, are all-NULL. + let missing_path_step = || match shredding_state.value_field() { + Some(_) => ShreddedPathStep::NotShredded, None => ShreddedPathStep::Missing, }; - let Some(typed_value) = typed_value else { - return Ok(missing_path_step(value)); + let Some(typed_value) = shredding_state.typed_value_field() else { + return Ok(missing_path_step()); }; match path_element { - VariantPathElement::Field { name } => match typed_value { - Cow::Borrowed(typed_value_ref) => { - let Some(nested) = walk_field_step(typed_value_ref, name)? else { - return Ok(missing_path_step(value)); - }; - // nested: &'a StructArray, so try_from produces ShreddingState<'a> with - // Cow::Borrowed views into `nested`. No Arc bumps on this path. - let next_state = ShreddingState::try_from(nested)?; - Ok(ShreddedPathStep::Success { - next_state, - consumed_typed_value: Cow::Borrowed(typed_value_ref), - }) - } - Cow::Owned(typed_value_arc) => { - let Some(nested) = walk_field_step(&typed_value_arc, name)? else { - return Ok(missing_path_step(value)); - }; - // `typed_value_arc` is the owned input and is about to move into the returned - // `consumed_typed_value`. We cannot also borrow into it (self-referential), so we - // clone `nested`'s columns to produce a `ShreddingState<'static>` (coerces to 'a). - let value_col = nested.column_by_name("value"); - if let Some(v) = value_col { - validate_binary_array(v.as_ref(), "value")?; - } - let typed_value_col = nested.column_by_name("typed_value"); - let next_state = ShreddingState::new(value_col.cloned(), typed_value_col.cloned()); - Ok(ShreddedPathStep::Success { - next_state, - consumed_typed_value: Cow::Owned(typed_value_arc), - }) - } - }, + VariantPathElement::Field { name } => { + // Try to step into the requested field name of a struct. + // First, try to downcast to StructArray + let Some(struct_array) = typed_value.as_any().downcast_ref::() else { + // Object field path step follows JSONPath semantics and returns missing path step (NotShredded/Missing) on non-struct path + return Ok(missing_path_step()); + }; + + // Now try to find the column - missing column in a present struct is just missing data + let Some(field) = struct_array.column_by_name(name) else { + // Missing column in a present struct is just missing, not wrong - return Ok + return Ok(missing_path_step()); + }; + + let struct_array = field.as_struct_opt().ok_or_else(|| { + // TODO: Should we blow up? Or just end the traversal and let the normal + // variant pathing code sort out the mess that it must anyway be + // prepared to handle? + ArrowError::InvalidArgumentError(format!( + "Expected Struct array while following path, got {}", + field.data_type(), + )) + })?; + + let state = ShreddingState::try_from(struct_array)?; + Ok(ShreddedPathStep::Success(state)) + } VariantPathElement::Index { .. } => { // TODO: Support array indexing. Among other things, it will require slicing not // only the array we have here, but also the corresponding metadata and null masks. @@ -195,19 +154,18 @@ fn shredded_get_path( // Peel away the prefix of path elements that traverses the shredded parts of this variant // column. Shredding will traverse the rest of the path on a per-row basis. - let mut shredding_state = input.shredding_state(); + let mut shredding_state = input.shredding_state().clone(); let mut accumulated_nulls = input.inner().nulls().cloned(); let mut path_index = 0; for path_element in path { - match follow_shredded_path_element(shredding_state, path_element, cast_options)? { - ShreddedPathStep::Success { - next_state, - consumed_typed_value, - } => { - // Union nulls from the typed_value we just descended through. - accumulated_nulls = - NullBuffer::union(accumulated_nulls.as_ref(), consumed_typed_value.nulls()); - shredding_state = next_state; + match follow_shredded_path_element(&shredding_state, path_element, cast_options)? { + ShreddedPathStep::Success(state) => { + // Union nulls from the typed_value we just accessed + if let Some(typed_value) = shredding_state.typed_value_field() { + accumulated_nulls = + NullBuffer::union(accumulated_nulls.as_ref(), typed_value.nulls()); + } + shredding_state = state; path_index += 1; continue; } @@ -219,9 +177,9 @@ fn shredded_get_path( }; return Ok(arr); } - ShreddedPathStep::NotShredded { consumed_value } => { + ShreddedPathStep::NotShredded => { let target = make_target_variant( - consumed_value.map(Cow::into_owned), + shredding_state.value_field().cloned(), None, accumulated_nulls, ); @@ -231,10 +189,9 @@ fn shredded_get_path( } // Path exhausted! Create a new `VariantArray` for the location we landed on. - let (final_value, final_typed_value) = shredding_state.into_fields(); let target = make_target_variant( - final_value.map(Cow::into_owned), - final_typed_value.map(Cow::into_owned), + shredding_state.value_field().cloned(), + shredding_state.typed_value_field().cloned(), accumulated_nulls, );