Skip to content

Commit a8d5016

Browse files
authored
Support remaining substrait time literal variations (#17707)
1 parent 95357ed commit a8d5016

File tree

2 files changed

+42
-4
lines changed

2 files changed

+42
-4
lines changed

datafusion/substrait/src/logical_plan/consumer/expr/literal.rs

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,9 @@ use crate::variation_const::{
2525
INTERVAL_MONTH_DAY_NANO_TYPE_REF, INTERVAL_YEAR_MONTH_TYPE_REF,
2626
LARGE_CONTAINER_TYPE_VARIATION_REF, TIMESTAMP_MICRO_TYPE_VARIATION_REF,
2727
TIMESTAMP_MILLI_TYPE_VARIATION_REF, TIMESTAMP_NANO_TYPE_VARIATION_REF,
28-
TIMESTAMP_SECOND_TYPE_VARIATION_REF, TIME_64_TYPE_VARIATION_REF,
29-
UNSIGNED_INTEGER_TYPE_VARIATION_REF, VIEW_CONTAINER_TYPE_VARIATION_REF,
28+
TIMESTAMP_SECOND_TYPE_VARIATION_REF, TIME_32_TYPE_VARIATION_REF,
29+
TIME_64_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
30+
VIEW_CONTAINER_TYPE_VARIATION_REF,
3031
};
3132
use datafusion::arrow::array::{new_empty_array, AsArray, MapArray};
3233
use datafusion::arrow::buffer::OffsetBuffer;
@@ -156,6 +157,22 @@ pub(crate) fn from_substrait_literal(
156157
},
157158
Some(LiteralType::Date(d)) => ScalarValue::Date32(Some(*d)),
158159
Some(LiteralType::PrecisionTime(pt)) => match pt.precision {
160+
0 => match lit.type_variation_reference {
161+
TIME_32_TYPE_VARIATION_REF => {
162+
ScalarValue::Time32Second(Some(pt.value as i32))
163+
}
164+
others => {
165+
return substrait_err!("Unknown type variation reference {others}");
166+
}
167+
},
168+
3 => match lit.type_variation_reference {
169+
TIME_32_TYPE_VARIATION_REF => {
170+
ScalarValue::Time32Millisecond(Some(pt.value as i32))
171+
}
172+
others => {
173+
return substrait_err!("Unknown type variation reference {others}");
174+
}
175+
},
159176
6 => match lit.type_variation_reference {
160177
TIME_64_TYPE_VARIATION_REF => {
161178
ScalarValue::Time64Microsecond(Some(pt.value))

datafusion/substrait/src/logical_plan/producer/expr/literal.rs

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,9 @@ use crate::logical_plan::producer::{to_substrait_type, SubstraitProducer};
1919
use crate::variation_const::{
2020
DATE_32_TYPE_VARIATION_REF, DECIMAL_128_TYPE_VARIATION_REF,
2121
DEFAULT_CONTAINER_TYPE_VARIATION_REF, DEFAULT_TYPE_VARIATION_REF,
22-
LARGE_CONTAINER_TYPE_VARIATION_REF, TIME_64_TYPE_VARIATION_REF,
23-
UNSIGNED_INTEGER_TYPE_VARIATION_REF, VIEW_CONTAINER_TYPE_VARIATION_REF,
22+
LARGE_CONTAINER_TYPE_VARIATION_REF, TIME_32_TYPE_VARIATION_REF,
23+
TIME_64_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
24+
VIEW_CONTAINER_TYPE_VARIATION_REF,
2425
};
2526
use datafusion::arrow::array::{Array, GenericListArray, OffsetSizeTrait};
2627
use datafusion::arrow::temporal_conversions::NANOSECONDS;
@@ -280,6 +281,20 @@ pub(crate) fn to_substrait_literal(
280281
};
281282
(map, DEFAULT_CONTAINER_TYPE_VARIATION_REF)
282283
}
284+
ScalarValue::Time32Second(Some(t)) => (
285+
LiteralType::PrecisionTime(PrecisionTime {
286+
precision: 0,
287+
value: *t as i64,
288+
}),
289+
TIME_32_TYPE_VARIATION_REF,
290+
),
291+
ScalarValue::Time32Millisecond(Some(t)) => (
292+
LiteralType::PrecisionTime(PrecisionTime {
293+
precision: 3,
294+
value: *t as i64,
295+
}),
296+
TIME_32_TYPE_VARIATION_REF,
297+
),
283298
ScalarValue::Time64Microsecond(Some(t)) => (
284299
LiteralType::PrecisionTime(PrecisionTime {
285300
precision: 6,
@@ -412,6 +427,12 @@ mod tests {
412427
round_trip_literal(ScalarValue::TimestampNanosecond(ts, tz))?;
413428
}
414429

430+
// Test Time32 literals
431+
round_trip_literal(ScalarValue::Time32Second(Some(45296)))?;
432+
round_trip_literal(ScalarValue::Time32Second(None))?;
433+
round_trip_literal(ScalarValue::Time32Millisecond(Some(45296789)))?;
434+
round_trip_literal(ScalarValue::Time32Millisecond(None))?;
435+
415436
// Test Time64 literals
416437
round_trip_literal(ScalarValue::Time64Microsecond(Some(45296789123)))?;
417438
round_trip_literal(ScalarValue::Time64Microsecond(None))?;

0 commit comments

Comments
 (0)