tustvold commented on code in PR #5764:
URL: https://github.com/apache/arrow-datafusion/pull/5764#discussion_r1151698714
##########
datafusion/physical-expr/src/expressions/datetime.rs:
##########
@@ -239,6 +251,542 @@ pub fn evaluate_array(
Ok(ColumnarValue::Array(ret))
}
+macro_rules! ts_sub_op {
+ ($lhs:ident, $rhs:ident, $lhs_tz:ident, $rhs_tz:ident, $coef:expr,
$caster:expr, $op:expr, $ts_unit:expr, $mode:expr, $type_in:ty, $type_out:ty)
=> {{
+ let prim_array_lhs = $caster(&$lhs)?;
+ let prim_array_rhs = $caster(&$rhs)?;
+ let ret = Arc::new(try_binary_op::<$type_in, $type_in, _, $type_out>(
+ prim_array_lhs,
+ prim_array_rhs,
+ |ts1, ts2| {
+ let (lhs_tz, rhs_tz) =
+ (parse_timezones($lhs_tz), parse_timezones($rhs_tz));
+ Ok($op(
+ $ts_unit(&with_timezone_to_naive_datetime::<$mode>(
+ ts1.mul_wrapping($coef),
+ &lhs_tz,
+ )?),
+ $ts_unit(&with_timezone_to_naive_datetime::<$mode>(
+ ts2.mul_wrapping($coef),
+ &rhs_tz,
+ )?),
+ ))
+ },
+ )?) as ArrayRef;
+ ret
+ }};
+}
+macro_rules! interval_op {
+ ($lhs:ident, $rhs:ident, $caster:expr, $op:expr, $sign:ident, $type_in:ty)
=> {{
+ let prim_array_lhs = $caster(&$lhs)?;
+ let prim_array_rhs = $caster(&$rhs)?;
+ let ret = Arc::new(binary::<$type_in, $type_in, _, $type_in>(
+ prim_array_lhs,
+ prim_array_rhs,
+ |interval1, interval2| $op(interval1, interval2, $sign),
+ )?) as ArrayRef;
+ ret
+ }};
+}
+macro_rules! interval_cross_op {
+ ($lhs:ident, $rhs:ident, $caster1:expr, $caster2:expr, $op:expr,
$sign:ident, $commute:ident, $type_in1:ty, $type_in2:ty) => {{
+ let prim_array_lhs = $caster1(&$lhs)?;
+ let prim_array_rhs = $caster2(&$rhs)?;
+ let ret = Arc::new(binary::<$type_in1, $type_in2, _,
IntervalMonthDayNanoType>(
+ prim_array_lhs,
+ prim_array_rhs,
+ |interval1, interval2| $op(interval1, interval2, $sign, $commute),
+ )?) as ArrayRef;
+ ret
+ }};
+}
+macro_rules! ts_interval_op {
+ ($lhs:ident, $rhs:ident, $caster1:expr, $caster2:expr, $op:expr,
$sign:ident, $type_in1:ty, $type_in2:ty) => {{
+ let prim_array_lhs = $caster1(&$lhs)?;
+ let prim_array_rhs = $caster2(&$rhs)?;
+ let ret = Arc::new(try_binary_op::<$type_in1, $type_in2, _, $type_in1>(
+ prim_array_lhs,
+ prim_array_rhs,
+ |ts, interval| Ok($op(ts, interval as i128, $sign)?),
+ )?) as ArrayRef;
+ ret
+ }};
+}
+// This function evaluates temporal array operations, such as timestamp -
timestamp, interval + interval,
+// timestamp + interval, and interval + timestamp. It takes two arrays as
input and an integer sign representing
+// the operation (+1 for addition and -1 for subtraction). It returns a
ColumnarValue as output, which can hold
+// either a scalar or an array.
+pub fn evaluate_temporal_arrays(
+ array_lhs: &ArrayRef,
+ sign: i32,
+ array_rhs: &ArrayRef,
+) -> Result<ColumnarValue> {
+ let ret = match (array_lhs.data_type(), array_rhs.data_type()) {
+ // Timestamp - Timestamp operations, operands of only the same types
are supported.
+ (DataType::Timestamp(_, _), DataType::Timestamp(_, _)) => {
+ ts_array_op(array_lhs, array_rhs)?
+ }
+ // Interval (+ , -) Interval operations
+ (DataType::Interval(_), DataType::Interval(_)) => {
+ interval_array_op(array_lhs, array_rhs, sign)?
+ }
+ // Timestamp (+ , -) Interval and Interval + Timestamp operations
+ // Interval - Timestamp operation is not rational hence not supported
+ (DataType::Timestamp(_, _), DataType::Interval(_)) => {
+ ts_interval_array_op(array_lhs, sign, array_rhs)?
+ }
+ (DataType::Interval(_), DataType::Timestamp(_, _)) if sign == 1 => {
+ ts_interval_array_op(array_rhs, sign, array_lhs)?
+ }
+ (_, _) => Err(DataFusionError::Execution(format!(
+ "Invalid array types for DateIntervalExpr: {:?} {} {:?}",
+ array_lhs.data_type(),
+ sign,
+ array_rhs.data_type()
+ )))?,
+ };
+ Ok(ColumnarValue::Array(ret))
+}
+
+#[inline]
+unsafe fn build_primitive_array<O: ArrowPrimitiveType>(
+ len: usize,
+ buffer: Buffer,
+ null_count: usize,
+ null_buffer: Option<Buffer>,
+) -> PrimitiveArray<O> {
+ PrimitiveArray::from(ArrayData::new_unchecked(
+ O::DATA_TYPE,
+ len,
+ Some(null_count),
+ null_buffer,
+ 0,
+ vec![buffer],
+ vec![],
+ ))
+}
+
+pub fn try_binary_op<A, B, F, O>(
+ a: &PrimitiveArray<A>,
+ b: &PrimitiveArray<B>,
+ op: F,
+) -> Result<PrimitiveArray<O>, ArrowError>
+where
+ A: ArrowPrimitiveType,
+ B: ArrowPrimitiveType,
+ O: ArrowPrimitiveType,
+ F: Fn(A::Native, B::Native) -> Result<O::Native, ArrowError>,
+{
+ if a.len() != b.len() {
+ return Err(ArrowError::ComputeError(
+ "Cannot perform binary operation on arrays of different
length".to_string(),
+ ));
+ }
+ let len = a.len();
+
+ if a.is_empty() {
+ return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)));
+ }
+
+ let null_buffer = combine_option_bitmap(&[a.data(), b.data()], len);
+ let null_count = null_buffer
+ .as_ref()
+ .map(|x| len - x.count_set_bits_offset(0, len))
+ .unwrap_or_default();
+
+ let values = a.values().iter().zip(b.values()).map(|(l, r)| op(*l, *r));
+ // JUSTIFICATION
+ // Benefit
+ // ~60% speedup
+ // Soundness
+ // `values` is an iterator with a known size from a PrimitiveArray
+ let buffer = unsafe { Buffer::try_from_trusted_len_iter(values) }?;
+
+ Ok(unsafe { build_primitive_array(len, buffer, null_count, null_buffer) })
+}
+
+/// Performs a timestamp subtraction operation on two arrays and returns the
resulting array.
+fn ts_array_op(array_lhs: &ArrayRef, array_rhs: &ArrayRef) -> Result<ArrayRef>
{
+ match (array_lhs.data_type(), array_rhs.data_type()) {
+ (
+ DataType::Timestamp(TimeUnit::Second, opt_tz_lhs),
+ DataType::Timestamp(TimeUnit::Second, opt_tz_rhs),
+ ) => Ok(ts_sub_op!(
+ array_lhs,
+ array_rhs,
+ opt_tz_lhs,
+ opt_tz_rhs,
+ 1000i64,
+ as_timestamp_second_array,
+ seconds_sub,
+ NaiveDateTime::timestamp,
+ MILLISECOND_MODE,
+ TimestampSecondType,
+ IntervalDayTimeType
+ )),
+ (
+ DataType::Timestamp(TimeUnit::Millisecond, opt_tz_lhs),
+ DataType::Timestamp(TimeUnit::Millisecond, opt_tz_rhs),
+ ) => Ok(ts_sub_op!(
+ array_lhs,
+ array_rhs,
+ opt_tz_lhs,
+ opt_tz_rhs,
+ 1i64,
+ as_timestamp_millisecond_array,
+ milliseconds_sub,
+ NaiveDateTime::timestamp_millis,
+ MILLISECOND_MODE,
+ TimestampMillisecondType,
+ IntervalDayTimeType
+ )),
+ (
+ DataType::Timestamp(TimeUnit::Microsecond, opt_tz_lhs),
+ DataType::Timestamp(TimeUnit::Microsecond, opt_tz_rhs),
+ ) => Ok(ts_sub_op!(
+ array_lhs,
+ array_rhs,
+ opt_tz_lhs,
+ opt_tz_rhs,
+ 1000i64,
+ as_timestamp_microsecond_array,
+ microseconds_sub,
+ NaiveDateTime::timestamp_micros,
+ NANOSECOND_MODE,
+ TimestampMicrosecondType,
+ IntervalMonthDayNanoType
+ )),
+ (
+ DataType::Timestamp(TimeUnit::Nanosecond, opt_tz_lhs),
+ DataType::Timestamp(TimeUnit::Nanosecond, opt_tz_rhs),
+ ) => Ok(ts_sub_op!(
+ array_lhs,
+ array_rhs,
+ opt_tz_lhs,
+ opt_tz_rhs,
+ 1i64,
+ as_timestamp_nanosecond_array,
+ nanoseconds_sub,
+ NaiveDateTime::timestamp_nanos,
+ NANOSECOND_MODE,
+ TimestampNanosecondType,
+ IntervalMonthDayNanoType
+ )),
+ (_, _) => Err(DataFusionError::Execution(format!(
+ "Invalid array types for Timestamp subtraction: {:?} - {:?}",
+ array_lhs.data_type(),
+ array_rhs.data_type()
+ ))),
+ }
+}
+/// Performs an interval operation on two arrays and returns the resulting
array.
+/// The operation sign determines whether to perform addition or subtraction.
+/// The data type and unit of the two input arrays must match the supported
combinations.
+fn interval_array_op(
+ array_lhs: &ArrayRef,
+ array_rhs: &ArrayRef,
+ sign: i32,
+) -> Result<ArrayRef> {
+ match (array_lhs.data_type(), array_rhs.data_type()) {
+ (
+ DataType::Interval(IntervalUnit::YearMonth),
+ DataType::Interval(IntervalUnit::YearMonth),
+ ) => Ok(interval_op!(
+ array_lhs,
+ array_rhs,
+ as_interval_ym_array,
+ op_ym,
+ sign,
+ IntervalYearMonthType
+ )),
+ (
+ DataType::Interval(IntervalUnit::YearMonth),
+ DataType::Interval(IntervalUnit::DayTime),
+ ) => Ok(interval_cross_op!(
+ array_lhs,
+ array_rhs,
+ as_interval_ym_array,
+ as_interval_dt_array,
+ op_ym_dt,
+ sign,
+ false,
+ IntervalYearMonthType,
+ IntervalDayTimeType
+ )),
+ (
+ DataType::Interval(IntervalUnit::YearMonth),
+ DataType::Interval(IntervalUnit::MonthDayNano),
+ ) => Ok(interval_cross_op!(
+ array_lhs,
+ array_rhs,
+ as_interval_ym_array,
+ as_interval_mdn_array,
+ op_ym_mdn,
+ sign,
+ false,
+ IntervalYearMonthType,
+ IntervalMonthDayNanoType
+ )),
+ (
+ DataType::Interval(IntervalUnit::DayTime),
+ DataType::Interval(IntervalUnit::YearMonth),
+ ) => Ok(interval_cross_op!(
+ array_rhs,
+ array_lhs,
+ as_interval_ym_array,
+ as_interval_dt_array,
+ op_ym_dt,
+ sign,
+ true,
+ IntervalYearMonthType,
+ IntervalDayTimeType
+ )),
+ (
+ DataType::Interval(IntervalUnit::DayTime),
+ DataType::Interval(IntervalUnit::DayTime),
+ ) => Ok(interval_op!(
+ array_lhs,
+ array_rhs,
+ as_interval_dt_array,
+ op_dt,
+ sign,
+ IntervalDayTimeType
+ )),
+ (
+ DataType::Interval(IntervalUnit::DayTime),
+ DataType::Interval(IntervalUnit::MonthDayNano),
+ ) => Ok(interval_cross_op!(
+ array_lhs,
+ array_rhs,
+ as_interval_dt_array,
+ as_interval_mdn_array,
+ op_dt_mdn,
+ sign,
+ false,
+ IntervalDayTimeType,
+ IntervalMonthDayNanoType
+ )),
+ (
+ DataType::Interval(IntervalUnit::MonthDayNano),
+ DataType::Interval(IntervalUnit::YearMonth),
+ ) => Ok(interval_cross_op!(
+ array_rhs,
+ array_lhs,
+ as_interval_ym_array,
+ as_interval_mdn_array,
+ op_ym_mdn,
+ sign,
+ true,
+ IntervalYearMonthType,
+ IntervalMonthDayNanoType
+ )),
+ (
+ DataType::Interval(IntervalUnit::MonthDayNano),
+ DataType::Interval(IntervalUnit::DayTime),
+ ) => Ok(interval_cross_op!(
+ array_rhs,
+ array_lhs,
+ as_interval_dt_array,
+ as_interval_mdn_array,
+ op_dt_mdn,
+ sign,
+ true,
+ IntervalDayTimeType,
+ IntervalMonthDayNanoType
+ )),
+ (
+ DataType::Interval(IntervalUnit::MonthDayNano),
+ DataType::Interval(IntervalUnit::MonthDayNano),
+ ) => Ok(interval_op!(
+ array_lhs,
+ array_rhs,
+ as_interval_mdn_array,
+ op_mdn,
+ sign,
+ IntervalMonthDayNanoType
+ )),
+ (_, _) => Err(DataFusionError::Execution(format!(
+ "Invalid array types for Interval operation: {:?} {} {:?}",
+ array_lhs.data_type(),
+ sign,
+ array_rhs.data_type()
+ ))),
+ }
+}
+/// Performs a timestamp/interval operation on two arrays and returns the
resulting array.
+/// The operation sign determines whether to perform addition or subtraction.
+/// The data type and unit of the two input arrays must match the supported
combinations.
+fn ts_interval_array_op(
Review Comment:
I think this discards the timezone of the input
--
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]
