From 0e874263f185fb68196dd0569c94fa00ceaedf5f Mon Sep 17 00:00:00 2001 From: Wendell Smith Date: Tue, 2 Jun 2026 15:46:46 -0400 Subject: [PATCH] feat!: add temporal extension literals --- API.md | 7 +- GRAMMAR.md | 19 +- src/extensions/args.rs | 367 ++++++++++++++++++++++++++++- src/extensions/mod.rs | 3 +- src/parser/expression_grammar.pest | 13 +- src/parser/extensions.rs | 80 +++++-- src/parser/literals.rs | 200 ++++++++++++++-- src/textify/extensions.rs | 15 +- src/textify/literals.rs | 139 +++++++++-- tests/extension_roundtrip.rs | 125 +++++++++- 10 files changed, 888 insertions(+), 80 deletions(-) diff --git a/API.md b/API.md index 6c90de27..8937096f 100644 --- a/API.md +++ b/API.md @@ -151,8 +151,11 @@ Untyped scalar extension literals such as `2` or `'path'` are represented as scalar `ExtensionValue` variants and render without expression type suffixes, even in verbose output. The same values can still be requested as `Expr` through `ArgsExtractor`, which widens them to default non-nullable Substrait literal -expressions. Typed literals, field references, function calls, and casts are -represented as expression values. +expressions. Temporal typed literals such as `'2024-01-01':date` or +`'2024-01-01T12:34:56.123456':precisiontimestamp<6>` are represented as +`ExtensionLiteral` values and can also be requested as `Expr`. Other typed +literals, field references, function calls, and casts are represented as +expression values. `ExtensionProtoConvert` converts between extension arguments and Substrait protobuf values in either direction, such as output columns and relation diff --git a/GRAMMAR.md b/GRAMMAR.md index 23326458..7066f456 100644 --- a/GRAMMAR.md +++ b/GRAMMAR.md @@ -239,11 +239,11 @@ A literal can be an integer, float, boolean, string, or null. Literals may inclu - **`null`**` := "null"` - Examples: `null:i64?`, `null:string?`, `null:date?` - A type annotation is required for `null` -- **`typed_literal`**` := string ":" type` - - String literals with type annotations for non-primitive types - - Examples: `'2023-01-01':date`, `'2023-12-25T14:30:45.123':timestamp` +- **typed literal syntax**` := (float / integer / boolean / string / "null") ":" type` + - Literals with explicit type annotations, including string literals for non-primitive types + - Examples: `'2023-01-01':date`, `'2023-12-25T14:30:45.123':timestamp`, `'2023-12-25T14:30:45.123456789':precisiontimestamp<9>` -All basic literal types (`integer`, `float`, `boolean`, and `string`) are supported, plus `date`, `time`, `timestamp`, and typed null literals. Other Substrait literal types (e.g., `interval_year`, `decimal`, `uuid`) are not yet implemented. +All basic literal types (`integer`, `float`, `boolean`, `string`, and typed `null`) are supported, plus `date`, `time`, `timestamp`, `precisiontime`, `precisiontimestamp`, and `precisiontimestamptz` typed literals. Other Substrait literal types (e.g., `interval_year`, `decimal`, `uuid`) are not yet implemented. ## Types @@ -309,6 +309,9 @@ Root[result] ### Compound Types Compound types follow the same syntax as standard Substrait parameterized types. +Precision temporal types use one integer parameter for precision: +`precisiontime`, `precisiontimestamp`, and `precisiontimestamptz`, +where `N` is between `0` and `12`. #### Examples @@ -959,9 +962,11 @@ Untyped scalar extension arguments such as `2`, `2.4`, `true`, and `'path'` are treated as extension scalar values and render without expression type suffixes, even in verbose output. They can still be consumed by extension handlers as expressions, in which case they widen to default non-nullable -Substrait literal expressions. Typed literals such as `2:i16` or -`'2024-01-01':date`, field references, function calls, and casts are expression -values. +Substrait literal expressions. Temporal typed literals such as +`'2024-01-01':date` and +`'2024-01-01T12:34:56.123456':precisiontimestamp<6>` are extension literal +values and can also widen to expression values. Other typed literals such as +`2:i16`, field references, function calls, and casts are expression values. #### Examples diff --git a/src/extensions/args.rs b/src/extensions/args.rs index d617b77c..6cabaeb6 100644 --- a/src/extensions/args.rs +++ b/src/extensions/args.rs @@ -29,10 +29,11 @@ use std::collections::HashSet; use std::fmt; use indexmap::IndexMap; +use prost_types::Timestamp as ProtoTimestamp; use substrait::proto; use substrait::proto::expression::field_reference::ReferenceType; -use substrait::proto::expression::literal::LiteralType; -use substrait::proto::expression::{RexType, reference_segment}; +use substrait::proto::expression::literal::{LiteralType, PrecisionTime, PrecisionTimestamp}; +use substrait::proto::expression::{Literal, RexType, reference_segment}; use super::ExtensionError; use crate::textify::expressions::Reference; @@ -109,8 +110,8 @@ impl From for Expr { } } -impl From for Expr { - fn from(literal: proto::expression::Literal) -> Self { +impl From for Expr { + fn from(literal: Literal) -> Self { proto::Expression { rex_type: Some(RexType::Literal(literal)), } @@ -132,7 +133,7 @@ impl From for proto::Expression { impl From for Expr { fn from(value: i64) -> Self { - proto::expression::Literal { + Literal { literal_type: Some(LiteralType::I64(value)), nullable: false, type_variation_reference: 0, @@ -143,7 +144,7 @@ impl From for Expr { impl From for Expr { fn from(value: f64) -> Self { - proto::expression::Literal { + Literal { literal_type: Some(LiteralType::Fp64(value)), nullable: false, type_variation_reference: 0, @@ -154,7 +155,7 @@ impl From for Expr { impl From for Expr { fn from(value: bool) -> Self { - proto::expression::Literal { + Literal { literal_type: Some(LiteralType::Boolean(value)), nullable: false, type_variation_reference: 0, @@ -165,7 +166,7 @@ impl From for Expr { impl From for Expr { fn from(value: String) -> Self { - proto::expression::Literal { + Literal { literal_type: Some(LiteralType::String(value)), nullable: false, type_variation_reference: 0, @@ -180,6 +181,196 @@ impl From<&str> for Expr { } } +/// A Substrait literal carried as an extension argument. +/// +/// Unlike [`Expr`], this represents literal argument syntax directly rather than +/// a full Substrait expression. +#[derive(Debug, Clone, PartialEq)] +pub struct ExtensionLiteral(Literal); + +impl ExtensionLiteral { + /// Borrow the underlying Substrait literal protobuf. + pub fn as_proto(&self) -> &Literal { + &self.0 + } + + /// Clone the underlying Substrait literal protobuf. + pub fn to_proto(&self) -> Literal { + self.0.clone() + } + + pub fn date_days(days: i32) -> Self { + Self(Literal { + literal_type: Some(LiteralType::Date(days)), + nullable: false, + type_variation_reference: 0, + }) + } + + #[allow(deprecated)] + pub fn time_micros(micros: i64) -> Self { + // TODO: Decide whether this raw-unit constructor should reject values + // outside one time-of-day. The text parser uses HH:MM:SS and cannot + // round-trip values that format as 24:00:00 or beyond. + Self(Literal { + literal_type: Some(LiteralType::Time(micros)), + nullable: false, + type_variation_reference: 0, + }) + } + + #[allow(deprecated)] + pub fn timestamp_micros(micros: i64) -> Self { + Self(Literal { + literal_type: Some(LiteralType::Timestamp(micros)), + nullable: false, + type_variation_reference: 0, + }) + } + + pub fn precision_time_units(precision: i32, value: i64) -> Result { + validate_precision(precision)?; + // TODO: Decide whether this raw-unit constructor should reject values + // outside one time-of-day. The text parser uses HH:MM:SS and cannot + // round-trip values that format as 24:00:00 or beyond. + Ok(Self(Literal { + literal_type: Some(LiteralType::PrecisionTime(PrecisionTime { + precision, + value, + })), + nullable: false, + type_variation_reference: 0, + })) + } + + pub fn precision_timestamp_units(precision: i32, value: i64) -> Result { + validate_precision(precision)?; + Ok(Self::precision_timestamp_units_unchecked(precision, value)) + } + + pub fn precision_timestamp_seconds(value: i64) -> Self { + Self::precision_timestamp_units_unchecked(0, value) + } + + pub fn precision_timestamp_millis(value: i64) -> Self { + Self::precision_timestamp_units_unchecked(3, value) + } + + pub fn precision_timestamp_micros(value: i64) -> Self { + Self::precision_timestamp_units_unchecked(6, value) + } + + pub fn precision_timestamp_nanos(value: i64) -> Self { + Self::precision_timestamp_units_unchecked(9, value) + } + + /// Build a timezone-less Substrait precision timestamp from a protobuf UTC + /// timestamp. The epoch value is preserved, but timezone semantics are not. + pub fn precision_timestamp( + precision: i32, + timestamp: ProtoTimestamp, + ) -> Result { + let value = timestamp_to_units(precision, timestamp)?; + Ok(Self::precision_timestamp_units_unchecked(precision, value)) + } + + /// Build a UTC Substrait precision timestamp from a protobuf UTC timestamp. + pub fn precision_timestamp_tz_utc( + precision: i32, + timestamp: ProtoTimestamp, + ) -> Result { + let value = timestamp_to_units(precision, timestamp)?; + Ok(Self(Literal { + literal_type: Some(LiteralType::PrecisionTimestampTz(PrecisionTimestamp { + precision, + value, + })), + nullable: false, + type_variation_reference: 0, + })) + } + + fn precision_timestamp_units_unchecked(precision: i32, value: i64) -> Self { + Self(Literal { + literal_type: Some(LiteralType::PrecisionTimestamp(PrecisionTimestamp { + precision, + value, + })), + nullable: false, + type_variation_reference: 0, + }) + } +} + +fn validate_precision(precision: i32) -> Result<(), ExtensionError> { + if (0..=12).contains(&precision) { + Ok(()) + } else { + Err(ExtensionError::InvalidArgument(format!( + "temporal precision must be between 0 and 12, got {precision}" + ))) + } +} + +fn timestamp_to_units(precision: i32, timestamp: ProtoTimestamp) -> Result { + validate_precision(precision)?; + if precision > 9 { + return Err(ExtensionError::InvalidArgument(format!( + "protobuf Timestamp can only represent precision 0 through 9, got {precision}" + ))); + } + let timestamp = timestamp.try_normalize().map_err(|original| { + ExtensionError::InvalidArgument(format!( + "timestamp out of range or overflow (seconds={}, nanos={})", + original.seconds, original.nanos + )) + })?; + let scale = 10_i64.pow(precision as u32); + let nanos_per_unit = 10_i32.pow((9 - precision) as u32); + if timestamp.nanos % nanos_per_unit != 0 { + return Err(ExtensionError::InvalidArgument(format!( + "timestamp nanos {} are not exactly representable at precision {precision}", + timestamp.nanos + ))); + } + let seconds_units = timestamp.seconds.checked_mul(scale).ok_or_else(|| { + ExtensionError::InvalidArgument("timestamp value overflows i64".to_string()) + })?; + seconds_units + .checked_add(i64::from(timestamp.nanos / nanos_per_unit)) + .ok_or_else(|| ExtensionError::InvalidArgument("timestamp value overflows i64".to_string())) +} + +impl TryFrom for ExtensionLiteral { + type Error = ExtensionError; + + #[allow(deprecated)] + fn try_from(literal: Literal) -> Result { + match literal.literal_type.as_ref() { + Some( + LiteralType::Date(_) + | LiteralType::Time(_) + | LiteralType::Timestamp(_) + | LiteralType::PrecisionTime(_) + | LiteralType::PrecisionTimestamp(_) + | LiteralType::PrecisionTimestampTz(_), + ) => Ok(Self(literal)), + Some(other) => Err(ExtensionError::InvalidArgument(format!( + "literal type {other:?} is not supported as an extension literal" + ))), + None => Err(ExtensionError::InvalidArgument( + "literal is missing literal_type".to_string(), + )), + } + } +} + +impl From for Literal { + fn from(literal: ExtensionLiteral) -> Self { + literal.0 + } +} + /// Represents extension arguments plus optional output columns. /// /// Named arguments are stored in an [`IndexMap`] whose iteration order @@ -371,6 +562,8 @@ pub enum ExtensionValue { /// Use `TryFrom<&ExtensionValue> for Expr` when a handler accepts either an /// expression or a scalar value widened into an expression. Expr(Expr), + /// Substrait literal value using typed literal syntax. + Literal(ExtensionLiteral), /// Enum value (e.g. &CORE, &Inner) — the string holds the identifier /// without the `&` prefix Enum(String), @@ -392,6 +585,7 @@ pub enum ExtensionValueKind { Enum, Tuple, Expression, + Literal, } impl fmt::Display for ExtensionValueKind { @@ -405,6 +599,7 @@ impl fmt::Display for ExtensionValueKind { ExtensionValueKind::Enum => write!(f, "enum"), ExtensionValueKind::Tuple => write!(f, "tuple"), ExtensionValueKind::Expression => write!(f, "expression"), + ExtensionValueKind::Literal => write!(f, "literal"), } } } @@ -418,6 +613,7 @@ impl ExtensionValue { ExtensionValue::Float(_) => ExtensionValueKind::Float, ExtensionValue::Boolean(_) => ExtensionValueKind::Boolean, ExtensionValue::Expr(_) => ExtensionValueKind::Expression, + ExtensionValue::Literal(_) => ExtensionValueKind::Literal, ExtensionValue::Enum(_) => ExtensionValueKind::Enum, ExtensionValue::Tuple(_) => ExtensionValueKind::Tuple, } @@ -430,14 +626,20 @@ impl From for ExtensionValue { } } +impl From for ExtensionValue { + fn from(literal: ExtensionLiteral) -> Self { + ExtensionValue::Literal(literal) + } +} + impl From for ExtensionValue { fn from(expr: proto::Expression) -> Self { Expr::from(expr).into() } } -impl From for ExtensionValue { - fn from(literal: proto::expression::Literal) -> Self { +impl From for ExtensionValue { + fn from(literal: Literal) -> Self { Expr::from(literal).into() } } @@ -529,6 +731,17 @@ impl<'a> TryFrom<&'a ExtensionValue> for &'a TupleValue { } } +impl TryFrom<&ExtensionValue> for ExtensionLiteral { + type Error = ExtensionError; + + fn try_from(value: &ExtensionValue) -> Result { + match value { + ExtensionValue::Literal(literal) => Ok(literal.clone()), + v => Err(invalid_type(ExtensionValueKind::Literal, v)), + } + } +} + impl TryFrom<&ExtensionValue> for i64 { type Error = ExtensionError; @@ -582,6 +795,7 @@ impl TryFrom<&ExtensionValue> for Expr { fn try_from(value: &ExtensionValue) -> Result { match value { ExtensionValue::Expr(e) => Ok(e.clone()), + ExtensionValue::Literal(literal) => Ok(Expr::from(literal.to_proto())), // Untyped extension scalars are intentionally expression-compatible: // `arg=2` carries no syntax that distinguishes "configuration // integer" from "i64 literal expression". Scalar-specific @@ -647,3 +861,136 @@ impl ExtensionArgs { ArgsExtractor::new(self) } } + +#[cfg(test)] +mod tests { + use prost_types::Timestamp as ProtoTimestamp; + use substrait::proto::expression::RexType; + use substrait::proto::expression::literal::LiteralType; + + use super::{Expr, ExtensionLiteral, ExtensionValue}; + + fn literal_type(literal: &ExtensionLiteral) -> &LiteralType { + literal + .as_proto() + .literal_type + .as_ref() + .expect("literal_type") + } + + #[test] + fn precision_timestamp_tz_utc_converts_protobuf_timestamp_units() { + let literal = ExtensionLiteral::precision_timestamp_tz_utc( + 9, + ProtoTimestamp { + seconds: 1, + nanos: 123_456_789, + }, + ) + .expect("timestamp literal"); + + match literal_type(&literal) { + LiteralType::PrecisionTimestampTz(value) => { + assert_eq!(value.precision, 9); + assert_eq!(value.value, 1_123_456_789); + } + other => panic!("Expected PrecisionTimestampTz, got {other:?}"), + } + } + + #[test] + fn precision_timestamp_converts_protobuf_timestamp_to_naive_literal() { + let literal = ExtensionLiteral::precision_timestamp( + 6, + ProtoTimestamp { + seconds: 1, + nanos: 123_456_000, + }, + ) + .expect("timestamp literal"); + + match literal_type(&literal) { + LiteralType::PrecisionTimestamp(value) => { + assert_eq!(value.precision, 6); + assert_eq!(value.value, 1_123_456); + } + other => panic!("Expected PrecisionTimestamp, got {other:?}"), + } + } + + #[test] + fn precision_timestamp_from_protobuf_rejects_lossy_precision() { + let err = ExtensionLiteral::precision_timestamp_tz_utc( + 6, + ProtoTimestamp { + seconds: 1, + nanos: 123_456_789, + }, + ) + .expect_err("lossy timestamp should be rejected"); + + assert!( + err.to_string() + .contains("not exactly representable at precision 6"), + "{err}" + ); + } + + #[test] + fn precision_timestamp_from_protobuf_rejects_precision_above_nanos() { + let err = ExtensionLiteral::precision_timestamp_tz_utc( + 10, + ProtoTimestamp { + seconds: 1, + nanos: 123_456_789, + }, + ) + .expect_err("protobuf timestamp cannot represent precision 10"); + + assert!( + err.to_string() + .contains("can only represent precision 0 through 9"), + "{err}" + ); + } + + #[test] + fn precision_timestamp_convenience_constructors_use_expected_precision() { + let cases = [ + (ExtensionLiteral::precision_timestamp_seconds(7), 0, 7), + (ExtensionLiteral::precision_timestamp_millis(7), 3, 7), + (ExtensionLiteral::precision_timestamp_micros(7), 6, 7), + (ExtensionLiteral::precision_timestamp_nanos(7), 9, 7), + ]; + + for (literal, expected_precision, expected_value) in cases { + match literal_type(&literal) { + LiteralType::PrecisionTimestamp(value) => { + assert_eq!(value.precision, expected_precision); + assert_eq!(value.value, expected_value); + } + other => panic!("Expected PrecisionTimestamp, got {other:?}"), + } + } + } + + #[test] + fn extension_literal_extracts_directly_and_widens_to_expr() { + let value = ExtensionValue::from(ExtensionLiteral::precision_timestamp_micros(42)); + + let literal = ExtensionLiteral::try_from(&value).expect("extension literal"); + match literal_type(&literal) { + LiteralType::PrecisionTimestamp(timestamp) => { + assert_eq!(timestamp.precision, 6); + assert_eq!(timestamp.value, 42); + } + other => panic!("Expected PrecisionTimestamp, got {other:?}"), + } + + let expr = Expr::try_from(&value).expect("expression widening"); + assert!(matches!( + expr.as_proto().rex_type.as_ref(), + Some(RexType::Literal(_)) + )); + } +} diff --git a/src/extensions/mod.rs b/src/extensions/mod.rs index b6d92aad..baab6a9c 100644 --- a/src/extensions/mod.rs +++ b/src/extensions/mod.rs @@ -19,7 +19,8 @@ pub mod examples; pub use any::{Any, AnyRef}; pub(crate) use args::AddendumKind; pub use args::{ - EnumValue, Expr, ExtensionArgs, ExtensionColumn, ExtensionValue, ExtensionValueKind, TupleValue, + EnumValue, Expr, ExtensionArgs, ExtensionColumn, ExtensionLiteral, ExtensionValue, + ExtensionValueKind, TupleValue, }; pub use registry::{ AnyConvertible, Explainable, Extension, ExtensionError, ExtensionProtoConvert, diff --git a/src/parser/expression_grammar.pest b/src/parser/expression_grammar.pest index 0ea91be5..24695fbd 100644 --- a/src/parser/expression_grammar.pest +++ b/src/parser/expression_grammar.pest @@ -55,7 +55,8 @@ function_signature = @{ identifier ~ (":" ~ argument_signature?)? } // Field reference reference = { "$" ~ integer } // Literal -literal = { (float | integer | boolean | string_literal | null) ~ (":" ~ sp ~ type)? } +literal_value = { float | integer | boolean | string_literal | null } +literal = { literal_value ~ (":" ~ sp ~ type)? } // -- Components for types and functions anchor = { "#" ~ sp ~ integer } @@ -353,10 +354,10 @@ arguments = { (empty | (extension_arguments ~ (sp ~ "," ~ sp ~ extension_named_a extension_arguments = { extension_argument ~ (sp ~ "," ~ sp ~ extension_argument)* } // Untyped scalar extension literals render independently of expression literal -// verbosity. Typed literals (for example, 2:i16 or '2024-01-01':date) fall -// through to expression parsing so their Substrait type information is -// preserved. -untyped_literal = { (float | integer | boolean | string_literal) ~ !(":") } +// verbosity. Literals with explicit type annotations (for example, 2:i16 or +// '2024-01-01':date) fall through to first-class extension literal parsing +// when their Substrait literal kind is supported; otherwise they become +// expression values. // Tuples follow the Python/Rust trailing-comma convention to disambiguate from parenthesised // expressions. @@ -370,7 +371,7 @@ tuple = { ~ (sp ~ "," ~ sp ~ extension_argument)+ ~ (sp ~ ",")? ~ sp ~ ")" } -extension_argument = { enum_value | untyped_literal | reference | expression | tuple } +extension_argument = { enum_value | literal | reference | expression | tuple } // Named arguments (name=value pairs) extension_named_arguments = { extension_named_argument ~ (sp ~ "," ~ sp ~ extension_named_argument)* } diff --git a/src/parser/extensions.rs b/src/parser/extensions.rs index ebb621d8..715298c7 100644 --- a/src/parser/extensions.rs +++ b/src/parser/extensions.rs @@ -1,18 +1,20 @@ use std::fmt; use std::str::FromStr; +use substrait::proto::expression::literal::LiteralType; use substrait::proto::{Expression, Type}; use thiserror::Error; use super::{ ErrorKind, ExpressionParser, MessageParseError, ParsePair, Rule, RuleIter, ScopedParsePair, - unescape_string, unwrap_single_pair, + unwrap_single_pair, }; use crate::extensions::simple::{self, ExtensionKind}; use crate::extensions::{ - AddendumKind, ExtensionArgs, ExtensionColumn, ExtensionValue, InsertError, SimpleExtensions, - TupleValue, + AddendumKind, ExtensionArgs, ExtensionColumn, ExtensionLiteral, ExtensionValue, InsertError, + SimpleExtensions, TupleValue, }; +use crate::parser::literals::parse_literal_pair; use crate::parser::structural::IndentedLine; #[derive(Debug, Clone, Error)] @@ -290,24 +292,7 @@ impl ScopedParsePair for ExtensionValue { let field_index = FieldIndex::parse_pair(inner); ExtensionValue::from(Reference(field_index.0)) } - Rule::untyped_literal => { - // Literal can contain integer, float, boolean, or string_literal - let value_pair = unwrap_single_pair(inner); - match value_pair.as_rule() { - Rule::string_literal => ExtensionValue::String(unescape_string(value_pair)), - Rule::integer => { - ExtensionValue::Integer(value_pair.as_str().parse::().unwrap()) - } - Rule::float => { - ExtensionValue::Float(value_pair.as_str().parse::().unwrap()) - } - Rule::boolean => ExtensionValue::Boolean(value_pair.as_str() == "true"), - _ => panic!( - "Unexpected extension scalar literal type: {:?}", - value_pair.as_rule() - ), - } - } + Rule::literal => parse_extension_literal(extensions, inner)?, Rule::tuple => { let tv = inner .into_inner() @@ -324,6 +309,37 @@ impl ScopedParsePair for ExtensionValue { } } +fn parse_extension_literal( + extensions: &SimpleExtensions, + literal: pest::iterators::Pair, +) -> Result { + assert_eq!(literal.as_rule(), Rule::literal); + let span = literal.as_span(); + let parsed = parse_literal_pair(extensions, literal)?; + if parsed.explicit_type.is_none() { + return match parsed.literal.literal_type { + Some(LiteralType::String(value)) => Ok(ExtensionValue::String(value)), + Some(LiteralType::I64(value)) => Ok(ExtensionValue::Integer(value)), + Some(LiteralType::Fp64(value)) => Ok(ExtensionValue::Float(value)), + Some(LiteralType::Boolean(value)) => Ok(ExtensionValue::Boolean(value)), + Some(other) => Err(MessageParseError::invalid( + "extension_scalar_literal", + span, + format!("Unexpected extension scalar literal type: {other:?}"), + )), + None => Err(MessageParseError::invalid( + "extension_scalar_literal", + span, + "Extension scalar literal is missing literal_type", + )), + }; + } + Ok(match ExtensionLiteral::try_from(parsed.literal.clone()) { + Ok(literal) => ExtensionValue::from(literal), + Err(_) => ExtensionValue::from(parsed.literal), + }) +} + impl ScopedParsePair for ExtensionColumn { fn rule() -> Rule { Rule::extension_column @@ -614,7 +630,7 @@ mod tests { use substrait::proto::expression::literal::LiteralType; use super::*; - use crate::extensions::{Expr, ExtensionValue}; + use crate::extensions::{Expr, ExtensionLiteral, ExtensionValue}; use crate::fixtures::TestContext; use crate::parser::Parser; use crate::parser::common::test_support::ScopedParse; @@ -1006,10 +1022,30 @@ Functions: assert_eq!(ctx.textify_no_errors(&expression), "42:i64"); } + #[test] + fn test_temporal_extension_literals_parse_as_extension_literals() { + let ctx = TestContext::new(); + for text in [ + "'2023-12-25':date", + "'14:30:45':time", + "'2023-01-01T12:00:00':timestamp", + "'14:30:45.123':precisiontime<3>", + "'2023-01-01T12:00:00.123456':precisiontimestamp<6>", + "'2023-01-01T12:00:00.123456789':precisiontimestamptz<9>", + ] { + let value = parse_extension_value(text); + let literal = ExtensionLiteral::try_from(&value).expect("temporal literal"); + assert_eq!(ctx.textify_no_errors(&literal), text); + let expr = Expr::try_from(&value).expect("temporal literal should widen to Expr"); + assert_eq!(ctx.textify_no_errors(&expr), text); + } + } + #[test] fn test_typed_extension_literal_parses_as_expression() { let value = parse_extension_value("42:i16"); assert!(i64::try_from(&value).is_err()); + assert!(ExtensionLiteral::try_from(&value).is_err()); let expr = Expr::try_from(&value).unwrap(); assert_eq!(ctx_text(&expr), "42:i16"); diff --git a/src/parser/literals.rs b/src/parser/literals.rs index 625439a6..c269c368 100644 --- a/src/parser/literals.rs +++ b/src/parser/literals.rs @@ -1,12 +1,17 @@ -use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime}; +use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Timelike}; use substrait::proto::Type; use substrait::proto::expression::Literal; -use substrait::proto::expression::literal::LiteralType; +use substrait::proto::expression::literal::{LiteralType, PrecisionTime, PrecisionTimestamp}; use substrait::proto::r#type::{Kind, Nullability}; -use super::{MessageParseError, Rule, ScopedParsePair, unescape_string}; +use super::{MessageParseError, Rule, ScopedParsePair, unescape_string, unwrap_single_pair}; use crate::extensions::SimpleExtensions; +pub(crate) struct ParsedLiteral { + pub(crate) literal: Literal, + pub(crate) explicit_type: Option, +} + #[derive(Debug)] enum LiteralSyntax<'i> { Integer { @@ -106,6 +111,9 @@ enum LiteralTarget { Date(LiteralAttrs), TimeMicros(LiteralAttrs), TimestampMicros(LiteralAttrs), + PrecisionTime { attrs: LiteralAttrs, precision: i32 }, + PrecisionTimestamp { attrs: LiteralAttrs, precision: i32 }, + PrecisionTimestampTz { attrs: LiteralAttrs, precision: i32 }, Null(Type), UnsupportedStringFallback, } @@ -152,6 +160,18 @@ impl LiteralTarget { Kind::Date(k) => { Self::Date(LiteralAttrs::new(k.nullability, k.type_variation_reference)) } + Kind::PrecisionTime(k) => Self::PrecisionTime { + attrs: LiteralAttrs::new(k.nullability, k.type_variation_reference), + precision: k.precision, + }, + Kind::PrecisionTimestamp(k) => Self::PrecisionTimestamp { + attrs: LiteralAttrs::new(k.nullability, k.type_variation_reference), + precision: k.precision, + }, + Kind::PrecisionTimestampTz(k) => Self::PrecisionTimestampTz { + attrs: LiteralAttrs::new(k.nullability, k.type_variation_reference), + precision: k.precision, + }, #[allow(deprecated)] Kind::Time(k) => { Self::TimeMicros(LiteralAttrs::new(k.nullability, k.type_variation_reference)) @@ -240,6 +260,27 @@ impl LiteralTarget { LiteralType::Timestamp(parse_timestamp_to_microseconds(&value, span)?), attrs, )), + Self::PrecisionTime { attrs, precision } => Ok(literal( + LiteralType::PrecisionTime(PrecisionTime { + precision, + value: parse_time_to_precision_units(&value, precision, span)?, + }), + attrs, + )), + Self::PrecisionTimestamp { attrs, precision } => Ok(literal( + LiteralType::PrecisionTimestamp(PrecisionTimestamp { + precision, + value: parse_timestamp_to_precision_units(&value, precision, span)?, + }), + attrs, + )), + Self::PrecisionTimestampTz { attrs, precision } => Ok(literal( + LiteralType::PrecisionTimestampTz(PrecisionTimestamp { + precision, + value: parse_timestamp_to_precision_units(&value, precision, span)?, + }), + attrs, + )), Self::UnsupportedStringFallback => Ok(literal( LiteralType::String(value), LiteralAttrs::required(), @@ -343,6 +384,119 @@ fn parse_timestamp_to_microseconds( )) } +fn precision_scale(precision: i32, span: pest::Span) -> Result { + if !(0..=12).contains(&precision) { + return Err(invalid_literal( + span, + format!("Invalid temporal precision {precision}; expected 0 through 12"), + )); + } + Ok(10_i64.pow(precision as u32)) +} + +fn parse_fraction_units( + fraction: Option<&str>, + precision: i32, + span: pest::Span, +) -> Result { + let Some(fraction) = fraction else { + return Ok(0); + }; + if !fraction.chars().all(|c| c.is_ascii_digit()) { + return Err(invalid_literal(span, "Fractional seconds must be digits")); + } + let precision = precision as usize; + if fraction.len() > precision { + return Err(invalid_literal( + span, + format!( + "Fractional seconds have precision {}, but literal type allows precision {precision}", + fraction.len() + ), + )); + } + let mut padded = fraction.to_owned(); + padded.extend(std::iter::repeat_n('0', precision - padded.len())); + Ok(if padded.is_empty() { + 0 + } else { + padded.parse::().unwrap() + }) +} + +fn parse_time_parts<'a>( + time_str: &'a str, + span: pest::Span, +) -> Result<(NaiveTime, Option<&'a str>), MessageParseError> { + let (base, fraction) = match time_str.split_once('.') { + Some((base, fraction)) => (base, Some(fraction)), + None => (time_str, None), + }; + let time = NaiveTime::parse_from_str(base, "%H:%M:%S").map_err(|_| { + invalid_literal( + span, + format!("Invalid time format: '{time_str}'. Expected HH:MM:SS or HH:MM:SS.fff"), + ) + })?; + Ok((time, fraction)) +} + +fn parse_time_to_precision_units( + time_str: &str, + precision: i32, + span: pest::Span, +) -> Result { + let scale = precision_scale(precision, span)?; + let (time, fraction) = parse_time_parts(time_str, span)?; + let seconds = i64::from(time.num_seconds_from_midnight()); + let units = seconds + .checked_mul(scale) + .ok_or_else(|| invalid_literal(span, "Time literal overflow"))?; + let fraction_units = parse_fraction_units(fraction, precision, span)?; + units + .checked_add(fraction_units) + .ok_or_else(|| invalid_literal(span, "Time literal overflow")) +} + +fn parse_timestamp_to_precision_units( + timestamp_str: &str, + precision: i32, + span: pest::Span, +) -> Result { + let scale = precision_scale(precision, span)?; + let (date_part, time_part) = timestamp_str + .split_once('T') + .or_else(|| timestamp_str.split_once(' ')) + .ok_or_else(|| { + invalid_literal( + span, + format!( + "Invalid timestamp format: '{timestamp_str}'. Expected YYYY-MM-DDTHH:MM:SS or YYYY-MM-DD HH:MM:SS" + ), + ) + })?; + let date = ["%Y-%m-%d", "%Y/%m/%d"] + .iter() + .find_map(|format| NaiveDate::parse_from_str(date_part, format).ok()) + .ok_or_else(|| { + invalid_literal( + span, + format!("Invalid date format: '{date_part}'. Expected YYYY-MM-DD or YYYY/MM/DD"), + ) + })?; + let (time, fraction) = parse_time_parts(time_part, span)?; + let datetime = date.and_time(time); + let epoch = DateTime::from_timestamp(0, 0).unwrap().naive_utc(); + let seconds = datetime.signed_duration_since(epoch).num_seconds(); + let units = seconds + .checked_mul(scale) + .ok_or_else(|| invalid_literal(span, "Timestamp literal overflow"))?; + let fraction_units = parse_fraction_units(fraction, precision, span)?; + units + .checked_add(fraction_units) + .ok_or_else(|| invalid_literal(span, "Timestamp literal overflow")) +} + impl ScopedParsePair for Literal { fn rule() -> Rule { Rule::literal @@ -357,17 +511,33 @@ impl ScopedParsePair for Literal { pair: pest::iterators::Pair, ) -> Result { assert_eq!(pair.as_rule(), Self::rule()); - let mut pairs = pair.into_inner(); - let value = pairs.next().unwrap(); - let typ = pairs.next(); - assert!(pairs.next().is_none()); - - let syntax = LiteralSyntax::from_pair(value); - let typ = match typ { - Some(t) => Some(Type::parse_pair(extensions, t)?), - None => None, - }; - let target = LiteralTarget::from_type(typ, &syntax)?; - target.parse(syntax) + parse_literal_pair(extensions, pair).map(|parsed| parsed.literal) + } +} + +pub(crate) fn parse_literal_pair( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, +) -> Result { + match pair.as_rule() { + Rule::literal => { + let mut pairs = pair.into_inner(); + let value = unwrap_single_pair(pairs.next().unwrap()); + let typ = pairs.next(); + assert!(pairs.next().is_none()); + + let syntax = LiteralSyntax::from_pair(value); + let explicit_type = match typ { + Some(t) => Some(Type::parse_pair(extensions, t)?), + None => None, + }; + let target = LiteralTarget::from_type(explicit_type.clone(), &syntax)?; + let literal = target.parse(syntax)?; + Ok(ParsedLiteral { + literal, + explicit_type, + }) + } + _ => unreachable!("Literal unexpected rule: {:?}", pair.as_rule()), } } diff --git a/src/textify/extensions.rs b/src/textify/extensions.rs index eb917ad8..1af2000f 100644 --- a/src/textify/extensions.rs +++ b/src/textify/extensions.rs @@ -7,7 +7,9 @@ use std::fmt; -use crate::extensions::{Expr, ExtensionArgs, ExtensionColumn, ExtensionValue, TupleValue}; +use crate::extensions::{ + Expr, ExtensionArgs, ExtensionColumn, ExtensionLiteral, ExtensionValue, TupleValue, +}; use crate::textify::foundation::{Scope, Textify}; use crate::textify::types::{Name, escaped}; @@ -40,12 +42,23 @@ impl Textify for ExtensionValue { ExtensionValue::Float(f) => write!(w, "{f}"), ExtensionValue::Boolean(b) => write!(w, "{b}"), ExtensionValue::Expr(expr) => expr.textify(ctx, w), + ExtensionValue::Literal(literal) => literal.textify(ctx, w), ExtensionValue::Enum(e) => write!(w, "&{e}"), ExtensionValue::Tuple(tv) => tv.textify(ctx, w), } } } +impl Textify for ExtensionLiteral { + fn name() -> &'static str { + "ExtensionLiteral" + } + + fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { + write!(w, "{}", ctx.display(self.as_proto())) + } +} + impl Textify for Expr { fn name() -> &'static str { "Expr" diff --git a/src/textify/literals.rs b/src/textify/literals.rs index 61ca08b8..de057a34 100644 --- a/src/textify/literals.rs +++ b/src/textify/literals.rs @@ -42,20 +42,42 @@ fn days_to_date_string(days: i32) -> String { date.format("%Y-%m-%d").to_string() } -fn microseconds_to_time_string(microseconds: i64) -> String { - let total_seconds = microseconds / 1_000_000; - let remaining_micros = microseconds % 1_000_000; +fn format_fraction(units: i64, precision: i32) -> String { + if precision == 0 || units == 0 { + return String::new(); + } + let mut fraction = format!("{:0width$}", units, width = precision as usize); + while fraction.ends_with('0') { + fraction.pop(); + } + if fraction.is_empty() { + String::new() + } else { + format!(".{fraction}") + } +} + +fn precision_scale(precision: i32) -> Option { + (0..=12) + .contains(&precision) + .then(|| 10_i64.pow(precision as u32)) +} + +fn precision_time_to_string(value: i64, precision: i32) -> Option { + let scale = precision_scale(precision)?; + let total_seconds = value.div_euclid(scale); + let fraction_units = value.rem_euclid(scale); let hours = total_seconds / 3600; let minutes = (total_seconds % 3600) / 60; let seconds = total_seconds % 60; - if remaining_micros == 0 { - format!("{hours:02}:{minutes:02}:{seconds:02}") - } else { - let fraction = format!("{remaining_micros:06}") - .trim_end_matches('0') - .to_string(); - format!("{hours:02}:{minutes:02}:{seconds:02}.{fraction}") - } + Some(format!( + "{hours:02}:{minutes:02}:{seconds:02}{}", + format_fraction(fraction_units, precision) + )) +} + +fn microseconds_to_time_string(microseconds: i64) -> String { + precision_time_to_string(microseconds, 6).expect("precision 6 is valid") } fn microseconds_to_timestamp_string(microseconds: i64) -> String { @@ -73,6 +95,18 @@ fn microseconds_to_timestamp_string(microseconds: i64) -> String { } } +fn precision_timestamp_to_string(value: i64, precision: i32) -> Option { + let scale = precision_scale(precision)?; + let seconds = value.div_euclid(scale); + let fraction_units = value.rem_euclid(scale); + let datetime = DateTime::from_timestamp(seconds, 0)?.naive_utc(); + Some(format!( + "{}{}", + datetime.format("%Y-%m-%dT%H:%M:%S"), + format_fraction(fraction_units, precision) + )) +} + fn write_literal_value( lit: &LiteralType, ctx: &S, @@ -110,10 +144,49 @@ fn write_literal_value( LiteralType::VarChar(_) => unimplemented_literal("VarChar", ctx, w), LiteralType::FixedBinary(_) => unimplemented_literal("FixedBinary", ctx, w), LiteralType::Decimal(_) => unimplemented_literal("Decimal", ctx, w), - LiteralType::PrecisionTime(_) => unimplemented_literal("PrecisionTime", ctx, w), - LiteralType::PrecisionTimestamp(_) => unimplemented_literal("PrecisionTimestamp", ctx, w), - LiteralType::PrecisionTimestampTz(_) => { - unimplemented_literal("PrecisionTimestampTz", ctx, w) + LiteralType::PrecisionTime(value) => { + let Some(formatted) = precision_time_to_string(value.value, value.precision) else { + return invalid_literal( + "PrecisionTime", + format!( + "invalid precision {} for precision time literal", + value.precision + ), + ctx, + w, + ); + }; + write!(w, "'{}'", escaped(&formatted)) + } + LiteralType::PrecisionTimestamp(value) => { + let Some(formatted) = precision_timestamp_to_string(value.value, value.precision) + else { + return invalid_literal( + "PrecisionTimestamp", + format!( + "invalid precision {} or out-of-range value {} for precision timestamp literal", + value.precision, value.value + ), + ctx, + w, + ); + }; + write!(w, "'{}'", escaped(&formatted)) + } + LiteralType::PrecisionTimestampTz(value) => { + let Some(formatted) = precision_timestamp_to_string(value.value, value.precision) + else { + return invalid_literal( + "PrecisionTimestampTz", + format!( + "invalid precision {} or out-of-range value {} for precision timestamp tz literal", + value.precision, value.value + ), + ctx, + w, + ); + }; + write!(w, "'{}'", escaped(&formatted)) } LiteralType::Struct(_) => unimplemented_literal("Struct", ctx, w), LiteralType::Map(_) => unimplemented_literal("Map", ctx, w), @@ -128,6 +201,23 @@ fn write_literal_value( } } +fn invalid_literal( + variant: &'static str, + description: impl Into>, + ctx: &S, + w: &mut W, +) -> fmt::Result { + write!( + w, + "{}", + ctx.failure(PlanError::invalid( + "LiteralType", + Some(variant), + description + )) + ) +} + fn literal_type(literal: &expr::Literal) -> Option { let lit = literal.literal_type.as_ref()?; let nullability = if literal.nullable { @@ -187,6 +277,25 @@ fn literal_type(literal: &expr::Literal) -> Option { nullability, type_variation_reference, }), + LiteralType::PrecisionTime(value) => Kind::PrecisionTime(ptype::PrecisionTime { + precision: value.precision, + nullability, + type_variation_reference, + }), + LiteralType::PrecisionTimestamp(value) => { + Kind::PrecisionTimestamp(ptype::PrecisionTimestamp { + precision: value.precision, + nullability, + type_variation_reference, + }) + } + LiteralType::PrecisionTimestampTz(value) => { + Kind::PrecisionTimestampTz(ptype::PrecisionTimestampTz { + precision: value.precision, + nullability, + type_variation_reference, + }) + } _ => return None, }; Some(Type { kind: Some(kind) }) diff --git a/tests/extension_roundtrip.rs b/tests/extension_roundtrip.rs index 302bead1..4456de57 100644 --- a/tests/extension_roundtrip.rs +++ b/tests/extension_roundtrip.rs @@ -4,13 +4,14 @@ mod common; use common::parse_type; use prost::{Message, Name}; +use prost_types::Timestamp; use substrait::proto; use substrait::proto::expression::RexType; use substrait::proto::expression::literal::LiteralType; use substrait::proto::r#type::Nullability; use substrait_explain::extensions::examples::PartitionHint; use substrait_explain::extensions::{ - EnumValue, Explainable, Expr, ExtensionArgs, ExtensionColumn, ExtensionError, + EnumValue, Explainable, Expr, ExtensionArgs, ExtensionColumn, ExtensionError, ExtensionLiteral, ExtensionProtoConvert, ExtensionRegistry, ExtensionValue, TupleValue, }; use substrait_explain::{Parser, format_with_registry}; @@ -297,6 +298,128 @@ Root[result] assert_eq!(formatted.trim(), plan_text.trim()); } +/// Test-only protobuf payload used to verify temporal literal argument +/// round-tripping for registered extensions. +#[derive(Clone, PartialEq, Message)] +pub struct TemporalLiteralConfig { + #[prost(int64, tag = "1")] + pub created_at_micros: i64, + #[prost(int64, tag = "2")] + pub watermark_seconds: i64, + #[prost(int32, tag = "3")] + pub watermark_nanos: i32, +} + +impl Name for TemporalLiteralConfig { + const NAME: &'static str = "TemporalLiteralConfig"; + const PACKAGE: &'static str = "test"; + + fn full_name() -> String { + "test.TemporalLiteralConfig".to_string() + } + + fn type_url() -> String { + "type.googleapis.com/test.TemporalLiteralConfig".to_string() + } +} + +impl Explainable for TemporalLiteralConfig { + fn name() -> &'static str { + "TemporalLiteralTest" + } + + fn from_args(args: &ExtensionArgs) -> Result { + let mut extractor = args.extractor(); + let created_at: ExtensionLiteral = extractor.expect_named_arg("created_at")?; + let watermark: ExtensionLiteral = extractor.expect_named_arg("watermark")?; + extractor.check_exhausted()?; + + #[allow(deprecated)] + let created_at_micros = match created_at.as_proto().literal_type.as_ref() { + Some(LiteralType::Timestamp(value)) => *value, + Some(other) => { + return Err(ExtensionError::InvalidArgument(format!( + "created_at must be a timestamp literal, got {other:?}" + ))); + } + None => { + return Err(ExtensionError::InvalidArgument( + "created_at must have literal_type".to_string(), + )); + } + }; + + let (watermark_seconds, watermark_nanos) = match watermark.as_proto().literal_type.as_ref() + { + Some(LiteralType::PrecisionTimestampTz(value)) if value.precision == 9 => { + let seconds = value.value.div_euclid(1_000_000_000); + let nanos = value.value.rem_euclid(1_000_000_000) as i32; + (seconds, nanos) + } + Some(other) => { + return Err(ExtensionError::InvalidArgument(format!( + "watermark must be a precisiontimestamptz<9> literal, got {other:?}" + ))); + } + None => { + return Err(ExtensionError::InvalidArgument( + "watermark must have literal_type".to_string(), + )); + } + }; + + Ok(TemporalLiteralConfig { + created_at_micros, + watermark_seconds, + watermark_nanos, + }) + } + + fn to_args(&self) -> Result { + let mut args = ExtensionArgs::default(); + args.insert( + "created_at", + ExtensionLiteral::timestamp_micros(self.created_at_micros), + ); + args.insert( + "watermark", + ExtensionLiteral::precision_timestamp_tz_utc( + 9, + Timestamp { + seconds: self.watermark_seconds, + nanos: self.watermark_nanos, + }, + )?, + ); + args.output_columns.push(ExtensionColumn::Named { + name: "value".to_string(), + r#type: parse_type("string"), + }); + Ok(args) + } +} + +#[test] +fn test_temporal_extension_literal_roundtrip() { + let mut registry = ExtensionRegistry::new(); + registry + .register_relation::() + .unwrap(); + + let plan_text = r#" +=== Plan +Root[result] + ExtensionLeaf:TemporalLiteralTest[created_at='2023-01-01T12:00:00':timestamp, watermark='2023-01-01T12:00:00.123456789':precisiontimestamptz<9> => value:string] +"#; + + let parser = Parser::new().with_extension_registry(registry.clone()); + let plan = parser.parse_plan(plan_text).expect("Failed to parse plan"); + + let (formatted, errors) = format_with_registry(&plan, &Default::default(), ®istry); + assert!(errors.is_empty(), "Unexpected errors: {errors:?}"); + assert_eq!(formatted.trim(), plan_text.trim()); +} + #[test] fn test_extension_unknown_arguments() { let mut registry = ExtensionRegistry::new();