Skip to content

feat: Support DATE_PART with intervals #182

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 1, 2025
Merged

feat: Support DATE_PART with intervals #182

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions datafusion/core/tests/sql/expr.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1124,11 +1124,19 @@ async fn test_extract_date_part() -> Result<()> {
"EXTRACT(year FROM to_timestamp('2020-09-08T12:00:00+00:00'))",
"2020"
);
test_expression!(
"date_part('YEAR', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
"1"
);
test_expression!("date_part('MONTH', CAST('2000-01-01' AS DATE))", "1");
test_expression!(
"EXTRACT(month FROM to_timestamp('2020-09-08T12:00:00+00:00'))",
"9"
);
test_expression!(
"date_part('MONTH', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
"2"
);
test_expression!("date_part('WEEK', CAST('2003-01-01' AS DATE))", "1");

// TODO Creating logical plan for 'SELECT EXTRACT(WEEK FROM to_timestamp('2020-09-08T12:00:00+00:00'))'
Expand All @@ -1144,11 +1152,19 @@ async fn test_extract_date_part() -> Result<()> {
"EXTRACT(day FROM to_timestamp('2020-09-08T12:00:00+00:00'))",
"8"
);
test_expression!(
"date_part('DAY', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
"3"
);
test_expression!("date_part('HOUR', CAST('2000-01-01' AS DATE))", "0");
test_expression!(
"EXTRACT(hour FROM to_timestamp('2020-09-08T12:03:03+00:00'))",
"12"
);
test_expression!(
"date_part('HOUR', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
"4"
);
test_expression!(
"EXTRACT(minute FROM to_timestamp('2020-09-08T12:12:00+00:00'))",
"12"
Expand All @@ -1157,6 +1173,10 @@ async fn test_extract_date_part() -> Result<()> {
"date_part('minute', to_timestamp('2020-09-08T12:12:00+00:00'))",
"12"
);
test_expression!(
"date_part('MINUTE', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
"5"
);
test_expression!(
"EXTRACT(second FROM to_timestamp('2020-09-08T12:00:12+00:00'))",
"12"
Expand All @@ -1165,6 +1185,10 @@ async fn test_extract_date_part() -> Result<()> {
"date_part('second', to_timestamp('2020-09-08T12:00:12+00:00'))",
"12"
);
test_expression!(
"date_part('SECOND', INTERVAL '1 year 2 month 3 day 4 hour 5 minute 6 second')",
"6"
);

// DOY
test_expression!(
Expand Down
272 changes: 270 additions & 2 deletions datafusion/cube_ext/src/temporal.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,13 @@
// specific language governing permissions and limitations
// under the License.

use arrow::array::{Array, Float64Array, Int32Array, Int32Builder, PrimitiveArray};
use arrow::compute::kernels::arity::unary;
use arrow::array::{
Array, Date32Array, Date64Array, Float64Array, Int32Array, Int32Builder,
IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray,
PrimitiveArray, TimestampMicrosecondArray, TimestampMillisecondArray,
TimestampNanosecondArray, TimestampSecondArray,
};
use arrow::compute::kernels::{arity::unary, temporal as arrow_temporal};
use arrow::datatypes::{
ArrowNumericType, ArrowPrimitiveType, ArrowTemporalType, DataType, Date32Type,
Date64Type, Float64Type, IntervalDayTimeType, IntervalMonthDayNanoType,
Expand Down Expand Up @@ -295,3 +300,266 @@ where

Ok(b.finish())
}

/// This macro will generate a trait `DatePartable` that is automatically implemented
/// for all Arrow temporal types.
macro_rules! date_partable {
($($gran:ident),*) => {
pub trait DatePartable: ArrowPrimitiveType + Sized {
$(
fn $gran(array: &PrimitiveArray<Self>) -> Result<Int32Array>;
)*
}

impl DatePartable for TimestampSecondType {
$(
fn $gran(array: &TimestampSecondArray) -> Result<Int32Array> {
arrow_temporal::$gran(array)
}
)*
}

impl DatePartable for TimestampMillisecondType {
$(
fn $gran(array: &TimestampMillisecondArray) -> Result<Int32Array> {
arrow_temporal::$gran(array)
}
)*
}

impl DatePartable for TimestampMicrosecondType {
$(
fn $gran(array: &TimestampMicrosecondArray) -> Result<Int32Array> {
arrow_temporal::$gran(array)
}
)*
}

impl DatePartable for TimestampNanosecondType {
$(
fn $gran(array: &TimestampNanosecondArray) -> Result<Int32Array> {
arrow_temporal::$gran(array)
}
)*
}

impl DatePartable for Date32Type {
$(
fn $gran(array: &Date32Array) -> Result<Int32Array> {
arrow_temporal::$gran(array)
}
)*
}

impl DatePartable for Date64Type {
$(
fn $gran(array: &Date64Array) -> Result<Int32Array> {
arrow_temporal::$gran(array)
}
)*
}

$(
pub fn $gran<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
where
T: DatePartable,
{
DatePartable::$gran(array)
}
)*
};
}

date_partable!(year, quarter, month, day, hour, minute, second);

fn interval_year_month_op(
array: &IntervalYearMonthArray,
op: fn(i32) -> Result<i32>,
) -> Result<Int32Array> {
let mut builder = Int32Builder::new(array.len());
for i in 0..array.len() {
if array.is_null(i) {
builder.append_null()?;
continue;
}
let value = array.value(i);
let result = op(value)?;
builder.append_value(result)?;
}
Ok(builder.finish())
}

fn interval_day_time_op(
array: &IntervalDayTimeArray,
op: fn(i64) -> Result<i32>,
) -> Result<Int32Array> {
let mut builder = Int32Builder::new(array.len());
for i in 0..array.len() {
if array.is_null(i) {
builder.append_null()?;
continue;
}
let value = array.value(i);
let result = op(value)?;
builder.append_value(result)?;
}
Ok(builder.finish())
}

fn interval_month_day_nano_op(
array: &IntervalMonthDayNanoArray,
op: fn(i128) -> Result<i32>,
) -> Result<Int32Array> {
let mut builder = Int32Builder::new(array.len());
for i in 0..array.len() {
if array.is_null(i) {
builder.append_null()?;
continue;
}
let value = array.value(i);
let result = op(value)?;
builder.append_value(result)?;
}
Ok(builder.finish())
}

impl DatePartable for IntervalYearMonthType {
fn year(array: &IntervalYearMonthArray) -> Result<Int32Array> {
interval_year_month_op(array, |v| Ok(v / 12))
}

fn quarter(array: &IntervalYearMonthArray) -> Result<Int32Array> {
interval_year_month_op(array, |v| Ok(v % 12 / 3 + 1))
}

fn month(array: &IntervalYearMonthArray) -> Result<Int32Array> {
interval_year_month_op(array, |v| Ok(v % 12))
}

fn day(array: &IntervalYearMonthArray) -> Result<Int32Array> {
interval_year_month_op(array, |_| Ok(0))
}

fn hour(array: &IntervalYearMonthArray) -> Result<Int32Array> {
interval_year_month_op(array, |_| Ok(0))
}

fn minute(array: &IntervalYearMonthArray) -> Result<Int32Array> {
interval_year_month_op(array, |_| Ok(0))
}

fn second(array: &IntervalYearMonthArray) -> Result<Int32Array> {
interval_year_month_op(array, |_| Ok(0))
}
}

impl DatePartable for IntervalDayTimeType {
fn year(array: &IntervalDayTimeArray) -> Result<Int32Array> {
interval_day_time_op(array, |_| Ok(0))
}

fn quarter(array: &IntervalDayTimeArray) -> Result<Int32Array> {
interval_day_time_op(array, |_| Ok(1))
}

fn month(array: &IntervalDayTimeArray) -> Result<Int32Array> {
interval_day_time_op(array, |_| Ok(0))
}

fn day(array: &IntervalDayTimeArray) -> Result<Int32Array> {
interval_day_time_op(array, |v| {
let (days, _) = IntervalDayTimeType::to_parts(v);
Ok(days)
})
}

fn hour(array: &IntervalDayTimeArray) -> Result<Int32Array> {
interval_day_time_op(array, |v| {
let (_, millis) = IntervalDayTimeType::to_parts(v);
Ok(millis / 3_600_000)
})
}

fn minute(array: &IntervalDayTimeArray) -> Result<Int32Array> {
interval_day_time_op(array, |v| {
let (_, millis) = IntervalDayTimeType::to_parts(v);
Ok(millis % 3_600_000 / 60_000)
})
}

fn second(array: &IntervalDayTimeArray) -> Result<Int32Array> {
// NOTE: this technically should return Float64 with millis in the decimal part,
// but we are returning Int32 for compatibility with the original implementation.
interval_day_time_op(array, |v| {
let (_, millis) = IntervalDayTimeType::to_parts(v);
Ok(millis % 60_000 / 1_000)
})
}
}

impl DatePartable for IntervalMonthDayNanoType {
fn year(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
interval_month_day_nano_op(array, |v| {
let (months, _, _) = IntervalMonthDayNanoType::to_parts(v);
Ok(months / 12)
})
}

fn quarter(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
interval_month_day_nano_op(array, |v| {
let (months, _, _) = IntervalMonthDayNanoType::to_parts(v);
Ok(months % 12 / 3 + 1)
})
}

fn month(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
interval_month_day_nano_op(array, |v| {
let (months, _, _) = IntervalMonthDayNanoType::to_parts(v);
Ok(months % 12)
})
}

fn day(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
interval_month_day_nano_op(array, |v| {
let (_, days, _) = IntervalMonthDayNanoType::to_parts(v);
Ok(days)
})
}

fn hour(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
interval_month_day_nano_op(array, |v| {
let (_, _, nanos) = IntervalMonthDayNanoType::to_parts(v);
(nanos / 3_600_000_000_000).try_into().map_err(|_| {
ArrowError::ComputeError("Unable to convert i64 nanos to i32".to_string())
})
})
}

fn minute(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
interval_month_day_nano_op(array, |v| {
let (_, _, nanos) = IntervalMonthDayNanoType::to_parts(v);
(nanos % 3_600_000_000_000 / 60_000_000_000)
.try_into()
.map_err(|_| {
ArrowError::ComputeError(
"Unable to convert i64 nanos to i32".to_string(),
)
})
})
}

fn second(array: &IntervalMonthDayNanoArray) -> Result<Int32Array> {
// NOTE: this technically should return Float64 with millis in the decimal part,
// but we are returning Int32 for compatibility with the original implementation.
interval_month_day_nano_op(array, |v| {
let (_, _, nanos) = IntervalMonthDayNanoType::to_parts(v);
(nanos % 60_000_000_000 / 1_000_000_000)
.try_into()
.map_err(|_| {
ArrowError::ComputeError(
"Unable to convert i64 nanos to i32".to_string(),
)
})
})
}
}
17 changes: 9 additions & 8 deletions datafusion/expr/src/function.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,14 +123,7 @@ pub fn return_type(
}
BuiltinScalarFunction::Concat => Ok(DataType::Utf8),
BuiltinScalarFunction::ConcatWithSeparator => Ok(DataType::Utf8),
BuiltinScalarFunction::DatePart => {
match &input_expr_types[1] {
// FIXME: DatePart should *always* return a numeric but this might break things
// so since interval wasn't supported in the first place, this is safe
DataType::Interval(_) => Ok(DataType::Float64),
_ => Ok(DataType::Int32),
}
}
BuiltinScalarFunction::DatePart => Ok(DataType::Float64),
BuiltinScalarFunction::DateTrunc => {
Ok(DataType::Timestamp(TimeUnit::Nanosecond, None))
}
Expand Down Expand Up @@ -442,6 +435,14 @@ pub fn signature(fun: &BuiltinScalarFunction) -> Signature {
DataType::Utf8,
DataType::Timestamp(TimeUnit::Nanosecond, None),
]),
TypeSignature::Exact(vec![
DataType::Utf8,
DataType::Interval(IntervalUnit::YearMonth),
]),
TypeSignature::Exact(vec![
DataType::Utf8,
DataType::Interval(IntervalUnit::DayTime),
]),
TypeSignature::Exact(vec![
DataType::Utf8,
DataType::Interval(IntervalUnit::MonthDayNano),
Expand Down
Loading
Loading