Skip to content

Commit f17dd4d

Browse files
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 <[email protected]>
1 parent b710a48 commit f17dd4d

File tree

1 file changed

+82
-0
lines changed

1 file changed

+82
-0
lines changed

crates/iceberg/src/transform/bucket.rs

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -167,6 +167,16 @@ impl TransformFunction for Bucket {
167167
.downcast_ref::<arrow_array::TimestampMicrosecondArray>()
168168
.unwrap()
169169
.unary(|v| self.bucket_timestamp(v)),
170+
DataType::Time64(TimeUnit::Nanosecond) => input
171+
.as_any()
172+
.downcast_ref::<arrow_array::Time64NanosecondArray>()
173+
.unwrap()
174+
.unary(|v| self.bucket_time(v / 1000)),
175+
DataType::Timestamp(TimeUnit::Nanosecond, _) => input
176+
.as_any()
177+
.downcast_ref::<arrow_array::TimestampNanosecondArray>()
178+
.unwrap()
179+
.unary(|v| self.bucket_timestamp(v / 1000)),
170180
DataType::Utf8 => arrow_array::Int32Array::from_iter(
171181
input
172182
.as_any()
@@ -228,6 +238,13 @@ impl TransformFunction for Bucket {
228238
(PrimitiveType::Date, PrimitiveLiteral::Int(v)) => self.bucket_date(*v),
229239
(PrimitiveType::Time, PrimitiveLiteral::Long(v)) => self.bucket_time(*v),
230240
(PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => self.bucket_timestamp(*v),
241+
(PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => self.bucket_timestamp(*v),
242+
(PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => {
243+
self.bucket_timestamp(*v / 1000)
244+
}
245+
(PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => {
246+
self.bucket_timestamp(*v / 1000)
247+
}
231248
(PrimitiveType::String, PrimitiveLiteral::String(v)) => self.bucket_str(v.as_str()),
232249
(PrimitiveType::Uuid, PrimitiveLiteral::UInt128(v)) => {
233250
self.bucket_bytes(uuid::Uuid::from_u128(*v).as_ref())
@@ -250,6 +267,9 @@ impl TransformFunction for Bucket {
250267

251268
#[cfg(test)]
252269
mod test {
270+
use std::sync::Arc;
271+
272+
use arrow_array::{ArrayRef, Int32Array, TimestampMicrosecondArray, TimestampNanosecondArray};
253273
use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime};
254274

255275
use super::Bucket;
@@ -888,4 +908,66 @@ mod test {
888908
Datum::int(32)
889909
);
890910
}
911+
912+
#[test]
913+
fn test_timestamptz_literal() {
914+
let bucket = Bucket::new(100);
915+
assert_eq!(
916+
bucket
917+
.transform_literal(&Datum::timestamptz_micros(1510871468000000))
918+
.unwrap()
919+
.unwrap(),
920+
Datum::int(7)
921+
);
922+
}
923+
924+
#[test]
925+
fn test_timestamp_ns_literal() {
926+
let bucket = Bucket::new(100);
927+
let ns_value = 1510871468000000i64 * 1000;
928+
assert_eq!(
929+
bucket
930+
.transform_literal(&Datum::timestamp_nanos(ns_value))
931+
.unwrap()
932+
.unwrap(),
933+
Datum::int(7)
934+
);
935+
}
936+
937+
#[test]
938+
fn test_timestamptz_ns_literal() {
939+
let bucket = Bucket::new(100);
940+
let ns_value = 1510871468000000i64 * 1000;
941+
assert_eq!(
942+
bucket
943+
.transform_literal(&Datum::timestamptz_nanos(ns_value))
944+
.unwrap()
945+
.unwrap(),
946+
Datum::int(7)
947+
);
948+
}
949+
950+
#[test]
951+
fn test_transform_timestamp_nanos_and_micros_array_equivalence() {
952+
let bucket = Bucket::new(100);
953+
let micros_value = 1510871468000000;
954+
let nanos_value = micros_value * 1000;
955+
956+
let micro_array = TimestampMicrosecondArray::from_iter_values(vec![micros_value]);
957+
let nano_array = TimestampNanosecondArray::from_iter_values(vec![nanos_value]);
958+
959+
let transformed_micro: ArrayRef = bucket.transform(Arc::new(micro_array)).unwrap();
960+
let transformed_nano: ArrayRef = bucket.transform(Arc::new(nano_array)).unwrap();
961+
962+
let micro_result = transformed_micro
963+
.as_any()
964+
.downcast_ref::<Int32Array>()
965+
.unwrap();
966+
let nano_result = transformed_nano
967+
.as_any()
968+
.downcast_ref::<Int32Array>()
969+
.unwrap();
970+
971+
assert_eq!(micro_result.value(0), nano_result.value(0));
972+
}
891973
}

0 commit comments

Comments
 (0)