fix: Used NumCast rather than manual

Author: Yuvraj-cyborg

datafusion/functions/src/math/power.rs

@@ -24,7 +24,7 @@ use crate::utils::{calculate_binary_decimal_math, calculate_binary_math};
 use arrow::array::{Array, ArrayRef};
 use arrow::datatypes::{
     ArrowNativeTypeOp, DataType, Decimal32Type, Decimal64Type, Decimal128Type,
-    Decimal256Type, Float64Type, Int64Type,
+    Decimal256Type, Float64Type, Int64Type, i256,
 };
 use arrow::error::ArrowError;
 use datafusion_common::types::{NativeType, logical_float64, logical_int64};
@@ -115,7 +115,7 @@ impl PowerFunc {
 ///   Scale it back to 1: 390625 / 10^4 = 39
 fn pow_decimal_int<T>(base: T, scale: i8, exp: i64) -> Result<T, ArrowError>
 where
-    T: From<i32> + ArrowNativeTypeOp + ToPrimitive,
+    T: From<i32> + ArrowNativeTypeOp + ToPrimitive + num_traits::NumCast,
 {
     if exp < 0 {
         return pow_decimal_float(base, scale, exp as f64);
@@ -128,7 +128,7 @@ where
     })?;
     if exp == 0 {
         // Edge case to provide 1 as result (10^s with scale)
-        let result: T = T::from(10).pow_checked(scale).map_err(|_| {
+        let result: T = <T as From<i32>>::from(10).pow_checked(scale).map_err(|_| {
             ArrowError::ArithmeticOverflow(format!(
                 "Cannot make unscale factor for {scale} and {exp}"
             ))
@@ -141,18 +141,20 @@ where
     let powered: T = base.pow_checked(exp).map_err(|_| {
         ArrowError::ArithmeticOverflow(format!("Cannot raise base {base:?} to exp {exp}"))
     })?;
-    let unscale_factor: T = T::from(10).pow_checked(scale * (exp - 1)).map_err(|_| {
-        ArrowError::ArithmeticOverflow(format!(
-            "Cannot make unscale factor for {scale} and {exp}"
-        ))
-    })?;
+    let unscale_factor: T = <T as From<i32>>::from(10)
+        .pow_checked(scale * (exp - 1))
+        .map_err(|_| {
+            ArrowError::ArithmeticOverflow(format!(
+                "Cannot make unscale factor for {scale} and {exp}"
+            ))
+        })?;
 
     powered.div_checked(unscale_factor)
 }
 
 fn pow_decimal_float<T>(base: T, scale: i8, exp: f64) -> Result<T, ArrowError>
 where
-    T: From<i32> + ArrowNativeTypeOp + ToPrimitive,
+    T: From<i32> + ArrowNativeTypeOp + ToPrimitive + num_traits::NumCast,
 {
     if exp.is_finite() && exp.trunc() == exp && exp >= 0f64 && exp < u32::MAX as f64 {
         return pow_decimal_int(base, scale, exp as i64);
@@ -167,23 +169,65 @@ where
     pow_decimal_float_fallback(base, scale, exp)
 }
 
-/// Fallback implementation using f64 for negative or non-integer exponents.
-/// This handles cases that cannot be computed using integer arithmetic.
-fn pow_decimal_float_fallback<T>(base: T, scale: i8, exp: f64) -> Result<T, ArrowError>
-where
-    T: From<i32> + ArrowNativeTypeOp + ToPrimitive,
-{
-    if scale < 0 {
-        return Err(ArrowError::NotYetImplemented(format!(
-            "Negative scale is not yet supported: {scale}"
+/// Decimal256 specialized float exponent version.
+fn pow_decimal256_float(base: i256, scale: i8, exp: f64) -> Result<i256, ArrowError> {
+    if exp.is_finite() && exp.trunc() == exp && exp >= 0f64 && exp < u32::MAX as f64 {
+        return pow_decimal256_int(base, scale, exp as i64);
+    }
+
+    if !exp.is_finite() {
+        return Err(ArrowError::ComputeError(format!(
+            "Cannot use non-finite exp: {exp}"
         )));
     }
 
-    let scale_factor = 10f64.powi(scale as i32);
-    let base_f64 = base.to_f64().ok_or_else(|| {
-        ArrowError::ComputeError("Cannot convert base to f64".to_string())
-    })? / scale_factor;
+    pow_decimal256_float_fallback(base, scale, exp)
+}
+
+/// Decimal256 specialized integer exponent version.
+fn pow_decimal256_int(base: i256, scale: i8, exp: i64) -> Result<i256, ArrowError> {
+    if exp < 0 {
+        return pow_decimal256_float(base, scale, exp as f64);
+    }
 
+    let scale: u32 = scale.try_into().map_err(|_| {
+        ArrowError::NotYetImplemented(format!(
+            "Negative scale is not yet supported value: {scale}"
+        ))
+    })?;
+    if exp == 0 {
+        let result: i256 = i256::from_i128(10).pow_checked(scale).map_err(|_| {
+            ArrowError::ArithmeticOverflow(format!(
+                "Cannot make unscale factor for {scale} and {exp}"
+            ))
+        })?;
+        return Ok(result);
+    }
+    let exp: u32 = exp.try_into().map_err(|_| {
+        ArrowError::ArithmeticOverflow(format!("Unsupported exp value: {exp}"))
+    })?;
+    let powered: i256 = base.pow_checked(exp).map_err(|_| {
+        ArrowError::ArithmeticOverflow(format!("Cannot raise base {base:?} to exp {exp}"))
+    })?;
+    let unscale_factor: i256 = i256::from_i128(10)
+        .pow_checked(scale * (exp - 1))
+        .map_err(|_| {
+            ArrowError::ArithmeticOverflow(format!(
+                "Cannot make unscale factor for {scale} and {exp}"
+            ))
+        })?;
+
+    powered.div_checked(unscale_factor)
+}
+
+/// Compute the f64 power result and scale it back.
+/// Returns the rounded i128 result for conversion to target type.
+#[inline]
+fn compute_pow_f64_result(
+    base_f64: f64,
+    scale: i8,
+    exp: f64,
+) -> Result<i128, ArrowError> {
     let result_f64 = base_f64.powf(exp);
 
     if !result_f64.is_finite() {
@@ -192,6 +236,7 @@ where
         )));
     }
 
+    let scale_factor = 10f64.powi(scale as i32);
     let result_scaled = result_f64 * scale_factor;
     let result_rounded = result_scaled.round();
 
@@ -201,64 +246,56 @@ where
         )));
     }
 
-    decimal_from_i128::<T>(result_rounded as i128)
+    Ok(result_rounded as i128)
 }
 
-fn decimal_from_i128<T>(value: i128) -> Result<T, ArrowError>
+/// Fallback implementation using f64 for negative or non-integer exponents.
+/// This handles cases that cannot be computed using integer arithmetic.
+fn pow_decimal_float_fallback<T>(base: T, scale: i8, exp: f64) -> Result<T, ArrowError>
 where
-    T: From<i32> + ArrowNativeTypeOp,
+    T: ToPrimitive + num_traits::NumCast + Copy,
 {
-    if value == 0 {
-        return Ok(T::from(0));
-    }
-
-    if value >= i32::MIN as i128 && value <= i32::MAX as i128 {
-        return Ok(T::from(value as i32));
+    if scale < 0 {
+        return Err(ArrowError::NotYetImplemented(format!(
+            "Negative scale is not yet supported: {scale}"
+        )));
     }
 
-    let is_negative = value < 0;
-    let abs_value = value.unsigned_abs();
-
-    let billion = 1_000_000_000u128;
-    let mut result = T::from(0);
-    let mut multiplier = T::from(1);
-    let billion_t = T::from(1_000_000_000);
+    let scale_factor = 10f64.powi(scale as i32);
+    let base_f64 = base.to_f64().ok_or_else(|| {
+        ArrowError::ComputeError("Cannot convert base to f64".to_string())
+    })? / scale_factor;
 
-    let mut remaining = abs_value;
-    while remaining > 0 {
-        let chunk = (remaining % billion) as i32;
-        remaining /= billion;
+    let result_i128 = compute_pow_f64_result(base_f64, scale, exp)?;
 
-        let chunk_value = T::from(chunk).mul_checked(multiplier).map_err(|_| {
-            ArrowError::ArithmeticOverflow(format!(
-                "Overflow while converting {value} to decimal type"
-            ))
-        })?;
-
-        result = result.add_checked(chunk_value).map_err(|_| {
-            ArrowError::ArithmeticOverflow(format!(
-                "Overflow while converting {value} to decimal type"
-            ))
-        })?;
+    num_traits::NumCast::from(result_i128).ok_or_else(|| {
+        ArrowError::ArithmeticOverflow(format!(
+            "Value {result_i128} is too large for the target decimal type"
+        ))
+    })
+}
 
-        if remaining > 0 {
-            multiplier = multiplier.mul_checked(billion_t).map_err(|_| {
-                ArrowError::ArithmeticOverflow(format!(
-                    "Overflow while converting {value} to decimal type"
-                ))
-            })?;
-        }
+/// Fallback implementation for Decimal256.
+fn pow_decimal256_float_fallback(
+    base: i256,
+    scale: i8,
+    exp: f64,
+) -> Result<i256, ArrowError> {
+    if scale < 0 {
+        return Err(ArrowError::NotYetImplemented(format!(
+            "Negative scale is not yet supported: {scale}"
+        )));
     }
 
-    if is_negative {
-        result = T::from(0).sub_checked(result).map_err(|_| {
-            ArrowError::ArithmeticOverflow(format!(
-                "Overflow while negating {value} in decimal type"
-            ))
-        })?;
-    }
+    let scale_factor = 10f64.powi(scale as i32);
+    let base_f64 = base.to_f64().ok_or_else(|| {
+        ArrowError::ComputeError("Cannot convert base to f64".to_string())
+    })? / scale_factor;
+
+    let result_i128 = compute_pow_f64_result(base_f64, scale, exp)?;
 
-    Ok(result)
+    // i256 can be constructed from i128 directly
+    Ok(i256::from_i128(result_i128))
 }
 
 impl ScalarUDFImpl for PowerFunc {
@@ -381,7 +418,7 @@ impl ScalarUDFImpl for PowerFunc {
                 >(
                     &base,
                     exponent,
-                    |b, e| pow_decimal_int(b, *scale, e),
+                    |b, e| pow_decimal256_int(b, *scale, e),
                     *precision,
                     *scale,
                 )?
@@ -395,7 +432,7 @@ impl ScalarUDFImpl for PowerFunc {
                 >(
                     &base,
                     exponent,
-                    |b, e| pow_decimal_float(b, *scale, e),
+                    |b, e| pow_decimal256_float(b, *scale, e),
                     *precision,
                     *scale,
                 )?