Support for extract(x from time) / date_part from time types

datafusion/common/src/cast.rs

@@ -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> {
@@ -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,

datafusion/functions/src/datetime/date_part.rs

@@ -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;
@@ -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,
             ),
@@ -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,
@@ -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())?;
@@ -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))