diff --git a/crates/adapters/src/format/avro/serializer.rs b/crates/adapters/src/format/avro/serializer.rs index 61cc2e4d030..fe46f2625e3 100644 --- a/crates/adapters/src/format/avro/serializer.rs +++ b/crates/adapters/src/format/avro/serializer.rs @@ -19,6 +19,7 @@ use super::schema::{OptionalField, schema_unwrap_optional}; pub fn avro_ser_config() -> SqlSerdeConfig { SqlSerdeConfig::default() .with_timestamp_format(TimestampFormat::MicrosSinceEpoch) + .with_timestamp_tz_format(TimestampFormat::MicrosSinceEpoch) .with_time_format(TimeFormat::Micros) .with_date_format(DateFormat::DaysSinceEpoch) .with_decimal_format(DecimalFormat::I128) diff --git a/crates/adapters/src/format/avro/test.rs b/crates/adapters/src/format/avro/test.rs index 98e945d8798..0381ba3b717 100644 --- a/crates/adapters/src/format/avro/test.rs +++ b/crates/adapters/src/format/avro/test.rs @@ -20,7 +20,7 @@ use dbsp::trace::BatchReaderFactories; use dbsp::typed_batch::{DynSpineSnapshot, SpineSnapshot as TypedSpineSnapshot, TypedBatch}; use dbsp::{DBData, IndexedZSetReader, OrdIndexedZSet, OrdZSet, ZWeight, utils::Tup2}; use feldera_adapterlib::transport::OutputBatchType; -use feldera_sqllib::{ByteArray, Uuid, Variant}; +use feldera_sqllib::{ByteArray, TimestampTz, Uuid, Variant}; use feldera_types::{ deserialize_table_record, format::avro::{AvroEncoderConfig, AvroEncoderKeyMode}, @@ -1063,6 +1063,177 @@ fn test_ms_time() { run_parser_test(vec![test]); } +#[derive( + Debug, + Default, + PartialEq, + Eq, + PartialOrd, + Ord, + Clone, + Hash, + SizeOf, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] +#[archive_attr(derive(Ord, Eq, PartialEq, PartialOrd))] +struct TestTimestampTz { + id: i64, + ts_tz: TimestampTz, + ts_tz_nullable: Option, +} + +impl TestTimestampTz { + pub fn avro_schema_micros() -> &'static str { + r#"{ + "type": "record", + "name": "TestTimestampTz", + "fields": [ + { "name": "id", "type": "long" }, + { "name": "ts_tz", "type": "long", "logicalType": "timestamp-micros" }, + { "name": "ts_tz_nullable", "type": ["null", { "type": "long", "logicalType": "timestamp-micros" }] } + ] + }"# + } + + pub fn avro_schema_millis() -> &'static str { + r#"{ + "type": "record", + "name": "TestTimestampTz", + "fields": [ + { "name": "id", "type": "long" }, + { "name": "ts_tz", "type": "long", "logicalType": "timestamp-millis" }, + { "name": "ts_tz_nullable", "type": ["null", { "type": "long", "logicalType": "timestamp-millis" }] } + ] + }"# + } + + pub fn relation_schema() -> Relation { + Relation { + name: SqlIdentifier::new("TestTimestampTz", false), + fields: vec![ + Field::new("id".into(), ColumnType::bigint(false)), + Field::new("ts_tz".into(), ColumnType::timestamp_tz(false)), + Field::new("ts_tz_nullable".into(), ColumnType::timestamp_tz(true)), + ], + materialized: false, + properties: BTreeMap::new(), + primary_key: None, + } + } + + /// Whole-millisecond values, so that both the `timestamp-micros` and the + /// `timestamp-millis` encodings round-trip exactly. + pub fn data() -> Vec { + vec![ + TestTimestampTz { + id: 1, + ts_tz: TimestampTz::from_milliseconds(1_713_597_703_123), + ts_tz_nullable: Some(TimestampTz::from_milliseconds(1_713_597_704_456)), + }, + TestTimestampTz { + id: 2, + ts_tz: TimestampTz::from_milliseconds(0), + ts_tz_nullable: None, + }, + ] + } +} + +serialize_table_record!(TestTimestampTz[3]{ + r#id["id"]: i64, + r#ts_tz["ts_tz"]: TimestampTz, + r#ts_tz_nullable["ts_tz_nullable"]: Option +}); + +deserialize_table_record!(TestTimestampTz["TestTimestampTz", Variant, 3] { + (r#id, "id", false, i64, |_| None), + (r#ts_tz, "ts_tz", false, TimestampTz, |_| None), + (r#ts_tz_nullable, "ts_tz_nullable", false, Option, |_| Some(None)) +}); + +/// TIMESTAMP WITH TIME ZONE columns encoded as `timestamp-micros`. +#[test] +fn test_timestamp_tz() { + let mut data = TestTimestampTz::data(); + // Sub-millisecond precision survives the micros encoding. + data.push(TestTimestampTz { + id: 3, + ts_tz: TimestampTz::from_microseconds(1_713_597_703_123_456), + ts_tz_nullable: Some(TimestampTz::from_microseconds(1_713_597_704_654_321)), + }); + + let test_case = gen_raw_parser_test( + &data, + &TestTimestampTz::relation_schema(), + TestTimestampTz::avro_schema_micros(), + ); + + run_parser_test(vec![test_case]); +} + +/// TIMESTAMP WITH TIME ZONE columns encoded as `timestamp-millis`. +#[test] +fn test_timestamp_tz_millis() { + let test_case = gen_raw_parser_test( + &TestTimestampTz::data(), + &TestTimestampTz::relation_schema(), + TestTimestampTz::avro_schema_millis(), + ); + + run_parser_test(vec![test_case]); +} + +/// Decode hand-built Avro values with hard-coded expected results. +#[test] +fn test_timestamp_tz_wire_format() { + fn test_case( + avro_schema_str: &str, + ts_tz: Value, + expected: TimestampTz, + ) -> TestCase { + let schema = AvroSchema::parse_str(avro_schema_str).unwrap(); + let record = Value::Record(vec![ + ("id".to_string(), Value::Long(1)), + ("ts_tz".to_string(), ts_tz), + ( + "ts_tz_nullable".to_string(), + Value::Union(0, Box::new(Value::Null)), + ), + ]); + + TestCase { + relation_schema: TestTimestampTz::relation_schema(), + config: AvroParserConfig { + update_format: AvroUpdateFormat::Raw, + schema: Some(avro_schema_str.to_string()), + skip_schema_id: false, + registry_config: Default::default(), + }, + input_batches: vec![(serialize_value(record, &schema), vec![])], + expected_output: vec![MockUpdate::Insert(TestTimestampTz { + id: 1, + ts_tz: expected, + ts_tz_nullable: None, + })], + } + } + + run_parser_test(vec![ + test_case( + TestTimestampTz::avro_schema_millis(), + Value::TimestampMillis(1_713_597_703_123), + TimestampTz::from_milliseconds(1_713_597_703_123), + ), + test_case( + TestTimestampTz::avro_schema_micros(), + Value::TimestampMicros(1_713_597_703_123_456), + TimestampTz::from_microseconds(1_713_597_703_123_456), + ), + ]); +} + /// Type used to serialize different integer types as Avro `int`. #[derive( Debug, diff --git a/crates/feldera-types/src/serde_with_context/serde_config.rs b/crates/feldera-types/src/serde_with_context/serde_config.rs index a9dbe279cf5..db71c202ba8 100644 --- a/crates/feldera-types/src/serde_with_context/serde_config.rs +++ b/crates/feldera-types/src/serde_with_context/serde_config.rs @@ -158,6 +158,11 @@ impl SqlSerdeConfig { self } + pub fn with_timestamp_tz_format(mut self, timestamp_tz_format: TimestampFormat) -> Self { + self.timestamp_tz_format = timestamp_tz_format; + self + } + pub fn with_decimal_format(mut self, decimal_format: DecimalFormat) -> Self { self.decimal_format = decimal_format; self diff --git a/python/tests/workloads/test_kafka_avro.py b/python/tests/workloads/test_kafka_avro.py index b121552aa48..b215e6dc24e 100644 --- a/python/tests/workloads/test_kafka_avro.py +++ b/python/tests/workloads/test_kafka_avro.py @@ -176,11 +176,12 @@ def sql_source_table(v: Variant) -> str: create table {v.source} ( id int, str varchar, - dec decimal, + dec decimal(10, 2), reall real, dbl double, booll boolean, tmestmp timestamp, + tmestmptz timestamp with time zone, datee date, tme time ) with ( @@ -260,11 +261,12 @@ def sql_loopback_table(v: Variant) -> str: create table {v.loopback} ( id int, str varchar, - dec decimal, + dec decimal(10, 2), reall real, dbl double, booll boolean, tmestmp timestamp, + tmestmptz timestamp with time zone, datee date, tme time ) with (