Skip to content

Commit 96c5f17

Browse files
MazterQyouMazterQyou
committed
feat: Support DATE_PART with intervals
Signed-off-by: Alex Qyoun-ae <[email protected]>
1 parent 454020d commit 96c5f17

File tree

4 files changed

+364
-38
lines changed

4 files changed

+364
-38
lines changed

‎datafusion/core/tests/sql/expr.rs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1124,11 +1124,19 @@ async fn test_extract_date_part() -> Result<()> {
11241124
"EXTRACT(year FROM to_timestamp('2020-09-08T12:00:00+00:00'))",
11251125
"2020"
11261126
);
1127+
test_expression!(
1128+
"date_part('YEAR', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
1129+
"1"
1130+
);
11271131
test_expression!("date_part('MONTH', CAST('2000-01-01' AS DATE))", "1");
11281132
test_expression!(
11291133
"EXTRACT(month FROM to_timestamp('2020-09-08T12:00:00+00:00'))",
11301134
"9"
11311135
);
1136+
test_expression!(
1137+
"date_part('MONTH', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
1138+
"2"
1139+
);
11321140
test_expression!("date_part('WEEK', CAST('2003-01-01' AS DATE))", "1");
11331141

11341142
// TODO Creating logical plan for 'SELECT EXTRACT(WEEK FROM to_timestamp('2020-09-08T12:00:00+00:00'))'
@@ -1144,11 +1152,19 @@ async fn test_extract_date_part() -> Result<()> {
11441152
"EXTRACT(day FROM to_timestamp('2020-09-08T12:00:00+00:00'))",
11451153
"8"
11461154
);
1155+
test_expression!(
1156+
"date_part('DAY', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
1157+
"3"
1158+
);
11471159
test_expression!("date_part('HOUR', CAST('2000-01-01' AS DATE))", "0");
11481160
test_expression!(
11491161
"EXTRACT(hour FROM to_timestamp('2020-09-08T12:03:03+00:00'))",
11501162
"12"
11511163
);
1164+
test_expression!(
1165+
"date_part('HOUR', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
1166+
"4"
1167+
);
11521168
test_expression!(
11531169
"EXTRACT(minute FROM to_timestamp('2020-09-08T12:12:00+00:00'))",
11541170
"12"
@@ -1157,6 +1173,10 @@ async fn test_extract_date_part() -> Result<()> {
11571173
"date_part('minute', to_timestamp('2020-09-08T12:12:00+00:00'))",
11581174
"12"
11591175
);
1176+
test_expression!(
1177+
"date_part('MINUTE', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
1178+
"5"
1179+
);
11601180
test_expression!(
11611181
"EXTRACT(second FROM to_timestamp('2020-09-08T12:00:12+00:00'))",
11621182
"12"
@@ -1165,6 +1185,10 @@ async fn test_extract_date_part() -> Result<()> {
11651185
"date_part('second', to_timestamp('2020-09-08T12:00:12+00:00'))",
11661186
"12"
11671187
);
1188+
test_expression!(
1189+
"date_part('SECOND', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
1190+
"6"
1191+
);
11681192

11691193
// DOY
11701194
test_expression!(

‎datafusion/cube_ext/src/temporal.rs‎

Lines changed: 270 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,13 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use arrow::array::{Array, Float64Array, Int32Array, Int32Builder, PrimitiveArray};
19-
use arrow::compute::kernels::arity::unary;
18+
use arrow::array::{
19+
Array, Date32Array, Date64Array, Float64Array, Int32Array, Int32Builder,
20+
IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray,
21+
PrimitiveArray, TimestampMicrosecondArray, TimestampMillisecondArray,
22+
TimestampNanosecondArray, TimestampSecondArray,
23+
};
24+
use arrow::compute::kernels::{arity::unary, temporal as arrow_temporal};
2025
use arrow::datatypes::{
2126
ArrowNumericType, ArrowPrimitiveType, ArrowTemporalType, DataType, Date32Type,
2227
Date64Type, Float64Type, IntervalDayTimeType, IntervalMonthDayNanoType,
@@ -295,3 +300,266 @@ where
295300

296301
Ok(b.finish())
297302
}
303+
304+
/// This macro will generate a trait `DatePartable` that is automatically implemented
305+
/// for all Arrow temporal types.
306+
macro_rules! date_partable {
307+
($($gran:ident),*) => {
308+
pub trait DatePartable: ArrowPrimitiveType + Sized {
309+
$(
310+
fn $gran(array: &PrimitiveArray<Self>) -> Result<Int32Array>;
311+
)*
312+
}
313+
314+
impl DatePartable for TimestampSecondType {
315+
$(
316+
fn $gran(array: &TimestampSecondArray) -> Result<Int32Array> {
317+
arrow_temporal::$gran(array)
318+
}
319+
)*
320+
}
321+
322+
impl DatePartable for TimestampMillisecondType {
323+
$(
324+
fn $gran(array: &TimestampMillisecondArray) -> Result<Int32Array> {
325+
arrow_temporal::$gran(array)
326+
}
327+
)*
328+
}
329+
330+
impl DatePartable for TimestampMicrosecondType {
331+
$(
332+
fn $gran(array: &TimestampMicrosecondArray) -> Result<Int32Array> {
333+
arrow_temporal::$gran(array)
334+
}
335+
)*
336+
}
337+
338+
impl DatePartable for TimestampNanosecondType {
339+
$(
340+
fn $gran(array: &TimestampNanosecondArray) -> Result<Int32Array> {
341+
arrow_temporal::$gran(array)
342+
}
343+
)*
344+
}
345+
346+
impl DatePartable for Date32Type {
347+
$(
348+
fn $gran(array: &Date32Array) -> Result<Int32Array> {
349+
arrow_temporal::$gran(array)
350+
}
351+
)*
352+
}
353+
354+
impl DatePartable for Date64Type {
355+
$(
356+
fn $gran(array: &Date64Array) -> Result<Int32Array> {
357+
arrow_temporal::$gran(array)
358+
}
359+
)*
360+
}
361+
362+
$(
363+
pub fn $gran<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
364+
where
365+
T: DatePartable,
366+
{
367+
DatePartable::$gran(array)
368+
}
369+
)*
370+
};
371+
}
372+
373+
date_partable!(year, quarter, month, day, hour, minute, second);
374+
375+
fn interval_year_month_op(
376+
array: &IntervalYearMonthArray,
377+
op: fn(i32) -> Result<i32>,
378+
) -> Result<Int32Array> {
379+
let mut builder = Int32Builder::new(array.len());
380+
for i in 0..array.len() {
381+
if array.is_null(i) {
382+
builder.append_null()?;
383+
continue;
384+
}
385+
let value = array.value(i);
386+
let result = op(value)?;
387+
builder.append_value(result)?;
388+
}
389+
Ok(builder.finish())
390+
}
391+
392+
fn interval_day_time_op(
393+
array: &IntervalDayTimeArray,
394+
op: fn(i64) -> Result<i32>,
395+
) -> Result<Int32Array> {
396+
let mut builder = Int32Builder::new(array.len());
397+
for i in 0..array.len() {
398+
if array.is_null(i) {
399+
builder.append_null()?;
400+
continue;
401+
}
402+
let value = array.value(i);
403+
let result = op(value)?;
404+
builder.append_value(result)?;
405+
}
406+
Ok(builder.finish())
407+
}
408+
409+
fn interval_month_day_nano_op(
410+
array: &IntervalMonthDayNanoArray,
411+
op: fn(i128) -> Result<i32>,
412+
) -> Result<Int32Array> {
413+
let mut builder = Int32Builder::new(array.len());
414+
for i in 0..array.len() {
415+
if array.is_null(i) {
416+
builder.append_null()?;
417+
continue;
418+
}
419+
let value = array.value(i);
420+
let result = op(value)?;
421+
builder.append_value(result)?;
422+
}
423+
Ok(builder.finish())
424+
}
425+
426+
impl DatePartable for IntervalYearMonthType {
427+
fn year(array: &IntervalYearMonthArray) -> Result<Int32Array> {
428+
interval_year_month_op(array, |v| Ok(v / 12))
429+
}
430+
431+
fn quarter(array: &IntervalYearMonthArray) -> Result<Int32Array> {
432+
interval_year_month_op(array, |v| Ok(v % 12 / 3 + 1))
433+
}
434+
435+
fn month(array: &IntervalYearMonthArray) -> Result<Int32Array> {
436+
interval_year_month_op(array, |v| Ok(v % 12))
437+
}
438+
439+
fn day(array: &IntervalYearMonthArray) -> Result<Int32Array> {
440+
interval_year_month_op(array, |_| Ok(0))
441+
}
442+
443+
fn hour(array: &IntervalYearMonthArray) -> Result<Int32Array> {
444+
interval_year_month_op(array, |_| Ok(0))
445+
}
446+
447+
fn minute(array: &IntervalYearMonthArray) -> Result<Int32Array> {
448+
interval_year_month_op(array, |_| Ok(0))
449+
}
450+
451+
fn second(array: &IntervalYearMonthArray) -> Result<Int32Array> {
452+
interval_year_month_op(array, |_| Ok(0))
453+
}
454+
}
455+
456+
impl DatePartable for IntervalDayTimeType {
457+
fn year(array: &IntervalDayTimeArray) -> Result<Int32Array> {
458+
interval_day_time_op(array, |_| Ok(0))
459+
}
460+
461+
fn quarter(array: &IntervalDayTimeArray) -> Result<Int32Array> {
462+
interval_day_time_op(array, |_| Ok(1))
463+
}
464+
465+
fn month(array: &IntervalDayTimeArray) -> Result<Int32Array> {
466+
interval_day_time_op(array, |_| Ok(0))
467+
}
468+
469+
fn day(array: &IntervalDayTimeArray) -> Result<Int32Array> {
470+
interval_day_time_op(array, |v| {
471+
let (days, _) = IntervalDayTimeType::to_parts(v);
472+
Ok(days)
473+
})
474+
}
475+
476+
fn hour(array: &IntervalDayTimeArray) -> Result<Int32Array> {
477+
interval_day_time_op(array, |v| {
478+
let (_, millis) = IntervalDayTimeType::to_parts(v);
479+
Ok(millis / 3_600_000)
480+
})
481+
}
482+
483+
fn minute(array: &IntervalDayTimeArray) -> Result<Int32Array> {
484+
interval_day_time_op(array, |v| {
485+
let (_, millis) = IntervalDayTimeType::to_parts(v);
486+
Ok(millis % 3_600_000 / 60_000)
487+
})
488+
}
489+
490+
fn second(array: &IntervalDayTimeArray) -> Result<Int32Array> {
491+
// NOTE: this technically should return Float64 with millis in the decimal part,
492+
// but we are returning Int32 for compatibility with the original implementation.
493+
interval_day_time_op(array, |v| {
494+
let (_, millis) = IntervalDayTimeType::to_parts(v);
495+
Ok(millis % 60_000 / 1_000)
496+
})
497+
}
498+
}
499+
500+
impl DatePartable for IntervalMonthDayNanoType {
501+
fn year(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
502+
interval_month_day_nano_op(array, |v| {
503+
let (months, _, _) = IntervalMonthDayNanoType::to_parts(v);
504+
Ok(months / 12)
505+
})
506+
}
507+
508+
fn quarter(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
509+
interval_month_day_nano_op(array, |v| {
510+
let (months, _, _) = IntervalMonthDayNanoType::to_parts(v);
511+
Ok(months % 12 / 3 + 1)
512+
})
513+
}
514+
515+
fn month(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
516+
interval_month_day_nano_op(array, |v| {
517+
let (months, _, _) = IntervalMonthDayNanoType::to_parts(v);
518+
Ok(months % 12)
519+
})
520+
}
521+
522+
fn day(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
523+
interval_month_day_nano_op(array, |v| {
524+
let (_, days, _) = IntervalMonthDayNanoType::to_parts(v);
525+
Ok(days)
526+
})
527+
}
528+
529+
fn hour(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
530+
interval_month_day_nano_op(array, |v| {
531+
let (_, _, nanos) = IntervalMonthDayNanoType::to_parts(v);
532+
(nanos / 3_600_000_000_000).try_into().map_err(|_| {
533+
ArrowError::ComputeError("Unable to convert i64 nanos to i32".to_string())
534+
})
535+
})
536+
}
537+
538+
fn minute(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
539+
interval_month_day_nano_op(array, |v| {
540+
let (_, _, nanos) = IntervalMonthDayNanoType::to_parts(v);
541+
(nanos % 3_600_000_000_000 / 60_000_000_000)
542+
.try_into()
543+
.map_err(|_| {
544+
ArrowError::ComputeError(
545+
"Unable to convert i64 nanos to i32".to_string(),
546+
)
547+
})
548+
})
549+
}
550+
551+
fn second(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
552+
// NOTE: this technically should return Float64 with millis in the decimal part,
553+
// but we are returning Int32 for compatibility with the original implementation.
554+
interval_month_day_nano_op(array, |v| {
555+
let (_, _, nanos) = IntervalMonthDayNanoType::to_parts(v);
556+
(nanos % 60_000_000_000 / 1_000_000_000)
557+
.try_into()
558+
.map_err(|_| {
559+
ArrowError::ComputeError(
560+
"Unable to convert i64 nanos to i32".to_string(),
561+
)
562+
})
563+
})
564+
}
565+
}

‎datafusion/expr/src/function.rs‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -123,14 +123,7 @@ pub fn return_type(
123123
}
124124
BuiltinScalarFunction::Concat => Ok(DataType::Utf8),
125125
BuiltinScalarFunction::ConcatWithSeparator => Ok(DataType::Utf8),
126-
BuiltinScalarFunction::DatePart => {
127-
match &input_expr_types[1] {
128-
// FIXME: DatePart should *always* return a numeric but this might break things
129-
// so since interval wasn't supported in the first place, this is safe
130-
DataType::Interval(_) => Ok(DataType::Float64),
131-
_ => Ok(DataType::Int32),
132-
}
133-
}
126+
BuiltinScalarFunction::DatePart => Ok(DataType::Float64),
134127
BuiltinScalarFunction::DateTrunc => {
135128
Ok(DataType::Timestamp(TimeUnit::Nanosecond, None))
136129
}
@@ -442,6 +435,14 @@ pub fn signature(fun: &BuiltinScalarFunction) -> Signature {
442435
DataType::Utf8,
443436
DataType::Timestamp(TimeUnit::Nanosecond, None),
444437
]),
438+
TypeSignature::Exact(vec![
439+
DataType::Utf8,
440+
DataType::Interval(IntervalUnit::YearMonth),
441+
]),
442+
TypeSignature::Exact(vec![
443+
DataType::Utf8,
444+
DataType::Interval(IntervalUnit::DayTime),
445+
]),
445446
TypeSignature::Exact(vec![
446447
DataType::Utf8,
447448
DataType::Interval(IntervalUnit::MonthDayNano),

0 commit comments

Comments
 (0)