Refactor some decimal-related code and tests in preparation for addin…

…g Decimal32 and Decimal64 support
apache · Feb 1, 2025 · 6988882 · 6988882
1 parent a6648fe
commit 6988882
Show file tree

Hide file tree

Showing 10 changed files with 467 additions and 754 deletions.
diff --git a/arrow-cast/src/cast/decimal.rs b/arrow-cast/src/cast/decimal.rs
@@ -25,6 +25,8 @@ pub(crate) trait DecimalCast: Sized {
     fn to_i256(self) -> Option<i256>;
 
     fn from_decimal<T: DecimalCast>(n: T) -> Option<Self>;
+
+    fn from_f64<T: DecimalCast>(n: f64) -> Option<Self>;
 }
 
 impl DecimalCast for i128 {
@@ -39,6 +41,10 @@ impl DecimalCast for i128 {
     fn from_decimal<T: DecimalCast>(n: T) -> Option<Self> {
         n.to_i128()
     }
+
+    fn from_f64<T: DecimalCast>(n: f64) -> Option<Self> {
+        n.to_i128()
+    }
 }
 
 impl DecimalCast for i256 {
@@ -53,6 +59,10 @@ impl DecimalCast for i256 {
     fn from_decimal<T: DecimalCast>(n: T) -> Option<Self> {
         n.to_i256()
     }
+
+    fn from_f64<T: DecimalCast>(n: f64) -> Option<Self> {
+        i256::from_f64(n)
+    }
 }
 
 pub(crate) fn cast_decimal_to_decimal_error<I, O>(
@@ -464,86 +474,41 @@ where
     Ok(Arc::new(result))
 }
 
-pub(crate) fn cast_floating_point_to_decimal128<T: ArrowPrimitiveType>(
+pub(crate) fn cast_floating_point_to_decimal<T: ArrowPrimitiveType, D, M>(
     array: &PrimitiveArray<T>,
     precision: u8,
     scale: i8,
     cast_options: &CastOptions,
 ) -> Result<ArrayRef, ArrowError>
 where
     <T as ArrowPrimitiveType>::Native: AsPrimitive<f64>,
+    D: DecimalType + ArrowPrimitiveType<Native = M>,
+    M: ArrowNativeTypeOp + DecimalCast,
 {
     let mul = 10_f64.powi(scale as i32);
 
     if cast_options.safe {
         array
-            .unary_opt::<_, Decimal128Type>(|v| {
-                (mul * v.as_())
-                    .round()
-                    .to_i128()
-                    .filter(|v| Decimal128Type::is_valid_decimal_precision(*v, precision))
+            .unary_opt::<_, D>(|v| {
+                M::from_f64::<M>((mul * v.as_()).round())
+                    .filter(|v| D::is_valid_decimal_precision(*v, precision))
             })
             .with_precision_and_scale(precision, scale)
             .map(|a| Arc::new(a) as ArrayRef)
     } else {
         array
-            .try_unary::<_, Decimal128Type, _>(|v| {
-                (mul * v.as_())
-                    .round()
-                    .to_i128()
+            .try_unary::<_, D, _>(|v| {
+                M::from_f64::<M>((mul * v.as_()).round())
                     .ok_or_else(|| {
                         ArrowError::CastError(format!(
                             "Cannot cast to {}({}, {}). Overflowing on {:?}",
-                            Decimal128Type::PREFIX,
+                            D::PREFIX,
                             precision,
                             scale,
                             v
                         ))
                     })
-                    .and_then(|v| {
-                        Decimal128Type::validate_decimal_precision(v, precision).map(|_| v)
-                    })
-            })?
-            .with_precision_and_scale(precision, scale)
-            .map(|a| Arc::new(a) as ArrayRef)
-    }
-}
-
-pub(crate) fn cast_floating_point_to_decimal256<T: ArrowPrimitiveType>(
-    array: &PrimitiveArray<T>,
-    precision: u8,
-    scale: i8,
-    cast_options: &CastOptions,
-) -> Result<ArrayRef, ArrowError>
-where
-    <T as ArrowPrimitiveType>::Native: AsPrimitive<f64>,
-{
-    let mul = 10_f64.powi(scale as i32);
-
-    if cast_options.safe {
-        array
-            .unary_opt::<_, Decimal256Type>(|v| {
-                i256::from_f64((v.as_() * mul).round())
-                    .filter(|v| Decimal256Type::is_valid_decimal_precision(*v, precision))
-            })
-            .with_precision_and_scale(precision, scale)
-            .map(|a| Arc::new(a) as ArrayRef)
-    } else {
-        array
-            .try_unary::<_, Decimal256Type, _>(|v| {
-                i256::from_f64((v.as_() * mul).round())
-                    .ok_or_else(|| {
-                        ArrowError::CastError(format!(
-                            "Cannot cast to {}({}, {}). Overflowing on {:?}",
-                            Decimal256Type::PREFIX,
-                            precision,
-                            scale,
-                            v
-                        ))
-                    })
-                    .and_then(|v| {
-                        Decimal256Type::validate_decimal_precision(v, precision).map(|_| v)
-                    })
+                    .and_then(|v| D::validate_decimal_precision(v, precision).map(|_| v))
             })?
             .with_precision_and_scale(precision, scale)
             .map(|a| Arc::new(a) as ArrayRef)

diff --git a/arrow-cast/src/cast/dictionary.rs b/arrow-cast/src/cast/dictionary.rs
@@ -214,50 +214,20 @@ pub(crate) fn cast_to_dictionary<K: ArrowDictionaryKeyType>(
         UInt16 => pack_numeric_to_dictionary::<K, UInt16Type>(array, dict_value_type, cast_options),
         UInt32 => pack_numeric_to_dictionary::<K, UInt32Type>(array, dict_value_type, cast_options),
         UInt64 => pack_numeric_to_dictionary::<K, UInt64Type>(array, dict_value_type, cast_options),
-        Decimal128(p, s) => {
-            let dict = pack_numeric_to_dictionary::<K, Decimal128Type>(
-                array,
-                dict_value_type,
-                cast_options,
-            )?;
-            let dict = dict
-                .as_dictionary::<K>()
-                .downcast_dict::<Decimal128Array>()
-                .ok_or_else(|| {
-                    ArrowError::ComputeError(
-                        "Internal Error: Cannot cast dict to Decimal128Array".to_string(),
-                    )
-                })?;
-            let value = dict.values().clone();
-            // Set correct precision/scale
-            let value = value.with_precision_and_scale(p, s)?;
-            Ok(Arc::new(DictionaryArray::<K>::try_new(
-                dict.keys().clone(),
-                Arc::new(value),
-            )?))
-        }
-        Decimal256(p, s) => {
-            let dict = pack_numeric_to_dictionary::<K, Decimal256Type>(
-                array,
-                dict_value_type,
-                cast_options,
-            )?;
-            let dict = dict
-                .as_dictionary::<K>()
-                .downcast_dict::<Decimal256Array>()
-                .ok_or_else(|| {
-                    ArrowError::ComputeError(
-                        "Internal Error: Cannot cast dict to Decimal256Array".to_string(),
-                    )
-                })?;
-            let value = dict.values().clone();
-            // Set correct precision/scale
-            let value = value.with_precision_and_scale(p, s)?;
-            Ok(Arc::new(DictionaryArray::<K>::try_new(
-                dict.keys().clone(),
-                Arc::new(value),
-            )?))
-        }
+        Decimal128(p, s) => pack_decimal_to_dictionary::<K, Decimal128Type, _>(
+            array,
+            dict_value_type,
+            p,
+            s,
+            cast_options,
+        ),
+        Decimal256(p, s) => pack_decimal_to_dictionary::<K, Decimal256Type, _>(
+            array,
+            dict_value_type,
+            p,
+            s,
+            cast_options,
+        ),
         Float16 => {
             pack_numeric_to_dictionary::<K, Float16Type>(array, dict_value_type, cast_options)
         }
@@ -359,6 +329,34 @@ where
     Ok(Arc::new(b.finish()))
 }
 
+pub(crate) fn pack_decimal_to_dictionary<K, D, M>(
+    array: &dyn Array,
+    dict_value_type: &DataType,
+    precision: u8,
+    scale: i8,
+    cast_options: &CastOptions,
+) -> Result<ArrayRef, ArrowError>
+where
+    K: ArrowDictionaryKeyType,
+    D: DecimalType + ArrowPrimitiveType<Native = M>,
+    M: ArrowNativeTypeOp + DecimalCast,
+{
+    let dict = pack_numeric_to_dictionary::<K, D>(array, dict_value_type, cast_options)?;
+    let dict = dict
+        .as_dictionary::<K>()
+        .downcast_dict::<PrimitiveArray<D>>()
+        .ok_or_else(|| {
+            ArrowError::ComputeError(format!("Internal Error: Cannot cast dict to {}", D::PREFIX))
+        })?;
+    let value = dict.values().clone();
+    // Set correct precision/scale
+    let value = value.with_precision_and_scale(precision, scale)?;
+    Ok(Arc::new(DictionaryArray::<K>::try_new(
+        dict.keys().clone(),
+        Arc::new(value),
+    )?))
+}
+
 pub(crate) fn string_view_to_dictionary<K, O: OffsetSizeTrait>(
     array: &dyn Array,
 ) -> Result<ArrayRef, ArrowError>