From 06084d0d7dd1de20fc26d5431a78684c96481208 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Sat, 4 Nov 2023 19:29:20 -0700
Subject: [PATCH] Cast from integer to timestamp

---
 arrow-cast/src/cast.rs | 178 ++++++++++++++++++++++++++++++++++++-----
 1 file changed, 160 insertions(+), 18 deletions(-)

diff --git a/arrow-cast/src/cast.rs b/arrow-cast/src/cast.rs
index 684e02b87e6c..b86ac5f8c8a6 100644
--- a/arrow-cast/src/cast.rs
+++ b/arrow-cast/src/cast.rs
@@ -227,8 +227,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
         (Time64(_), Time32(to_unit)) => {
             matches!(to_unit, Second | Millisecond)
         }
-        (Timestamp(_, _), Int64) => true,
-        (Int64, Timestamp(_, _)) => true,
+        (Timestamp(_, _), _) if to_type.is_integer() => true,
+        (_, Timestamp(_, _)) if from_type.is_integer() => true,
         (Date64, Timestamp(_, None)) => true,
         (Date32, Timestamp(_, None)) => true,
         (
@@ -1621,24 +1621,104 @@ pub fn cast_with_options(
                 .unary::<_, Time64MicrosecondType>(|x| x / (NANOSECONDS / MICROSECONDS)),
         )),
 
-        (Timestamp(TimeUnit::Second, _), Int64) => {
-            cast_reinterpret_arrays::<TimestampSecondType, Int64Type>(array)
+        (Timestamp(TimeUnit::Second, _), _) if to_type.is_integer() => {
+            let array = cast_reinterpret_arrays::<TimestampSecondType, Int64Type>(array)?;
+
+            match to_type {
+                Int8 => cast_numeric_arrays::<Int64Type, Int8Type>(&array, cast_options),
+                Int16 => cast_numeric_arrays::<Int64Type, Int16Type>(&array, cast_options),
+                Int32 => cast_numeric_arrays::<Int64Type, Int32Type>(&array, cast_options),
+                Int64 => Ok(array),
+                UInt8 => cast_numeric_arrays::<Int64Type, UInt8Type>(&array, cast_options),
+                UInt16 => cast_numeric_arrays::<Int64Type, UInt16Type>(&array, cast_options),
+                UInt32 => cast_numeric_arrays::<Int64Type, UInt32Type>(&array, cast_options),
+                UInt64 => cast_numeric_arrays::<Int64Type, UInt64Type>(&array, cast_options),
+                _ => Err(ArrowError::CastError(format!(
+                    "Casting from {:?} to {:?} not supported",
+                    from_type, to_type
+                ))),
+            }
         }
-        (Timestamp(TimeUnit::Millisecond, _), Int64) => {
-            cast_reinterpret_arrays::<TimestampMillisecondType, Int64Type>(array)
+        (Timestamp(TimeUnit::Millisecond, _), _) if to_type.is_integer() => {
+            let array = cast_reinterpret_arrays::<TimestampMillisecondType, Int64Type>(array)?;
+
+            match to_type {
+                Int8 => cast_numeric_arrays::<Int64Type, Int8Type>(&array, cast_options),
+                Int16 => cast_numeric_arrays::<Int64Type, Int16Type>(&array, cast_options),
+                Int32 => cast_numeric_arrays::<Int64Type, Int32Type>(&array, cast_options),
+                Int64 => Ok(array),
+                UInt8 => cast_numeric_arrays::<Int64Type, UInt8Type>(&array, cast_options),
+                UInt16 => cast_numeric_arrays::<Int64Type, UInt16Type>(&array, cast_options),
+                UInt32 => cast_numeric_arrays::<Int64Type, UInt32Type>(&array, cast_options),
+                UInt64 => cast_numeric_arrays::<Int64Type, UInt64Type>(&array, cast_options),
+                _ => Err(ArrowError::CastError(format!(
+                    "Casting from {:?} to {:?} not supported",
+                    from_type, to_type
+                ))),
+            }
         }
-        (Timestamp(TimeUnit::Microsecond, _), Int64) => {
-            cast_reinterpret_arrays::<TimestampMicrosecondType, Int64Type>(array)
+        (Timestamp(TimeUnit::Microsecond, _), _) if to_type.is_integer() => {
+            let array = cast_reinterpret_arrays::<TimestampMicrosecondType, Int64Type>(array)?;
+
+            match to_type {
+                Int8 => cast_numeric_arrays::<Int64Type, Int8Type>(&array, cast_options),
+                Int16 => cast_numeric_arrays::<Int64Type, Int16Type>(&array, cast_options),
+                Int32 => cast_numeric_arrays::<Int64Type, Int32Type>(&array, cast_options),
+                Int64 => Ok(array),
+                UInt8 => cast_numeric_arrays::<Int64Type, UInt8Type>(&array, cast_options),
+                UInt16 => cast_numeric_arrays::<Int64Type, UInt16Type>(&array, cast_options),
+                UInt32 => cast_numeric_arrays::<Int64Type, UInt32Type>(&array, cast_options),
+                UInt64 => cast_numeric_arrays::<Int64Type, UInt64Type>(&array, cast_options),
+                _ => Err(ArrowError::CastError(format!(
+                    "Casting from {:?} to {:?} not supported",
+                    from_type, to_type
+                ))),
+            }
         }
-        (Timestamp(TimeUnit::Nanosecond, _), Int64) => {
-            cast_reinterpret_arrays::<TimestampNanosecondType, Int64Type>(array)
+        (Timestamp(TimeUnit::Nanosecond, _), _) if to_type.is_integer() => {
+            let array = cast_reinterpret_arrays::<TimestampNanosecondType, Int64Type>(array)?;
+
+            match to_type {
+                Int8 => cast_numeric_arrays::<Int64Type, Int8Type>(&array, cast_options),
+                Int16 => cast_numeric_arrays::<Int64Type, Int16Type>(&array, cast_options),
+                Int32 => cast_numeric_arrays::<Int64Type, Int32Type>(&array, cast_options),
+                Int64 => Ok(array),
+                UInt8 => cast_numeric_arrays::<Int64Type, UInt8Type>(&array, cast_options),
+                UInt16 => cast_numeric_arrays::<Int64Type, UInt16Type>(&array, cast_options),
+                UInt32 => cast_numeric_arrays::<Int64Type, UInt32Type>(&array, cast_options),
+                UInt64 => cast_numeric_arrays::<Int64Type, UInt64Type>(&array, cast_options),
+                _ => Err(ArrowError::CastError(format!(
+                    "Casting from {:?} to {:?} not supported",
+                    from_type, to_type
+                ))),
+            }
         }
 
-        (Int64, Timestamp(unit, tz)) => Ok(make_timestamp_array(
-            array.as_primitive(),
-            unit.clone(),
-            tz.clone(),
-        )),
+        (_, Timestamp(unit, tz)) if from_type.is_integer() => {
+            let array = match from_type {
+                Int8 => cast_numeric_arrays::<Int8Type, Int64Type>(array, cast_options),
+                Int16 => cast_numeric_arrays::<Int16Type, Int64Type>(array, cast_options),
+                Int32 => cast_numeric_arrays::<Int32Type, Int64Type>(array, cast_options),
+                Int64 => {
+                    return Ok(make_timestamp_array(
+                        array.as_primitive(),
+                        unit.clone(),
+                        tz.clone(),
+                    ))
+                }
+                UInt8 => cast_numeric_arrays::<UInt8Type, Int64Type>(array, cast_options),
+                UInt16 => cast_numeric_arrays::<UInt16Type, Int64Type>(array, cast_options),
+                UInt32 => cast_numeric_arrays::<UInt32Type, Int64Type>(array, cast_options),
+                UInt64 => cast_numeric_arrays::<UInt64Type, Int64Type>(array, cast_options),
+                _ => unreachable!(),
+            }?;
+
+            Ok(make_timestamp_array(
+                array.as_primitive(),
+                unit.clone(),
+                tz.clone(),
+            ))
+        }
 
         (Timestamp(from_unit, from_tz), Timestamp(to_unit, to_tz)) => {
             let array = cast_with_options(array, &Int64, cast_options)?;
@@ -4559,10 +4639,72 @@ mod tests {
     }
 
     #[test]
-    #[should_panic(expected = "Casting from Int32 to Timestamp(Microsecond, None) not supported")]
-    fn test_cast_int32_to_timestamp() {
+    fn test_cast_integer_to_timestamp() {
+        let array = Int64Array::from(vec![Some(2), Some(10), None]);
+        let expected = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        let array = Int8Array::from(vec![Some(2), Some(10), None]);
+        let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        assert_eq!(&actual, &expected);
+
+        let array = Int16Array::from(vec![Some(2), Some(10), None]);
+        let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        assert_eq!(&actual, &expected);
+
         let array = Int32Array::from(vec![Some(2), Some(10), None]);
-        cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+        let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        assert_eq!(&actual, &expected);
+
+        let array = UInt8Array::from(vec![Some(2), Some(10), None]);
+        let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        assert_eq!(&actual, &expected);
+
+        let array = UInt16Array::from(vec![Some(2), Some(10), None]);
+        let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        assert_eq!(&actual, &expected);
+
+        let array = UInt32Array::from(vec![Some(2), Some(10), None]);
+        let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        assert_eq!(&actual, &expected);
+
+        let array = UInt64Array::from(vec![Some(2), Some(10), None]);
+        let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
+
+        assert_eq!(&actual, &expected);
+    }
+
+    #[test]
+    fn test_cast_timestamp_to_integer() {
+        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::Int8).unwrap(), &DataType::Int64).unwrap();
+        assert_eq!(&actual, &expected);
+
+        let actual = cast(&cast(&array, &DataType::Int16).unwrap(), &DataType::Int64).unwrap();
+        assert_eq!(&actual, &expected);
+
+        let actual = cast(&cast(&array, &DataType::Int32).unwrap(), &DataType::Int64).unwrap();
+        assert_eq!(&actual, &expected);
+
+        let actual = cast(&cast(&array, &DataType::UInt8).unwrap(), &DataType::Int64).unwrap();
+        assert_eq!(&actual, &expected);
+
+        let actual = cast(&cast(&array, &DataType::UInt16).unwrap(), &DataType::Int64).unwrap();
+        assert_eq!(&actual, &expected);
+
+        let actual = cast(&cast(&array, &DataType::UInt32).unwrap(), &DataType::Int64).unwrap();
+        assert_eq!(&actual, &expected);
+
+        let actual = cast(&cast(&array, &DataType::UInt64).unwrap(), &DataType::Int64).unwrap();
+        assert_eq!(&actual, &expected);
     }
 
     #[test]