starting to work

lib/explorer/polars_backend/native.ex

@@ -341,7 +341,7 @@ defmodule Explorer.PolarsBackend.Native do
   def s_from_list_bool(_name, _val), do: err()
   def s_from_list_date(_name, _val), do: err()
   def s_from_list_time(_name, _val), do: err()
-  def s_from_list_datetime(_name, _val, _precision), do: err()
+  def s_from_list_datetime(_name, _val, _precision, _tz), do: err()
   def s_from_list_duration(_name, _val, _precision), do: err()
   def s_from_list_f32(_name, _val), do: err()
   def s_from_list_f64(_name, _val), do: err()

lib/explorer/polars_backend/shared.ex

@@ -170,27 +170,67 @@ defmodule Explorer.PolarsBackend.Shared do
   def from_list(list, dtype, name) when is_list(list) do
     case dtype do
       # Signed integers
-      {:s, 8} -> Native.s_from_list_s8(name, list)
-      {:s, 16} -> Native.s_from_list_s16(name, list)
-      {:s, 32} -> Native.s_from_list_s32(name, list)
-      {:s, 64} -> Native.s_from_list_s64(name, list)
+      {:s, 8} ->
+        Native.s_from_list_s8(name, list)
+
+      {:s, 16} ->
+        Native.s_from_list_s16(name, list)
+
+      {:s, 32} ->
+        Native.s_from_list_s32(name, list)
+
+      {:s, 64} ->
+        Native.s_from_list_s64(name, list)
+
       # Unsigned integers
-      {:u, 8} -> Native.s_from_list_u8(name, list)
-      {:u, 16} -> Native.s_from_list_u16(name, list)
-      {:u, 32} -> Native.s_from_list_u32(name, list)
-      {:u, 64} -> Native.s_from_list_u64(name, list)
+      {:u, 8} ->
+        Native.s_from_list_u8(name, list)
+
+      {:u, 16} ->
+        Native.s_from_list_u16(name, list)
+
+      {:u, 32} ->
+        Native.s_from_list_u32(name, list)
+
+      {:u, 64} ->
+        Native.s_from_list_u64(name, list)
+
       # Floats
-      {:f, 32} -> Native.s_from_list_f32(name, list)
-      {:f, 64} -> Native.s_from_list_f64(name, list)
-      :boolean -> Native.s_from_list_bool(name, list)
-      :string -> Native.s_from_list_str(name, list)
-      :category -> Native.s_from_list_categories(name, list)
-      :date -> Native.s_from_list_date(name, list)
-      :time -> Native.s_from_list_time(name, list)
-      {:datetime, precision} -> Native.s_from_list_datetime(name, list, Atom.to_string(precision))
-      {:duration, precision} -> Native.s_from_list_duration(name, list, Atom.to_string(precision))
-      :binary -> Native.s_from_list_binary(name, list)
-      :null -> Native.s_from_list_null(name, length(list))
+      {:f, 32} ->
+        Native.s_from_list_f32(name, list)
+
+      {:f, 64} ->
+        Native.s_from_list_f64(name, list)
+
+      :boolean ->
+        Native.s_from_list_bool(name, list)
+
+      :string ->
+        Native.s_from_list_str(name, list)
+
+      :category ->
+        Native.s_from_list_categories(name, list)
+
+      :date ->
+        Native.s_from_list_date(name, list)
+
+      :time ->
+        Native.s_from_list_time(name, list)
+
+      {:datetime, precision} ->
+        Native.s_from_list_datetime(name, list, Atom.to_string(precision), nil)
+
+      {:datetime, precision, tz} ->
+        Native.s_from_list_datetime(name, list, Atom.to_string(precision), tz)
+
+      {:duration, precision} ->
+        Native.s_from_list_duration(name, list, Atom.to_string(precision))
+
+      :binary ->
+        Native.s_from_list_binary(name, list)
+
+      :null ->
+        Native.s_from_list_null(name, length(list))
     end
   end
 

lib/explorer/shared.ex

@@ -311,8 +311,10 @@ defmodule Explorer.Shared do
 
   defp infer_type(%Date{} = _item), do: :date
   defp infer_type(%Time{} = _item), do: :time
+  defp infer_type(%DateTime{time_zone: tz} = _item), do: {:datetime, :microsecond, tz}
   defp infer_type(%NaiveDateTime{} = _item), do: {:datetime, :microsecond}
   defp infer_type(%Explorer.Duration{precision: precision} = _item), do: {:duration, precision}
+  defp infer_type(%_{} = item), do: raise(ArgumentError, "unsupported datatype: #{inspect(item)}")
   defp infer_type(item) when is_integer(item), do: {:s, 64}
   defp infer_type(item) when is_float(item) or item in @non_finite, do: {:f, 64}
   defp infer_type(item) when is_boolean(item), do: :boolean
@@ -559,19 +561,20 @@ defmodule Explorer.Shared do
   @doc """
   Converts dtype to its string representation.
   """
-  def dtype_to_string({:datetime, :millisecond}), do: "datetime[ms]"
-  def dtype_to_string({:datetime, :microsecond}), do: "datetime[μs]"
-  def dtype_to_string({:datetime, :nanosecond}), do: "datetime[ns]"
-  def dtype_to_string({:duration, :millisecond}), do: "duration[ms]"
-  def dtype_to_string({:duration, :microsecond}), do: "duration[μs]"
-  def dtype_to_string({:duration, :nanosecond}), do: "duration[ns]"
+  def dtype_to_string({:datetime, p}), do: "datetime[#{precision(p)}]"
+  def dtype_to_string({:datetime, p, tz}), do: "datetime[#{precision(p)}, #{tz}]"
+  def dtype_to_string({:duration, p}), do: "duration[#{precision(p)}]"
   def dtype_to_string({:list, dtype}), do: "list[" <> dtype_to_string(dtype) <> "]"
   def dtype_to_string({:struct, fields}), do: "struct[#{length(fields)}]"
   def dtype_to_string({:f, size}), do: "f" <> Integer.to_string(size)
   def dtype_to_string({:s, size}), do: "s" <> Integer.to_string(size)
   def dtype_to_string({:u, size}), do: "u" <> Integer.to_string(size)
   def dtype_to_string(other) when is_atom(other), do: Atom.to_string(other)
 
+  defp precision(:millisecond), do: "ms"
+  defp precision(:microsecond), do: "μs"
+  defp precision(:nanosecond), do: "ns"
+
   @threshold 0.77
   @max_suggestions 5
 

native/explorer/src/datatypes.rs

@@ -17,6 +17,9 @@ use std::str::FromStr;
 #[cfg(feature = "aws")]
 use polars::prelude::cloud::AmazonS3ConfigKey as S3Key;
 
+#[cfg(feature = "timezones")]
+use chrono_tz::Tz;
+
 pub use ex_dtypes::*;
 
 pub struct ExDataFrameRef(pub DataFrame);
@@ -257,7 +260,7 @@ fn time_unit_of_ex_duration(duration: &ExDuration) -> TimeUnit {
 
 #[derive(NifStruct, Copy, Clone, Debug)]
 #[module = "NaiveDateTime"]
-pub struct ExDateTime {
+pub struct ExNaiveDateTime {
     pub calendar: Atom,
     pub day: u32,
     pub month: u32,
@@ -268,11 +271,23 @@ pub struct ExDateTime {
     pub microsecond: (u32, u32),
 }
 
+#[derive(NifStruct, Copy, Clone, Debug)]
+#[module = "DateTime"]
+pub struct ExDateTime<'a> {
+    pub calendar: Atom,
+    pub day: u32,
+    pub hour: u32,
+    pub microsecond: (u32, u32),
+    pub minute: u32,
+    pub month: u32,
+    pub second: u32,
+    pub time_zone: &'a str,
+    pub year: i32,
+}
+
 pub use polars::export::arrow::temporal_conversions::date32_to_date as days_to_date;
 
-/// Converts a microsecond i64 to a `NaiveDateTime`.
-/// This is because when getting a timestamp, it might have negative values.
-pub fn timestamp_to_datetime(microseconds: i64) -> NaiveDateTime {
+pub fn timestamp_to_datetime_utc(microseconds: i64) -> DateTime<Utc> {
     let sign = microseconds.signum();
     let seconds = match sign {
         -1 => microseconds / 1_000_000 - 1,
@@ -283,9 +298,13 @@ pub fn timestamp_to_datetime(microseconds: i64) -> NaiveDateTime {
         _ => microseconds % 1_000_000,
     };
     let nanoseconds = remainder.abs() * 1_000;
-    DateTime::<Utc>::from_timestamp(seconds, nanoseconds.try_into().unwrap())
-        .expect("construct a UTC")
-        .naive_utc()
+    DateTime::<Utc>::from_timestamp(seconds, nanoseconds.try_into().unwrap()).expect("construct a UTC")
+}
+
+/// Converts a microsecond i64 to a `NaiveDateTime`.
+/// This is because when getting a timestamp, it might have negative values.
+pub fn timestamp_to_datetime(microseconds: i64) -> NaiveDateTime {
+    timestamp_to_datetime_utc(microseconds).naive_utc()
 }
 
 // Limit the number of digits in the microsecond part of a timestamp to 6.
@@ -299,14 +318,14 @@ fn microseconds_six_digits(microseconds: u32) -> u32 {
     }
 }
 
-impl From<i64> for ExDateTime {
+impl From<i64> for ExNaiveDateTime {
     fn from(microseconds: i64) -> Self {
         timestamp_to_datetime(microseconds).into()
     }
 }
 
-impl From<ExDateTime> for i64 {
-    fn from(dt: ExDateTime) -> i64 {
+impl From<ExNaiveDateTime> for i64 {
+    fn from(dt: ExNaiveDateTime) -> i64 {
         let duration = NaiveDate::from_ymd_opt(dt.year, dt.month, dt.day)
             .unwrap()
             .and_hms_micro_opt(dt.hour, dt.minute, dt.second, dt.microsecond.0)
@@ -325,18 +344,18 @@ impl From<ExDateTime> for i64 {
     }
 }
 
-impl From<ExDateTime> for NaiveDateTime {
-    fn from(dt: ExDateTime) -> NaiveDateTime {
+impl From<ExNaiveDateTime> for NaiveDateTime {
+    fn from(dt: ExNaiveDateTime) -> NaiveDateTime {
         NaiveDate::from_ymd_opt(dt.year, dt.month, dt.day)
             .unwrap()
             .and_hms_micro_opt(dt.hour, dt.minute, dt.second, dt.microsecond.0)
             .unwrap()
     }
 }
 
-impl From<NaiveDateTime> for ExDateTime {
+impl From<NaiveDateTime> for ExNaiveDateTime {
     fn from(dt: NaiveDateTime) -> Self {
-        ExDateTime {
+        ExNaiveDateTime {
             calendar: atoms::calendar_iso_module(),
             day: dt.day(),
             month: dt.month(),
@@ -352,12 +371,75 @@ impl From<NaiveDateTime> for ExDateTime {
     }
 }
 
-impl Literal for ExDateTime {
+impl Literal for ExNaiveDateTime {
     fn lit(self) -> Expr {
         NaiveDateTime::from(self).lit()
     }
 }
 
+impl From<i64> for ExDateTime<'_> {
+    fn from(microseconds: i64) -> Self {
+        timestamp_to_datetime_utc(microseconds).into()
+    }
+}
+
+impl From<ExDateTime<'_>> for i64 {
+    fn from(dt: ExDateTime<'_>) -> i64 {
+        let duration = NaiveDate::from_ymd_opt(dt.year, dt.month, dt.day)
+            .unwrap()
+            .and_hms_micro_opt(dt.hour, dt.minute, dt.second, dt.microsecond.0)
+            .unwrap()
+            .signed_duration_since(
+                NaiveDate::from_ymd_opt(1970, 1, 1)
+                    .unwrap()
+                    .and_hms_opt(0, 0, 0)
+                    .unwrap(),
+            );
+
+        match duration.num_microseconds() {
+            Some(us) => us,
+            None => duration.num_milliseconds() * 1_000,
+        }
+    }
+}
+
+impl From<chrono::DateTime<chrono::Utc>> for ExDateTime<'static> {
+    fn from(dt: chrono::DateTime<chrono::Utc>) -> ExDateTime<'static> {
+        dt.timestamp_micros().into()
+    }
+}
+
+// impl<Tz: chrono::TimeZone + std::str::FromStr> From<ExDateTime<'_>> for DateTime<Tz> {
+//     fn from(dt: ExDateTime<'_>) -> DateTime<Tz> where <Tz as FromStr>::Err: core::fmt::Debug {
+//         let tz: Tz = dt.time_zone.parse().unwrap();
+//         tz.ymd(dt.year, dt.month, dt.day)
+//             .and_hms_micro_opt(dt.hour, dt.minute, dt.second, dt.microsecond.0)
+//             .unwrap()
+//     }
+// }
+
+// impl From<DateTime<T=Tz>> for ExDateTime<'_> {
+//     fn from(dt: DateTime<T=Tz>) -> Self {
+//         ExDateTime {
+//             calendar: atoms::calendar_iso_module(),
+//             day: dt.day(),
+//             hour: dt.hour(),
+//             microsecond: (microseconds_six_digits(dt.and_utc().timestamp_subsec_micros()), 6),
+//             minute: dt.minute(),
+//             month: dt.month(),
+//             second: dt.second(),
+//             time_zone: dt.time_zone().to_string(),
+//             year: dt.year(),
+//         }
+//     }
+// }
+
+// impl Literal for ExDateTime<'_> {
+//     fn lit(self) -> Expr {
+//         DateTime::from(self).lit()
+//     }
+// }
+
 #[derive(NifStruct, Copy, Clone, Debug)]
 #[module = "Time"]
 pub struct ExTime {
@@ -433,7 +515,7 @@ pub enum ExValidValue<'a> {
     Str(&'a str),
     Date(ExDate),
     Time(ExTime),
-    DateTime(ExDateTime),
+    DateTime(ExNaiveDateTime),
     Duration(ExDuration),
 }
 
@@ -476,7 +558,7 @@ impl<'a> rustler::Decoder<'a> for ExValidValue<'a> {
                     Ok(ExValidValue::Date(date))
                 } else if let Ok(time) = term.decode::<ExTime>() {
                     Ok(ExValidValue::Time(time))
-                } else if let Ok(datetime) = term.decode::<ExDateTime>() {
+                } else if let Ok(datetime) = term.decode::<ExNaiveDateTime>() {
                     Ok(ExValidValue::DateTime(datetime))
                 } else if let Ok(duration) = term.decode::<ExDuration>() {
                     Ok(ExValidValue::Duration(duration))

native/explorer/src/expressions.rs

@@ -10,7 +10,7 @@ use polars::prelude::{GetOutput, IntoSeries, Utf8JsonPathImpl};
 use polars::series::Series;
 
 use crate::datatypes::{
-    ExCorrelationMethod, ExDate, ExDateTime, ExDuration, ExRankMethod, ExSeriesDtype, ExValidValue,
+    ExCorrelationMethod, ExDate, ExNaiveDateTime, ExDuration, ExRankMethod, ExSeriesDtype, ExValidValue,
 };
 use crate::series::{cast_str_to_f64, ewm_opts, rolling_opts};
 use crate::{ExDataFrame, ExExpr, ExSeries};
@@ -71,7 +71,7 @@ pub fn expr_date(date: ExDate) -> ExExpr {
 }
 
 #[rustler::nif]
-pub fn expr_datetime(datetime: ExDateTime) -> ExExpr {
+pub fn expr_datetime(datetime: ExNaiveDateTime) -> ExExpr {
     ExExpr::new(datetime.lit())
 }