theirix commented on code in PR #18032:
URL: https://github.com/apache/datafusion/pull/18032#discussion_r2487699585
##########
datafusion/functions/src/math/power.rs:
##########
@@ -104,45 +225,161 @@ impl ScalarUDFImpl for PowerFunc {
fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
let args = ColumnarValue::values_to_arrays(&args.args)?;
- let arr: ArrayRef = match args[0].data_type() {
+ let base = &args[0];
+ let exponent = ColumnarValue::Array(Arc::clone(&args[1]));
+
+ let arr: ArrayRef = match base.data_type() {
+ DataType::Float32 => {
+ calculate_binary_math::<Float32Type, Float32Type, Float32Type,
_>(
+ base,
+ &exponent,
+ |b, e| Ok(f32::powf(b, e)),
+ )?
+ }
DataType::Float64 => {
- let bases = args[0].as_primitive::<Float64Type>();
- let exponents = args[1].as_primitive::<Float64Type>();
- let result = arrow::compute::binary::<_, _, _, Float64Type>(
- bases,
- exponents,
- f64::powf,
- )?;
- Arc::new(result) as _
+ calculate_binary_math::<Float64Type, Float64Type, Float64Type,
_>(
+ &base,
+ &exponent,
+ |b, e| Ok(f64::powf(b, e)),
+ )?
+ }
+ DataType::Int32 => {
+ calculate_binary_math::<Int32Type, Int32Type, Int32Type, _>(
+ &base,
+ &exponent,
+ |b, e| match e.try_into() {
+ Ok(exp_u32) => b.pow_checked(exp_u32),
+ Err(_) => Err(ArrowError::ArithmeticOverflow(format!(
+ "Exponent {e} in integer computation is out of
bounds."
+ ))),
+ },
+ )?
}
DataType::Int64 => {
- let bases = downcast_named_arg!(&args[0], "base", Int64Array);
- let exponents = downcast_named_arg!(&args[1], "exponent",
Int64Array);
- bases
- .iter()
- .zip(exponents.iter())
- .map(|(base, exp)| match (base, exp) {
- (Some(base), Some(exp)) => Ok(Some(base.pow_checked(
- exp.try_into().map_err(|_| {
- exec_datafusion_err!(
- "Can't use negative exponents: {exp} in
integer computation, please use Float."
- )
- })?,
- ).map_err(|e| arrow_datafusion_err!(e))?)),
- _ => Ok(None),
- })
- .collect::<Result<Int64Array>>()
- .map(Arc::new)? as _
+ calculate_binary_math::<Int64Type, Int64Type, Int64Type, _>(
+ &base,
+ &exponent,
+ |b, e| match e.try_into() {
+ Ok(exp_u32) => b.pow_checked(exp_u32),
+ Err(_) => Err(ArrowError::ArithmeticOverflow(format!(
+ "Exponent {e} in integer computation is out of
bounds."
+ ))),
+ },
+ )?
+ }
+ DataType::Decimal32(_precision, scale) => {
+ let array = match exponent.data_type() {
+ DataType::Int64 => {
+ calculate_binary_math::<Decimal32Type, Int64Type,
Decimal32Type, _>(
+ &base,
+ &exponent,
+ |b, e| pow_decimal32_int(b, *scale, e),
+ )?
+ }
+ DataType::Float64 => {
+ calculate_binary_math::<
+ Decimal32Type,
+ Float64Type,
+ Decimal32Type,
+ _,
+ >(&base, &exponent, |b, e| {
+ pow_decimal32_float(b, *scale, e)
+ })?
+ }
+ other => {
+ return exec_err!("Unsupported data type {other:?} for
exponent")
+ }
+ };
+ Arc::new(reset_decimal_scale(array.as_ref())?)
+ }
+ DataType::Decimal64(_precision, scale) => {
+ let array = match exponent.data_type() {
+ DataType::Int64 => {
+ calculate_binary_math::<Decimal64Type, Int64Type,
Decimal64Type, _>(
+ &base,
+ &exponent,
+ |b, e| pow_decimal64_int(b, *scale, e),
+ )?
+ }
+ DataType::Float64 => {
+ calculate_binary_math::<
+ Decimal64Type,
+ Float64Type,
+ Decimal64Type,
+ _,
+ >(&base, &exponent, |b, e| {
+ pow_decimal64_float(b, *scale, e)
+ })?
+ }
+ other => {
+ return exec_err!("Unsupported data type {other:?} for
exponent")
+ }
+ };
+ Arc::new(reset_decimal_scale(array.as_ref())?)
+ }
+ DataType::Decimal128(_precision, scale) => {
+ let array = match exponent.data_type() {
+ DataType::Int64 => {
+ calculate_binary_math::<
+ Decimal128Type,
+ Int64Type,
+ Decimal128Type,
+ _,
+ >(&base, &exponent, |b, e| {
+ pow_decimal128_int(b, *scale, e)
+ })?
+ }
+ DataType::Float64 => {
+ calculate_binary_math::<
+ Decimal128Type,
+ Float64Type,
+ Decimal128Type,
+ _,
+ >(&base, &exponent, |b, e| {
+ pow_decimal128_float(b, *scale, e)
+ })?
+ }
+ other => {
+ return exec_err!("Unsupported data type {other:?} for
exponent")
+ }
+ };
+ Arc::new(reset_decimal_scale(array.as_ref())?)
+ }
+ DataType::Decimal256(_precision, scale) => {
+ let array = match exponent.data_type() {
+ DataType::Int64 => {
+ calculate_binary_math::<
+ Decimal256Type,
+ Int64Type,
+ Decimal256Type,
+ _,
+ >(&base, &exponent, |b, e| {
+ pow_decimal256_int(b, *scale, e)
+ })?
+ }
+ DataType::Float64 => {
+ calculate_binary_math::<
+ Decimal256Type,
+ Float64Type,
+ Decimal256Type,
+ _,
+ >(&base, &exponent, |b, e| {
+ pow_decimal256_float(b, *scale, e)
+ })?
+ }
+ other => {
+ return exec_err!("Unsupported data type {other:?} for
exponent")
+ }
+ };
+ Arc::new(reset_decimal_scale(array.as_ref())?)
Review Comment:
You raised a good point. Initially, my PR addressed only zero-scale numbers,
as I was unable to specify a class of decimals for a type. In recent commits,
I've implemented proper support for scaled integers.
Functions like `abs`, `round` could work flawlessly with decimals because
they don't change type (precision and scale) and value size. However, when we
start performing operations on decimals using arrow-compute with
`arrow::compute::binary`, we operate on the underlying type (f64, i128, i256)
without knowing the precision or scale. The function `calculate_binary_math` is
a generic wrapper around arrow-compute which is not aware of a scale concept.
So we implement handling scaled numbers in a kernel - `pow_decimal128_int`.
Also, the raw i128 result should be cast to the expected precision and scale
using a simple helper around `with_precision_and_scale`.
A function like `power` is quite different, and it requires handling a scale
to raise it to a given power. I've added documentation to rust-docs,
unit/integration tests for the new behaviour. Please tell if I need to address
more in the code.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
