Skip to content

Commit c2ea6b3

Browse files
maxburkealamb
andauthored
Implement min/max for interval types (#11015)
* Implement min/max for interval types * Add sqllogictests for min/max intervals * Add tests for interval min/max * update sql logic tests --------- Co-authored-by: Andrew Lamb <[email protected]>
1 parent 8d8dd90 commit c2ea6b3

File tree

2 files changed

+289
-157
lines changed

2 files changed

+289
-157
lines changed

datafusion/physical-expr/src/aggregate/min_max.rs

Lines changed: 131 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,18 +24,20 @@ use crate::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator;
2424
use crate::{AggregateExpr, PhysicalExpr};
2525
use arrow::compute;
2626
use arrow::datatypes::{
27-
DataType, Date32Type, Date64Type, Time32MillisecondType, Time32SecondType,
28-
Time64MicrosecondType, Time64NanosecondType, TimeUnit, TimestampMicrosecondType,
29-
TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
27+
DataType, Date32Type, Date64Type, IntervalUnit, Time32MillisecondType,
28+
Time32SecondType, Time64MicrosecondType, Time64NanosecondType, TimeUnit,
29+
TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
30+
TimestampSecondType,
3031
};
3132
use arrow::{
3233
array::{
3334
ArrayRef, BinaryArray, BooleanArray, Date32Array, Date64Array, Float32Array,
34-
Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray,
35-
LargeStringArray, StringArray, Time32MillisecondArray, Time32SecondArray,
36-
Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray,
37-
TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
38-
UInt16Array, UInt32Array, UInt64Array, UInt8Array,
35+
Float64Array, Int16Array, Int32Array, Int64Array, Int8Array,
36+
IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray,
37+
LargeBinaryArray, LargeStringArray, StringArray, Time32MillisecondArray,
38+
Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
39+
TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
40+
TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
3941
},
4042
datatypes::Field,
4143
};
@@ -408,6 +410,25 @@ macro_rules! min_max_batch {
408410
$OP
409411
)
410412
}
413+
DataType::Interval(IntervalUnit::YearMonth) => {
414+
typed_min_max_batch!(
415+
$VALUES,
416+
IntervalYearMonthArray,
417+
IntervalYearMonth,
418+
$OP
419+
)
420+
}
421+
DataType::Interval(IntervalUnit::DayTime) => {
422+
typed_min_max_batch!($VALUES, IntervalDayTimeArray, IntervalDayTime, $OP)
423+
}
424+
DataType::Interval(IntervalUnit::MonthDayNano) => {
425+
typed_min_max_batch!(
426+
$VALUES,
427+
IntervalMonthDayNanoArray,
428+
IntervalMonthDayNano,
429+
$OP
430+
)
431+
}
411432
other => {
412433
// This should have been handled before
413434
return internal_err!(
@@ -1121,6 +1142,108 @@ impl Accumulator for SlidingMinAccumulator {
11211142
#[cfg(test)]
11221143
mod tests {
11231144
use super::*;
1145+
use arrow::datatypes::{
1146+
IntervalDayTimeType, IntervalMonthDayNanoType, IntervalYearMonthType,
1147+
};
1148+
1149+
#[test]
1150+
fn interval_min_max() {
1151+
// IntervalYearMonth
1152+
let b = IntervalYearMonthArray::from(vec![
1153+
IntervalYearMonthType::make_value(0, 1),
1154+
IntervalYearMonthType::make_value(5, 34),
1155+
IntervalYearMonthType::make_value(-2, 4),
1156+
IntervalYearMonthType::make_value(7, -4),
1157+
IntervalYearMonthType::make_value(0, 1),
1158+
]);
1159+
let b: ArrayRef = Arc::new(b);
1160+
1161+
let mut min =
1162+
MinAccumulator::try_new(&DataType::Interval(IntervalUnit::YearMonth))
1163+
.unwrap();
1164+
min.update_batch(&[b.clone()]).unwrap();
1165+
let min_res = min.evaluate().unwrap();
1166+
assert_eq!(
1167+
min_res,
1168+
ScalarValue::IntervalYearMonth(Some(IntervalYearMonthType::make_value(
1169+
-2, 4
1170+
)))
1171+
);
1172+
1173+
let mut max =
1174+
MaxAccumulator::try_new(&DataType::Interval(IntervalUnit::YearMonth))
1175+
.unwrap();
1176+
max.update_batch(&[b.clone()]).unwrap();
1177+
let max_res = max.evaluate().unwrap();
1178+
assert_eq!(
1179+
max_res,
1180+
ScalarValue::IntervalYearMonth(Some(IntervalYearMonthType::make_value(
1181+
5, 34
1182+
)))
1183+
);
1184+
1185+
// IntervalDayTime
1186+
let b = IntervalDayTimeArray::from(vec![
1187+
IntervalDayTimeType::make_value(0, 0),
1188+
IntervalDayTimeType::make_value(5, 454000),
1189+
IntervalDayTimeType::make_value(-34, 0),
1190+
IntervalDayTimeType::make_value(7, -4000),
1191+
IntervalDayTimeType::make_value(1, 0),
1192+
]);
1193+
let b: ArrayRef = Arc::new(b);
1194+
1195+
let mut min =
1196+
MinAccumulator::try_new(&DataType::Interval(IntervalUnit::DayTime)).unwrap();
1197+
min.update_batch(&[b.clone()]).unwrap();
1198+
let min_res = min.evaluate().unwrap();
1199+
assert_eq!(
1200+
min_res,
1201+
ScalarValue::IntervalDayTime(Some(IntervalDayTimeType::make_value(-34, 0)))
1202+
);
1203+
1204+
let mut max =
1205+
MaxAccumulator::try_new(&DataType::Interval(IntervalUnit::DayTime)).unwrap();
1206+
max.update_batch(&[b.clone()]).unwrap();
1207+
let max_res = max.evaluate().unwrap();
1208+
assert_eq!(
1209+
max_res,
1210+
ScalarValue::IntervalDayTime(Some(IntervalDayTimeType::make_value(7, -4000)))
1211+
);
1212+
1213+
// IntervalMonthDayNano
1214+
let b = IntervalMonthDayNanoArray::from(vec![
1215+
IntervalMonthDayNanoType::make_value(1, 0, 0),
1216+
IntervalMonthDayNanoType::make_value(344, 34, -43_000_000_000),
1217+
IntervalMonthDayNanoType::make_value(-593, -33, 13_000_000_000),
1218+
IntervalMonthDayNanoType::make_value(5, 2, 493_000_000_000),
1219+
IntervalMonthDayNanoType::make_value(1, 0, 0),
1220+
]);
1221+
let b: ArrayRef = Arc::new(b);
1222+
1223+
let mut min =
1224+
MinAccumulator::try_new(&DataType::Interval(IntervalUnit::MonthDayNano))
1225+
.unwrap();
1226+
min.update_batch(&[b.clone()]).unwrap();
1227+
let min_res = min.evaluate().unwrap();
1228+
assert_eq!(
1229+
min_res,
1230+
ScalarValue::IntervalMonthDayNano(Some(
1231+
IntervalMonthDayNanoType::make_value(-593, -33, 13_000_000_000)
1232+
))
1233+
);
1234+
1235+
let mut max =
1236+
MaxAccumulator::try_new(&DataType::Interval(IntervalUnit::MonthDayNano))
1237+
.unwrap();
1238+
max.update_batch(&[b.clone()]).unwrap();
1239+
let max_res = max.evaluate().unwrap();
1240+
assert_eq!(
1241+
max_res,
1242+
ScalarValue::IntervalMonthDayNano(Some(
1243+
IntervalMonthDayNanoType::make_value(344, 34, -43_000_000_000)
1244+
))
1245+
);
1246+
}
11241247

11251248
#[test]
11261249
fn float_min_max_with_nans() {

0 commit comments

Comments
 (0)