From 1522ac5a7d737bdfe74242af53768672f23b2db5 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Sat, 25 Nov 2023 09:49:50 +0800
Subject: [PATCH] Casting between floating and timestamp

---
 arrow-cast/src/cast.rs | 50 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 43 insertions(+), 7 deletions(-)

diff --git a/arrow-cast/src/cast.rs b/arrow-cast/src/cast.rs
index dd3e271afb0d..0a9cfc203dd4 100644
--- a/arrow-cast/src/cast.rs
+++ b/arrow-cast/src/cast.rs
@@ -231,8 +231,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
         (Time64(_), Time32(to_unit)) => {
             matches!(to_unit, Second | Millisecond)
         }
-        (Timestamp(_, _), _) if to_type.is_integer() => true,
-        (_, Timestamp(_, _)) if from_type.is_integer() => true,
+        (Timestamp(_, _), _) if to_type.is_integer() || to_type.is_floating() => true,
+        (_, Timestamp(_, _)) if from_type.is_integer() || from_type.is_floating() => true,
         (Date64, Timestamp(_, None)) => true,
         (Date32, Timestamp(_, None)) => true,
         (
@@ -1634,24 +1634,31 @@ pub fn cast_with_options(
                 .unary::<_, Time64MicrosecondType>(|x| x / (NANOSECONDS / MICROSECONDS)),
         )),
 
-        (Timestamp(TimeUnit::Second, _), _) if to_type.is_integer() => {
+        // Timestamp to integer/floating
+        (Timestamp(TimeUnit::Second, _), _) if to_type.is_integer() || to_type.is_floating() => {
             let array = cast_reinterpret_arrays::<TimestampSecondType, Int64Type>(array)?;
             cast_with_options(&array, to_type, cast_options)
         }
-        (Timestamp(TimeUnit::Millisecond, _), _) if to_type.is_integer() => {
+        (Timestamp(TimeUnit::Millisecond, _), _)
+            if to_type.is_integer() || to_type.is_floating() =>
+        {
             let array = cast_reinterpret_arrays::<TimestampMillisecondType, Int64Type>(array)?;
             cast_with_options(&array, to_type, cast_options)
         }
-        (Timestamp(TimeUnit::Microsecond, _), _) if to_type.is_integer() => {
+        (Timestamp(TimeUnit::Microsecond, _), _)
+            if to_type.is_integer() || to_type.is_floating() =>
+        {
             let array = cast_reinterpret_arrays::<TimestampMicrosecondType, Int64Type>(array)?;
             cast_with_options(&array, to_type, cast_options)
         }
-        (Timestamp(TimeUnit::Nanosecond, _), _) if to_type.is_integer() => {
+        (Timestamp(TimeUnit::Nanosecond, _), _)
+            if to_type.is_integer() || to_type.is_floating() =>
+        {
             let array = cast_reinterpret_arrays::<TimestampNanosecondType, Int64Type>(array)?;
             cast_with_options(&array, to_type, cast_options)
         }
 
-        (_, Timestamp(unit, tz)) if from_type.is_integer() => {
+        (_, Timestamp(unit, tz)) if from_type.is_integer() || from_type.is_floating() => {
             let array = cast_with_options(array, &Int64, cast_options)?;
             Ok(make_timestamp_array(
                 array.as_primitive(),
@@ -4719,6 +4726,35 @@ mod tests {
         assert_eq!(&actual, &expected);
     }
 
+    #[test]
+    fn test_cast_floating_to_timestamp() {
+        let array = Int64Array::from(vec![Some(2), Some(10), None]);
+        let expected = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        let array = Float32Array::from(vec![Some(2.0), Some(10.6), None]);
+        let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        assert_eq!(&actual, &expected);
+
+        let array = Float64Array::from(vec![Some(2.1), Some(10.2), None]);
+        let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        assert_eq!(&actual, &expected);
+    }
+
+    #[test]
+    fn test_cast_timestamp_to_floating() {
+        let array = TimestampMillisecondArray::from(vec![Some(5), Some(1), None])
+            .with_timezone("UTC".to_string());
+        let expected = cast(&array, &DataType::Int64).unwrap();
+
+        let actual = cast(&cast(&array, &DataType::Float32).unwrap(), &DataType::Int64).unwrap();
+        assert_eq!(&actual, &expected);
+
+        let actual = cast(&cast(&array, &DataType::Float64).unwrap(), &DataType::Int64).unwrap();
+        assert_eq!(&actual, &expected);
+    }
+
     #[test]
     fn test_cast_list_i32_to_list_u16() {
         let value_data = Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 100000000]).into_data();