(iceberg-rust) branch main updated: feat: Support `TimestampNs` and TimestampTzNs` in bucket transform (#1150)

[liurenjie1024]

This is an automated email from the ASF dual-hosted git repository.

liurenjie1024 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git


The following commit(s) were added to refs/heads/main by this push:
     new f17dd4d2 feat: Support `TimestampNs` and TimestampTzNs` in bucket 
transform (#1150)
f17dd4d2 is described below

commit f17dd4d230116e16401733857e474083924a634c
Author: Jonathan Chen <chenleejonat...@gmail.com>
AuthorDate: Tue Apr 1 23:35:02 2025 -0400

    feat: Support `TimestampNs` and TimestampTzNs` in bucket transform (#1150)
    
    ## Which issue does this PR close?
    
    - Closes #1110 .
    
    ## What changes are included in this PR?
    Add bucket transforms for `TimestampNs` and TimestampTzNs`
    
    
    ## Are these changes tested?
    Unit tests
    
    ---------
    
    Co-authored-by: Renjie Liu <liurenjie2...@gmail.com>
---
 crates/iceberg/src/transform/bucket.rs | 82 ++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)

diff --git a/crates/iceberg/src/transform/bucket.rs 
b/crates/iceberg/src/transform/bucket.rs
index ce39826b..b64a4631 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::<arrow_array::TimestampMicrosecondArray>()
                 .unwrap()
                 .unary(|v| self.bucket_timestamp(v)),
+            DataType::Time64(TimeUnit::Nanosecond) => input
+                .as_any()
+                .downcast_ref::<arrow_array::Time64NanosecondArray>()
+                .unwrap()
+                .unary(|v| self.bucket_time(v / 1000)),
+            DataType::Timestamp(TimeUnit::Nanosecond, _) => input
+                .as_any()
+                .downcast_ref::<arrow_array::TimestampNanosecondArray>()
+                .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::<Int32Array>()
+            .unwrap();
+        let nano_result = transformed_nano
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap();
+
+        assert_eq!(micro_result.value(0), nano_result.value(0));
+    }
 }