Skip to content

Commit

Permalink
Support for extract(x from time) / date_part from time types (#8693)
Browse files Browse the repository at this point in the history
* Initial support for `extract(x from time)`

* Update function docs

* Add extract tests
  • Loading branch information
Jefffrey authored Mar 21, 2024
1 parent 1d8a41b commit dc373a3
Show file tree
Hide file tree
Showing 4 changed files with 345 additions and 37 deletions.
37 changes: 29 additions & 8 deletions datafusion/common/src/cast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,17 +24,18 @@ use crate::{downcast_value, DataFusionError, Result};
use arrow::{
array::{
Array, BinaryArray, BooleanArray, Date32Array, Date64Array, Decimal128Array,
DictionaryArray, FixedSizeBinaryArray, FixedSizeListArray, Float32Array,
Float64Array, GenericBinaryArray, GenericListArray, GenericStringArray,
Int32Array, Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray,
IntervalYearMonthArray, LargeListArray, ListArray, MapArray, NullArray,
OffsetSizeTrait, PrimitiveArray, StringArray, StructArray,
TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
TimestampSecondArray, UInt32Array, UInt64Array, UInt8Array, UnionArray,
Decimal256Array, DictionaryArray, FixedSizeBinaryArray, FixedSizeListArray,
Float32Array, Float64Array, GenericBinaryArray, GenericListArray,
GenericStringArray, Int32Array, Int64Array, IntervalDayTimeArray,
IntervalMonthDayNanoArray, IntervalYearMonthArray, LargeListArray, ListArray,
MapArray, NullArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray,
Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray,
Time64NanosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray,
TimestampNanosecondArray, TimestampSecondArray, UInt32Array, UInt64Array,
UInt8Array, UnionArray,
},
datatypes::{ArrowDictionaryKeyType, ArrowPrimitiveType},
};
use arrow_array::Decimal256Array;

// Downcast ArrayRef to Date32Array
pub fn as_date32_array(array: &dyn Array) -> Result<&Date32Array> {
Expand Down Expand Up @@ -154,6 +155,26 @@ pub fn as_union_array(array: &dyn Array) -> Result<&UnionArray> {
Ok(downcast_value!(array, UnionArray))
}

// Downcast ArrayRef to Time32SecondArray
pub fn as_time32_second_array(array: &dyn Array) -> Result<&Time32SecondArray> {
Ok(downcast_value!(array, Time32SecondArray))
}

// Downcast ArrayRef to Time32MillisecondArray
pub fn as_time32_millisecond_array(array: &dyn Array) -> Result<&Time32MillisecondArray> {
Ok(downcast_value!(array, Time32MillisecondArray))
}

// Downcast ArrayRef to Time64MicrosecondArray
pub fn as_time64_microsecond_array(array: &dyn Array) -> Result<&Time64MicrosecondArray> {
Ok(downcast_value!(array, Time64MicrosecondArray))
}

// Downcast ArrayRef to Time64NanosecondArray
pub fn as_time64_nanosecond_array(array: &dyn Array) -> Result<&Time64NanosecondArray> {
Ok(downcast_value!(array, Time64NanosecondArray))
}

// Downcast ArrayRef to TimestampNanosecondArray
pub fn as_timestamp_nanosecond_array(
array: &dyn Array,
Expand Down
31 changes: 23 additions & 8 deletions datafusion/functions/src/datetime/date_part.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,14 +20,17 @@ use std::sync::Arc;

use arrow::array::{Array, ArrayRef, Float64Array};
use arrow::compute::{binary, cast, date_part, DatePart};
use arrow::datatypes::DataType::{Date32, Date64, Float64, Timestamp, Utf8};
use arrow::datatypes::DataType::{
Date32, Date64, Float64, Time32, Time64, Timestamp, Utf8,
};
use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second};
use arrow::datatypes::{DataType, TimeUnit};

use datafusion_common::cast::{
as_date32_array, as_date64_array, as_int32_array, as_timestamp_microsecond_array,
as_timestamp_millisecond_array, as_timestamp_nanosecond_array,
as_timestamp_second_array,
as_date32_array, as_date64_array, as_int32_array, as_time32_millisecond_array,
as_time32_second_array, as_time64_microsecond_array, as_time64_nanosecond_array,
as_timestamp_microsecond_array, as_timestamp_millisecond_array,
as_timestamp_nanosecond_array, as_timestamp_second_array,
};
use datafusion_common::{exec_err, Result, ScalarValue};
use datafusion_expr::TypeSignature::Exact;
Expand Down Expand Up @@ -68,6 +71,10 @@ impl DatePartFunc {
]),
Exact(vec![Utf8, Date64]),
Exact(vec![Utf8, Date32]),
Exact(vec![Utf8, Time32(Second)]),
Exact(vec![Utf8, Time32(Millisecond)]),
Exact(vec![Utf8, Time64(Microsecond)]),
Exact(vec![Utf8, Time64(Nanosecond)]),
],
Volatility::Immutable,
),
Expand Down Expand Up @@ -149,12 +156,9 @@ fn date_part_f64(array: &dyn Array, part: DatePart) -> Result<ArrayRef> {
Ok(cast(date_part(array, part)?.as_ref(), &Float64)?)
}

/// invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the
/// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the
/// result to a total number of seconds, milliseconds, microseconds or
/// nanoseconds
///
/// # Panics
/// If `array` is not a temporal type such as Timestamp or Date32
fn seconds(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
let sf = match unit {
Second => 1_f64,
Expand All @@ -163,6 +167,7 @@ fn seconds(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
Nanosecond => 1_000_000_000_f64,
};
let secs = date_part(array, DatePart::Second)?;
// This assumes array is primitive and not a dictionary
let secs = as_int32_array(secs.as_ref())?;
let subsecs = date_part(array, DatePart::Nanosecond)?;
let subsecs = as_int32_array(subsecs.as_ref())?;
Expand All @@ -189,6 +194,16 @@ fn epoch(array: &dyn Array) -> Result<ArrayRef> {
}
Date32 => as_date32_array(array)?.unary(|x| x as f64 * SECONDS_IN_A_DAY),
Date64 => as_date64_array(array)?.unary(|x| x as f64 / 1_000_f64),
Time32(Second) => as_time32_second_array(array)?.unary(|x| x as f64),
Time32(Millisecond) => {
as_time32_millisecond_array(array)?.unary(|x| x as f64 / 1_000_f64)
}
Time64(Microsecond) => {
as_time64_microsecond_array(array)?.unary(|x| x as f64 / 1_000_000_f64)
}
Time64(Nanosecond) => {
as_time64_nanosecond_array(array)?.unary(|x| x as f64 / 1_000_000_000_f64)
}
d => return exec_err!("Can not convert {d:?} to epoch"),
};
Ok(Arc::new(f))
Expand Down
Loading

0 comments on commit dc373a3

Please sign in to comment.