findepi commented on code in PR #18360:
URL: https://github.com/apache/datafusion/pull/18360#discussion_r2475557450
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -493,6 +515,21 @@ fn general_date_trunc_array_fine_granularity<T:
ArrowTimestampType>(
}
}
+/// Applies the unary operation in place if possible, or cloning the array if
not
+fn try_unary_mut_or_clone<F>(
+ array: PrimitiveArray<Int64Type>,
+ op: F,
+) -> Result<PrimitiveArray<Int64Type>>
+where
+ F: Fn(i64) -> Result<i64>,
Review Comment:
not really date_trunc specific. can this be made more flexible with a more
generous use of generics?
perhaps it could even be in arrow-rs. it makes `try_unary_mut` significantly
more approachable
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -455,36 +456,57 @@ fn general_date_trunc_array_fine_granularity<T:
ArrowTimestampType>(
array: &PrimitiveArray<T>,
granularity: &str,
) -> Result<ArrayRef> {
- let unit = match (tu, granularity) {
- (Second, "minute") => Some(Int64Array::new_scalar(60)),
- (Second, "hour") => Some(Int64Array::new_scalar(3600)),
- (Second, "day") => Some(Int64Array::new_scalar(86400)),
-
- (Millisecond, "second") => Some(Int64Array::new_scalar(1_000)),
- (Millisecond, "minute") => Some(Int64Array::new_scalar(60_000)),
- (Millisecond, "hour") => Some(Int64Array::new_scalar(3_600_000)),
- (Millisecond, "day") => Some(Int64Array::new_scalar(86_400_000)),
-
- (Microsecond, "millisecond") => Some(Int64Array::new_scalar(1_000)),
- (Microsecond, "second") => Some(Int64Array::new_scalar(1_000_000)),
- (Microsecond, "minute") => Some(Int64Array::new_scalar(60_000_000)),
- (Microsecond, "hour") => Some(Int64Array::new_scalar(3_600_000_000)),
- (Microsecond, "day") => Some(Int64Array::new_scalar(86_400_000_000)),
-
- (Nanosecond, "microsecond") => Some(Int64Array::new_scalar(1_000)),
- (Nanosecond, "millisecond") => Some(Int64Array::new_scalar(1_000_000)),
- (Nanosecond, "second") => Some(Int64Array::new_scalar(1_000_000_000)),
- (Nanosecond, "minute") => Some(Int64Array::new_scalar(60_000_000_000)),
- (Nanosecond, "hour") =>
Some(Int64Array::new_scalar(3_600_000_000_000)),
- (Nanosecond, "day") =>
Some(Int64Array::new_scalar(86_400_000_000_000)),
+ let unit: Option<i64> = match (tu, granularity) {
+ (Second, "minute") => Some(60),
+ (Second, "hour") => Some(3600),
+ (Second, "day") => Some(86400),
+
+ (Millisecond, "second") => Some(1_000),
+ (Millisecond, "minute") => Some(60_000),
+ (Millisecond, "hour") => Some(3_600_000),
+ (Millisecond, "day") => Some(86_400_000),
+
+ (Microsecond, "millisecond") => Some(1_000),
+ (Microsecond, "second") => Some(1_000_000),
+ (Microsecond, "minute") => Some(60_000_000),
+ (Microsecond, "hour") => Some(3_600_000_000),
+ (Microsecond, "day") => Some(86_400_000_000),
+
+ (Nanosecond, "microsecond") => Some(1_000),
+ (Nanosecond, "millisecond") => Some(1_000_000),
+ (Nanosecond, "second") => Some(1_000_000_000),
+ (Nanosecond, "minute") => Some(60_000_000_000),
+ (Nanosecond, "hour") => Some(3_600_000_000_000),
+ (Nanosecond, "day") => Some(86_400_000_000_000),
_ => None,
};
if let Some(unit) = unit {
let original_type = array.data_type();
- let array = arrow::compute::cast(array, &DataType::Int64)?;
- let array = arrow::compute::kernels::numeric::div(&array, &unit)?;
- let array = arrow::compute::kernels::numeric::mul(&array, &unit)?;
+ let input = arrow::compute::cast(array, &DataType::Int64)?;
+ // Optimize performance by doing operations in place if possible
+ let array = input.as_primitive::<Int64Type>().clone();
+ drop(input); // ensure the input reference is dropped (so we can reuse
the memory if possible)
+ let array = try_unary_mut_or_clone(array, |i| {
+ i.checked_div(unit)
+ .ok_or_else(|| exec_datafusion_err!("division overflow"))
+ })?;
+ let array = try_unary_mut_or_clone(array, |i| {
Review Comment:
technically speaking, only the first `try_unary_mut_or_clone` is needed
on the second transformation, we're guaranteed to be the pointer into the
array, and the so taking the `or_clone` path would be an error
--
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]
