// SPDX-License-Identifier: Apache-2.0

// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership

// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements

use crate::utils::{calculate_binary_decimal_math_cast, calculate_binary_math};

use arrow::array::{Array, ArrayRef, AsArray};

use arrow::datatypes::DataType::{

Decimal32, Decimal64, Decimal128, Decimal256, Float32, Float64, Int8, Int16, Int32,

Int64, UInt8, UInt16, UInt32, UInt64,

};

use arrow::datatypes::{

ArrowNativeTypeOp, ArrowPrimitiveType, DataType, Decimal32Type, Decimal64Type,

Decimal128Type, Decimal256Type, DecimalType, Float32Type, Float64Type, Int8Type,

Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type,

};

use arrow::datatypes::{Field, FieldRef};

use arrow::error::ArrowError;

use datafusion_common::types::{

NativeType, logical_float32, logical_float64, logical_int32,

};

use datafusion_common::{Result, ScalarValue, exec_err, internal_err};

use datafusion_expr::sort_properties::{ExprProperties, SortProperties};

use datafusion_expr::{

Coercion, ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs,

ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility,

};

use datafusion_macros::user_doc;

use num_traits::{PrimInt, Signed, cast, checked_pow};

use std::sync::Arc;

fn output_scale_for_decimal(precision: u8, input_scale: i8, decimal_places: i32) -> i8 {

// `decimal_places` controls the maximum output scale, but scale cannot exceed the input scale.

//

// For negative-scale decimals, allow further scale reduction to match negative `decimal_places`

// (e.g. scale -2 rounded to -3 becomes scale -3). This preserves fixed precision by

// representing the rounded result at a coarser scale.

if input_scale < 0 {

// Decimal scales must be within [-precision, precision] and fit in i8. For negative-scale

// decimals, allow rounding to move the output scale further negative, but cap it at

// `-precision` (beyond that, the rounded result is always 0).

let min_scale = -i32::from(precision);

let new_scale = i32::from(input_scale).min(decimal_places).max(min_scale);

return new_scale as i8;

}

// The `min` ensures the result is always within i8 range because `input_scale` is i8.

let decimal_places = decimal_places.max(0);

i32::from(input_scale).min(decimal_places) as i8

}

fn normalize_decimal_places_for_decimal(

decimal_places: i32,

precision: u8,

scale: i8,

) -> Option<i32> {

if decimal_places >= 0 {

return Some(decimal_places);

}

// For fixed precision decimals, the absolute value is strictly less than 10^(precision - scale).

// If the rounding position is beyond that (abs(decimal_places) > precision - scale), the

// rounded result is always 0, and we can avoid overflow in intermediate 10^n computations.

let max_rounding_pow10 = i64::from(precision) - i64::from(scale);

if max_rounding_pow10 <= 0 {

return None;

}

let abs_decimal_places = i64::from(decimal_places.unsigned_abs());

(abs_decimal_places <= max_rounding_pow10).then_some(decimal_places)

}

fn validate_decimal_precision<T: DecimalType>(

value: T::Native,

precision: u8,

scale: i8,

) -> Result<T::Native, ArrowError> {

T::validate_decimal_precision(value, precision, scale).map_err(|e| {

ArrowError::ComputeError(format!(

"Decimal overflow: rounded value exceeds precision {precision}: {e}"

))

})?;

Ok(value)

}

fn calculate_new_precision_scale<T: DecimalType>(

precision: u8,

scale: i8,

decimal_places: Option<i32>,

) -> Result<DataType> {

if let Some(decimal_places) = decimal_places {

let new_scale = output_scale_for_decimal(precision, scale, decimal_places);

// When rounding an integer decimal (scale == 0) to a negative `decimal_places`, a carry can

// add an extra digit to the integer part (e.g. 99 -> 100 when rounding to -1). This can

// only happen when the rounding position is within the existing precision.

let abs_decimal_places = decimal_places.unsigned_abs();

let new_precision = if scale == 0

&& decimal_places < 0

&& abs_decimal_places <= u32::from(precision)

{

precision.saturating_add(1).min(T::MAX_PRECISION)

} else {

precision

};

Ok(T::TYPE_CONSTRUCTOR(new_precision, new_scale))

} else {

let new_precision = precision.saturating_add(1).min(T::MAX_PRECISION);

Ok(T::TYPE_CONSTRUCTOR(new_precision, scale))

}

}

fn decimal_places_from_scalar(scalar: &ScalarValue) -> Result<i32> {

let out_of_range = |value: String| {

datafusion_common::DataFusionError::Execution(format!(

"round decimal_places {value} is out of supported i32 range"

))

};

match scalar {

ScalarValue::Int8(Some(v)) => Ok(i32::from(*v)),

ScalarValue::Int16(Some(v)) => Ok(i32::from(*v)),

ScalarValue::Int32(Some(v)) => Ok(*v),

ScalarValue::Int64(Some(v)) => {

i32::try_from(*v).map_err(|_| out_of_range(v.to_string()))

}

ScalarValue::UInt8(Some(v)) => Ok(i32::from(*v)),

ScalarValue::UInt16(Some(v)) => Ok(i32::from(*v)),

ScalarValue::UInt32(Some(v)) => {

i32::try_from(*v).map_err(|_| out_of_range(v.to_string()))

}

ScalarValue::UInt64(Some(v)) => {

i32::try_from(*v).map_err(|_| out_of_range(v.to_string()))

}

other => exec_err!(

"Unexpected datatype for decimal_places: {}",

other.data_type()

),

}

}

#[user_doc(

doc_section(label = "Math Functions"),

description = "Rounds a number to the nearest integer.",

syntax_example = "round(numeric_expression[, decimal_places])",

standard_argument(name = "numeric_expression", prefix = "Numeric"),

argument(

name = "decimal_places",

description = "Optional. The number of decimal places to round to. Defaults to 0."

),

sql_example = r#"```sql

> SELECT round(3.14159);

+----------------+

| round(3.14159) |

+----------------+

| 3.0 |

+----------------+

```"#

)]

#[derive(Debug, PartialEq, Eq, Hash)]

pub struct RoundFunc {

signature: Signature,

}

impl Default for RoundFunc {

fn default() -> Self {

RoundFunc::new()

}

}

impl RoundFunc {

pub fn new() -> Self {

let decimal = Coercion::new_exact(TypeSignatureClass::Decimal);

let decimal_places = Coercion::new_implicit(

TypeSignatureClass::Native(logical_int32()),

vec![TypeSignatureClass::Integer],

NativeType::Int32,

);

let integer = Coercion::new_exact(TypeSignatureClass::Integer);

let float32 = Coercion::new_exact(TypeSignatureClass::Native(logical_float32()));

let float64 = Coercion::new_implicit(

TypeSignatureClass::Native(logical_float64()),

vec![TypeSignatureClass::Numeric],

NativeType::Float64,

);

Self {

signature: Signature::one_of(

vec![

TypeSignature::Coercible(vec![

decimal.clone(),

decimal_places.clone(),

]),

TypeSignature::Coercible(vec![decimal]),

TypeSignature::Coercible(vec![

integer.clone(),

decimal_places.clone(),

]),

TypeSignature::Coercible(vec![integer]),

TypeSignature::Coercible(vec![

float32.clone(),

decimal_places.clone(),

]),

TypeSignature::Coercible(vec![float32]),

TypeSignature::Coercible(vec![float64.clone(), decimal_places]),

TypeSignature::Coercible(vec![float64]),

],

Volatility::Immutable,

),

}

}

}

impl ScalarUDFImpl for RoundFunc {

fn name(&self) -> &str {

"round"

}

fn is_strict(&self) -> bool {

true

}

fn signature(&self) -> &Signature {

&self.signature

}

fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {

let input_field = &args.arg_fields[0];

let input_type = input_field.data_type();

// If decimal_places is a scalar literal, we can incorporate it into the output type

// (scale reduction). Otherwise, keep the input scale as we can't pick a per-row scale.

let decimal_places: Option<i32> = match args.scalar_arguments.get(1) {

None => Some(0), // No dp argument means default to 0

Some(None) => None, // dp is not a literal (e.g. column)

Some(Some(scalar)) if scalar.is_null() => Some(0), // null dp => default to 0

Some(Some(scalar)) => Some(decimal_places_from_scalar(scalar)?),

};

// Calculate return type based on input type

// For decimals: reduce scale to decimal_places (reclaims precision for integer part)

// This matches Spark/DuckDB behavior where ROUND adjusts the scale

// BUT only if dp is a scalar literal - otherwise keep original scale and add

// extra precision to accommodate potential carry-over.

let return_type =

match input_type {

input_type if input_type.is_integer() => input_type.clone(),

Float32 => Float32,

Decimal32(precision, scale) => calculate_new_precision_scale::<

Decimal32Type,

>(

*precision, *scale, decimal_places

)?,

Decimal64(precision, scale) => calculate_new_precision_scale::<

Decimal64Type,

>(

*precision, *scale, decimal_places

)?,

Decimal128(precision, scale) => calculate_new_precision_scale::<

Decimal128Type,

>(

*precision, *scale, decimal_places

)?,

Decimal256(precision, scale) => calculate_new_precision_scale::<

Decimal256Type,

>(

*precision, *scale, decimal_places

)?,

_ => Float64,

};

let nullable = args.arg_fields.iter().any(|f| f.is_nullable());

Ok(Arc::new(Field::new(self.name(), return_type, nullable)))

}

fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {

internal_err!("use return_field_from_args instead")

}

fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {

if args.arg_fields.iter().any(|a| a.data_type().is_null()) {

return ColumnarValue::Scalar(ScalarValue::Null)

.cast_to(args.return_type(), None);

}

let default_decimal_places = ColumnarValue::Scalar(ScalarValue::Int32(Some(0)));

let decimal_places = if args.args.len() == 2 {

&args.args[1]

} else {

&default_decimal_places

};

if let (ColumnarValue::Scalar(value_scalar), ColumnarValue::Scalar(dp_scalar)) =

(&args.args[0], decimal_places)

{

if value_scalar.is_null() || dp_scalar.is_null() {

return ColumnarValue::Scalar(ScalarValue::Null)

.cast_to(args.return_type(), None);

}

let dp = if let ScalarValue::Int32(Some(dp)) = dp_scalar {

*dp

} else {

return internal_err!(

"Unexpected datatype for decimal_places: {}",

dp_scalar.data_type()

);

};

match (value_scalar, args.return_type()) {

(value_scalar, return_type) if return_type.is_integer() => {

round_integer_scalar(value_scalar, return_type, dp)

}

(ScalarValue::Float32(Some(v)), _) => {

let rounded = round_float(*v, dp)?;

Ok(ColumnarValue::Scalar(ScalarValue::from(rounded)))

}

(ScalarValue::Float64(Some(v)), _) => {

let rounded = round_float(*v, dp)?;

Ok(ColumnarValue::Scalar(ScalarValue::from(rounded)))

}

(

ScalarValue::Decimal32(Some(v), in_precision, scale),

Decimal32(out_precision, out_scale),

) => {

let rounded =

round_decimal_or_zero(*v, *in_precision, *scale, *out_scale, dp)?;

let rounded = if *out_precision == Decimal32Type::MAX_PRECISION

&& *scale == 0

&& dp < 0

{

// With scale == 0 and negative dp, rounding can carry into an additional

// digit (e.g. 99 -> 100). If we're already at max precision we can't widen

// the type, so validate and error rather than producing an invalid decimal.

validate_decimal_precision::<Decimal32Type>(

rounded,

*out_precision,

*out_scale,

)

} else {

Ok(rounded)

}?;

let scalar =

ScalarValue::Decimal32(Some(rounded), *out_precision, *out_scale);

Ok(ColumnarValue::Scalar(scalar))

}

(

ScalarValue::Decimal64(Some(v), in_precision, scale),

Decimal64(out_precision, out_scale),

) => {

let rounded =

round_decimal_or_zero(*v, *in_precision, *scale, *out_scale, dp)?;

let rounded = if *out_precision == Decimal64Type::MAX_PRECISION

&& *scale == 0

&& dp < 0

{

// See Decimal32 branch for details.

validate_decimal_precision::<Decimal64Type>(

rounded,

*out_precision,

*out_scale,

)

} else {

Ok(rounded)

}?;

let scalar =

ScalarValue::Decimal64(Some(rounded), *out_precision, *out_scale);

Ok(ColumnarValue::Scalar(scalar))

}

(

ScalarValue::Decimal128(Some(v), in_precision, scale),

Decimal128(out_precision, out_scale),

) => {

let rounded =

round_decimal_or_zero(*v, *in_precision, *scale, *out_scale, dp)?;

let rounded = if *out_precision == Decimal128Type::MAX_PRECISION

&& *scale == 0

&& dp < 0

{

// See Decimal32 branch for details.

validate_decimal_precision::<Decimal128Type>(

rounded,

*out_precision,

*out_scale,

)

} else {

Ok(rounded)

}?;

let scalar = ScalarValue::Decimal128(

Some(rounded),

*out_precision,

*out_scale,

);

Ok(ColumnarValue::Scalar(scalar))

}

(

ScalarValue::Decimal256(Some(v), in_precision, scale),

Decimal256(out_precision, out_scale),

) => {

let rounded =

round_decimal_or_zero(*v, *in_precision, *scale, *out_scale, dp)?;

let rounded = if *out_precision == Decimal256Type::MAX_PRECISION

&& *scale == 0

&& dp < 0

{

// See Decimal32 branch for details.

validate_decimal_precision::<Decimal256Type>(

rounded,

*out_precision,

*out_scale,

)

} else {

Ok(rounded)

}?;

let scalar = ScalarValue::Decimal256(

Some(rounded),

*out_precision,

*out_scale,

);

Ok(ColumnarValue::Scalar(scalar))

}

(ScalarValue::Null, _) => ColumnarValue::Scalar(ScalarValue::Null)

.cast_to(args.return_type(), None),

(value_scalar, return_type) => {

internal_err!(

"Unexpected datatype for round(value, decimal_places): value {}, return type {}",

value_scalar.data_type(),

return_type

)

}

}

} else {

round_columnar(

&args.args[0],

decimal_places,

args.number_rows,

args.return_type(),

)

}

}

fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {

// round preserves the order of the first argument

let value = &input[0];

let precision = input.get(1);

if precision

.map(|r| r.sort_properties.eq(&SortProperties::Singleton))

.unwrap_or(true)

{

Ok(value.sort_properties)

} else {

Ok(SortProperties::Unordered)

}

}

fn documentation(&self) -> Option<&Documentation> {

self.doc()

}

}

fn round_columnar(

value: &ColumnarValue,

decimal_places: &ColumnarValue,

number_rows: usize,

return_type: &DataType,

) -> Result<ColumnarValue> {

let value_array = value.to_array(number_rows)?;

let both_scalars = matches!(value, ColumnarValue::Scalar(_))

&& matches!(decimal_places, ColumnarValue::Scalar(_));

let decimal_places_is_array = matches!(decimal_places, ColumnarValue::Array(_));

let arr: ArrayRef = match (value_array.data_type(), return_type) {

(input_type, return_type)

if input_type == return_type && return_type.is_integer() =>

{

match decimal_places {

ColumnarValue::Scalar(ScalarValue::Int32(Some(dp))) if *dp >= 0 => {

value_array

}

_ => round_integer_array(

value_array.as_ref(),

decimal_places,

return_type,

)?,

}

}

(Float64, _) => round_float_column::<Float64Type>(&value_array, decimal_places)?,

(Float32, _) => round_float_column::<Float32Type>(&value_array, decimal_places)?,

(Decimal32(input_precision, scale), Decimal32(precision, new_scale)) => {

// reduce scale to reclaim integer precision

let result = calculate_binary_decimal_math_cast::<

Decimal32Type,

Int32Type,

Decimal32Type,

_,

>(

value_array.as_ref(),

decimal_places,

|v, dp| {

let rounded = round_decimal_or_zero(

v,

*input_precision,

*scale,

*new_scale,

dp,

)?;

if *precision == Decimal32Type::MAX_PRECISION

&& (decimal_places_is_array || (*scale == 0 && dp < 0))

{

// If we're already at max precision, we can't widen the result type. For

// dp arrays, or for scale == 0 with negative dp, rounding can overflow the

// fixed-precision type. Validate per-row and return an error instead of

// producing an invalid decimal that Arrow may display incorrectly.

validate_decimal_precision::<Decimal32Type>(

rounded, *precision, *new_scale,

)

} else {

Ok(rounded)

}

},

*precision,

*new_scale,

&Int32,

)?;

result as _

}

(Decimal64(input_precision, scale), Decimal64(precision, new_scale)) => {

let result = calculate_binary_decimal_math_cast::<

Decimal64Type,

Int32Type,

Decimal64Type,

_,

>(

value_array.as_ref(),

decimal_places,

|v, dp| {

let rounded = round_decimal_or_zero(

v,

*input_precision,

*scale,

*new_scale,

dp,

)?;

if *precision == Decimal64Type::MAX_PRECISION

&& (decimal_places_is_array || (*scale == 0 && dp < 0))

{

// See Decimal32 branch for details.

validate_decimal_precision::<Decimal64Type>(

rounded, *precision, *new_scale,

)

} else {

Ok(rounded)

}

},

*precision,

*new_scale,

&Int32,

)?;

result as _

}

(Decimal128(input_precision, scale), Decimal128(precision, new_scale)) => {

let result = calculate_binary_decimal_math_cast::<

Decimal128Type,

Int32Type,

Decimal128Type,

_,

>(

value_array.as_ref(),

decimal_places,

|v, dp| {

let rounded = round_decimal_or_zero(

v,

*input_precision,

*scale,

*new_scale,

dp,

)?;

if *precision == Decimal128Type::MAX_PRECISION

&& (decimal_places_is_array || (*scale == 0 && dp < 0))

{

// See Decimal32 branch for details.

validate_decimal_precision::<Decimal128Type>(

rounded, *precision, *new_scale,

)

} else {

Ok(rounded)

}

},

*precision,

*new_scale,

&Int32,

)?;

result as _

}

(Decimal256(input_precision, scale), Decimal256(precision, new_scale)) => {

let result = calculate_binary_decimal_math_cast::<

Decimal256Type,

Int32Type,

Decimal256Type,

_,

>(

value_array.as_ref(),

decimal_places,

|v, dp| {

let rounded = round_decimal_or_zero(

v,

*input_precision,

*scale,

*new_scale,

dp,

)?;

if *precision == Decimal256Type::MAX_PRECISION

&& (decimal_places_is_array || (*scale == 0 && dp < 0))

{

// See Decimal32 branch for details.

validate_decimal_precision::<Decimal256Type>(

rounded, *precision, *new_scale,

)

} else {

Ok(rounded)

}

},

*precision,

*new_scale,

&Int32,

)?;

result as _

}

(other, _) => exec_err!("Unsupported data type {other:?} for function round")?,

};

if both_scalars {

ScalarValue::try_from_array(&arr, 0).map(ColumnarValue::Scalar)

} else {

Ok(ColumnarValue::Array(arr))

}

}

fn round_signed_integer<T>(

value: T,

decimal_places: i32,

type_name: &str,

) -> Result<T, ArrowError>

where

T: PrimInt + Signed,

{

if decimal_places >= 0 || value == T::zero() {

return Ok(value);

}

let ten = cast::<_, T>(10).expect("10 fits in all integer types");

let Some(factor) = checked_pow(ten, decimal_places.unsigned_abs() as usize) else {

return Ok(T::zero());

};

let two = cast::<_, T>(2).expect("2 fits in all integer types");

let one = T::one();

let threshold = factor / two;

let mut quotient = value / factor;

let remainder = value % factor;

if remainder >= threshold {

quotient = quotient.checked_add(&one).ok_or_else(|| {

ArrowError::ComputeError(format!("Overflow while rounding {type_name}"))

})?;

} else if remainder <= -threshold {

quotient = quotient.checked_sub(&one).ok_or_else(|| {

ArrowError::ComputeError(format!("Overflow while rounding {type_name}"))

})?;

}

quotient.checked_mul(&factor).ok_or_else(|| {

ArrowError::ComputeError(format!("Overflow while rounding {type_name}"))

})

}

fn round_unsigned_integer<T>(

value: T,

decimal_places: i32,

type_name: &str,

) -> Result<T, ArrowError>

where

T: PrimInt,

{

if decimal_places >= 0 || value == T::zero() {

return Ok(value);

}

let ten = cast::<_, T>(10).expect("10 fits in all integer types");

let Some(factor) = checked_pow(ten, decimal_places.unsigned_abs() as usize) else {

return Ok(T::zero());

};

let two = cast::<_, T>(2).expect("2 fits in all integer types");

let one = T::one();

let threshold = factor / two;

let mut quotient = value / factor;

let remainder = value % factor;

if remainder >= threshold {

quotient = quotient.checked_add(&one).ok_or_else(|| {

ArrowError::ComputeError(format!("Overflow while rounding {type_name}"))

})?;

}

quotient.checked_mul(&factor).ok_or_else(|| {

ArrowError::ComputeError(format!("Overflow while rounding {type_name}"))

})

}

fn round_integer_scalar(

value: &ScalarValue,

return_type: &DataType,

decimal_places: i32,

) -> Result<ColumnarValue> {

match (value, return_type) {

(ScalarValue::Int8(Some(v)), Int8) => Ok(ColumnarValue::Scalar(

ScalarValue::Int8(Some(round_signed_integer(*v, decimal_places, "Int8")?)),

)),

(ScalarValue::Int16(Some(v)), Int16) => Ok(ColumnarValue::Scalar(

ScalarValue::Int16(Some(round_signed_integer(*v, decimal_places, "Int16")?)),

)),

(ScalarValue::Int32(Some(v)), Int32) => Ok(ColumnarValue::Scalar(

ScalarValue::Int32(Some(round_signed_integer(*v, decimal_places, "Int32")?)),

)),

(ScalarValue::Int64(Some(v)), Int64) => Ok(ColumnarValue::Scalar(

ScalarValue::Int64(Some(round_signed_integer(*v, decimal_places, "Int64")?)),

)),

(ScalarValue::UInt8(Some(v)), UInt8) => {

Ok(ColumnarValue::Scalar(ScalarValue::UInt8(Some(

round_unsigned_integer(*v, decimal_places, "UInt8")?,

))))

}

(ScalarValue::UInt16(Some(v)), UInt16) => {

Ok(ColumnarValue::Scalar(ScalarValue::UInt16(Some(

round_unsigned_integer(*v, decimal_places, "UInt16")?,

))))

}

(ScalarValue::UInt32(Some(v)), UInt32) => {

Ok(ColumnarValue::Scalar(ScalarValue::UInt32(Some(

round_unsigned_integer(*v, decimal_places, "UInt32")?,

))))

}

(ScalarValue::UInt64(Some(v)), UInt64) => {

Ok(ColumnarValue::Scalar(ScalarValue::UInt64(Some(

round_unsigned_integer(*v, decimal_places, "UInt64")?,

))))

}

_ => internal_err!(

"Unexpected integer round input/output types: {} -> {}",

value.data_type(),

return_type

),

}

}

macro_rules! round_integer_array {

($ARRAY:expr, $DP:expr, $ARRAY_TYPE:ty, $ROUND_FN:ident, $TYPE_NAME:expr) => {{

let array = $ARRAY.as_primitive::<$ARRAY_TYPE>();

let result = calculate_binary_math::<$ARRAY_TYPE, Int32Type, $ARRAY_TYPE, _>(

array,

$DP,

|v, dp| $ROUND_FN(v, dp, $TYPE_NAME),

)?;

Ok(result as ArrayRef)

}};

}

fn round_integer_array(

value_array: &dyn Array,

decimal_places: &ColumnarValue,

return_type: &DataType,

) -> Result<ArrayRef> {

match return_type {

Int8 => round_integer_array!(

value_array,

decimal_places,

Int8Type,

round_signed_integer,

"Int8"

),

Int16 => round_integer_array!(

value_array,

decimal_places,

Int16Type,

round_signed_integer,

"Int16"

),

Int32 => round_integer_array!(

value_array,

decimal_places,

Int32Type,

round_signed_integer,

"Int32"

),

Int64 => round_integer_array!(

value_array,

decimal_places,

Int64Type,

round_signed_integer,

"Int64"

),

UInt8 => round_integer_array!(

value_array,

decimal_places,

UInt8Type,

round_unsigned_integer,

"UInt8"

),

UInt16 => round_integer_array!(

value_array,

decimal_places,

UInt16Type,

round_unsigned_integer,

"UInt16"

),

UInt32 => round_integer_array!(

value_array,

decimal_places,

UInt32Type,

round_unsigned_integer,

"UInt32"

),

UInt64 => round_integer_array!(

value_array,

decimal_places,

UInt64Type,

round_unsigned_integer,

"UInt64"

),

_ => internal_err!("Unexpected return type for integer round: {return_type}"),

}

}

/// Rounds a float array to `decimal_places`.

///

/// The shared `calculate_binary_math` kernel routes through `try_unary` and

/// re-evaluates `round_float` (including `10f64.powi(decimal_places)` and a

/// `Result` check) for every element. When `decimal_places` is a non-null

/// scalar, the scaling factor can instead be hoisted out of the loop and the

/// infallible `unary` kernel used, which the compiler can autovectorize.

/// `unary` also computes over null slots, but it carries the input null buffer

/// through to the output, so those values stay masked.

fn round_float_column<PT>(

value_array: &ArrayRef,

decimal_places: &ColumnarValue,

) -> Result<ArrayRef>

where

PT: ArrowPrimitiveType,

PT::Native: num_traits::Float,

{

// Bring `Float` into scope so `.round()` resolves on the `PT::Native`

// projection below.

use num_traits::Float;

if let ColumnarValue::Scalar(ScalarValue::Int32(Some(decimal_places))) =

decimal_places

{

let factor = round_factor::<PT::Native>(*decimal_places)?;

let result = value_array

.as_primitive::<PT>()

.unary::<_, PT>(|value| (value * factor).round() / factor);

return Ok(Arc::new(result) as ArrayRef);

}

let result = calculate_binary_math::<PT, Int32Type, PT, _>(

value_array.as_ref(),

decimal_places,

round_float::<PT::Native>,

)?;

Ok(result as _)

}

/// Computes the power-of-ten scaling factor used to round to `decimal_places`.

fn round_factor<T: num_traits::Float>(decimal_places: i32) -> Result<T, ArrowError> {

T::from(10_f64.powi(decimal_places)).ok_or_else(|| {

ArrowError::ComputeError(format!(

"Invalid value for decimal places: {decimal_places}"

))

})

}

fn round_float<T>(value: T, decimal_places: i32) -> Result<T, ArrowError>

where

T: num_traits::Float,

{

let factor = round_factor::<T>(decimal_places)?;

Ok((value * factor).round() / factor)

}

fn round_decimal<V: ArrowNativeTypeOp>(

value: V,

input_scale: i8,

output_scale: i8,

decimal_places: i32,

) -> Result<V, ArrowError> {

let diff = i64::from(input_scale) - i64::from(decimal_places);

if diff <= 0 {

return Ok(value);

}

debug_assert!(diff <= i64::from(u32::MAX));

let diff = diff as u32;

let one = V::ONE;

let two = V::from_usize(2).ok_or_else(|| {

ArrowError::ComputeError("Internal error: could not create constant 2".into())

})?;

let ten = V::from_usize(10).ok_or_else(|| {

ArrowError::ComputeError("Internal error: could not create constant 10".into())

})?;

let factor = ten.pow_checked(diff).map_err(|_| {

ArrowError::ComputeError(format!(

"Overflow while rounding decimal with scale {input_scale} and decimal places {decimal_places}"

))

})?;

let mut quotient = value.div_wrapping(factor);

let remainder = value.mod_wrapping(factor);

// `factor` is an even number (10^n, n > 0), so `factor / 2` is the tie threshold

let threshold = factor.div_wrapping(two);

if remainder >= threshold {

quotient = quotient.add_checked(one).map_err(|_| {

ArrowError::ComputeError("Overflow while rounding decimal".into())

})?;

} else if remainder <= threshold.neg_wrapping() {

quotient = quotient.sub_checked(one).map_err(|_| {

ArrowError::ComputeError("Overflow while rounding decimal".into())

})?;

}

// `quotient` is the rounded value at scale `decimal_places`. Rescale to the desired

// `output_scale` (which is always >= `decimal_places` in cases where diff > 0).

let scale_shift = i64::from(output_scale) - i64::from(decimal_places);

if scale_shift == 0 {

return Ok(quotient);

}

debug_assert!(scale_shift > 0);

debug_assert!(scale_shift <= i64::from(u32::MAX));

let scale_shift = scale_shift as u32;

let shift_factor = ten.pow_checked(scale_shift).map_err(|_| {

ArrowError::ComputeError(format!(

"Overflow while rounding decimal with scale {input_scale} and decimal places {decimal_places}"

))

})?;

quotient

.mul_checked(shift_factor)

.map_err(|_| ArrowError::ComputeError("Overflow while rounding decimal".into()))

}

fn round_decimal_or_zero<V: ArrowNativeTypeOp>(

value: V,

precision: u8,

input_scale: i8,

output_scale: i8,

decimal_places: i32,

) -> Result<V, ArrowError> {

if let Some(dp) =

normalize_decimal_places_for_decimal(decimal_places, precision, input_scale)

{

round_decimal(value, input_scale, output_scale, dp)

} else {

V::from_usize(0).ok_or_else(|| {

ArrowError::ComputeError("Internal error: could not create constant 0".into())

})

}

}

#[cfg(test)]

mod test {

use std::sync::Arc;

use arrow::array::{ArrayRef, Float32Array, Float64Array, Int64Array};

use arrow::datatypes::DataType;

use datafusion_common::DataFusionError;

use datafusion_common::ScalarValue;

use datafusion_common::cast::{as_float32_array, as_float64_array};

use datafusion_expr::ColumnarValue;

fn round_arrays(

value: ArrayRef,

decimal_places: Option<ArrayRef>,

) -> Result<ArrayRef, DataFusionError> {

let number_rows = value.len();

// NOTE: For decimal inputs, the actual ROUND return type can differ from the

// input type (scale reduction for literal `decimal_places`). These unit tests

// only exercise Float32/Float64 behavior.

let return_type = value.data_type().clone();

let value = ColumnarValue::Array(value);

let decimal_places = decimal_places

.map(ColumnarValue::Array)

.unwrap_or_else(|| ColumnarValue::Scalar(ScalarValue::Int32(Some(0))));

let result =

super::round_columnar(&value, &decimal_places, number_rows, &return_type)?;

match result {