Add Struct json_decode/3 for decoding json from string

lib/explorer/backend/lazy_series.ex

@@ -143,7 +143,8 @@ defmodule Explorer.Backend.LazySeries do
     lengths: 1,
     member: 3,
     # Struct functions
-    field: 2
+    field: 2,
+    json_decode: 3
   ]
 
   @comparison_operations [:equal, :not_equal, :greater, :greater_equal, :less, :less_equal]
@@ -1094,6 +1095,13 @@ defmodule Explorer.Backend.LazySeries do
     Backend.Series.new(data, dtype)
   end
 
+  @impl true
+  def json_decode(series, dtype, infer_schema_length) do
+    data = new(:json_decode, [lazy_series!(series), dtype, infer_schema_length], dtype)
+
+    Backend.Series.new(data, dtype)
+  end
+
   @remaining_non_lazy_operations [
     at: 2,
     at_every: 2,

lib/explorer/backend/series.ex

@@ -305,6 +305,7 @@ defmodule Explorer.Backend.Series do
 
   # Struct
   @callback field(s, String.t()) :: s
+  @callback json_decode(s, option(dtype()), option(non_neg_integer())) :: s
 
   # Functions
 

lib/explorer/polars_backend/expression.ex

@@ -142,7 +142,8 @@ defmodule Explorer.PolarsBackend.Expression do
     member: 3,
 
     # Structs
-    field: 2
+    field: 2,
+    json_decode: 3
   ]
 
   @custom_expressions [

lib/explorer/polars_backend/native.ex

@@ -449,6 +449,7 @@ defmodule Explorer.PolarsBackend.Native do
   def s_member(_s, _value, _inner_dtype), do: err()
 
   def s_field(_s, _name), do: err()
+  def s_json_decode(_s, _dtype, _infer_schema_length), do: err()
 
   defp err, do: :erlang.nif_error(:nif_not_loaded)
 end

lib/explorer/polars_backend/series.ex

@@ -738,6 +738,10 @@ defmodule Explorer.PolarsBackend.Series do
   def field(%Series{dtype: {:struct, _inner_dtype}} = series, name),
     do: Shared.apply_series(series, :s_field, [name])
 
+  @impl true
+  def json_decode(series, dtype, infer_schema_length),
+    do: Shared.apply_series(series, :s_json_decode, [dtype, infer_schema_length])
+
   # Polars specific functions
 
   def name(series), do: Shared.apply_series(series, :s_name)

lib/explorer/polars_backend/shared.ex

@@ -42,36 +42,55 @@ defmodule Explorer.PolarsBackend.Shared do
   def apply_dataframe(%DataFrame{} = df, %DataFrame{} = out_df, fun, args) do
     case apply(Native, fun, [df.data | args]) do
       {:ok, %module{} = new_df} when module in @polars_df ->
-        if @check_frames do
-          # We need to collect here, because the lazy frame may not have
-          # the full picture of the result yet.
-          check_df =
-            if match?(%PolarsLazyFrame{}, new_df) do
-              {:ok, new_df} = Native.lf_collect(new_df)
-              create_dataframe(new_df)
-            else
-              create_dataframe(new_df)
+        {struct?, dtypes} =
+          if @check_frames do
+            # We need to collect here, because the lazy frame may not have
+            # the full picture of the result yet.
+            check_df =
+              if match?(%PolarsLazyFrame{}, new_df) do
+                {:ok, new_df} = Native.lf_collect(new_df)
+                create_dataframe(new_df)
+              else
+                create_dataframe(new_df)
+              end
+
+            # When dealing with structs in mutate, we may not know dtype of struct series.
+            # We have to accept the dtype returned by polars, else we will have mismatch error.
+            {struct?, out_dtypes} =
+              if fun == :df_mutate_with_exprs do
+                Enum.reduce(check_df.dtypes, {false, out_df.dtypes}, fn
+                  {key, {:struct, _} = dtype}, {_, dtypes} -> {true, Map.put(dtypes, key, dtype)}
+                  _, acc -> acc
+                end)
+              else
+                {false, out_df.dtypes}
+              end
+
+            if Enum.sort(out_df.names) != Enum.sort(check_df.names) or
+                 out_dtypes != check_df.dtypes do
+              raise """
+              DataFrame mismatch.
+
+              expected:
+
+                  names: #{inspect(out_df.names)}
+                  dtypes: #{inspect(out_df.dtypes)}
+
+              got:
+
+                  names: #{inspect(check_df.names)}
+                  dtypes: #{inspect(check_df.dtypes)}
+              """
             end
 
-          if Enum.sort(out_df.names) != Enum.sort(check_df.names) or
-               out_df.dtypes != check_df.dtypes do
-            raise """
-            DataFrame mismatch.
-
-            expected:
-
-                names: #{inspect(out_df.names)}
-                dtypes: #{inspect(out_df.dtypes)}
-
-            got:
-
-                names: #{inspect(check_df.names)}
-                dtypes: #{inspect(check_df.dtypes)}
-            """
+            {struct?, out_dtypes}
           end
-        end
 
-        %{out_df | data: new_df}
+        if struct? do
+          %{out_df | data: new_df, dtypes: dtypes}
+        else
+          %{out_df | data: new_df}
+        end
 
       {:error, error} ->
         raise runtime_error(error)

lib/explorer/series.ex

@@ -6043,6 +6043,29 @@ defmodule Explorer.Series do
     end
   end
 
+  @doc """
+  Decode json from string.
+
+  ## Examples
+
+      iex> s = Series.from_list(["{\\"a\\":1}"])
+      iex> Series.json_decode(s)
+      #Explorer.Series<
+        Polars[1]
+        struct[1] [%{"a" => 1}]
+      >
+
+  Will raise `RuntimeError` for invalid json.
+  """
+  @doc type: :struct_wise
+  @spec json_decode(Series.t(), Keyword.t()) :: Series.t()
+  def json_decode(%Series{dtype: :string} = series, opts \\ []) do
+    opts = Keyword.validate!(opts, [:dtype, :infer_schema_length])
+    dtype = if opts[:dtype], do: Shared.normalise_dtype!(opts[:dtype])
+
+    apply_series(series, :json_decode, [dtype, opts[:infer_schema_length]])
+  end
+
   # Helpers
 
   defp apply_series(series, fun, args \\ []) do

native/explorer/Cargo.toml

@@ -36,6 +36,9 @@ object_store = { version = "0.8", default-features = false, optional = true }
 [target.'cfg(not(any(all(windows, target_env = "gnu"), all(target_os = "linux", target_env = "musl"))))'.dependencies]
 mimalloc = { version = "*", default-features = false }
 
+[patch.crates-io]
+jsonpath_lib = { version = "0.3", git = "https://github.com/ritchie46/jsonpath", branch = "improve_compiled" }
+
 [dependencies.polars]
 version = "0.36"
 default-features = false
@@ -79,6 +82,7 @@ features = [
   "moment",
   "rank",
   "propagate_nans",
+  "extract_jsonpath"
 ]
 
 [dependencies.polars-ops]

native/explorer/src/expressions.rs

@@ -1065,3 +1065,17 @@ pub fn expr_field(expr: ExExpr, name: &str) -> ExExpr {
     let expr = expr.clone_inner().struct_().field_by_name(name);
     ExExpr::new(expr)
 }
+
+#[rustler::nif]
+pub fn expr_json_decode(
+    expr: ExExpr,
+    ex_dtype: Option<ExSeriesDtype>,
+    infer_schema_length: Option<usize>,
+) -> ExExpr {
+    let dtype = ex_dtype.map(|x| DataType::try_from(&x).unwrap());
+    let expr = expr
+        .clone_inner()
+        .str()
+        .json_decode(dtype, infer_schema_length);
+    ExExpr::new(expr)
+}

native/explorer/src/lib.rs

@@ -280,6 +280,7 @@ rustler::init!(
         expr_member,
         // struct expressions
         expr_field,
+        expr_json_decode,
         // lazyframe
         lf_collect,
         lf_describe_plan,
@@ -480,6 +481,7 @@ rustler::init!(
         s_lengths,
         s_member,
         s_field,
+        s_json_decode,
     ],
     load = on_load
 );

native/explorer/src/series.rs

@@ -1824,3 +1824,24 @@ pub fn s_field(s: ExSeries, name: &str) -> Result<ExSeries, ExplorerError> {
         .clone();
     Ok(ExSeries::new(s2))
 }
+
+#[rustler::nif]
+pub fn s_json_decode(
+    s: ExSeries,
+    ex_dtype: Option<ExSeriesDtype>,
+    infer_schema_length: Option<usize>,
+) -> Result<ExSeries, ExplorerError> {
+    let dtype = ex_dtype.map(|x| DataType::try_from(&x).unwrap());
+    let s2 = s
+        .clone_inner()
+        .into_frame()
+        .lazy()
+        .select([col(s.name())
+            .str()
+            .json_decode(dtype, infer_schema_length)
+            .alias(s.name())])
+        .collect()?
+        .column(s.name())?
+        .clone();
+    Ok(ExSeries::new(s2))
+}

test/explorer/data_frame_test.exs

@@ -1870,6 +1870,32 @@ defmodule Explorer.DataFrameTest do
                member?: [true, false]
              }
     end
+
+    test "extracts struct from json - json_decode" do
+      df = DF.new([%{a: "{\"n\": 1}"}])
+      dfj = DF.mutate(df, aj: json_decode(a, dtype: {:struct, %{"n" => {:s, 64}}}))
+      assert dfj.dtypes == %{"a" => :string, "aj" => {:struct, %{"n" => {:s, 64}}}}
+      assert DF.to_rows(dfj) == [%{"a" => "{\"n\": 1}", "aj" => %{"n" => 1}}]
+    end
+
+    test "extracts struct from json - json_decode with dtype" do
+      df = DF.new([%{a: "{\"n\": 1}"}])
+      dfj = DF.mutate(df, aj: json_decode(a, dtype: {:struct, %{"n" => {:f, 64}}}))
+      assert dfj.dtypes == %{"a" => :string, "aj" => {:struct, %{"n" => {:f, 64}}}}
+      assert DF.to_rows(dfj) == [%{"a" => "{\"n\": 1}", "aj" => %{"n" => 1.0}}]
+    end
+
+    test "extracts struct from json - json_decode with infer_schema_length" do
+      df = DF.new([%{a: "{\"n\": 1}"}])
+
+      dfj =
+        DF.mutate(df,
+          aj: json_decode(a, infer_schema_length: 100)
+        )
+
+      assert dfj.dtypes == %{"a" => :string, "aj" => {:struct, %{"n" => {:s, 64}}}}
+      assert DF.to_rows(dfj) == [%{"a" => "{\"n\": 1}", "aj" => %{"n" => 1}}]
+    end
   end
 
   describe "sort_by/3" do