get datetimes back out (incorrectly)

native/explorer/src/encoding.rs

@@ -6,7 +6,7 @@ use std::{mem, slice};
 
 use crate::atoms::{
     self, calendar, day, hour, infinity, microsecond, millisecond, minute, month, nan, nanosecond,
-    neg_infinity, precision, second, value, year,
+    neg_infinity, precision, second, std_offset, time_zone, utc_offset, value, year, zone_abbr,
 };
 use crate::datatypes::{
     days_to_date, time64ns_to_time, timestamp_to_datetime, ExSeries, ExSeriesRef,
@@ -103,7 +103,7 @@ fn date_series_to_list<'b>(s: &Series, env: Env<'b>) -> Result<Term<'b>, Explore
     ))
 }
 
-macro_rules! unsafe_encode_datetime {
+macro_rules! unsafe_encode_naive_datetime {
     ($v: expr, $naive_datetime_struct_keys: ident, $calendar_iso_module: ident, $naive_datetime_module: ident, $env: ident) => {{
         let dt = timestamp_to_datetime($v);
         let microseconds = dt.and_utc().timestamp_subsec_micros();
@@ -158,7 +158,7 @@ fn naive_datetime_struct_keys(env: Env) -> [NIF_TERM; 9] {
 }
 
 #[inline]
-fn datetime_to_microseconds(v: i64, time_unit: TimeUnit) -> i64 {
+fn naive_datetime_to_microseconds(v: i64, time_unit: TimeUnit) -> i64 {
     match time_unit {
         TimeUnit::Milliseconds => v * 1000,
         TimeUnit::Microseconds => v,
@@ -167,13 +167,13 @@ fn datetime_to_microseconds(v: i64, time_unit: TimeUnit) -> i64 {
 }
 
 #[inline]
-pub fn encode_datetime(v: i64, time_unit: TimeUnit, env: Env) -> Result<Term, ExplorerError> {
+pub fn encode_naive_datetime(v: i64, time_unit: TimeUnit, env: Env) -> Result<Term, ExplorerError> {
     let naive_datetime_struct_keys = &naive_datetime_struct_keys(env);
     let calendar_iso_module = atoms::calendar_iso_module().encode(env).as_c_arg();
     let naive_datetime_module = atoms::naive_datetime_module().encode(env).as_c_arg();
-    let microseconds_time = datetime_to_microseconds(v, time_unit);
+    let microseconds_time = naive_datetime_to_microseconds(v, time_unit);
 
-    Ok(unsafe_encode_datetime!(
+    Ok(unsafe_encode_naive_datetime!(
         microseconds_time,
         naive_datetime_struct_keys,
         calendar_iso_module,
@@ -183,7 +183,7 @@ pub fn encode_datetime(v: i64, time_unit: TimeUnit, env: Env) -> Result<Term, Ex
 }
 
 #[inline]
-fn datetime_series_to_list<'b>(
+fn naive_datetime_series_to_list<'b>(
     s: &Series,
     time_unit: TimeUnit,
     env: Env<'b>,
@@ -196,9 +196,9 @@ fn datetime_series_to_list<'b>(
         env,
         s.datetime()?.into_iter().map(|option| option
             .map(|v| {
-                let microseconds_time = datetime_to_microseconds(v, time_unit);
+                let microseconds_time = naive_datetime_to_microseconds(v, time_unit);
 
-                unsafe_encode_datetime!(
+                unsafe_encode_naive_datetime!(
                     microseconds_time,
                     naive_datetime_struct_keys,
                     calendar_iso_module,
@@ -210,6 +210,129 @@ fn datetime_series_to_list<'b>(
     ))
 }
 
+macro_rules! unsafe_encode_datetime {
+    (
+        $v: expr,
+        $time_zone: expr,
+        $datetime_struct_keys: ident,
+        $calendar_iso_module: ident,
+        $datetime_module: ident,
+        $env: ident
+    ) => {{
+        let dt = timestamp_to_datetime($v);
+        let microseconds = dt.and_utc().timestamp_subsec_micros();
+
+        // Limit the number of digits in the microsecond part of a timestamp to 6.
+        // This is necessary because the microsecond part of Elixir is only 6 digits.
+        let limited_ms = if microseconds > 999_999 {
+            999_999
+        } else {
+            microseconds
+        };
+
+        unsafe {
+            Term::new(
+                $env,
+                map::make_map_from_arrays(
+                    $env.as_c_arg(),
+                    $datetime_struct_keys,
+                    &[
+                        $datetime_module,
+                        $calendar_iso_module,
+                        dt.day().encode($env).as_c_arg(),
+                        dt.hour().encode($env).as_c_arg(),
+                        (limited_ms, 6).encode($env).as_c_arg(),
+                        dt.minute().encode($env).as_c_arg(),
+                        dt.month().encode($env).as_c_arg(),
+                        dt.second().encode($env).as_c_arg(),
+                        // std_offset
+                        0.encode($env).as_c_arg(),
+                        $time_zone.to_string().encode($env).as_c_arg(),
+                        // utc_offset
+                        0.encode($env).as_c_arg(),
+                        dt.year().encode($env).as_c_arg(),
+                        "zone_abbr".encode($env).as_c_arg(),
+                    ],
+                )
+                .unwrap(),
+            )
+        }
+    }};
+}
+
+// Here we build the DateTime struct manually, as it's much faster than using NifStruct
+// This is because we already have the keys (we know this at compile time), and the types,
+// so we can build the struct directly.
+fn datetime_struct_keys(env: Env) -> [NIF_TERM; 13] {
+    return [
+        atom::__struct__().encode(env).as_c_arg(),
+        calendar().encode(env).as_c_arg(),
+        day().encode(env).as_c_arg(),
+        hour().encode(env).as_c_arg(),
+        microsecond().encode(env).as_c_arg(),
+        minute().encode(env).as_c_arg(),
+        month().encode(env).as_c_arg(),
+        second().encode(env).as_c_arg(),
+        std_offset().encode(env).as_c_arg(),
+        time_zone().encode(env).as_c_arg(),
+        utc_offset().encode(env).as_c_arg(),
+        year().encode(env).as_c_arg(),
+        zone_abbr().encode(env).as_c_arg(),
+    ];
+}
+
+#[inline]
+pub fn encode_datetime(
+    v: i64,
+    time_unit: TimeUnit,
+    time_zone: String,
+    env: Env,
+) -> Result<Term, ExplorerError> {
+    let datetime_struct_keys = &datetime_struct_keys(env);
+    let calendar_iso_module = atoms::calendar_iso_module().encode(env).as_c_arg();
+    let datetime_module = atoms::datetime_module().encode(env).as_c_arg();
+    let microseconds_time = naive_datetime_to_microseconds(v, time_unit);
+
+    Ok(unsafe_encode_datetime!(
+        microseconds_time,
+        time_zone,
+        datetime_struct_keys,
+        calendar_iso_module,
+        datetime_module,
+        env
+    ))
+}
+
+#[inline]
+fn datetime_series_to_list<'b>(
+    s: &Series,
+    time_unit: TimeUnit,
+    time_zone: String,
+    env: Env<'b>,
+) -> Result<Term<'b>, ExplorerError> {
+    let datetime_struct_keys = &datetime_struct_keys(env);
+    let calendar_iso_module = atoms::calendar_iso_module().encode(env).as_c_arg();
+    let datetime_module = atoms::datetime_module().encode(env).as_c_arg();
+
+    Ok(unsafe_iterator_series_to_list!(
+        env,
+        s.datetime()?.into_iter().map(|option| option
+            .map(|v| {
+                let microseconds_time = naive_datetime_to_microseconds(v, time_unit);
+
+                unsafe_encode_datetime!(
+                    microseconds_time,
+                    time_zone,
+                    datetime_struct_keys,
+                    calendar_iso_module,
+                    datetime_module,
+                    env
+                )
+            })
+            .encode(env))
+    ))
+}
+
 fn time_unit_to_atom(time_unit: TimeUnit) -> atom::Atom {
     match time_unit {
         TimeUnit::Milliseconds => millisecond(),
@@ -565,7 +688,10 @@ pub fn term_from_value<'b>(v: AnyValue, env: Env<'b>) -> Result<Term<'b>, Explor
         AnyValue::Float64(v) => Ok(Some(term_from_float64(v, env)).encode(env)),
         AnyValue::Date(v) => encode_date(v, env),
         AnyValue::Time(v) => encode_time(v, env),
-        AnyValue::Datetime(v, time_unit, None) => encode_datetime(v, time_unit, env),
+        AnyValue::Datetime(v, time_unit, None) => encode_naive_datetime(v, time_unit, env),
+        AnyValue::Datetime(v, time_unit, Some(time_zone)) => {
+            encode_datetime(v, time_unit, time_zone.to_string(), env)
+        }
         AnyValue::Duration(v, time_unit) => encode_duration(v, time_unit, env),
         AnyValue::Categorical(idx, mapping, _) => Ok(mapping.get(idx).encode(env)),
         AnyValue::List(series) => list_from_series(ExSeries::new(series), env),
@@ -599,7 +725,10 @@ pub fn list_from_series(s: ExSeries, env: Env) -> Result<Term, ExplorerError> {
 
         DataType::Date => date_series_to_list(&s, env),
         DataType::Time => time_series_to_list(&s, env),
-        DataType::Datetime(time_unit, None) => datetime_series_to_list(&s, *time_unit, env),
+        DataType::Datetime(time_unit, None) => naive_datetime_series_to_list(&s, *time_unit, env),
+        DataType::Datetime(time_unit, Some(time_zone)) => {
+            datetime_series_to_list(&s, *time_unit, time_zone.clone().to_string(), env)
+        }
         DataType::Duration(time_unit) => duration_series_to_list(&s, *time_unit, env),
         DataType::Binary => generic_binary_series_to_list(&s.resource, &s, env),
         DataType::String => generic_string_series_to_list(&s, env),

native/explorer/src/lib.rs

@@ -50,6 +50,7 @@ mod atoms {
     rustler::atoms! {
         calendar_iso_module = "Elixir.Calendar.ISO",
         date_module = "Elixir.Date",
+        datetime_module = "Elixir.DateTime",
         duration_module = "Elixir.Explorer.Duration",
         naive_datetime_module = "Elixir.NaiveDateTime",
         time_module = "Elixir.Time",
@@ -67,7 +68,11 @@ mod atoms {
         calendar,
         nan,
         infinity,
-        neg_infinity
+        neg_infinity,
+        std_offset,
+        time_zone,
+        utc_offset,
+        zone_abbr,
     }
 }
 

native/explorer/src/series.rs

@@ -7,7 +7,8 @@ use crate::{
     encoding, ExDataFrame, ExSeries, ExplorerError,
 };
 
-use encoding::encode_datetime;
+use encoding::encode_naive_datetime;
+// use encoding::encode_datetime;
 
 use polars::prelude::*;
 use polars_ops::chunked_array::cov::{cov, pearson_corr};
@@ -1007,7 +1008,7 @@ pub fn s_min(env: Env, s: ExSeries) -> Result<Term, ExplorerError> {
         DataType::Time => Ok(s.min::<i64>()?.map(ExTime::from).encode(env)),
         DataType::Datetime(unit, _) => Ok(s
             .min::<i64>()?
-            .map(|v| encode_datetime(v, *unit, env).unwrap())
+            .map(|v| encode_naive_datetime(v, *unit, env).unwrap())
             .encode(env)),
         dt => panic!("min/1 not implemented for {dt:?}"),
     }
@@ -1029,7 +1030,7 @@ pub fn s_max(env: Env, s: ExSeries) -> Result<Term, ExplorerError> {
         DataType::Time => Ok(s.max::<i64>()?.map(ExTime::from).encode(env)),
         DataType::Datetime(unit, _) => Ok(s
             .max::<i64>()?
-            .map(|v| encode_datetime(v, *unit, env).unwrap())
+            .map(|v| encode_naive_datetime(v, *unit, env).unwrap())
             .encode(env)),
         dt => panic!("max/1 not implemented for {dt:?}"),
     }
@@ -1230,7 +1231,7 @@ pub fn s_quantile<'a>(
         },
         DataType::Datetime(unit, None) => match s.datetime()?.quantile(quantile, strategy)? {
             None => Ok(None::<ExNaiveDateTime>.encode(env)),
-            Some(time) => Ok(encode_datetime(time as i64, *unit, env)
+            Some(time) => Ok(encode_naive_datetime(time as i64, *unit, env)
                 .unwrap()
                 .encode(env)),
         },

test/explorer/series/datetime_test.exs

@@ -50,13 +50,18 @@ defmodule Explorer.Series.DateTimeTest do
 
   describe "timezones" do
     test "UTC" do
-      datetimes = [~U[2024-01-01T12:00:00Z], ~U[2024-01-01T13:00:00Z], ~U[2024-01-01T14:00:00Z]]
+      datetimes_in = [
+        ~U[2024-01-01T12:00:00.000000Z],
+        ~U[2024-01-01T13:00:00.000000Z],
+        ~U[2024-01-01T14:00:00.000000Z]
+      ]
 
-      datetimes
-      |> Series.from_list()
-      |> IO.inspect()
-      |> Series.to_list()
-      |> IO.inspect()
+      datetimes_out =
+        datetimes_in
+        |> Series.from_list()
+        |> Series.to_list()
+
+      assert datetimes_out == datetimes_in
     end
   end
 end