Add support for Arrow Time types in Substrait (#16558)

* Add support for Arrow Time types in Substrait

This commit adds support for Arrow Time types Time32 and Time64 in
Substrait plans.

Resolves #16296
Resolves #16275

* Clean up test

Author: jkosh44

datafusion/substrait/src/logical_plan/consumer/types.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use super::utils::{next_struct_field_name, DEFAULT_TIMEZONE};
+use super::utils::{from_substrait_precision, next_struct_field_name, DEFAULT_TIMEZONE};
 use super::SubstraitConsumer;
 #[allow(deprecated)]
 use crate::variation_const::{
@@ -27,7 +27,8 @@ use crate::variation_const::{
     INTERVAL_MONTH_DAY_NANO_TYPE_REF, INTERVAL_YEAR_MONTH_TYPE_REF,
     LARGE_CONTAINER_TYPE_VARIATION_REF, TIMESTAMP_MICRO_TYPE_VARIATION_REF,
     TIMESTAMP_MILLI_TYPE_VARIATION_REF, TIMESTAMP_NANO_TYPE_VARIATION_REF,
-    TIMESTAMP_SECOND_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
+    TIMESTAMP_SECOND_TYPE_VARIATION_REF, TIME_32_TYPE_VARIATION_REF,
+    TIME_64_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
     VIEW_CONTAINER_TYPE_VARIATION_REF,
 };
 use datafusion::arrow::datatypes::{
@@ -107,29 +108,24 @@ pub fn from_substrait_type(
                 }
             }
             r#type::Kind::PrecisionTimestamp(pts) => {
-                let unit = match pts.precision {
-                    0 => Ok(TimeUnit::Second),
-                    3 => Ok(TimeUnit::Millisecond),
-                    6 => Ok(TimeUnit::Microsecond),
-                    9 => Ok(TimeUnit::Nanosecond),
-                    p => not_impl_err!(
-                        "Unsupported Substrait precision {p} for PrecisionTimestamp"
-                    ),
-                }?;
+                let unit = from_substrait_precision(pts.precision, "PrecisionTimestamp")?;
                 Ok(DataType::Timestamp(unit, None))
             }
             r#type::Kind::PrecisionTimestampTz(pts) => {
-                let unit = match pts.precision {
-                    0 => Ok(TimeUnit::Second),
-                    3 => Ok(TimeUnit::Millisecond),
-                    6 => Ok(TimeUnit::Microsecond),
-                    9 => Ok(TimeUnit::Nanosecond),
-                    p => not_impl_err!(
-                        "Unsupported Substrait precision {p} for PrecisionTimestampTz"
-                    ),
-                }?;
+                let unit =
+                    from_substrait_precision(pts.precision, "PrecisionTimestampTz")?;
                 Ok(DataType::Timestamp(unit, Some(DEFAULT_TIMEZONE.into())))
             }
+            r#type::Kind::PrecisionTime(pt) => {
+                let time_unit = from_substrait_precision(pt.precision, "PrecisionTime")?;
+                match pt.type_variation_reference {
+                    TIME_32_TYPE_VARIATION_REF => Ok(DataType::Time32(time_unit)),
+                    TIME_64_TYPE_VARIATION_REF => Ok(DataType::Time64(time_unit)),
+                    v => not_impl_err!(
+                        "Unsupported Substrait type variation {v} of type {s_kind:?}"
+                    ),
+                }
+            }
             r#type::Kind::Date(date) => match date.type_variation_reference {
                 DATE_32_TYPE_VARIATION_REF => Ok(DataType::Date32),
                 DATE_64_TYPE_VARIATION_REF => Ok(DataType::Date64),
@@ -220,14 +216,9 @@ pub fn from_substrait_type(
                 }
                 DURATION_INTERVAL_DAY_TYPE_VARIATION_REF => {
                     let duration_unit = match i.precision {
-                        Some(0) => Ok(TimeUnit::Second),
-                        Some(3) => Ok(TimeUnit::Millisecond),
-                        Some(6) => Ok(TimeUnit::Microsecond),
-                        Some(9) => Ok(TimeUnit::Nanosecond),
-                        p => {
-                            not_impl_err!(
-                                "Unsupported Substrait precision {p:?} for Duration"
-                            )
+                        Some(p) => from_substrait_precision(p, "Duration"),
+                        None => {
+                            not_impl_err!("Missing Substrait precision for Duration")
                         }
                     }?;
                     Ok(DataType::Duration(duration_unit))

datafusion/substrait/src/logical_plan/consumer/utils.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::logical_plan::consumer::SubstraitConsumer;
-use datafusion::arrow::datatypes::{DataType, Field, Schema, UnionFields};
+use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UnionFields};
 use datafusion::common::{
     exec_err, not_impl_err, substrait_datafusion_err, substrait_err, DFSchema,
     DFSchemaRef, TableReference,
@@ -476,6 +476,21 @@ pub async fn from_substrait_sorts(
     Ok(sorts)
 }
 
+pub(crate) fn from_substrait_precision(
+    precision: i32,
+    type_name: &str,
+) -> datafusion::common::Result<TimeUnit> {
+    match precision {
+        0 => Ok(TimeUnit::Second),
+        3 => Ok(TimeUnit::Millisecond),
+        6 => Ok(TimeUnit::Microsecond),
+        9 => Ok(TimeUnit::Nanosecond),
+        precision => {
+            not_impl_err!("Unsupported Substrait precision {precision}, for {type_name}")
+        }
+    }
+}
+
 #[cfg(test)]
 pub(crate) mod tests {
     use super::make_renamed_schema;

datafusion/substrait/src/logical_plan/producer/types.rs

@@ -15,16 +15,18 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::logical_plan::producer::to_substrait_precision;
 use crate::logical_plan::producer::utils::flatten_names;
 use crate::variation_const::{
     DATE_32_TYPE_VARIATION_REF, DATE_64_TYPE_VARIATION_REF,
     DECIMAL_128_TYPE_VARIATION_REF, DECIMAL_256_TYPE_VARIATION_REF,
     DEFAULT_CONTAINER_TYPE_VARIATION_REF, DEFAULT_INTERVAL_DAY_TYPE_VARIATION_REF,
     DEFAULT_TYPE_VARIATION_REF, DURATION_INTERVAL_DAY_TYPE_VARIATION_REF,
-    LARGE_CONTAINER_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
+    LARGE_CONTAINER_TYPE_VARIATION_REF, TIME_32_TYPE_VARIATION_REF,
+    TIME_64_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
     VIEW_CONTAINER_TYPE_VARIATION_REF,
 };
-use datafusion::arrow::datatypes::{DataType, IntervalUnit, TimeUnit};
+use datafusion::arrow::datatypes::{DataType, IntervalUnit};
 use datafusion::common::{internal_err, not_impl_err, plan_err, DFSchemaRef};
 use substrait::proto::{r#type, NamedStruct};
 
@@ -107,12 +109,7 @@ pub(crate) fn to_substrait_type(
             })),
         }),
         DataType::Timestamp(unit, tz) => {
-            let precision = match unit {
-                TimeUnit::Second => 0,
-                TimeUnit::Millisecond => 3,
-                TimeUnit::Microsecond => 6,
-                TimeUnit::Nanosecond => 9,
-            };
+            let precision = to_substrait_precision(unit);
             let kind = match tz {
                 None => r#type::Kind::PrecisionTimestamp(r#type::PrecisionTimestamp {
                     type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
@@ -132,6 +129,26 @@ pub(crate) fn to_substrait_type(
             };
             Ok(substrait::proto::Type { kind: Some(kind) })
         }
+        DataType::Time32(unit) => {
+            let precision = to_substrait_precision(unit);
+            Ok(substrait::proto::Type {
+                kind: Some(r#type::Kind::PrecisionTime(r#type::PrecisionTime {
+                    precision,
+                    type_variation_reference: TIME_32_TYPE_VARIATION_REF,
+                    nullability,
+                })),
+            })
+        }
+        DataType::Time64(unit) => {
+            let precision = to_substrait_precision(unit);
+            Ok(substrait::proto::Type {
+                kind: Some(r#type::Kind::PrecisionTime(r#type::PrecisionTime {
+                    precision,
+                    type_variation_reference: TIME_64_TYPE_VARIATION_REF,
+                    nullability,
+                })),
+            })
+        }
         DataType::Date32 => Ok(substrait::proto::Type {
             kind: Some(r#type::Kind::Date(r#type::Date {
                 type_variation_reference: DATE_32_TYPE_VARIATION_REF,
@@ -173,12 +190,7 @@ pub(crate) fn to_substrait_type(
             }
         }
         DataType::Duration(duration_unit) => {
-            let precision = match duration_unit {
-                TimeUnit::Second => 0,
-                TimeUnit::Millisecond => 3,
-                TimeUnit::Microsecond => 6,
-                TimeUnit::Nanosecond => 9,
-            };
+            let precision = to_substrait_precision(duration_unit);
             Ok(substrait::proto::Type {
                 kind: Some(r#type::Kind::IntervalDay(r#type::IntervalDay {
                     type_variation_reference: DURATION_INTERVAL_DAY_TYPE_VARIATION_REF,
@@ -335,7 +347,7 @@ mod tests {
     use crate::logical_plan::consumer::{
         from_substrait_named_struct, from_substrait_type_without_names,
     };
-    use datafusion::arrow::datatypes::{Field, Fields, Schema};
+    use datafusion::arrow::datatypes::{Field, Fields, Schema, TimeUnit};
     use datafusion::common::{DFSchema, Result};
     use std::sync::Arc;
 
@@ -360,6 +372,10 @@ mod tests {
             round_trip_type(DataType::Timestamp(TimeUnit::Nanosecond, tz))?;
         }
 
+        round_trip_type(DataType::Time32(TimeUnit::Second))?;
+        round_trip_type(DataType::Time32(TimeUnit::Millisecond))?;
+        round_trip_type(DataType::Time64(TimeUnit::Microsecond))?;
+        round_trip_type(DataType::Time64(TimeUnit::Nanosecond))?;
         round_trip_type(DataType::Date32)?;
         round_trip_type(DataType::Date64)?;
         round_trip_type(DataType::Binary)?;

datafusion/substrait/src/logical_plan/producer/utils.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::logical_plan::producer::SubstraitProducer;
-use datafusion::arrow::datatypes::{DataType, Field};
+use datafusion::arrow::datatypes::{DataType, Field, TimeUnit};
 use datafusion::common::{plan_err, DFSchemaRef};
 use datafusion::logical_expr::SortExpr;
 use substrait::proto::sort_field::{SortDirection, SortKind};
@@ -76,3 +76,12 @@ pub(crate) fn substrait_sort_field(
         sort_kind: Some(SortKind::Direction(d as i32)),
     })
 }
+
+pub(crate) fn to_substrait_precision(time_unit: &TimeUnit) -> i32 {
+    match time_unit {
+        TimeUnit::Second => 0,
+        TimeUnit::Millisecond => 3,
+        TimeUnit::Microsecond => 6,
+        TimeUnit::Nanosecond => 9,
+    }
+}

datafusion/substrait/src/variation_const.rs

@@ -50,6 +50,8 @@ pub const TIMESTAMP_NANO_TYPE_VARIATION_REF: u32 = 3;
 
 pub const DATE_32_TYPE_VARIATION_REF: u32 = 0;
 pub const DATE_64_TYPE_VARIATION_REF: u32 = 1;
+pub const TIME_32_TYPE_VARIATION_REF: u32 = 0;
+pub const TIME_64_TYPE_VARIATION_REF: u32 = 1;
 pub const DEFAULT_CONTAINER_TYPE_VARIATION_REF: u32 = 0;
 pub const LARGE_CONTAINER_TYPE_VARIATION_REF: u32 = 1;
 pub const VIEW_CONTAINER_TYPE_VARIATION_REF: u32 = 2;