add datetime literals to expressions

lib/explorer/polars_backend/expression.ex

@@ -334,7 +334,7 @@ defmodule Explorer.PolarsBackend.Expression do
   def to_expr(number) when is_float(number), do: Native.expr_float(number)
   def to_expr(%Date{} = date), do: Native.expr_date(date)
   def to_expr(%NaiveDateTime{} = naive_datetime), do: Native.expr_naive_datetime(naive_datetime)
-  # def to_expr(%DateTime{} = datetime), do: Native.expr_datetime(datetime)
+  def to_expr(%DateTime{} = datetime), do: Native.expr_datetime(datetime)
   def to_expr(%Explorer.Duration{} = duration), do: Native.expr_duration(duration)
   def to_expr(%PolarsSeries{} = polars_series), do: Native.expr_series(polars_series)
 

lib/explorer/polars_backend/native.ex

@@ -194,7 +194,7 @@ defmodule Explorer.PolarsBackend.Native do
   def expr_boolean(_bool), do: err()
   def expr_date(_date), do: err()
   def expr_naive_datetime(_datetime), do: err()
-  # def expr_datetime(_datetime), do: err()
+  def expr_datetime(_datetime), do: err()
   def expr_duration(_duration), do: err()
   def expr_describe_filter_plan(_df, _expr), do: err()
   def expr_float(_number), do: err()

native/explorer/src/datatypes.rs

@@ -467,12 +467,27 @@ impl From<ExDateTime<'_>> for DateTime<Tz> {
     }
 }
 
-// TODO: Polars doesn't provide a default `Literal` impl. Find out why.
-// impl Literal for ExDateTime<'_> {
-//     fn lit(self) -> Expr {
-//         DateTime::from(self).lit()
-//     }
-// }
+impl<'tz> From<ExDateTime<'tz>> for NaiveDateTime {
+    fn from(dt: ExDateTime) -> 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<'tz> Literal for ExDateTime<'tz> {
+    fn lit(self) -> Expr {
+        let ndt = NaiveDateTime::from(self);
+        let time_zone = self.time_zone.to_string();
+
+        Expr::Literal(LiteralValue::DateTime(
+            ndt.and_utc().timestamp_micros(),
+            TimeUnit::Microseconds,
+            Some(time_zone),
+        ))
+    }
+}
 
 #[derive(NifStruct, Copy, Clone, Debug)]
 #[module = "Time"]

native/explorer/src/expressions.rs

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

native/explorer/src/lib.rs

@@ -141,7 +141,7 @@ rustler::init!(
         expr_column,
         expr_date,
         expr_naive_datetime,
-        // expr_datetime,
+        expr_datetime,
         expr_duration,
         expr_day_of_week,
         expr_day_of_year,

test/explorer/polars_backend/expression_test.exs

@@ -50,7 +50,7 @@ defmodule Explorer.PolarsBackend.ExpressionTest do
       assert %Expression{} = Expression.to_expr(lazy)
     end
 
-    test "with datetime value" do
+    test "with naive datetime value" do
       lazy = %LazySeries{
         op: :equal,
         args: [%LazySeries{op: :column, args: ["col_b"]}, ~N[2022-07-07 18:09:17.824019]]
@@ -59,6 +59,15 @@ defmodule Explorer.PolarsBackend.ExpressionTest do
       assert %Expression{} = Expression.to_expr(lazy)
     end
 
+    test "with datetime value" do
+      lazy = %LazySeries{
+        op: :equal,
+        args: [%LazySeries{op: :column, args: ["col_b"]}, ~U[2022-07-07 18:09:17.824019Z]]
+      }
+
+      assert %Expression{} = Expression.to_expr(lazy)
+    end
+
     test "with series" do
       lazy = %LazySeries{
         op: :equal,