Merge remote-tracking branch 'apache/main' into variant-value-builder

Author: alamb

arrow-cast/src/cast/decimal.rs

@@ -81,7 +81,9 @@ impl DecimalCast for i64 {
     }
 
     fn from_f64(n: f64) -> Option<Self> {
-        n.to_i64()
+        // Call implementation explicitly otherwise this resolves to `to_i64`
+        // in arrow-buffer that behaves differently.
+        num::traits::ToPrimitive::to_i64(&n)
     }
 }
 

arrow-cast/src/cast/mod.rs

@@ -3101,6 +3101,35 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_cast_floating_to_decimals() {
+        for output_type in [
+            DataType::Decimal32(9, 3),
+            DataType::Decimal64(9, 3),
+            DataType::Decimal128(9, 3),
+            DataType::Decimal256(9, 3),
+        ] {
+            let input_type = DataType::Float64;
+            assert!(can_cast_types(&input_type, &output_type));
+
+            let array = vec![Some(1.1_f64)];
+            let array = PrimitiveArray::<Float64Type>::from_iter(array);
+            let result = cast_with_options(
+                &array,
+                &output_type,
+                &CastOptions {
+                    safe: false,
+                    format_options: FormatOptions::default(),
+                },
+            );
+            assert!(
+                result.is_ok(),
+                "Failed to cast to {output_type} with: {}",
+                result.unwrap_err()
+            );
+        }
+    }
+
     #[test]
     fn test_cast_decimal128_to_decimal128_overflow() {
         let input_type = DataType::Decimal128(38, 3);

parquet-testing

@@ -129,7 +129,7 @@ impl VariantArray {
         Ok(Self {
             inner: inner.clone(),
             metadata: metadata.clone(),
-            shredding_state: ShreddingState::try_new(metadata.clone(), value, typed_value)?,
+            shredding_state: ShreddingState::try_new(value, typed_value)?,
         })
     }
 
@@ -154,8 +154,7 @@ impl VariantArray {
         // This would be a lot simpler if ShreddingState were just a pair of Option... we already
         // have everything we need.
         let inner = builder.build();
-        let shredding_state =
-            ShreddingState::try_new(metadata.clone(), value, typed_value).unwrap(); // valid by construction
+        let shredding_state = ShreddingState::try_new(value, typed_value).unwrap(); // valid by construction
         Self {
             inner,
             metadata,
@@ -222,7 +221,7 @@ impl VariantArray {
                     typed_value_to_variant(typed_value, index)
                 }
             }
-            ShreddingState::AllNull { .. } => {
+            ShreddingState::AllNull => {
                 // AllNull case: neither value nor typed_value fields exist
                 // NOTE: This handles the case where neither value nor typed_value fields exist.
                 // For top-level variants, this returns Variant::Null (JSON null).
@@ -325,14 +324,11 @@ impl ShreddedVariantFieldArray {
             .and_then(|col| col.as_binary_view_opt().cloned());
         let typed_value = inner_struct.column_by_name("typed_value").cloned();
 
-        // Use a dummy metadata for the constructor (ShreddedVariantFieldArray doesn't have metadata)
-        let dummy_metadata = arrow::array::BinaryViewArray::new_null(inner_struct.len());
-
         // Note this clone is cheap, it just bumps the ref count
         let inner = inner_struct.clone();
         Ok(Self {
             inner: inner.clone(),
-            shredding_state: ShreddingState::try_new(dummy_metadata, value, typed_value)?,
+            shredding_state: ShreddingState::try_new(value, typed_value)?,
         })
     }
 
@@ -432,16 +428,10 @@ impl Array for ShreddedVariantFieldArray {
 #[derive(Debug)]
 pub enum ShreddingState {
     /// This variant has no typed_value field
-    Unshredded {
-        metadata: BinaryViewArray,
-        value: BinaryViewArray,
-    },
+    Unshredded { value: BinaryViewArray },
     /// This variant has a typed_value field and no value field
     /// meaning it is the shredded type
-    Typed {
-        metadata: BinaryViewArray,
-        typed_value: ArrayRef,
-    },
+    Typed { typed_value: ArrayRef },
     /// Imperfectly shredded: Shredded values reside in `typed_value` while those that failed to
     /// shred reside in `value`. Missing field values are NULL in both columns, while NULL primitive
     /// values have NULL `typed_value` and `Variant::Null` in `value`.
@@ -453,7 +443,6 @@ pub enum ShreddingState {
     /// the `value` is a variant object containing the subset of fields for which shredding was
     /// not even attempted.
     PartiallyShredded {
-        metadata: BinaryViewArray,
         value: BinaryViewArray,
         typed_value: ArrayRef,
     },
@@ -463,38 +452,20 @@ pub enum ShreddingState {
     /// Note: By strict spec interpretation, this should only be valid for shredded object fields,
     /// not top-level variants. However, we allow it and treat as Variant::Null for pragmatic
     /// handling of missing data.
-    AllNull { metadata: BinaryViewArray },
+    AllNull,
 }
 
 impl ShreddingState {
     /// try to create a new `ShreddingState` from the given fields
     pub fn try_new(
-        metadata: BinaryViewArray,
         value: Option<BinaryViewArray>,
         typed_value: Option<ArrayRef>,
     ) -> Result<Self, ArrowError> {
-        match (metadata, value, typed_value) {
-            (metadata, Some(value), Some(typed_value)) => Ok(Self::PartiallyShredded {
-                metadata,
-                value,
-                typed_value,
-            }),
-            (metadata, Some(value), None) => Ok(Self::Unshredded { metadata, value }),
-            (metadata, None, Some(typed_value)) => Ok(Self::Typed {
-                metadata,
-                typed_value,
-            }),
-            (metadata, None, None) => Ok(Self::AllNull { metadata }),
-        }
-    }
-
-    /// Return a reference to the metadata field
-    pub fn metadata_field(&self) -> &BinaryViewArray {
-        match self {
-            ShreddingState::Unshredded { metadata, .. } => metadata,
-            ShreddingState::Typed { metadata, .. } => metadata,
-            ShreddingState::PartiallyShredded { metadata, .. } => metadata,
-            ShreddingState::AllNull { metadata } => metadata,
+        match (value, typed_value) {
+            (Some(value), Some(typed_value)) => Ok(Self::PartiallyShredded { value, typed_value }),
+            (Some(value), None) => Ok(Self::Unshredded { value }),
+            (None, Some(typed_value)) => Ok(Self::Typed { typed_value }),
+            (None, None) => Ok(Self::AllNull),
         }
     }
 
@@ -504,7 +475,7 @@ impl ShreddingState {
             ShreddingState::Unshredded { value, .. } => Some(value),
             ShreddingState::Typed { .. } => None,
             ShreddingState::PartiallyShredded { value, .. } => Some(value),
-            ShreddingState::AllNull { .. } => None,
+            ShreddingState::AllNull => None,
         }
     }
 
@@ -514,36 +485,26 @@ impl ShreddingState {
             ShreddingState::Unshredded { .. } => None,
             ShreddingState::Typed { typed_value, .. } => Some(typed_value),
             ShreddingState::PartiallyShredded { typed_value, .. } => Some(typed_value),
-            ShreddingState::AllNull { .. } => None,
+            ShreddingState::AllNull => None,
         }
     }
 
     /// Slice all the underlying arrays
     pub fn slice(&self, offset: usize, length: usize) -> Self {
         match self {
-            ShreddingState::Unshredded { metadata, value } => ShreddingState::Unshredded {
-                metadata: metadata.slice(offset, length),
+            ShreddingState::Unshredded { value } => ShreddingState::Unshredded {
                 value: value.slice(offset, length),
             },
-            ShreddingState::Typed {
-                metadata,
-                typed_value,
-            } => ShreddingState::Typed {
-                metadata: metadata.slice(offset, length),
-                typed_value: typed_value.slice(offset, length),
-            },
-            ShreddingState::PartiallyShredded {
-                metadata,
-                value,
-                typed_value,
-            } => ShreddingState::PartiallyShredded {
-                metadata: metadata.slice(offset, length),
-                value: value.slice(offset, length),
+            ShreddingState::Typed { typed_value } => ShreddingState::Typed {
                 typed_value: typed_value.slice(offset, length),
             },
-            ShreddingState::AllNull { metadata } => ShreddingState::AllNull {
-                metadata: metadata.slice(offset, length),
-            },
+            ShreddingState::PartiallyShredded { value, typed_value } => {
+                ShreddingState::PartiallyShredded {
+                    value: value.slice(offset, length),
+                    typed_value: typed_value.slice(offset, length),
+                }
+            }
+            ShreddingState::AllNull => ShreddingState::AllNull,
         }
     }
 }
@@ -744,7 +705,7 @@ mod test {
         // Verify the shredding state is AllNull
         assert!(matches!(
             variant_array.shredding_state(),
-            ShreddingState::AllNull { .. }
+            ShreddingState::AllNull
         ));
 
         // Verify that value() returns Variant::Null (compensating for spec violation)
@@ -801,17 +762,10 @@ mod test {
 
     #[test]
     fn all_null_shredding_state() {
-        let metadata = BinaryViewArray::from(vec![b"test" as &[u8]]);
-        let shredding_state = ShreddingState::try_new(metadata.clone(), None, None).unwrap();
+        let shredding_state = ShreddingState::try_new(None, None).unwrap();
 
         // Verify the shredding state is AllNull
-        assert!(matches!(shredding_state, ShreddingState::AllNull { .. }));
-
-        // Verify metadata is preserved correctly
-        if let ShreddingState::AllNull { metadata: m } = shredding_state {
-            assert_eq!(m.len(), metadata.len());
-            assert_eq!(m.value(0), metadata.value(0));
-        }
+        assert!(matches!(shredding_state, ShreddingState::AllNull));
     }
 
     #[test]
@@ -827,7 +781,7 @@ mod test {
         // Verify the shredding state is AllNull
         assert!(matches!(
             variant_array.shredding_state(),
-            ShreddingState::AllNull { .. }
+            ShreddingState::AllNull
         ));
 
         // Verify all values are null

parquet-variant-compute/src/variant_array.rs

@@ -135,7 +135,8 @@ fn shredded_get_path(
     let shred_basic_variant =
         |target: VariantArray, path: VariantPath<'_>, as_field: Option<&Field>| {
             let as_type = as_field.map(|f| f.data_type());
-            let mut builder = make_variant_to_arrow_row_builder(path, as_type, cast_options)?;
+            let mut builder =
+                make_variant_to_arrow_row_builder(path, as_type, cast_options, target.len())?;
             for i in 0..target.len() {
                 if target.is_null(i) {
                     builder.append_null()?;

parquet-variant-compute/src/variant_get.rs

@@ -76,7 +76,7 @@ impl<'a> VariantToArrowRowBuilder<'a> {
         }
     }
 
-    pub fn finish(&mut self) -> Result<ArrayRef> {
+    pub fn finish(self) -> Result<ArrayRef> {
         use VariantToArrowRowBuilder::*;
         match self {
             Int8(b) => b.finish(),
@@ -97,19 +97,41 @@ pub(crate) fn make_variant_to_arrow_row_builder<'a>(
     path: VariantPath<'a>,
     data_type: Option<&'a DataType>,
     cast_options: &'a CastOptions,
+    capacity: usize,
 ) -> Result<VariantToArrowRowBuilder<'a>> {
     use VariantToArrowRowBuilder::*;
 
     let mut builder = match data_type {
         // If no data type was requested, build an unshredded VariantArray.
-        None => BinaryVariant(VariantToBinaryVariantArrowRowBuilder::new(16)),
-        Some(DataType::Int8) => Int8(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Int16) => Int16(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Int32) => Int32(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Int64) => Int64(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Float16) => Float16(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Float32) => Float32(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Float64) => Float64(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
+        None => BinaryVariant(VariantToBinaryVariantArrowRowBuilder::new(capacity)),
+        Some(DataType::Int8) => Int8(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Int16) => Int16(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Int32) => Int32(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Int64) => Int64(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Float16) => Float16(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Float32) => Float32(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Float64) => Float64(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
         _ => {
             return Err(ArrowError::NotYetImplemented(format!(
                 "variant_get with path={:?} and data_type={:?} not yet implemented",
@@ -150,7 +172,7 @@ impl<'a> VariantPathRowBuilder<'a> {
         }
     }
 
-    fn finish(&mut self) -> Result<ArrayRef> {
+    fn finish(self) -> Result<ArrayRef> {
         self.builder.finish()
     }
 }
@@ -180,9 +202,9 @@ pub(crate) struct VariantToPrimitiveArrowRowBuilder<'a, T: ArrowPrimitiveType> {
 }
 
 impl<'a, T: ArrowPrimitiveType> VariantToPrimitiveArrowRowBuilder<'a, T> {
-    fn new(cast_options: &'a CastOptions<'a>) -> Self {
+    fn new(cast_options: &'a CastOptions<'a>, capacity: usize) -> Self {
         Self {
-            builder: PrimitiveBuilder::<T>::new(),
+            builder: PrimitiveBuilder::<T>::with_capacity(capacity),
             cast_options,
         }
     }
@@ -217,7 +239,7 @@ where
         }
     }
 
-    fn finish(&mut self) -> Result<ArrayRef> {
+    fn finish(mut self) -> Result<ArrayRef> {
         Ok(Arc::new(self.builder.finish()))
     }
 }
@@ -253,9 +275,7 @@ impl VariantToBinaryVariantArrowRowBuilder {
         Ok(true)
     }
 
-    fn finish(&mut self) -> Result<ArrayRef> {
-        // VariantArrayBuilder::build takes ownership, so we need to replace it
-        let builder = std::mem::replace(&mut self.builder, VariantArrayBuilder::new(0));
-        Ok(Arc::new(builder.build()))
+    fn finish(self) -> Result<ArrayRef> {
+        Ok(Arc::new(self.builder.build()))
     }
 }