From 1ff988401a654aa709e59e63d70f40c8267e4dd8 Mon Sep 17 00:00:00 2001 From: Colin Marc Date: Thu, 31 Jul 2025 12:53:39 +0200 Subject: [PATCH] fix(iceberg-datafusion): handle more types of timestamp conversion This adds support for converting date values, which helps with predicate expressions such as `date > DATE_TRUNC('day', ts)`. --- crates/iceberg/src/spec/values/datum.rs | 94 +++++- crates/iceberg/src/spec/values/tests.rs | 393 ++++++++++++------------ 2 files changed, 269 insertions(+), 218 deletions(-) diff --git a/crates/iceberg/src/spec/values/datum.rs b/crates/iceberg/src/spec/values/datum.rs index 71de11beb3..38782e7415 100644 --- a/crates/iceberg/src/spec/values/datum.rs +++ b/crates/iceberg/src/spec/values/datum.rs @@ -1132,11 +1132,12 @@ impl Datum { match target_type { Type::Primitive(target_primitive_type) => { match (&self.literal, &self.r#type, target_primitive_type) { - (PrimitiveLiteral::Int(val), _, PrimitiveType::Int) => Ok(Datum::int(*val)), - (PrimitiveLiteral::Int(val), _, PrimitiveType::Date) => Ok(Datum::date(*val)), - (PrimitiveLiteral::Int(val), _, PrimitiveType::Long) => Ok(Datum::long(*val)), - (PrimitiveLiteral::Long(val), _, PrimitiveType::Int) => { - Ok(Datum::i64_to_i32(*val)) + (_, self_type, target_type) if self_type == target_type => Ok(self), + (PrimitiveLiteral::Int(val), source_type, target_type) => { + convert_int(*val, source_type, target_type) + } + (PrimitiveLiteral::Long(val), source_type, target_type) => { + convert_long(*val, source_type, target_type) } (PrimitiveLiteral::Double(val), _, PrimitiveType::Float) => { Ok(Datum::f64_to_f32(val.0)) @@ -1144,13 +1145,6 @@ impl Datum { (PrimitiveLiteral::Float(val), _, PrimitiveType::Double) => { Ok(Datum::double(val.0 as f64)) } - (PrimitiveLiteral::Long(val), _, PrimitiveType::Timestamp) => { - Ok(Datum::timestamp_micros(*val)) - } - (PrimitiveLiteral::Long(val), _, PrimitiveType::Timestamptz) => { - Ok(Datum::timestamptz_micros(*val)) - } - // Let's wait with nano's until this clears up: https://github.com/apache/iceberg/pull/11775 (PrimitiveLiteral::Int128(val), _, PrimitiveType::Long) => { Ok(Datum::i128_to_i64(*val)) } @@ -1199,7 +1193,6 @@ impl Datum { } // TODO: implement more type conversions - (_, self_type, target_type) if self_type == target_type => Ok(self), _ => Err(invalid_data!( "Can't convert datum from {} type to {} type.", self.r#type, @@ -1312,3 +1305,78 @@ fn human_float(value: T) -> String { format!("{sign}{}.{fraction}E{exponent}", &digits[..1]) } } + +const MICROS_PER_DAY: i64 = 24 * 60 * 60 * 1_000_000; +const NANOS_PER_MICRO: i64 = 1000; +const NANOS_PER_DAY: i64 = NANOS_PER_MICRO * MICROS_PER_DAY; + +/// Converts an int literal between two [PrimitiveType]s. +fn convert_int( + val: i32, + source_type: &PrimitiveType, + target_type: &PrimitiveType, +) -> Result { + let datum = match (source_type, target_type) { + (_, PrimitiveType::Int) => Datum::int(val), + (_, PrimitiveType::Date) => Datum::date(val), + (_, PrimitiveType::Long) => Datum::long(val), + (PrimitiveType::Date, PrimitiveType::Timestamp) => Datum::timestamp_micros( + (val as i64) + .checked_mul(MICROS_PER_DAY) + .ok_or_else(overflow)?, + ), + (PrimitiveType::Date, PrimitiveType::TimestampNs) => Datum::timestamp_nanos( + (val as i64) + .checked_mul(NANOS_PER_DAY) + .ok_or_else(overflow)?, + ), + _ => { + return Err(invalid_data!( + "Can't convert datum from {source_type} type to {target_type} type." + )); + } + }; + + Ok(datum) +} + +/// Converts a long literal between two [PrimitiveType]s. +fn convert_long( + val: i64, + source_type: &PrimitiveType, + target_type: &PrimitiveType, +) -> Result { + let datum = match (source_type, target_type) { + (_, PrimitiveType::Int) => Datum::i64_to_i32(val), + (PrimitiveType::Timestamp | PrimitiveType::Timestamptz, PrimitiveType::Timestamp) => { + Datum::timestamp_micros(val) + } + (PrimitiveType::Timestamp | PrimitiveType::Timestamptz, PrimitiveType::Timestamptz) => { + Datum::timestamptz_micros(val) + } + (PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs, PrimitiveType::TimestampNs) => { + Datum::timestamp_nanos(val) + } + ( + PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs, + PrimitiveType::TimestamptzNs, + ) => Datum::timestamptz_nanos(val), + (PrimitiveType::Timestamp | PrimitiveType::Timestamptz, PrimitiveType::TimestampNs) => { + Datum::timestamp_nanos(val.checked_mul(NANOS_PER_MICRO).ok_or_else(overflow)?) + } + (PrimitiveType::Timestamp | PrimitiveType::Timestamptz, PrimitiveType::TimestamptzNs) => { + Datum::timestamptz_nanos(val.checked_mul(NANOS_PER_MICRO).ok_or_else(overflow)?) + } + _ => { + return Err(invalid_data!( + "Can't convert datum from {source_type} type to {target_type} type." + )); + } + }; + + Ok(datum) +} + +fn overflow() -> Error { + invalid_data!("integer overflow") +} diff --git a/crates/iceberg/src/spec/values/tests.rs b/crates/iceberg/src/spec/values/tests.rs index 2de6e2916a..dbfb4c88da 100644 --- a/crates/iceberg/src/spec/values/tests.rs +++ b/crates/iceberg/src/spec/values/tests.rs @@ -21,6 +21,7 @@ use std::collections::HashSet; use apache_avro::to_value; use apache_avro::types::Value; +use chrono::DateTime; use ordered_float::OrderedFloat; use serde_bytes::ByteBuf; use serde_json::Value as JsonValue; @@ -1511,213 +1512,195 @@ fn test_datum_ser_deser() { test_fn(datum); } -#[test] -fn test_datum_date_convert_to_int() { - let datum_date = Datum::date(12345); - - let result = datum_date.to(&Primitive(PrimitiveType::Int)).unwrap(); - - let expected = Datum::int(12345); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_int_convert_to_date() { - let datum_int = Datum::int(12345); - - let result = datum_int.to(&Primitive(PrimitiveType::Date)).unwrap(); - - let expected = Datum::date(12345); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_long_convert_to_int() { - let datum = Datum::long(12345); - - let result = datum.to(&Primitive(PrimitiveType::Int)).unwrap(); - - let expected = Datum::int(12345); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_long_convert_to_int_above_max() { - let datum = Datum::long(INT_MAX as i64 + 1); - - let result = datum.to(&Primitive(PrimitiveType::Int)).unwrap(); - - let expected = Datum::new(PrimitiveType::Int, PrimitiveLiteral::AboveMax); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_long_convert_to_int_below_min() { - let datum = Datum::long(INT_MIN as i64 - 1); - - let result = datum.to(&Primitive(PrimitiveType::Int)).unwrap(); - - let expected = Datum::new(PrimitiveType::Int, PrimitiveLiteral::BelowMin); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_double_convert_to_float() { - let datum = Datum::double(2.5); - - let result = datum.to(&Primitive(PrimitiveType::Float)).unwrap(); - - let expected = Datum::float(2.5_f32); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_double_convert_to_float_above_max() { - let datum = Datum::double(1e39); - - let result = datum.to(&Primitive(PrimitiveType::Float)).unwrap(); - - let expected = Datum::new(PrimitiveType::Float, PrimitiveLiteral::AboveMax); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_double_convert_to_float_below_min() { - let datum = Datum::double(-1e39); - - let result = datum.to(&Primitive(PrimitiveType::Float)).unwrap(); - - let expected = Datum::new(PrimitiveType::Float, PrimitiveLiteral::BelowMin); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_float_convert_to_double() { - let datum = Datum::float(2.5_f32); - - let result = datum.to(&Primitive(PrimitiveType::Double)).unwrap(); - - let expected = Datum::double(2.5); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_long_convert_to_timestamp() { - let datum = Datum::long(12345); - - let result = datum.to(&Primitive(PrimitiveType::Timestamp)).unwrap(); - - let expected = Datum::timestamp_micros(12345); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_long_convert_to_timestamptz() { - let datum = Datum::long(12345); +macro_rules! datum_convert_tests { + ($($name:ident: ($source:expr => $target_type:ident, $expected:expr),)*) => { + $( + #[test] + fn $name() { + let source = $source; + let target_type = PrimitiveType::$target_type; + let expected = $expected; - let result = datum.to(&Primitive(PrimitiveType::Timestamptz)).unwrap(); - - let expected = Datum::timestamptz_micros(12345); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_decimal_convert_to_long() { - let datum = Datum::decimal(decimal_new(12345, 0)).unwrap(); - - let result = datum.to(&Primitive(PrimitiveType::Long)).unwrap(); - - let expected = Datum::long(12345); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_decimal_convert_to_long_above_max() { - let datum = Datum::decimal(decimal_from_i128_with_scale(LONG_MAX as i128 + 1, 0)).unwrap(); - - let result = datum.to(&Primitive(PrimitiveType::Long)).unwrap(); - - let expected = Datum::new(PrimitiveType::Long, PrimitiveLiteral::AboveMax); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_decimal_convert_to_long_below_min() { - let datum = Datum::decimal(decimal_from_i128_with_scale(LONG_MIN as i128 - 1, 0)).unwrap(); - - let result = datum.to(&Primitive(PrimitiveType::Long)).unwrap(); - - let expected = Datum::new(PrimitiveType::Long, PrimitiveLiteral::BelowMin); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_string_convert_to_boolean() { - let datum = Datum::string("true"); - - let result = datum.to(&Primitive(PrimitiveType::Boolean)).unwrap(); - - let expected = Datum::bool(true); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_string_convert_to_int() { - let datum = Datum::string("12345"); - - let result = datum.to(&Primitive(PrimitiveType::Int)).unwrap(); - - let expected = Datum::int(12345); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_string_convert_to_long() { - let datum = Datum::string("12345"); - - let result = datum.to(&Primitive(PrimitiveType::Long)).unwrap(); - - let expected = Datum::long(12345); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_string_convert_to_timestamp() { - let datum = Datum::string("1925-05-20T19:25:00.000"); - - let result = datum.to(&Primitive(PrimitiveType::Timestamp)).unwrap(); - - let expected = Datum::timestamp_micros(-1407990900000000); - - assert_eq!(result, expected); -} - -#[test] -fn test_datum_string_convert_to_timestamptz() { - let datum = Datum::string("1925-05-20T19:25:00.000 UTC"); - - let result = datum.to(&Primitive(PrimitiveType::Timestamptz)).unwrap(); - - let expected = Datum::timestamptz_micros(-1407990900000000); + let result = source.to(&Primitive(target_type)).unwrap(); + assert_eq!(result, expected); + } + )* + } + } - assert_eq!(result, expected); +datum_convert_tests! { + // Date conversions + convert_date_to_int: ( + Datum::date(12345) => Int, + Datum::int(12345) + ), + convert_int_to_date: ( + Datum::int(12345) => Date, + Datum::date(12345) + ), + convert_date_to_timestamp: ( + Datum::date(1) => Timestamp, + Datum::timestamp_from_datetime( + DateTime::parse_from_rfc3339("1970-01-02T00:00:00Z") + .unwrap() + .naive_utc() + ) + ), + convert_date_to_timestamp_nanos: ( + Datum::date(1) => TimestampNs, + Datum::timestamp_from_datetime( + DateTime::parse_from_rfc3339("1970-01-02T00:00:00Z") + .unwrap() + .naive_utc() + ) + .to(&Primitive(PrimitiveType::TimestampNs)) + .unwrap() + ), + convert_date_negative_to_timestamp: ( + Datum::date(-1) => Timestamp, + Datum::timestamp_from_datetime( + DateTime::parse_from_rfc3339("1969-12-31T00:00:00Z") + .unwrap() + .naive_utc() + ) + ), + convert_date_to_timestamp_year_9999: ( + Datum::date(2932896) => Timestamp, + Datum::timestamp_from_datetime( + DateTime::parse_from_rfc3339("9999-12-31T00:00:00Z") + .unwrap() + .naive_utc() + ) + ), + // Long conversions + convert_long_to_int: ( + Datum::long(12345) => Int, + Datum::int(12345) + ), + convert_long_to_int_above_max: ( + Datum::long(INT_MAX as i64 + 1) => Int, + Datum::new(PrimitiveType::Int, PrimitiveLiteral::AboveMax) + ), + convert_long_to_int_below_min: ( + Datum::long(INT_MIN as i64 - 1) => Int, + Datum::new(PrimitiveType::Int, PrimitiveLiteral::BelowMin) + ), + // Decimal conversions + convert_decimal_to_long: ( + Datum::decimal(decimal_new(12345, 0)).unwrap() => Long, + Datum::long(12345) + ), + convert_decimal_to_long_above_max: ( + Datum::decimal(decimal_from_i128_with_scale(LONG_MAX as i128 + 1, 0)).unwrap() => Long, + Datum::new(PrimitiveType::Long, PrimitiveLiteral::AboveMax) + ), + convert_decimal_to_long_below_min: ( + Datum::decimal(decimal_from_i128_with_scale(LONG_MIN as i128 - 1, 0)).unwrap() => Long, + Datum::new(PrimitiveType::Long, PrimitiveLiteral::BelowMin) + ), + // Float conversions + convert_double_to_float: ( + Datum::double(2.5) => Float, + Datum::float(2.5_f32) + ), + convert_double_to_float_above_max: ( + Datum::double(1e39) => Float, + Datum::new(PrimitiveType::Float, PrimitiveLiteral::AboveMax) + ), + convert_double_to_float_below_min: ( + Datum::double(-1e39) => Float, + Datum::new(PrimitiveType::Float, PrimitiveLiteral::BelowMin) + ), + convert_float_to_double: ( + Datum::float(2.5_f32) => Double, + Datum::double(2.5) + ), + // String conversions + convert_string_to_boolean: ( + Datum::string("true") => Boolean, + Datum::bool(true) + ), + convert_string_to_int: ( + Datum::string("12345") => Int, + Datum::int(12345) + ), + convert_string_to_long: ( + Datum::string("12345") => Long, + Datum::long(12345) + ), + convert_string_to_timestamp: ( + Datum::string("1925-05-20T19:25:00.000") => Timestamp, + Datum::timestamp_micros(-1407990900000000) + ), + convert_string_to_timestamptz: ( + Datum::string("1925-05-20T19:25:00.000 UTC") => Timestamptz, + Datum::timestamptz_micros(-1407990900000000) + ), + // Timestamp conversions (micros to nanos) + convert_timestamp_micros_to_timestamp_nanos: ( + Datum::timestamp_micros(12345) => TimestampNs, + Datum::timestamp_nanos(12345000) + ), + convert_timestamp_micros_to_timestamptz_nanos: ( + Datum::timestamp_micros(12345) => TimestamptzNs, + Datum::timestamptz_nanos(12345000) + ), + convert_timestamptz_micros_to_timestamp_nanos: ( + Datum::timestamptz_micros(12345) => TimestampNs, + Datum::timestamp_nanos(12345000) + ), + convert_timestamptz_micros_to_timestamptz_nanos: ( + Datum::timestamptz_micros(12345) => TimestamptzNs, + Datum::timestamptz_nanos(12345000) + ), + // Timestamp conversions (nanos to nanos) + convert_timestamp_nanos_to_timestamp_nanos: ( + Datum::timestamp_nanos(12345) => TimestampNs, + Datum::timestamp_nanos(12345) + ), + convert_timestamp_nanos_to_timestamptz_nanos: ( + Datum::timestamp_nanos(12345) => TimestamptzNs, + Datum::timestamptz_nanos(12345) + ), + convert_timestamptz_nanos_to_timestamp_nanos: ( + Datum::timestamptz_nanos(12345) => TimestampNs, + Datum::timestamp_nanos(12345) + ), + convert_timestamptz_nanos_to_timestamptz_nanos: ( + Datum::timestamptz_nanos(12345) => TimestamptzNs, + Datum::timestamptz_nanos(12345) + ), + // Same-type conversions + convert_time_to_time: ( + Datum::time_micros(1_000_000).unwrap() => Time, + Datum::time_micros(1_000_000).unwrap() + ), +} + +#[test] +fn test_convert_date_to_timestamp_overflow() { + let date = Datum::date(i32::MAX / 2); + + let result = date.clone().to(&Primitive(PrimitiveType::Timestamp)); + let result = result.expect_err("should overflow"); + assert!(result.to_string().contains("integer overflow")); + + let result = date.clone().to(&Primitive(PrimitiveType::TimestampNs)); + let result = result.expect_err("should overflow"); + assert!(result.to_string().contains("integer overflow")); +} + +#[test] +fn test_convert_timestamp_to_timestamp_overflow() { + let ts = Datum::timestamp_micros(i64::MAX / 10); + let result = ts.clone().to(&Primitive(PrimitiveType::TimestampNs)); + let result = result.expect_err("should overflow"); + assert!(result.to_string().contains("integer overflow")); + + let tstz = Datum::timestamptz_micros(i64::MAX / 10); + let result = tstz.clone().to(&Primitive(PrimitiveType::TimestamptzNs)); + let result = result.expect_err("should overflow"); + assert!(result.to_string().contains("integer overflow")); } #[test]