From fa511d80da82a7b2fc9f9abd84ab9476844b87c4 Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Thu, 18 Sep 2025 21:08:24 -0400 Subject: [PATCH] feat(substrait): add time literal support Adds support for `ScalarValue::Time64Microsecond` and `ScalarValue::Time64Nanosecond` to be converted to and from Substrait literals. This includes the `PrecisionTime` literal type and specific `TIME_64_TYPE_VARIATION_REF` for 6-digit (microseconds) and 9-digit (nanoseconds) precision. --- .../src/logical_plan/consumer/expr/literal.rs | 27 +++++++++++++++++-- .../src/logical_plan/producer/expr/literal.rs | 26 +++++++++++++++--- 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs b/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs index d054e5267554..b8181f3d05cb 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs @@ -25,8 +25,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, - VIEW_CONTAINER_TYPE_VARIATION_REF, + TIMESTAMP_SECOND_TYPE_VARIATION_REF, TIME_64_TYPE_VARIATION_REF, + UNSIGNED_INTEGER_TYPE_VARIATION_REF, VIEW_CONTAINER_TYPE_VARIATION_REF, }; use datafusion::arrow::array::{new_empty_array, AsArray, MapArray}; use datafusion::arrow::buffer::OffsetBuffer; @@ -155,6 +155,29 @@ pub(crate) fn from_substrait_literal( } }, Some(LiteralType::Date(d)) => ScalarValue::Date32(Some(*d)), + Some(LiteralType::PrecisionTime(pt)) => match pt.precision { + 6 => match lit.type_variation_reference { + TIME_64_TYPE_VARIATION_REF => { + ScalarValue::Time64Microsecond(Some(pt.value)) + } + others => { + return substrait_err!("Unknown type variation reference {others}"); + } + }, + 9 => match lit.type_variation_reference { + TIME_64_TYPE_VARIATION_REF => { + ScalarValue::Time64Nanosecond(Some(pt.value)) + } + others => { + return substrait_err!("Unknown type variation reference {others}"); + } + }, + p => { + return not_impl_err!( + "Unsupported Substrait precision {p} for PrecisionTime" + ); + } + }, Some(LiteralType::String(s)) => match lit.type_variation_reference { DEFAULT_CONTAINER_TYPE_VARIATION_REF => ScalarValue::Utf8(Some(s.clone())), LARGE_CONTAINER_TYPE_VARIATION_REF => ScalarValue::LargeUtf8(Some(s.clone())), diff --git a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs index 31f4866bdc85..62ad7ac1f34e 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs @@ -19,8 +19,8 @@ use crate::logical_plan::producer::{to_substrait_type, SubstraitProducer}; use crate::variation_const::{ DATE_32_TYPE_VARIATION_REF, DECIMAL_128_TYPE_VARIATION_REF, DEFAULT_CONTAINER_TYPE_VARIATION_REF, DEFAULT_TYPE_VARIATION_REF, - LARGE_CONTAINER_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF, - VIEW_CONTAINER_TYPE_VARIATION_REF, + LARGE_CONTAINER_TYPE_VARIATION_REF, TIME_64_TYPE_VARIATION_REF, + UNSIGNED_INTEGER_TYPE_VARIATION_REF, VIEW_CONTAINER_TYPE_VARIATION_REF, }; use datafusion::arrow::array::{Array, GenericListArray, OffsetSizeTrait}; use datafusion::arrow::temporal_conversions::NANOSECONDS; @@ -29,7 +29,7 @@ use substrait::proto::expression::literal::interval_day_to_second::PrecisionMode use substrait::proto::expression::literal::map::KeyValue; use substrait::proto::expression::literal::{ Decimal, IntervalCompound, IntervalDayToSecond, IntervalYearToMonth, List, - LiteralType, Map, PrecisionTimestamp, Struct, + LiteralType, Map, PrecisionTime, PrecisionTimestamp, Struct, }; use substrait::proto::expression::{Literal, RexType}; use substrait::proto::{r#type, Expression}; @@ -280,6 +280,20 @@ pub(crate) fn to_substrait_literal( }; (map, DEFAULT_CONTAINER_TYPE_VARIATION_REF) } + ScalarValue::Time64Microsecond(Some(t)) => ( + LiteralType::PrecisionTime(PrecisionTime { + precision: 6, + value: *t, + }), + TIME_64_TYPE_VARIATION_REF, + ), + ScalarValue::Time64Nanosecond(Some(t)) => ( + LiteralType::PrecisionTime(PrecisionTime { + precision: 9, + value: *t, + }), + TIME_64_TYPE_VARIATION_REF, + ), ScalarValue::Struct(s) => ( LiteralType::Struct(Struct { fields: s @@ -398,6 +412,12 @@ mod tests { round_trip_literal(ScalarValue::TimestampNanosecond(ts, tz))?; } + // Test Time64 literals + round_trip_literal(ScalarValue::Time64Microsecond(Some(45296789123)))?; + round_trip_literal(ScalarValue::Time64Microsecond(None))?; + round_trip_literal(ScalarValue::Time64Nanosecond(Some(45296789123000)))?; + round_trip_literal(ScalarValue::Time64Nanosecond(None))?; + round_trip_literal(ScalarValue::List(ScalarValue::new_list_nullable( &[ScalarValue::Float32(Some(1.0))], &DataType::Float32,