diff --git a/crates/iceberg/src/transform/bucket.rs b/crates/iceberg/src/transform/bucket.rs index ce39826bb3..b64a4631d0 100644 --- a/crates/iceberg/src/transform/bucket.rs +++ b/crates/iceberg/src/transform/bucket.rs @@ -167,6 +167,16 @@ impl TransformFunction for Bucket { .downcast_ref::() .unwrap() .unary(|v| self.bucket_timestamp(v)), + DataType::Time64(TimeUnit::Nanosecond) => input + .as_any() + .downcast_ref::() + .unwrap() + .unary(|v| self.bucket_time(v / 1000)), + DataType::Timestamp(TimeUnit::Nanosecond, _) => input + .as_any() + .downcast_ref::() + .unwrap() + .unary(|v| self.bucket_timestamp(v / 1000)), DataType::Utf8 => arrow_array::Int32Array::from_iter( input .as_any() @@ -228,6 +238,13 @@ impl TransformFunction for Bucket { (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => self.bucket_date(*v), (PrimitiveType::Time, PrimitiveLiteral::Long(v)) => self.bucket_time(*v), (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => self.bucket_timestamp(*v), + (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => self.bucket_timestamp(*v), + (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => { + self.bucket_timestamp(*v / 1000) + } + (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => { + self.bucket_timestamp(*v / 1000) + } (PrimitiveType::String, PrimitiveLiteral::String(v)) => self.bucket_str(v.as_str()), (PrimitiveType::Uuid, PrimitiveLiteral::UInt128(v)) => { self.bucket_bytes(uuid::Uuid::from_u128(*v).as_ref()) @@ -250,6 +267,9 @@ impl TransformFunction for Bucket { #[cfg(test)] mod test { + use std::sync::Arc; + + use arrow_array::{ArrayRef, Int32Array, TimestampMicrosecondArray, TimestampNanosecondArray}; use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime}; use super::Bucket; @@ -888,4 +908,66 @@ mod test { Datum::int(32) ); } + + #[test] + fn test_timestamptz_literal() { + let bucket = Bucket::new(100); + assert_eq!( + bucket + .transform_literal(&Datum::timestamptz_micros(1510871468000000)) + .unwrap() + .unwrap(), + Datum::int(7) + ); + } + + #[test] + fn test_timestamp_ns_literal() { + let bucket = Bucket::new(100); + let ns_value = 1510871468000000i64 * 1000; + assert_eq!( + bucket + .transform_literal(&Datum::timestamp_nanos(ns_value)) + .unwrap() + .unwrap(), + Datum::int(7) + ); + } + + #[test] + fn test_timestamptz_ns_literal() { + let bucket = Bucket::new(100); + let ns_value = 1510871468000000i64 * 1000; + assert_eq!( + bucket + .transform_literal(&Datum::timestamptz_nanos(ns_value)) + .unwrap() + .unwrap(), + Datum::int(7) + ); + } + + #[test] + fn test_transform_timestamp_nanos_and_micros_array_equivalence() { + let bucket = Bucket::new(100); + let micros_value = 1510871468000000; + let nanos_value = micros_value * 1000; + + let micro_array = TimestampMicrosecondArray::from_iter_values(vec![micros_value]); + let nano_array = TimestampNanosecondArray::from_iter_values(vec![nanos_value]); + + let transformed_micro: ArrayRef = bucket.transform(Arc::new(micro_array)).unwrap(); + let transformed_nano: ArrayRef = bucket.transform(Arc::new(nano_array)).unwrap(); + + let micro_result = transformed_micro + .as_any() + .downcast_ref::() + .unwrap(); + let nano_result = transformed_nano + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(micro_result.value(0), nano_result.value(0)); + } }