Jefffrey commented on code in PR #18032:
URL: https://github.com/apache/datafusion/pull/18032#discussion_r2453694039
##########
datafusion/functions/src/math/power.rs:
##########
@@ -68,8 +77,44 @@ impl PowerFunc {
Self {
signature: Signature::one_of(
vec![
+ TypeSignature::Exact(vec![Int32, Int32]),
TypeSignature::Exact(vec![Int64, Int64]),
+ TypeSignature::Exact(vec![Float32, Float32]),
TypeSignature::Exact(vec![Float64, Float64]),
+ // Extra signatures for decimals to avoid casting them to
floats
+ TypeSignature::Exact(vec![
+ Decimal32(DECIMAL32_MAX_PRECISION, 0),
+ Int64,
+ ]),
+ TypeSignature::Exact(vec![
+ Decimal32(DECIMAL32_MAX_PRECISION, 0),
+ Float64,
+ ]),
+ TypeSignature::Exact(vec![
+ Decimal32(DECIMAL64_MAX_PRECISION, 0),
+ Int64,
+ ]),
+ TypeSignature::Exact(vec![
+ Decimal32(DECIMAL64_MAX_PRECISION, 0),
+ Float64,
+ ]),
+ TypeSignature::Exact(vec![
+ Decimal128(DECIMAL128_MAX_PRECISION, 0),
+ Int64,
+ ]),
+ TypeSignature::Exact(vec![
+ Decimal128(DECIMAL128_MAX_PRECISION, 0),
+ Float64,
+ ]),
+ TypeSignature::Exact(vec![
+ Decimal256(DECIMAL256_MAX_PRECISION, 0),
+ Int64,
+ ]),
+ TypeSignature::Exact(vec![
+ Decimal256(DECIMAL256_MAX_PRECISION, 0),
+ Float64,
+ ]),
+ Numeric(2), // Catch-all for all decimals
Review Comment:
Yeah this signature is a bit hard with decimals; it might be better off
doing this as user_defined type and use `coerce_types` until we can get better
support for arbitrary decimals in the type signature 🤔
- I recall another PR running into a similar problem:
https://github.com/apache/datafusion/pull/17113/files#r2282718716
Especially this last `Numeric(2)` which I think allows having a decimal type
as the exponent which I don't think the implementation supports?
##########
datafusion/functions/src/math/power.rs:
##########
@@ -92,8 +195,26 @@ impl ScalarUDFImpl for PowerFunc {
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
Review Comment:
Can probably simplify to:
```rust
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
Ok(arg_types[0].clone())
}
```
And assume the signature has done the checking for us
##########
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)?;
Review Comment:
We could probably refactor this like so:
```rust
let base = args.args[0].to_array(args.number_rows)?;
let exponent = &args.args[1];
```
So we don't needlessly convert `args.arg[1]` from ColumnarValue (which may
be Scalar) to Array back to ColumnarValue (now Array)
##########
datafusion/functions/src/math/power.rs:
##########
@@ -279,4 +522,364 @@ mod tests {
}
}
}
+
+ #[test]
Review Comment:
Would some of these tests be better off in as SLTs?
##########
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:
Could you help me understand why we set scale to 0? You mentioned this:
> Also, I discovered that a default-constructed decimal type always has
scale 10, which is not what we want for operations on
`ArrowPrimitiveType::Native`. Added descaling to zero, leaving underlying data
as-is.
Could you explain why we don't want that?
--
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]
