Jefffrey commented on code in PR #22697:
URL: https://github.com/apache/datafusion/pull/22697#discussion_r3377961775
##########
datafusion/spark/src/function/math/round.rs:
##########
@@ -462,18 +462,22 @@ fn spark_round(args: &[ColumnarValue], enable_ansi_mode:
bool) -> Result<Columna
impl_integer_array_round!(array, UInt32Type, scale,
enable_ansi_mode)
}
DataType::UInt64 => {
- let array = array.as_primitive::<UInt64Type>();
- let result: PrimitiveArray<UInt64Type> = array.try_unary(|x| {
- let v_i64 = i64::try_from(x).map_err(|_| {
- (exec_err!(
- "round: UInt64 value {x} exceeds i64::MAX and
cannot be rounded"
- ) as Result<(), _>)
- .unwrap_err()
Review Comment:
i wonder if we should be having this codepath in the first place,
considering spark doesnt have a u64 type 🤔
##########
datafusion/functions/src/math/round.rs:
##########
@@ -630,6 +663,180 @@ fn round_columnar(
}
}
+fn round_integer_value(value: i128, decimal_places: i32) -> Result<i128,
ArrowError> {
+ if decimal_places >= 0 {
+ return Ok(value);
+ }
+
+ let Some(factor) = 10_i128.checked_pow(decimal_places.unsigned_abs()) else
{
+ return Ok(0);
+ };
+
+ let remainder = value % factor;
+ let threshold = factor / 2;
+
+ if remainder >= threshold {
+ value
+ .checked_sub(remainder)
+ .and_then(|v| v.checked_add(factor))
+ .ok_or_else(|| {
+ ArrowError::ComputeError("Overflow while rounding
integer".to_string())
+ })
+ } else if remainder <= -threshold {
+ value
+ .checked_sub(remainder)
+ .and_then(|v| v.checked_sub(factor))
+ .ok_or_else(|| {
+ ArrowError::ComputeError("Overflow while rounding
integer".to_string())
+ })
+ } else {
+ value.checked_sub(remainder).ok_or_else(|| {
+ ArrowError::ComputeError("Overflow while rounding
integer".to_string())
+ })
+ }
+}
+
+fn round_integer_scalar(
+ value: &ScalarValue,
+ return_type: &DataType,
+ decimal_places: i32,
+) -> Result<ColumnarValue> {
+ let rounded = match value {
+ ScalarValue::Int8(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::Int16(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::Int32(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::Int64(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::UInt8(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::UInt16(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::UInt32(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::UInt64(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ _ => {
+ return internal_err!(
+ "Unexpected datatype for integer round: {}",
+ value.data_type()
+ );
+ }
+ };
+
+ let scalar = match return_type {
Review Comment:
it feels like we can combine these matches; the input and output types are
always going to be the same, right?
##########
datafusion/functions/src/math/round.rs:
##########
@@ -630,6 +663,180 @@ fn round_columnar(
}
}
+fn round_integer_value(value: i128, decimal_places: i32) -> Result<i128,
ArrowError> {
Review Comment:
is converting all ints to `i128` strictly needed for correctness, or its for
ergonomics?
##########
datafusion/functions/src/math/round.rs:
##########
@@ -120,6 +122,13 @@ fn calculate_new_precision_scale<T: DecimalType>(
}
}
+fn is_integer_data_type(data_type: &DataType) -> bool {
Review Comment:
can already use
[`DataType::is_integer`](https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html#method.is_integer)
##########
datafusion/functions/src/math/round.rs:
##########
@@ -630,6 +663,180 @@ fn round_columnar(
}
}
+fn round_integer_value(value: i128, decimal_places: i32) -> Result<i128,
ArrowError> {
+ if decimal_places >= 0 {
+ return Ok(value);
+ }
+
+ let Some(factor) = 10_i128.checked_pow(decimal_places.unsigned_abs()) else
{
+ return Ok(0);
+ };
+
+ let remainder = value % factor;
+ let threshold = factor / 2;
+
+ if remainder >= threshold {
+ value
+ .checked_sub(remainder)
+ .and_then(|v| v.checked_add(factor))
+ .ok_or_else(|| {
+ ArrowError::ComputeError("Overflow while rounding
integer".to_string())
+ })
+ } else if remainder <= -threshold {
+ value
+ .checked_sub(remainder)
+ .and_then(|v| v.checked_sub(factor))
+ .ok_or_else(|| {
+ ArrowError::ComputeError("Overflow while rounding
integer".to_string())
+ })
+ } else {
+ value.checked_sub(remainder).ok_or_else(|| {
+ ArrowError::ComputeError("Overflow while rounding
integer".to_string())
+ })
+ }
+}
+
+fn round_integer_scalar(
+ value: &ScalarValue,
+ return_type: &DataType,
+ decimal_places: i32,
+) -> Result<ColumnarValue> {
+ let rounded = match value {
+ ScalarValue::Int8(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::Int16(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::Int32(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::Int64(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::UInt8(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::UInt16(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::UInt32(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ ScalarValue::UInt64(Some(v)) => {
+ round_integer_value(i128::from(*v), decimal_places)?
+ }
+ _ => {
+ return internal_err!(
+ "Unexpected datatype for integer round: {}",
+ value.data_type()
+ );
+ }
+ };
+
+ let scalar = match return_type {
+ Int8 => ScalarValue::Int8(Some(i8::try_from(rounded).map_err(|_| {
+ ArrowError::ComputeError("Overflow while rounding
Int8".to_string())
Review Comment:
these dont need to be `ArrowErrors` since we're return a datafusion error in
this function
--
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]
