simonvandel commented on code in PR #12909:
URL: https://github.com/apache/datafusion/pull/12909#discussion_r1798350152
##########
datafusion/functions/src/math/trunc.rs:
##########
@@ -111,44 +111,66 @@ fn trunc(args: &[ArrayRef]) -> Result<ArrayRef> {
);
}
- //if only one arg then invoke toolchain trunc(num) and precision = 0 by
default
- //or then invoke the compute_truncate method to process precision
+ // If only one arg then invoke toolchain trunc(num) and precision = 0 by
default
+ // or then invoke the compute_truncate method to process precision
let num = &args[0];
let precision = if args.len() == 1 {
ColumnarValue::Scalar(Int64(Some(0)))
} else {
ColumnarValue::Array(Arc::clone(&args[1]))
};
- match args[0].data_type() {
+ match num.data_type() {
Float64 => match precision {
- ColumnarValue::Scalar(Int64(Some(0))) => Ok(Arc::new(
- make_function_scalar_inputs!(num, "num", Float64Array, {
f64::trunc }),
- ) as ArrayRef),
- ColumnarValue::Array(precision) =>
Ok(Arc::new(make_function_inputs2!(
- num,
- precision,
- "x",
- "y",
- Float64Array,
- Int64Array,
- { compute_truncate64 }
- )) as ArrayRef),
+ ColumnarValue::Scalar(Int64(Some(0))) => {
+ Ok(Arc::new(
+ args[0]
+ .as_primitive::<Float64Type>()
+ .unary::<_, Float64Type>(|x: f64| {
+ if x == 0_f64 {
+ 0_f64
+ } else {
+ x.trunc()
+ }
+ }),
+ ) as ArrayRef)
+ }
+ ColumnarValue::Array(precision) => {
+ let num_array = num.as_primitive::<Float64Type>();
+ let precision_array = precision.as_primitive::<Int64Type>();
+ let result: PrimitiveArray<Float64Type> =
+ arrow_arith::arity::binary(num_array, precision_array, |x,
y| {
+ compute_truncate64(x, y)
+ })?;
+
+ Ok(Arc::new(result) as ArrayRef)
+ }
_ => exec_err!("trunc function requires a scalar or array for
precision"),
},
Float32 => match precision {
- ColumnarValue::Scalar(Int64(Some(0))) => Ok(Arc::new(
- make_function_scalar_inputs!(num, "num", Float32Array, {
f32::trunc }),
- ) as ArrayRef),
- ColumnarValue::Array(precision) =>
Ok(Arc::new(make_function_inputs2!(
- num,
- precision,
- "x",
- "y",
- Float32Array,
- Int64Array,
- { compute_truncate32 }
- )) as ArrayRef),
+ ColumnarValue::Scalar(Int64(Some(0))) => {
+ Ok(Arc::new(
+ args[0]
+ .as_primitive::<Float32Type>()
+ .unary::<_, Float32Type>(|x: f32| {
+ if x == 0_f32 {
+ 0_f32
+ } else {
+ x.trunc()
+ }
+ }),
+ ) as ArrayRef)
+ }
+ ColumnarValue::Array(precision) => {
+ let num_array = num.as_primitive::<Float32Type>();
+ let precision_array = precision.as_primitive::<Int64Type>();
+ let result: PrimitiveArray<Float32Type> =
+ arrow_arith::arity::binary(num_array, precision_array, |x,
y| {
Review Comment:
Is this the same function as this?
https://docs.rs/arrow/latest/arrow/compute/fn.binary.html
If so you could remove the new dependency (although it is in the dependency
closure anyway already)
--
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]
