refactor: Make json_path_match expr non-anonymous (#15682)

pola-rs · Apr 16, 2024 · b338076 · b338076
1 parent 0e68e75
commit b338076
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 12 deletions.
diff --git a/crates/polars-plan/src/dsl/function_expr/strings.rs b/crates/polars-plan/src/dsl/function_expr/strings.rs
@@ -62,6 +62,8 @@ pub enum StringFunction {
         dtype: Option<DataType>,
         infer_schema_len: Option<usize>,
     },
+    #[cfg(feature = "extract_jsonpath")]
+    JsonPathMatch(String),
     #[cfg(feature = "regex")]
     Replace {
         // negative is replace all
@@ -146,6 +148,8 @@ impl StringFunction {
             Find { .. } => mapper.with_dtype(DataType::UInt32),
             #[cfg(feature = "extract_jsonpath")]
             JsonDecode { dtype, .. } => mapper.with_opt_dtype(dtype.clone()),
+            #[cfg(feature = "extract_jsonpath")]
+            JsonPathMatch(_) => mapper.with_dtype(DataType::String),
             LenBytes => mapper.with_dtype(DataType::UInt32),
             LenChars => mapper.with_dtype(DataType::UInt32),
             #[cfg(feature = "regex")]
@@ -216,6 +220,8 @@ impl Display for StringFunction {
             Tail { .. } => "tail",
             #[cfg(feature = "extract_jsonpath")]
             JsonDecode { .. } => "json_decode",
+            #[cfg(feature = "extract_jsonpath")]
+            JsonPathMatch(_) => "json_path_match",
             LenBytes => "len_bytes",
             Lowercase => "lowercase",
             LenChars => "len_chars",
@@ -367,6 +373,8 @@ impl From<StringFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
                 dtype,
                 infer_schema_len,
             } => map!(strings::json_decode, dtype.clone(), infer_schema_len),
+            #[cfg(feature = "extract_jsonpath")]
+            JsonPathMatch(pat) => map!(strings::json_path_match, &pat),
             #[cfg(feature = "find_many")]
             ContainsMany {
                 ascii_case_insensitive,
@@ -984,3 +992,9 @@ pub(super) fn json_decode(
     let ca = s.str()?;
     ca.json_decode(dtype, infer_schema_len)
 }
+
+#[cfg(feature = "extract_jsonpath")]
+pub(super) fn json_path_match(s: &Series, pat: &str) -> PolarsResult<Series> {
+    let ca = s.str()?;
+    Ok(ca.json_path_match(pat)?.into_series())
+}
diff --git a/crates/polars-plan/src/dsl/string.rs b/crates/polars-plan/src/dsl/string.rs
@@ -560,4 +560,10 @@ impl StringNameSpace {
                 infer_schema_len,
             }))
     }
+
+    #[cfg(feature = "extract_jsonpath")]
+    pub fn json_path_match(self, pat: String) -> Expr {
+        self.0
+            .map_private(FunctionExpr::StringExpr(StringFunction::JsonPathMatch(pat)))
+    }
 }
diff --git a/py-polars/src/expr/string.rs b/py-polars/src/expr/string.rs
@@ -238,18 +238,7 @@ impl PyExpr {
 
     #[cfg(feature = "extract_jsonpath")]
     fn str_json_path_match(&self, pat: String) -> Self {
-        let function = move |s: Series| {
-            let ca = s.str()?;
-            match ca.json_path_match(&pat) {
-                Ok(ca) => Ok(Some(ca.into_series())),
-                Err(e) => Err(PolarsError::ComputeError(format!("{e:?}").into())),
-            }
-        };
-        self.inner
-            .clone()
-            .map(function, GetOutput::from_type(DataType::String))
-            .with_fmt("str.json_path_match")
-            .into()
+        self.inner.clone().str().json_path_match(pat).into()
     }
 
     fn str_extract(&self, pat: Self, group_index: usize) -> Self {