Jefffrey commented on code in PR #19628: URL: https://github.com/apache/datafusion/pull/19628#discussion_r2659502642
########## datafusion/spark/src/function/math/decimal_div.rs: ########## @@ -0,0 +1,434 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spark-compatible decimal division functions. +//! +//! This module implements Spark's decimal division semantics, which require +//! special handling for precision and scale that differs from standard SQL. +//! +//! # Scale Expansion +//! +//! For Decimal(p1, s1) / Decimal(p2, s2) = Decimal(p3, s3): +//! The dividend needs to be scaled to s2 + s3 + 1 to get correct precision. +//! This can exceed Decimal128's maximum scale (38), requiring BigInt fallback. + +use arrow::array::{Array, ArrayRef, AsArray, Decimal128Array}; +use arrow::datatypes::{DECIMAL128_MAX_PRECISION, DataType, Decimal128Type}; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + Volatility, +}; +use num::{BigInt, Signed, ToPrimitive}; +use std::any::Any; +use std::sync::Arc; + +/// Extract precision and scale from a Decimal128 DataType. +fn get_precision_scale(data_type: &DataType) -> (u8, i8) { + match data_type { + DataType::Decimal128(p, s) => (*p, *s), + _ => unreachable!("Expected Decimal128 type"), + } +} + +/// Internal implementation for both regular and integral decimal division. +/// +/// # Arguments +/// * `args` - Two ColumnarValue arguments (dividend and divisor) +/// * `result_precision` - The precision of the result type +/// * `result_scale` - The scale of the result type +/// * `is_integral_div` - If true, performs integer division (truncates result) +fn spark_decimal_div_internal( + args: &[ColumnarValue], + result_precision: u8, + result_scale: i8, + is_integral_div: bool, +) -> Result<ColumnarValue> { + assert_eq_or_internal_err!( + args.len(), + 2, + "decimal division expects exactly two arguments" + ); + + let left = &args[0]; + let right = &args[1]; + + let (left, right): (ArrayRef, ArrayRef) = match (left, right) { + (ColumnarValue::Array(l), ColumnarValue::Array(r)) => { + (Arc::clone(l), Arc::clone(r)) + } + (ColumnarValue::Scalar(l), ColumnarValue::Array(r)) => { + (l.to_array_of_size(r.len())?, Arc::clone(r)) + } + (ColumnarValue::Array(l), ColumnarValue::Scalar(r)) => { + (Arc::clone(l), r.to_array_of_size(l.len())?) + } + (ColumnarValue::Scalar(l), ColumnarValue::Scalar(r)) => { + (l.to_array()?, r.to_array()?) + } + }; + + let left = left.as_primitive::<Decimal128Type>(); + let right = right.as_primitive::<Decimal128Type>(); + let (p1, s1) = get_precision_scale(left.data_type()); + let (p2, s2) = get_precision_scale(right.data_type()); + + // Calculate the scale expansion needed + // To get Decimal(p3, s3) from p1/p2, we need to widen s1 to s2 + s3 + 1 + let l_exp = ((s2 + result_scale + 1) as u32).saturating_sub(s1 as u32); + let r_exp = (s1 as u32).saturating_sub((s2 + result_scale + 1) as u32); + + let result: Decimal128Array = if p1 as u32 + l_exp > DECIMAL128_MAX_PRECISION as u32 + || p2 as u32 + r_exp > DECIMAL128_MAX_PRECISION as u32 + { + // Use BigInt for high precision calculations that would overflow i128 + let ten = BigInt::from(10); + let l_mul = ten.pow(l_exp); + let r_mul = ten.pow(r_exp); + let five = BigInt::from(5); + let zero = BigInt::from(0); + + arrow::compute::kernels::arity::try_binary(left, right, |l, r| { + let l = BigInt::from(l) * &l_mul; + let r = BigInt::from(r) * &r_mul; + // Legacy mode: divide by zero returns 0 + let div = if r.eq(&zero) { zero.clone() } else { &l / &r }; + let res = if is_integral_div { + div + } else if div.is_negative() { + div - &five + } else { + div + &five + } / &ten; + Ok(res.to_i128().unwrap_or(i128::MAX)) + })? + } else { + // Standard i128 calculation when precision is within bounds + let l_mul = 10_i128.pow(l_exp); + let r_mul = 10_i128.pow(r_exp); + + arrow::compute::kernels::arity::try_binary(left, right, |l, r| { + let l = l * l_mul; + let r = r * r_mul; + // Legacy mode: divide by zero returns 0 + let div = if r == 0 { 0 } else { l / r }; + let res = if is_integral_div { + div + } else if div.is_negative() { + div - 5 + } else { + div + 5 + } / 10; + Ok(res) + })? + }; + + let result = + result.with_data_type(DataType::Decimal128(result_precision, result_scale)); + Ok(ColumnarValue::Array(Arc::new(result))) +} + +/// Spark-compatible decimal division function. +/// +/// Performs division with Spark's rounding behavior (round half away from zero). +pub fn spark_decimal_div( + args: &[ColumnarValue], + result_precision: u8, + result_scale: i8, +) -> Result<ColumnarValue> { + spark_decimal_div_internal(args, result_precision, result_scale, false) +} + +/// Spark-compatible integral decimal division function. +/// +/// Performs integer division (truncates toward zero). +pub fn spark_decimal_integral_div( + args: &[ColumnarValue], + result_precision: u8, + result_scale: i8, +) -> Result<ColumnarValue> { + spark_decimal_div_internal(args, result_precision, result_scale, true) +} + +/// SparkDecimalDiv implements the Spark-compatible decimal division function. +/// +/// This UDF takes the result precision and scale as part of its configuration, +/// since Spark determines these at query planning time. +#[derive(Debug)] +pub struct SparkDecimalDiv { + signature: Signature, + result_precision: u8, + result_scale: i8, +} + +impl PartialEq for SparkDecimalDiv { Review Comment: Could we remove these manual impls and use derive? ########## datafusion/spark/src/function/math/decimal_div.rs: ########## @@ -0,0 +1,434 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spark-compatible decimal division functions. +//! +//! This module implements Spark's decimal division semantics, which require +//! special handling for precision and scale that differs from standard SQL. +//! +//! # Scale Expansion +//! +//! For Decimal(p1, s1) / Decimal(p2, s2) = Decimal(p3, s3): +//! The dividend needs to be scaled to s2 + s3 + 1 to get correct precision. +//! This can exceed Decimal128's maximum scale (38), requiring BigInt fallback. Review Comment: Should we put a note here that these UDFs are only meant to be used at the physical level instead of the logical? Since it doesn't handle things like type coercion to supported types (decimal128). ########## datafusion/spark/src/function/math/decimal_div.rs: ########## @@ -0,0 +1,434 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spark-compatible decimal division functions. +//! +//! This module implements Spark's decimal division semantics, which require +//! special handling for precision and scale that differs from standard SQL. +//! +//! # Scale Expansion +//! +//! For Decimal(p1, s1) / Decimal(p2, s2) = Decimal(p3, s3): +//! The dividend needs to be scaled to s2 + s3 + 1 to get correct precision. +//! This can exceed Decimal128's maximum scale (38), requiring BigInt fallback. + +use arrow::array::{Array, ArrayRef, AsArray, Decimal128Array}; +use arrow::datatypes::{DECIMAL128_MAX_PRECISION, DataType, Decimal128Type}; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + Volatility, +}; +use num::{BigInt, Signed, ToPrimitive}; +use std::any::Any; +use std::sync::Arc; + +/// Extract precision and scale from a Decimal128 DataType. +fn get_precision_scale(data_type: &DataType) -> (u8, i8) { + match data_type { + DataType::Decimal128(p, s) => (*p, *s), + _ => unreachable!("Expected Decimal128 type"), + } +} + +/// Internal implementation for both regular and integral decimal division. +/// +/// # Arguments +/// * `args` - Two ColumnarValue arguments (dividend and divisor) +/// * `result_precision` - The precision of the result type +/// * `result_scale` - The scale of the result type +/// * `is_integral_div` - If true, performs integer division (truncates result) +fn spark_decimal_div_internal( + args: &[ColumnarValue], + result_precision: u8, + result_scale: i8, + is_integral_div: bool, +) -> Result<ColumnarValue> { Review Comment: ```suggestion fn spark_decimal_div_internal<const IS_INTEGRAL_DIV: bool>( left: ArrayRef, right: ArrayRef, result_precision: u8, result_scale: i8, ) -> Result<ColumnarValue> { ``` So callers can call like so: ```rust let arrays = ColumnarValue::values_to_arrays(args.args)?; let (left, right) = take_function_args("fn_name", arrays)?; spark_decimal_div_internal::<TRUE>(left, right, precision, scale) ``` So can remove code inside related to checking args length, conversion to arrays, etc. -- 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]
