martin-g commented on code in PR #19628: URL: https://github.com/apache/datafusion/pull/19628#discussion_r2668571415
########## 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)) Review Comment: If `res` is a big negative number that is smaller than i128::MIN then unwrap_or() should return i128::MIN, not MAX -- 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]
