Merge branch 'main' into issue-9645-ree-optimization

Author: liamzwbao

arrow-arith/src/aggregate.rs

@@ -20,7 +20,7 @@
 use arrow_array::cast::*;
 use arrow_array::iterator::ArrayIter;
 use arrow_array::*;
-use arrow_buffer::{ArrowNativeType, NullBuffer};
+use arrow_buffer::NullBuffer;
 use arrow_data::bit_iterator::try_for_each_valid_idx;
 use arrow_schema::*;
 use std::borrow::BorrowMut;
@@ -541,11 +541,9 @@ pub fn min_string_view(array: &StringViewArray) -> Option<&str> {
 ///
 /// This doesn't detect overflow. Once overflowing, the result will wrap around.
 /// For an overflow-checking variant, use [`sum_array_checked`] instead.
-pub fn sum_array<T, A: ArrayAccessor<Item = T::Native>>(array: A) -> Option<T::Native>
-where
-    T: ArrowNumericType,
-    T::Native: ArrowNativeTypeOp,
-{
+pub fn sum_array<T: ArrowNumericType, A: ArrayAccessor<Item = T::Native>>(
+    array: A,
+) -> Option<T::Native> {
     match array.data_type() {
         DataType::Dictionary(_, _) => {
             let null_count = array.null_count();
@@ -583,13 +581,9 @@ where
 /// use [`sum_array`] instead.
 /// Additionally returns an `Err` on run-end-encoded arrays with a provided
 /// values type parameter that is incorrect.
-pub fn sum_array_checked<T, A: ArrayAccessor<Item = T::Native>>(
+pub fn sum_array_checked<T: ArrowNumericType, A: ArrayAccessor<Item = T::Native>>(
     array: A,
-) -> Result<Option<T::Native>, ArrowError>
-where
-    T: ArrowNumericType,
-    T::Native: ArrowNativeTypeOp,
-{
+) -> Result<Option<T::Native>, ArrowError> {
     match array.data_type() {
         DataType::Dictionary(_, _) => {
             let null_count = array.null_count();
@@ -717,21 +711,17 @@ mod ree {
 
 /// Returns the min of values in the array of `ArrowNumericType` type, or dictionary
 /// array with value of `ArrowNumericType` type.
-pub fn min_array<T, A: ArrayAccessor<Item = T::Native>>(array: A) -> Option<T::Native>
-where
-    T: ArrowNumericType,
-    T::Native: ArrowNativeType,
-{
+pub fn min_array<T: ArrowNumericType, A: ArrayAccessor<Item = T::Native>>(
+    array: A,
+) -> Option<T::Native> {
     min_max_array_helper::<T, A, _, _>(array, |a, b| a.is_gt(*b), min)
 }
 
 /// Returns the max of values in the array of `ArrowNumericType` type, or dictionary
 /// array with value of `ArrowNumericType` type.
-pub fn max_array<T, A: ArrayAccessor<Item = T::Native>>(array: A) -> Option<T::Native>
-where
-    T: ArrowNumericType,
-    T::Native: ArrowNativeTypeOp,
-{
+pub fn max_array<T: ArrowNumericType, A: ArrayAccessor<Item = T::Native>>(
+    array: A,
+) -> Option<T::Native> {
     min_max_array_helper::<T, A, _, _>(array, |a, b| a.is_lt(*b), max)
 }
 
@@ -874,11 +864,9 @@ pub fn bool_or(array: &BooleanArray) -> Option<bool> {
 ///
 /// This detects overflow and returns an `Err` for that. For an non-overflow-checking variant,
 /// use [`sum`] instead.
-pub fn sum_checked<T>(array: &PrimitiveArray<T>) -> Result<Option<T::Native>, ArrowError>
-where
-    T: ArrowNumericType,
-    T::Native: ArrowNativeTypeOp,
-{
+pub fn sum_checked<T: ArrowNumericType>(
+    array: &PrimitiveArray<T>,
+) -> Result<Option<T::Native>, ArrowError> {
     let null_count = array.null_count();
 
     if null_count == array.len() {
@@ -922,10 +910,7 @@ where
 ///
 /// This doesn't detect overflow in release mode by default. Once overflowing, the result will
 /// wrap around. For an overflow-checking variant, use [`sum_checked`] instead.
-pub fn sum<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native>
-where
-    T::Native: ArrowNativeTypeOp,
-{
+pub fn sum<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native> {
     aggregate::<T::Native, T, SumAccumulator<T::Native>>(array)
 }
 
@@ -940,10 +925,7 @@ where
 /// let result = min(&array);
 /// assert_eq!(result, Some(2));
 /// ```
-pub fn min<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native>
-where
-    T::Native: PartialOrd,
-{
+pub fn min<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native> {
     aggregate::<T::Native, T, MinAccumulator<T::Native>>(array)
 }
 
@@ -958,10 +940,7 @@ where
 /// let result = max(&array);
 /// assert_eq!(result, Some(8));
 /// ```
-pub fn max<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native>
-where
-    T::Native: PartialOrd,
-{
+pub fn max<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native> {
     aggregate::<T::Native, T, MaxAccumulator<T::Native>>(array)
 }
 

arrow-arith/src/arithmetic.rs

@@ -170,11 +170,7 @@ pub fn multiply_fixed_point(
 }
 
 /// Divide a decimal native value by given divisor and round the result.
-fn divide_and_round<I>(input: I::Native, div: I::Native) -> I::Native
-where
-    I: DecimalType,
-    I::Native: ArrowNativeTypeOp,
-{
+fn divide_and_round<I: DecimalType>(input: I::Native, div: I::Native) -> I::Native {
     let d = input.div_wrapping(div);
     let r = input.mod_wrapping(div);
 

arrow-array/src/array/byte_view_array.rs

@@ -670,6 +670,37 @@ impl<T: ByteViewType + ?Sized> GenericByteViewArray<T> {
         }
     }
 
+    /// Returns the total number of bytes of all non-null values in this array.
+    ///
+    /// Unlike [`Self::total_buffer_bytes_used`], this method includes inlined strings
+    /// (those with length ≤ [`MAX_INLINE_VIEW_LEN`]), making it suitable as a
+    /// capacity hint when pre-allocating output buffers.
+    ///
+    /// Null values are excluded from the sum.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// # use arrow_array::StringViewArray;
+    /// let array = StringViewArray::from_iter(vec![
+    ///     Some("hello"),   // 5 bytes, inlined
+    ///     None,            // excluded
+    ///     Some("large payload over 12 bytes"),  // 27 bytes, non-inlined
+    /// ]);
+    /// assert_eq!(array.total_bytes_len(), 5 + 27);
+    /// ```
+    pub fn total_bytes_len(&self) -> usize {
+        match self.nulls() {
+            None => self.views().iter().map(|v| (*v as u32) as usize).sum(),
+            Some(nulls) => self
+                .views()
+                .iter()
+                .zip(nulls.iter())
+                .map(|(v, is_valid)| if is_valid { (*v as u32) as usize } else { 0 })
+                .sum(),
+        }
+    }
+
     /// Returns the total number of bytes used by all non inlined views in all
     /// buffers.
     ///
@@ -1809,4 +1840,41 @@ mod tests {
         assert!(from_utf8(array.value(2)).is_ok());
         array
     }
+
+    #[test]
+    fn test_total_bytes_len() {
+        // inlined: "hello"=5, "world"=5, "lulu"=4 → 14
+        // non-inlined: "large payload over 12 bytes"=27
+        // null: should not count
+        let mut builder = StringViewBuilder::new();
+        builder.append_value("hello");
+        builder.append_value("world");
+        builder.append_value("lulu");
+        builder.append_null();
+        builder.append_value("large payload over 12 bytes");
+        let array = builder.finish();
+        assert_eq!(array.total_bytes_len(), 5 + 5 + 4 + 27);
+    }
+
+    #[test]
+    fn test_total_bytes_len_empty() {
+        let array = StringViewArray::from_iter::<Vec<Option<&str>>>(vec![]);
+        assert_eq!(array.total_bytes_len(), 0);
+    }
+
+    #[test]
+    fn test_total_bytes_len_all_nulls() {
+        let array = StringViewArray::new_null(5);
+        assert_eq!(array.total_bytes_len(), 0);
+    }
+
+    #[test]
+    fn test_total_bytes_len_binary_view() {
+        let array = BinaryViewArray::from_iter(vec![
+            Some(b"hi".as_ref()),
+            None,
+            Some(b"large payload over 12 bytes".as_ref()),
+        ]);
+        assert_eq!(array.total_bytes_len(), 2 + 27);
+    }
 }

arrow-cast/src/cast/decimal.rs

@@ -812,7 +812,7 @@ where
     T: ArrowPrimitiveType,
     <T as ArrowPrimitiveType>::Native: NumCast,
     D: DecimalType + ArrowPrimitiveType,
-    <D as ArrowPrimitiveType>::Native: ArrowNativeTypeOp + ToPrimitive,
+    <D as ArrowPrimitiveType>::Native: ToPrimitive,
 {
     let array = array.as_primitive::<D>();
 

arrow-cast/src/cast/mod.rs

@@ -43,13 +43,15 @@ mod list;
 mod map;
 mod run_array;
 mod string;
+mod union;
 
 use crate::cast::decimal::*;
 use crate::cast::dictionary::*;
 use crate::cast::list::*;
 use crate::cast::map::*;
 use crate::cast::run_array::*;
 use crate::cast::string::*;
+pub use crate::cast::union::*;
 
 use arrow_buffer::IntervalMonthDayNano;
 use arrow_data::ByteView;
@@ -71,6 +73,7 @@ use arrow_select::take::take;
 use num_traits::{NumCast, ToPrimitive, cast::AsPrimitive};
 
 pub use decimal::{DecimalCast, rescale_decimal};
+pub use string::cast_single_string_to_boolean_default;
 
 /// CastOptions provides a way to override the default cast behaviors
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -108,6 +111,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
             can_cast_types(from_value_type, to_value_type)
         }
         (Dictionary(_, value_type), _) => can_cast_types(value_type, to_type),
+        (Union(fields, _), _) => union::resolve_child_array(fields, to_type).is_some(),
+        (_, Union(_, _)) => false,
         (RunEndEncoded(_, value_type), _) => can_cast_types(value_type.data_type(), to_type),
         (_, RunEndEncoded(_, value_type)) => can_cast_types(from_type, value_type.data_type()),
         (_, Dictionary(_, value_type)) => can_cast_types(from_type, value_type),
@@ -230,7 +235,6 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
         }
         (Struct(_), _) => false,
         (_, Struct(_)) => false,
-
         (_, Boolean) => from_type.is_integer() || from_type.is_floating() || from_type.is_string(),
         (Boolean, _) => to_type.is_integer() || to_type.is_floating() || to_type.is_string(),
 
@@ -781,6 +785,14 @@ pub fn cast_with_options(
                 ))),
             }
         }
+        (Union(_, _), _) => union_extract_by_type(
+            array.as_any().downcast_ref::<UnionArray>().unwrap(),
+            to_type,
+            cast_options,
+        ),
+        (_, Union(_, _)) => Err(ArrowError::CastError(format!(
+            "Casting from {from_type} to {to_type} not supported"
+        ))),
         (Dictionary(index_type, _), _) => match **index_type {
             Int8 => dictionary_cast::<Int8Type>(array, to_type, cast_options),
             Int16 => dictionary_cast::<Int16Type>(array, to_type, cast_options),
@@ -2287,7 +2299,7 @@ fn cast_from_decimal<D, F>(
 ) -> Result<ArrayRef, ArrowError>
 where
     D: DecimalType + ArrowPrimitiveType,
-    <D as ArrowPrimitiveType>::Native: ArrowNativeTypeOp + ToPrimitive,
+    <D as ArrowPrimitiveType>::Native: ToPrimitive,
     F: Fn(D::Native) -> f64,
 {
     use DataType::*;
@@ -2464,7 +2476,7 @@ where
     R::Native: NumCast,
 {
     from.try_unary(|value| {
-        num_traits::cast::cast::<T::Native, R::Native>(value).ok_or_else(|| {
+        num_cast::<T::Native, R::Native>(value).ok_or_else(|| {
             ArrowError::CastError(format!(
                 "Can't cast value {:?} to type {}",
                 value,
@@ -2474,6 +2486,17 @@ where
     })
 }
 
+/// Natural cast between numeric types
+/// Return None if the input `value` can't be casted to type `O`.
+#[inline]
+pub fn num_cast<I, O>(value: I) -> Option<O>
+where
+    I: NumCast,
+    O: NumCast,
+{
+    num_traits::cast::cast::<I, O>(value)
+}
+
 // Natural cast between numeric types
 // If the value of T can't be casted to R, it will be converted to null
 fn numeric_cast<T, R>(from: &PrimitiveArray<T>) -> PrimitiveArray<R>
@@ -2483,7 +2506,7 @@ where
     T::Native: NumCast,
     R::Native: NumCast,
 {
-    from.unary_opt::<_, R>(num_traits::cast::cast::<T::Native, R::Native>)
+    from.unary_opt::<_, R>(num_cast::<T::Native, R::Native>)
 }
 
 fn cast_numeric_to_binary<FROM: ArrowPrimitiveType, O: OffsetSizeTrait>(
@@ -2540,16 +2563,23 @@ where
     for i in 0..from.len() {
         if from.is_null(i) {
             b.append_null();
-        } else if from.value(i) != T::default_value() {
-            b.append_value(true);
         } else {
-            b.append_value(false);
+            b.append_value(cast_num_to_bool::<T::Native>(from.value(i)));
         }
     }
 
     Ok(b.finish())
 }
 
+/// Cast numeric types to boolean
+#[inline]
+pub fn cast_num_to_bool<I>(value: I) -> bool
+where
+    I: Default + PartialEq,
+{
+    value != I::default()
+}
+
 /// Cast Boolean types to numeric
 ///
 /// `false` returns 0 while `true` returns 1
@@ -2575,11 +2605,8 @@ where
     let iter = (0..from.len()).map(|i| {
         if from.is_null(i) {
             None
-        } else if from.value(i) {
-            // a workaround to cast a primitive to T::Native, infallible
-            num_traits::cast::cast(1)
         } else {
-            Some(T::default_value())
+            single_bool_to_numeric::<T::Native>(from.value(i))
         }
     });
     // Benefit:
@@ -2589,6 +2616,20 @@ where
     unsafe { PrimitiveArray::<T>::from_trusted_len_iter(iter) }
 }
 
+/// Cast single bool value to numeric value.
+#[inline]
+pub fn single_bool_to_numeric<O>(value: bool) -> Option<O>
+where
+    O: num_traits::NumCast + Default,
+{
+    if value {
+        // a workaround to cast a primitive to type O, infallible
+        num_traits::cast::cast(1)
+    } else {
+        Some(O::default())
+    }
+}
+
 /// Helper function to cast from one `BinaryArray` or 'LargeBinaryArray' to 'FixedSizeBinaryArray'.
 fn cast_binary_to_fixed_size_binary<O: OffsetSizeTrait>(
     array: &dyn Array,