Jefffrey commented on code in PR #19384:
URL: https://github.com/apache/datafusion/pull/19384#discussion_r2638821457
##########
datafusion/functions/src/math/round.rs:
##########
@@ -123,119 +165,260 @@ impl ScalarUDFImpl for RoundFunc {
}
}
-/// Round SQL function
-fn round(args: &[ArrayRef]) -> Result<ArrayRef> {
- if args.len() != 1 && args.len() != 2 {
- return exec_err!(
- "round function requires one or two arguments, got {}",
- args.len()
- );
+fn round_columnar(
+ value: &ColumnarValue,
+ decimal_places: &ColumnarValue,
+ number_rows: usize,
+) -> Result<ColumnarValue> {
+ let value_array = value.to_array(number_rows)?;
+ let both_scalars = matches!(value, ColumnarValue::Scalar(_))
+ && matches!(decimal_places, ColumnarValue::Scalar(_));
+
+ let arr: ArrayRef = match value_array.data_type() {
+ Float64 => {
+ let result = calculate_binary_math::<Float64Type, Int32Type,
Float64Type, _>(
+ value_array.as_ref(),
+ decimal_places,
+ round_float::<f64>,
+ )?;
+ result as _
+ }
+ Float32 => {
+ let result = calculate_binary_math::<Float32Type, Int32Type,
Float32Type, _>(
+ value_array.as_ref(),
+ decimal_places,
+ round_float::<f32>,
+ )?;
+ result as _
+ }
+ Decimal32(precision, scale) => {
+ let result = calculate_binary_decimal_math::<
+ Decimal32Type,
+ Int32Type,
+ Decimal32Type,
+ _,
+ >(
+ value_array.as_ref(),
+ decimal_places,
+ |v, dp| round_decimal32(v, *scale, dp),
+ *precision,
+ *scale,
+ )?;
+ result as _
+ }
+ Decimal64(precision, scale) => {
+ let result = calculate_binary_decimal_math::<
+ Decimal64Type,
+ Int32Type,
+ Decimal64Type,
+ _,
+ >(
+ value_array.as_ref(),
+ decimal_places,
+ |v, dp| round_decimal64(v, *scale, dp),
+ *precision,
+ *scale,
+ )?;
+ result as _
+ }
+ Decimal128(precision, scale) => {
+ let result = calculate_binary_decimal_math::<
+ Decimal128Type,
+ Int32Type,
+ Decimal128Type,
+ _,
+ >(
+ value_array.as_ref(),
+ decimal_places,
+ |v, dp| round_decimal128(v, *scale, dp),
+ *precision,
+ *scale,
+ )?;
+ result as _
+ }
+ Decimal256(precision, scale) => {
+ let result = calculate_binary_decimal_math::<
+ Decimal256Type,
+ Int32Type,
+ Decimal256Type,
+ _,
+ >(
+ value_array.as_ref(),
+ decimal_places,
+ |v, dp| round_decimal256(v, *scale, dp),
+ *precision,
+ *scale,
+ )?;
+ 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_float<T>(value: T, decimal_places: i32) -> Result<T, ArrowError>
+where
+ T: num_traits::Float,
+{
+ let factor = T::from(10_f64.powi(decimal_places)).ok_or_else(|| {
+ ArrowError::ComputeError(format!(
+ "Invalid value for decimal places: {decimal_places}"
+ ))
+ })?;
+ Ok((value * factor).round() / factor)
+}
+
+fn round_decimal32(
+ value: i32,
+ scale: i8,
+ decimal_places: i32,
+) -> Result<i32, ArrowError> {
+ let rounded = round_decimal_i128(i128::from(value), scale,
decimal_places)?;
+ i32::try_from(rounded)
+ .map_err(|_| ArrowError::ComputeError("Overflow while rounding
decimal32".into()))
+}
+
+fn round_decimal64(
+ value: i64,
+ scale: i8,
+ decimal_places: i32,
+) -> Result<i64, ArrowError> {
+ let rounded = round_decimal_i128(i128::from(value), scale,
decimal_places)?;
+ i64::try_from(rounded)
+ .map_err(|_| ArrowError::ComputeError("Overflow while rounding
decimal64".into()))
+}
+
+fn round_decimal128(
+ value: i128,
+ scale: i8,
+ decimal_places: i32,
+) -> Result<i128, ArrowError> {
+ round_decimal_i128(value, scale, decimal_places)
+}
+
+fn round_decimal256(
+ value: i256,
+ scale: i8,
+ decimal_places: i32,
+) -> Result<i256, ArrowError> {
+ round_decimal_i256(value, scale, decimal_places)
+}
+
+fn round_decimal_i128(
Review Comment:
We can replace these decimal implementations with a single generic version:
```rust
fn round_decimal<V: ArrowNativeTypeOp>(
value: V,
scale: i8,
decimal_places: i32,
) -> Result<V, ArrowError> {
let diff = scale as i32 - decimal_places;
if diff <= 0 {
return Ok(value);
}
let diff = diff as u32; // safe since we check sign above
// these can't truncate/overflow
let one = V::ONE;
let two = V::from_usize(2).unwrap();
let ten = V::from_usize(10).unwrap();
let factor = ten.pow_checked(diff).map_err(|_| {
ArrowError::ComputeError(format!(
"Overflow while rounding decimal with scale {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
.mul_checked(factor)
.map_err(|_| ArrowError::ComputeError("Overflow while rounding
decimal".into()))
}
```
Then can call them directly like so:
```rust
Decimal128(precision, scale) => {
let result = calculate_binary_decimal_math::<
Decimal128Type,
Int32Type,
Decimal128Type,
_,
>(
value_array.as_ref(),
decimal_places,
|v, dp| round_decimal(v, *scale, dp), // <--- here
*precision,
*scale,
)?;
result as _
}
```
--
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]
