add json_path_match/2

lib/explorer/backend/lazy_series.ex

@@ -127,6 +127,7 @@ defmodule Explorer.Backend.LazySeries do
     substring: 3,
     split: 2,
     json_decode: 2,
+    json_path_match: 2,
     # Float round
     round: 2,
     floor: 1,
@@ -1103,6 +1104,13 @@ defmodule Explorer.Backend.LazySeries do
     Backend.Series.new(data, dtype)
   end
 
+  @impl true
+  def json_path_match(series, json_path) do
+    data = new(:json_path_match, [lazy_series!(series), json_path], :string)
+
+    Backend.Series.new(data, :string)
+  end
+
   @remaining_non_lazy_operations [
     at: 2,
     at_every: 2,

lib/explorer/backend/series.ex

@@ -287,6 +287,7 @@ defmodule Explorer.Backend.Series do
   @callback substring(s, integer(), non_neg_integer() | nil) :: s
   @callback split(s, String.t()) :: s
   @callback json_decode(s, dtype()) :: s
+  @callback json_path_match(s, String.t()) :: s
 
   # Date / DateTime
 

lib/explorer/polars_backend/expression.ex

@@ -137,6 +137,7 @@ defmodule Explorer.PolarsBackend.Expression do
     substring: 3,
     split: 2,
     json_decode: 2,
+    json_path_match: 2,
 
     # Lists
     join: 2,

lib/explorer/polars_backend/native.ex

@@ -451,6 +451,7 @@ defmodule Explorer.PolarsBackend.Native do
 
   def s_field(_s, _name), do: err()
   def s_json_decode(_s, _dtype), do: err()
+  def s_json_path_match(_s, _json_path), do: err()
 
   defp err, do: :erlang.nif_error(:nif_not_loaded)
 end

lib/explorer/polars_backend/series.ex

@@ -743,6 +743,10 @@ defmodule Explorer.PolarsBackend.Series do
   def json_decode(series, dtype),
     do: Shared.apply_series(series, :s_json_decode, [dtype])
 
+  @impl true
+  def json_path_match(series, json_path),
+    do: Shared.apply_series(series, :s_json_path_match, [json_path])
+
   # Polars specific functions
 
   def name(series), do: Shared.apply_series(series, :s_name)

lib/explorer/series.ex

@@ -6118,6 +6118,36 @@ defmodule Explorer.Series do
     apply_series(series, :json_decode, [dtype])
   end
 
+  @doc """
+  Extracts a string series using a [`json_path`](https://goessner.net/articles/JsonPath/) from a series. 
+
+  ## Examples
+
+      iex> s = Series.from_list(["{\\"a\\":1}", "{\\"a\\":2}"])
+      iex> Series.json_path_match(s, "$.a")
+      #Explorer.Series<
+        Polars[2]
+        string [\"1\", \"2\"]
+      >
+
+  If `json_path` is not found or if the string is invalid JSON or `nil`, 
+  nil is returned for the given entry:
+
+      iex> s = Series.from_list(["{\\"a\\":1}", nil, "{\\"a\\":2}"])
+      iex> Series.json_path_match(s, "$.b")
+      #Explorer.Series<
+        Polars[3]
+        string [nil, nil, nil]
+      >
+
+  It raises an exception if the `json_path` is invalid.
+  """
+  @doc type: :string_wise
+  @spec json_path_match(Series.t(), String.t()) :: Series.t()
+  def json_path_match(%Series{dtype: :string} = series, json_path) do
+    apply_series(series, :json_path_match, [json_path])
+  end
+
   # Helpers
 
   defp apply_series(series, fun, args \\ []) do

native/explorer/src/expressions.rs

@@ -4,6 +4,11 @@
 // or an expression and returns an expression that is
 // wrapped in an Elixir struct.
 
+use polars::error::PolarsError;
+
+use polars::prelude::{GetOutput, IntoSeries, Utf8JsonPathImpl};
+use polars::series::Series;
+
 use crate::datatypes::{
     ExCorrelationMethod, ExDate, ExDateTime, ExDuration, ExRankMethod, ExSeriesDtype, ExValidValue,
 };
@@ -1080,9 +1085,26 @@ pub fn expr_json_decode(expr: ExExpr, ex_dtype: ExSeriesDtype) -> ExExpr {
     ExExpr::new(expr)
 }
 
+#[rustler::nif]
+pub fn expr_json_path_match(expr: ExExpr, json_path: &str) -> ExExpr {
+    let p = json_path.to_owned();
+    let function = move |s: Series| {
+        let ca = s.str()?;
+        match ca.json_path_match(&p) {
+            Ok(ca) => Ok(Some(ca.into_series())),
+            Err(e) => Err(PolarsError::ComputeError(format!("{e:?}").into())),
+        }
+    };
+    let expr = expr
+        .clone_inner()
+        .map(function, GetOutput::from_type(DataType::String));
+    ExExpr::new(expr)
+}
+
 #[rustler::nif]
 pub fn expr_struct(ex_exprs: Vec<ExExpr>) -> ExExpr {
     let exprs = ex_exprs.iter().map(|e| e.clone_inner()).collect();
     let expr = dsl::as_struct(exprs);
+
     ExExpr::new(expr)
 }

native/explorer/src/lib.rs

@@ -271,6 +271,7 @@ rustler::init!(
         expr_rstrip,
         expr_substring,
         expr_replace,
+        expr_json_path_match,
         // float round expressions
         expr_round,
         expr_floor,
@@ -484,6 +485,7 @@ rustler::init!(
         s_member,
         s_field,
         s_json_decode,
+        s_json_path_match
     ],
     load = on_load
 );

native/explorer/src/series.rs

@@ -1841,3 +1841,26 @@ pub fn s_json_decode(s: ExSeries, ex_dtype: ExSeriesDtype) -> Result<ExSeries, E
         .clone();
     Ok(ExSeries::new(s2))
 }
+
+#[rustler::nif]
+pub fn s_json_path_match(s: ExSeries, json_path: &str) -> Result<ExSeries, ExplorerError> {
+    let p = json_path.to_owned();
+    let function = move |s: Series| {
+        let ca = s.str()?;
+        match ca.json_path_match(&p) {
+            Ok(ca) => Ok(Some(ca.into_series())),
+            Err(e) => Err(PolarsError::ComputeError(format!("{e:?}").into())),
+        }
+    };
+    let s2 = s
+        .clone_inner()
+        .into_frame()
+        .lazy()
+        .select([col(s.name())
+            .map(function, GetOutput::from_type(DataType::String))
+            .alias(s.name())])
+        .collect()?
+        .column(s.name())?
+        .clone();
+    Ok(ExSeries::new(s2))
+}

test/explorer/data_frame_test.exs

@@ -725,6 +725,16 @@ defmodule Explorer.DataFrameTest do
                "b" => :string,
                "c" => {:struct, [{"s", {:struct, [{"a", {:s, 64}}, {"b", :string}]}}]}
              }
+
+      assert DF.to_columns(df1, atom_keys: true) == %{
+               a: [1, nil, 3],
+               b: ["a", "b", nil],
+               c: [
+                 %{"s" => %{"a" => 1, "b" => "a"}},
+                 %{"s" => %{"a" => nil, "b" => "b"}},
+                 %{"s" => %{"a" => 3, "b" => nil}}
+               ]
+             }
     end
 
     test "add eager series to dataframe" do
@@ -739,6 +749,58 @@ defmodule Explorer.DataFrameTest do
                "s" => :string,
                "a1" => {:s, 64}
              }
+
+      assert DF.to_columns(df1, atom_keys: true) == %{
+               a: [1, 2, 3],
+               a1: [2, 3, 4],
+               s: ["x", "y", "z"]
+             }
+    end
+
+    test "json_path_match/2 - extract value from a valid json string using json path" do
+      df = DF.new([%{a: ~s({"n": 1})}, %{a: ~s({"m": 1})}])
+      df1 = DF.mutate(df, m: json_path_match(a, "$.m"), n: json_path_match(a, "$.n"))
+      assert df1.names == ["a", "m", "n"]
+
+      assert df1.dtypes == %{
+               "a" => :string,
+               "n" => :string,
+               "m" => :string
+             }
+
+      assert DF.to_columns(df1, atom_keys: true) == %{
+               a: ["{\"n\": 1}", "{\"m\": 1}"],
+               n: ["1", nil],
+               m: [nil, "1"]
+             }
+    end
+
+    test "json_path_match/2 - extracts nil from invalidjson using json path" do
+      df = DF.new([%{a: ~s({"n": 1)}, %{a: ~s({"m": 1})}])
+      df1 = DF.mutate(df, m: json_path_match(a, "$.m"), n: json_path_match(a, "$.n"))
+      assert df1.names == ["a", "m", "n"]
+
+      assert df1.dtypes == %{
+               "a" => :string,
+               "n" => :string,
+               "m" => :string
+             }
+
+      assert DF.to_columns(df1, atom_keys: true) == %{
+               a: ["{\"n\": 1", "{\"m\": 1}"],
+               m: [nil, "1"],
+               n: [nil, nil]
+             }
+    end
+
+    test "json_path_match/2 - raises for invalid json path" do
+      df = DF.new([%{a: ~s({"n": 1})}, %{a: ~s({"m": 1})}])
+
+      assert_raise RuntimeError,
+                   "Polars Error: ComputeError(ErrString(\"error compiling JSONpath expression path error: \\nEof\\n\"))",
+                   fn ->
+                     DF.mutate(df, n: json_path_match(a, "$."))
+                   end
     end
 
     test "adds new columns" do