Skip to content

Commit 61ea9f2

Browse files
authored
Add datetime/interval/duration into dyn scalar comparison (#3730)
* Add datatime/interval/duration into comparison * Add some tests
1 parent 18388b2 commit 61ea9f2

File tree

1 file changed

+219
-0
lines changed

1 file changed

+219
-0
lines changed

arrow-ord/src/comparison.rs

Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -523,6 +523,91 @@ macro_rules! dyn_compare_scalar {
523523
let left = as_primitive_array::<Decimal128Type>($LEFT);
524524
$OP::<Decimal128Type>(left, right)
525525
}
526+
DataType::Date32 => {
527+
let right = try_to_type!($RIGHT, to_i32)?;
528+
let left = as_primitive_array::<Date32Type>($LEFT);
529+
$OP::<Date32Type>(left, right)
530+
}
531+
DataType::Date64 => {
532+
let right = try_to_type!($RIGHT, to_i64)?;
533+
let left = as_primitive_array::<Date64Type>($LEFT);
534+
$OP::<Date64Type>(left, right)
535+
}
536+
DataType::Timestamp(TimeUnit::Nanosecond, _) => {
537+
let right = try_to_type!($RIGHT, to_i64)?;
538+
let left = as_primitive_array::<TimestampNanosecondType>($LEFT);
539+
$OP::<TimestampNanosecondType>(left, right)
540+
}
541+
DataType::Timestamp(TimeUnit::Microsecond, _) => {
542+
let right = try_to_type!($RIGHT, to_i64)?;
543+
let left = as_primitive_array::<TimestampMicrosecondType>($LEFT);
544+
$OP::<TimestampMicrosecondType>(left, right)
545+
}
546+
DataType::Timestamp(TimeUnit::Millisecond, _) => {
547+
let right = try_to_type!($RIGHT, to_i64)?;
548+
let left = as_primitive_array::<TimestampMillisecondType>($LEFT);
549+
$OP::<TimestampMillisecondType>(left, right)
550+
}
551+
DataType::Timestamp(TimeUnit::Second, _) => {
552+
let right = try_to_type!($RIGHT, to_i64)?;
553+
let left = as_primitive_array::<TimestampSecondType>($LEFT);
554+
$OP::<TimestampSecondType>(left, right)
555+
}
556+
DataType::Time32(TimeUnit::Second) => {
557+
let right = try_to_type!($RIGHT, to_i32)?;
558+
let left = as_primitive_array::<Time32SecondType>($LEFT);
559+
$OP::<Time32SecondType>(left, right)
560+
}
561+
DataType::Time32(TimeUnit::Millisecond) => {
562+
let right = try_to_type!($RIGHT, to_i32)?;
563+
let left = as_primitive_array::<Time32MillisecondType>($LEFT);
564+
$OP::<Time32MillisecondType>(left, right)
565+
}
566+
DataType::Time64(TimeUnit::Microsecond) => {
567+
let right = try_to_type!($RIGHT, to_i64)?;
568+
let left = as_primitive_array::<Time64MicrosecondType>($LEFT);
569+
$OP::<Time64MicrosecondType>(left, right)
570+
}
571+
DataType::Time64(TimeUnit::Nanosecond) => {
572+
let right = try_to_type!($RIGHT, to_i64)?;
573+
let left = as_primitive_array::<Time64NanosecondType>($LEFT);
574+
$OP::<Time64NanosecondType>(left, right)
575+
}
576+
DataType::Interval(IntervalUnit::YearMonth) => {
577+
let right = try_to_type!($RIGHT, to_i32)?;
578+
let left = as_primitive_array::<IntervalYearMonthType>($LEFT);
579+
$OP::<IntervalYearMonthType>(left, right)
580+
}
581+
DataType::Interval(IntervalUnit::DayTime) => {
582+
let right = try_to_type!($RIGHT, to_i64)?;
583+
let left = as_primitive_array::<IntervalDayTimeType>($LEFT);
584+
$OP::<IntervalDayTimeType>(left, right)
585+
}
586+
DataType::Interval(IntervalUnit::MonthDayNano) => {
587+
let right = try_to_type!($RIGHT, to_i128)?;
588+
let left = as_primitive_array::<IntervalMonthDayNanoType>($LEFT);
589+
$OP::<IntervalMonthDayNanoType>(left, right)
590+
}
591+
DataType::Duration(TimeUnit::Second) => {
592+
let right = try_to_type!($RIGHT, to_i64)?;
593+
let left = as_primitive_array::<DurationSecondType>($LEFT);
594+
$OP::<DurationSecondType>(left, right)
595+
}
596+
DataType::Duration(TimeUnit::Millisecond) => {
597+
let right = try_to_type!($RIGHT, to_i64)?;
598+
let left = as_primitive_array::<DurationMillisecondType>($LEFT);
599+
$OP::<DurationMillisecondType>(left, right)
600+
}
601+
DataType::Duration(TimeUnit::Microsecond) => {
602+
let right = try_to_type!($RIGHT, to_i64)?;
603+
let left = as_primitive_array::<DurationMicrosecondType>($LEFT);
604+
$OP::<DurationMicrosecondType>(left, right)
605+
}
606+
DataType::Duration(TimeUnit::Nanosecond) => {
607+
let right = try_to_type!($RIGHT, to_i64)?;
608+
let left = as_primitive_array::<DurationNanosecondType>($LEFT);
609+
$OP::<DurationNanosecondType>(left, right)
610+
}
526611
_ => Err(ArrowError::ComputeError(format!(
527612
"Unsupported data type {:?} for comparison {} with {:?}",
528613
$LEFT.data_type(),
@@ -1707,6 +1792,22 @@ macro_rules! typed_compares {
17071792
DataType::Interval(IntervalUnit::MonthDayNano),
17081793
DataType::Interval(IntervalUnit::MonthDayNano),
17091794
) => cmp_primitive_array::<IntervalMonthDayNanoType, _>($LEFT, $RIGHT, $OP),
1795+
(
1796+
DataType::Duration(TimeUnit::Second),
1797+
DataType::Duration(TimeUnit::Second),
1798+
) => cmp_primitive_array::<DurationSecondType, _>($LEFT, $RIGHT, $OP),
1799+
(
1800+
DataType::Duration(TimeUnit::Millisecond),
1801+
DataType::Duration(TimeUnit::Millisecond),
1802+
) => cmp_primitive_array::<DurationMillisecondType, _>($LEFT, $RIGHT, $OP),
1803+
(
1804+
DataType::Duration(TimeUnit::Microsecond),
1805+
DataType::Duration(TimeUnit::Microsecond),
1806+
) => cmp_primitive_array::<DurationMicrosecondType, _>($LEFT, $RIGHT, $OP),
1807+
(
1808+
DataType::Duration(TimeUnit::Nanosecond),
1809+
DataType::Duration(TimeUnit::Nanosecond),
1810+
) => cmp_primitive_array::<DurationNanosecondType, _>($LEFT, $RIGHT, $OP),
17101811
(t1, t2) if t1 == t2 => Err(ArrowError::NotYetImplemented(format!(
17111812
"Comparing arrays of type {} is not yet implemented",
17121813
t1
@@ -4006,6 +4107,124 @@ mod tests {
40064107
);
40074108
}
40084109

4110+
fn test_primitive_dyn_scalar<T: ArrowPrimitiveType>(array: PrimitiveArray<T>) {
4111+
let a_eq = eq_dyn_scalar(&array, 8).unwrap();
4112+
assert_eq!(
4113+
a_eq,
4114+
BooleanArray::from(vec![Some(false), None, Some(true), None, Some(false)])
4115+
);
4116+
4117+
let a_eq = gt_eq_dyn_scalar(&array, 8).unwrap();
4118+
assert_eq!(
4119+
a_eq,
4120+
BooleanArray::from(vec![Some(false), None, Some(true), None, Some(true)])
4121+
);
4122+
4123+
let a_eq = gt_dyn_scalar(&array, 8).unwrap();
4124+
assert_eq!(
4125+
a_eq,
4126+
BooleanArray::from(vec![Some(false), None, Some(false), None, Some(true)])
4127+
);
4128+
4129+
let a_eq = lt_eq_dyn_scalar(&array, 8).unwrap();
4130+
assert_eq!(
4131+
a_eq,
4132+
BooleanArray::from(vec![Some(true), None, Some(true), None, Some(false)])
4133+
);
4134+
4135+
let a_eq = lt_dyn_scalar(&array, 8).unwrap();
4136+
assert_eq!(
4137+
a_eq,
4138+
BooleanArray::from(vec![Some(true), None, Some(false), None, Some(false)])
4139+
);
4140+
}
4141+
4142+
#[test]
4143+
fn test_timestamp_dyn_scalar() {
4144+
let array =
4145+
TimestampSecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4146+
test_primitive_dyn_scalar(array);
4147+
4148+
let array =
4149+
TimestampMicrosecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4150+
test_primitive_dyn_scalar(array);
4151+
4152+
let array =
4153+
TimestampMicrosecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4154+
test_primitive_dyn_scalar(array);
4155+
4156+
let array =
4157+
TimestampNanosecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4158+
test_primitive_dyn_scalar(array);
4159+
}
4160+
4161+
#[test]
4162+
fn test_date32_dyn_scalar() {
4163+
let array = Date32Array::from(vec![Some(1), None, Some(8), None, Some(10)]);
4164+
test_primitive_dyn_scalar(array);
4165+
}
4166+
4167+
#[test]
4168+
fn test_date64_dyn_scalar() {
4169+
let array = Date64Array::from(vec![Some(1), None, Some(8), None, Some(10)]);
4170+
test_primitive_dyn_scalar(array);
4171+
}
4172+
4173+
#[test]
4174+
fn test_time32_dyn_scalar() {
4175+
let array = Time32SecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4176+
test_primitive_dyn_scalar(array);
4177+
4178+
let array =
4179+
Time32MillisecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4180+
test_primitive_dyn_scalar(array);
4181+
}
4182+
4183+
#[test]
4184+
fn test_time64_dyn_scalar() {
4185+
let array =
4186+
Time64MicrosecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4187+
test_primitive_dyn_scalar(array);
4188+
4189+
let array =
4190+
Time64NanosecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4191+
test_primitive_dyn_scalar(array);
4192+
}
4193+
4194+
#[test]
4195+
fn test_interval_dyn_scalar() {
4196+
let array =
4197+
IntervalDayTimeArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4198+
test_primitive_dyn_scalar(array);
4199+
4200+
let array =
4201+
IntervalMonthDayNanoArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4202+
test_primitive_dyn_scalar(array);
4203+
4204+
let array =
4205+
IntervalYearMonthArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4206+
test_primitive_dyn_scalar(array);
4207+
}
4208+
4209+
#[test]
4210+
fn test_duration_dyn_scalar() {
4211+
let array =
4212+
DurationSecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4213+
test_primitive_dyn_scalar(array);
4214+
4215+
let array =
4216+
DurationMicrosecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4217+
test_primitive_dyn_scalar(array);
4218+
4219+
let array =
4220+
DurationMillisecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4221+
test_primitive_dyn_scalar(array);
4222+
4223+
let array =
4224+
DurationNanosecondArray::from(vec![Some(1), None, Some(8), None, Some(10)]);
4225+
test_primitive_dyn_scalar(array);
4226+
}
4227+
40094228
#[test]
40104229
fn test_lt_eq_dyn_scalar_with_dict() {
40114230
let mut builder =

0 commit comments

Comments
 (0)