[Proof of concept] Expr serialize

lib/explorer/polars_backend/expression.ex

@@ -356,6 +356,18 @@ defmodule Explorer.PolarsBackend.Expression do
     Native.expr_describe_filter_plan(polars_df, expression)
   end
 
+  def to_json(%__MODULE__{} = expression) do
+    expression
+    |> Native.expr_to_json()
+    |> Jason.decode!()
+  end
+
+  def from_json(%{} = json_map) do
+    json_map
+    |> Jason.encode!()
+    |> Native.expr_from_json()
+  end
+
   defp dtype(%LazySeries{dtype: dtype}), do: dtype
 
   defp dtype(%PolarsSeries{} = polars_series) do

lib/explorer/polars_backend/native.ex

@@ -196,6 +196,8 @@ defmodule Explorer.PolarsBackend.Native do
   def expr_datetime(_datetime), do: err()
   def expr_duration(_duration), do: err()
   def expr_describe_filter_plan(_df, _expr), do: err()
+  def expr_to_json(_expr), do: err()
+  def expr_from_json(_expr), do: err()
   def expr_float(_number), do: err()
   def expr_integer(_number), do: err()
   def expr_int_range(_start, _end, _step, _dtype), do: err()

mix.exs

@@ -43,6 +43,8 @@ defmodule Explorer.MixProject do
       {:table_rex, "~> 3.1.1 or ~> 4.0.0"},
       {:castore, "~> 1.0", optional: true},
       {:adbc, "~> 0.1", optional: true},
+      # DELETEME!
+      {:jason, "~> 1.4"},
 
       ## Optional
       {:rustler, "~> 0.32.0", optional: not (@dev? or @force_build?)},

native/explorer/Cargo.toml

@@ -19,6 +19,7 @@ rand = { version = "0.8", features = ["alloc"] }
 rand_pcg = "0.3"
 rustler = { version = "0.32", default-features = false, features = ["derive"] }
 thiserror = "1"
+serde_json = "1"
 smartstring = "1"
 either = "1"
 
@@ -76,6 +77,8 @@ features = [
   "rolling_window",
   "round_series",
   "rows",
+  "serde",
+  "serde-lazy",
   "simd",
   "streaming",
   "strings",

native/explorer/src/expressions.rs

@@ -830,6 +830,17 @@ pub fn expr_describe_filter_plan(data: ExDataFrame, expr: ExExpr) -> String {
     df.lazy().filter(expressions).describe_plan()
 }
 
+#[rustler::nif]
+pub fn expr_to_json(expr: ExExpr) -> String {
+    serde_json::to_string(&expr.clone_inner()).unwrap()
+}
+
+#[rustler::nif]
+pub fn expr_from_json(expr_json: String) -> ExExpr {
+    let expr: Expr = serde_json::from_str(&expr_json).unwrap();
+    ExExpr::new(expr)
+}
+
 #[rustler::nif]
 pub fn expr_contains(expr: ExExpr, pattern: &str) -> ExExpr {
     let expr = expr.clone_inner();

native/explorer/src/lib.rs

@@ -258,6 +258,8 @@ rustler::init!(
         expr_ewm_variance,
         // inspect expressions
         expr_describe_filter_plan,
+        expr_to_json,
+        expr_from_json,
         // string expressions
         expr_contains,
         expr_re_contains,

test/explorer/polars_backend/expression_test.exs

@@ -89,4 +89,67 @@ defmodule Explorer.PolarsBackend.ExpressionTest do
                """)
     end
   end
+
+  describe "json" do
+    test "can convert exprs to/from json" do
+      lazy = %LazySeries{op: :column, args: ["a"]}
+      expr1 = Expression.to_expr(lazy)
+      json = Expression.to_json(expr1)
+      expr2 = Expression.from_json(json)
+
+      assert json == %{"Column" => "a"}
+      assert %Expression{} = expr2
+    end
+
+    test "can perform an unsupported operation via json-derived exprs" do
+      # Built in Python from:
+      # `pl.col("list_col_unsorted").list.sort().meta.serialize()`
+      list_col_sorted_expr_json = %{
+        "Function" => %{
+          "input" => [%{"Column" => "list_col_unsorted"}],
+          "function" => %{
+            "ListExpr" => %{
+              "Sort" => %{
+                "descending" => false,
+                "nulls_last" => false,
+                "multithreaded" => true,
+                "maintain_order" => false
+              }
+            }
+          },
+          "options" => %{
+            "collect_groups" => "ElementWise",
+            "fmt_str" => "",
+            "input_wildcard_expansion" => false,
+            "returns_scalar" => false,
+            "cast_to_supertypes" => false,
+            "allow_rename" => false,
+            "pass_name_to_apply" => false,
+            "changes_length" => false,
+            "check_lengths" => true,
+            "allow_group_aware" => true
+          }
+        }
+      }
+
+      mutate_with_json = fn df, name, json ->
+        expr =
+          json
+          |> Explorer.PolarsBackend.Expression.from_json()
+          |> Explorer.PolarsBackend.Expression.alias_expr(name)
+
+        ldf = Explorer.DataFrame.lazy(df)
+        {:ok, lpdf_new} = Explorer.PolarsBackend.Native.lf_mutate_with(ldf.data, [expr])
+        {:ok, pdf_new} = Explorer.PolarsBackend.Native.lf_collect(lpdf_new)
+        Explorer.PolarsBackend.Shared.create_dataframe(pdf_new)
+      end
+
+      df = Explorer.DataFrame.new(%{list_col_unsorted: [[1, 5, 3], [1, 1, 2, 0]]})
+      df_new = mutate_with_json.(df, "list_col_sorted", list_col_sorted_expr_json)
+
+      series = Explorer.DataFrame.to_series(df_new)
+      assert Explorer.Series.to_list(series["list_col_unsorted"]) == [[1, 5, 3], [1, 1, 2, 0]]
+      assert Explorer.Series.to_list(series["list_col_sorted"]) == [[1, 3, 5], [0, 1, 1, 2]]
+    end
+  end
 end