@@ -167,6 +167,16 @@ impl TransformFunction for Bucket {
167
167
. downcast_ref :: < arrow_array:: TimestampMicrosecondArray > ( )
168
168
. unwrap ( )
169
169
. 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 ) ) ,
170
180
DataType :: Utf8 => arrow_array:: Int32Array :: from_iter (
171
181
input
172
182
. as_any ( )
@@ -228,6 +238,13 @@ impl TransformFunction for Bucket {
228
238
( PrimitiveType :: Date , PrimitiveLiteral :: Int ( v) ) => self . bucket_date ( * v) ,
229
239
( PrimitiveType :: Time , PrimitiveLiteral :: Long ( v) ) => self . bucket_time ( * v) ,
230
240
( 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
+ }
231
248
( PrimitiveType :: String , PrimitiveLiteral :: String ( v) ) => self . bucket_str ( v. as_str ( ) ) ,
232
249
( PrimitiveType :: Uuid , PrimitiveLiteral :: UInt128 ( v) ) => {
233
250
self . bucket_bytes ( uuid:: Uuid :: from_u128 ( * v) . as_ref ( ) )
@@ -250,6 +267,9 @@ impl TransformFunction for Bucket {
250
267
251
268
#[ cfg( test) ]
252
269
mod test {
270
+ use std:: sync:: Arc ;
271
+
272
+ use arrow_array:: { ArrayRef , Int32Array , TimestampMicrosecondArray , TimestampNanosecondArray } ;
253
273
use chrono:: { DateTime , NaiveDate , NaiveDateTime , NaiveTime } ;
254
274
255
275
use super :: Bucket ;
@@ -888,4 +908,66 @@ mod test {
888
908
Datum :: int( 32 )
889
909
) ;
890
910
}
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
+ }
891
973
}
0 commit comments