1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

use arrow::array::PrimitiveArray;
use arrow::compute::arity::unary;
use arrow::datatypes::{DataType as ArrowDataType, TimeUnit};
use arrow::temporal_conversions::{
    timestamp_ms_to_datetime, timestamp_ns_to_datetime, timestamp_us_to_datetime,
};

use crate::prelude::ArrayRef;

pub fn cast_timezone(
    arr: &PrimitiveArray<i64>,
    tu: TimeUnit,
    from: chrono_tz::Tz,
    to: chrono_tz::Tz,
) -> ArrayRef {
    use chrono::TimeZone;

    match tu {
        TimeUnit::Millisecond => Box::new(unary(
            arr,
            |value| {
                let ndt = timestamp_ms_to_datetime(value);
                let tz_aware = from.from_local_datetime(&ndt).unwrap();
                let new_tz_aware = tz_aware.with_timezone(&to);
                new_tz_aware.naive_local().timestamp_millis()
            },
            ArrowDataType::Int64,
        )),
        TimeUnit::Microsecond => Box::new(unary(
            arr,
            |value| {
                let ndt = timestamp_us_to_datetime(value);
                let tz_aware = from.from_local_datetime(&ndt).unwrap();
                let new_tz_aware = tz_aware.with_timezone(&to);
                new_tz_aware.naive_local().timestamp_micros()
            },
            ArrowDataType::Int64,
        )),
        TimeUnit::Nanosecond => Box::new(unary(
            arr,
            |value| {
                let ndt = timestamp_ns_to_datetime(value);
                let tz_aware = from.from_local_datetime(&ndt).unwrap();
                let new_tz_aware = tz_aware.with_timezone(&to);
                new_tz_aware.naive_local().timestamp_nanos()
            },
            ArrowDataType::Int64,
        )),
        _ => unreachable!(),
    }
}