feat: Added escape_regex operation to the str namespace and as a …

…global function (#19257)

Cargo.toml

@@ -73,6 +73,7 @@ raw-cpuid = "11"
 rayon = "1.9"
 recursive = "0.1"
 regex = "1.9"
+regex-syntax = "0.8.5"
 reqwest = { version = "0.12", default-features = false }
 ryu = "1.0.13"
 serde = { version = "1.0.188", features = ["derive", "rc"] }

crates/polars-ops/Cargo.toml

@@ -34,6 +34,7 @@ rand = { workspace = true, optional = true, features = ["small_rng", "std"] }
 rand_distr = { workspace = true, optional = true }
 rayon = { workspace = true }
 regex = { workspace = true }
+regex-syntax = { workspace = true }
 serde = { workspace = true, optional = true }
 serde_json = { workspace = true, optional = true }
 unicode-reverse = { workspace = true, optional = true }

crates/polars-ops/src/chunked_array/strings/escape_regex.rs

@@ -0,0 +1,21 @@
+use polars_core::prelude::{StringChunked, StringChunkedBuilder};
+
+#[inline]
+pub fn escape_regex_str(s: &str) -> String {
+    regex_syntax::escape(s)
+}
+
+pub fn escape_regex(ca: &StringChunked) -> StringChunked {
+    let mut buffer = String::new();
+    let mut builder = StringChunkedBuilder::new(ca.name().clone(), ca.len());
+    for opt_s in ca.iter() {
+        if let Some(s) = opt_s {
+            buffer.clear();
+            regex_syntax::escape_into(s, &mut buffer);
+            builder.append_value(&buffer);
+        } else {
+            builder.append_null();
+        }
+    }
+    builder.finish()
+}

crates/polars-ops/src/chunked_array/strings/mod.rs

@@ -3,6 +3,8 @@ mod case;
 #[cfg(feature = "strings")]
 mod concat;
 #[cfg(feature = "strings")]
+mod escape_regex;
+#[cfg(feature = "strings")]
 mod extract;
 #[cfg(feature = "find_many")]
 mod find_many;
@@ -20,12 +22,13 @@ mod split;
 mod strip;
 #[cfg(feature = "strings")]
 mod substring;
-
 #[cfg(all(not(feature = "nightly"), feature = "strings"))]
 mod unicode_internals;
 
 #[cfg(feature = "strings")]
 pub use concat::*;
+#[cfg(feature = "strings")]
+pub use escape_regex::*;
 #[cfg(feature = "find_many")]
 pub use find_many::*;
 #[cfg(feature = "extract_jsonpath")]

crates/polars-ops/src/chunked_array/strings/namespace.rs

@@ -640,6 +640,12 @@ pub trait StringNameSpaceImpl: AsString {
 
         substring::tail(ca, n.i64()?)
     }
+    #[cfg(feature = "strings")]
+    /// Escapes all regular expression meta characters in the string.
+    fn str_escape_regex(&self) -> StringChunked {
+        let ca = self.as_string();
+        escape_regex::escape_regex(ca)
+    }
 }
 
 impl StringNameSpaceImpl for StringChunked {}

crates/polars-plan/src/dsl/function_expr/strings.rs

@@ -130,6 +130,8 @@ pub enum StringFunction {
         ascii_case_insensitive: bool,
         overlapping: bool,
     },
+    #[cfg(feature = "regex")]
+    EscapeRegex,
 }
 
 impl StringFunction {
@@ -197,6 +199,8 @@ impl StringFunction {
             ReplaceMany { .. } => mapper.with_same_dtype(),
             #[cfg(feature = "find_many")]
             ExtractMany { .. } => mapper.with_dtype(DataType::List(Box::new(DataType::String))),
+            #[cfg(feature = "regex")]
+            EscapeRegex => mapper.with_same_dtype(),
         }
     }
 }
@@ -285,6 +289,8 @@ impl Display for StringFunction {
             ReplaceMany { .. } => "replace_many",
             #[cfg(feature = "find_many")]
             ExtractMany { .. } => "extract_many",
+            #[cfg(feature = "regex")]
+            EscapeRegex => "escape_regex",
         };
         write!(f, "str.{s}")
     }
@@ -400,6 +406,8 @@ impl From<StringFunction> for SpecialEq<Arc<dyn ColumnsUdf>> {
             } => {
                 map_as_slice!(extract_many, ascii_case_insensitive, overlapping)
             },
+            #[cfg(feature = "regex")]
+            EscapeRegex => map!(escape_regex),
         }
     }
 }
@@ -1033,3 +1041,9 @@ pub(super) fn json_path_match(s: &[Column]) -> PolarsResult<Column> {
     let pat = s[1].str()?;
     Ok(ca.json_path_match(pat)?.into_column())
 }
+
+#[cfg(feature = "regex")]
+pub(super) fn escape_regex(s: &Column) -> PolarsResult<Column> {
+    let ca = s.str()?;
+    Ok(ca.str_escape_regex().into_column())
+}

crates/polars-plan/src/dsl/string.rs

@@ -592,4 +592,14 @@ impl StringNameSpace {
             None,
         )
     }
+
+    #[cfg(feature = "strings")]
+    pub fn escape_regex(self) -> Expr {
+        self.0.map_many_private(
+            FunctionExpr::StringExpr(StringFunction::EscapeRegex),
+            &[],
+            false,
+            None,
+        )
+    }
 }

crates/polars-python/src/expr/string.rs

@@ -339,4 +339,9 @@ impl PyExpr {
             .extract_many(patterns.inner, ascii_case_insensitive, overlapping)
             .into()
     }
+
+    #[cfg(feature = "regex")]
+    fn str_escape_regex(&self) -> Self {
+        self.inner.clone().str().escape_regex().into()
+    }
 }

crates/polars-python/src/functions/mod.rs

@@ -8,6 +8,7 @@ mod misc;
 mod random;
 mod range;
 mod string_cache;
+mod strings;
 mod whenthen;
 
 pub use aggregation::*;
@@ -20,4 +21,5 @@ pub use misc::*;
 pub use random::*;
 pub use range::*;
 pub use string_cache::*;
+pub use strings::*;
 pub use whenthen::*;

crates/polars-python/src/functions/strings.rs

@@ -0,0 +1,7 @@
+use pyo3::prelude::*;
+
+#[pyfunction]
+pub fn escape_regex(s: &str) -> PyResult<String> {
+    let escaped_s = polars_ops::chunked_array::strings::escape_regex_str(s);
+    Ok(escaped_s)
+}

crates/polars-python/src/lazyframe/visitor/expr_nodes.rs

@@ -174,6 +174,7 @@ pub enum PyStringFunction {
     ZFill,
     ContainsMany,
     ReplaceMany,
+    EscapeRegex,
 }
 
 #[pymethods]
@@ -953,6 +954,9 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
                     StringFunction::ExtractMany { .. } => {
                         return Err(PyNotImplementedError::new_err("extract_many"))
                     },
+                    StringFunction::EscapeRegex => {
+                        (PyStringFunction::EscapeRegex.into_py(py),).to_object(py)
+                    },
                 },
                 FunctionExpr::StructExpr(_) => {
                     return Err(PyNotImplementedError::new_err("struct expr"))

py-polars/docs/source/reference/expressions/string.rst

@@ -16,6 +16,7 @@ The following methods are available under the `expr.str` attribute.
     Expr.str.decode
     Expr.str.encode
     Expr.str.ends_with
+    Expr.str.escape_regex
     Expr.str.explode
     Expr.str.extract
     Expr.str.extract_all

py-polars/docs/source/reference/functions.rst

@@ -25,6 +25,7 @@ Miscellaneous
 
     align_frames
     concat
+    escape_regex
 
 Parallelization
 ~~~~~~~~~~~~~~~

py-polars/polars/__init__.py

@@ -104,6 +104,7 @@
     datetime_ranges,
     duration,
     element,
+    escape_regex,
     exclude,
     field,
     first,
@@ -303,6 +304,7 @@
     "time_range",
     "time_ranges",
     "zeros",
+    "escape_regex",
     # polars.functions.aggregation
     "all",
     "all_horizontal",

py-polars/polars/expr/string.py

@@ -2781,6 +2781,28 @@ def concat(
             delimiter = "-"
         return self.join(delimiter, ignore_nulls=ignore_nulls)
 
+    def escape_regex(self) -> Expr:
+        r"""
+        Returns string values with all regular expression meta characters escaped.
+
+        Examples
+        --------
+        >>> df = pl.DataFrame({"text": ["abc", "def", None, "abc(\\w+)"]})
+        >>> df.with_columns(pl.col("text").str.escape_regex().alias("escaped"))
+         shape: (4, 2)
+        ┌──────────┬──────────────┐
+        │ text     ┆ escaped      │
+        │ ---      ┆ ---          │
+        │ str      ┆ str          │
+        ╞══════════╪══════════════╡
+        │ abc      ┆ abc          │
+        │ def      ┆ def          │
+        │ null     ┆ null         │
+        │ abc(\w+) ┆ abc\(\\w\+\) │
+        └──────────┴──────────────┘
+        """
+        return wrap_expr(self._pyexpr.str_escape_regex())
+
 
 def _validate_format_argument(format: str | None) -> None:
     if format is not None and ".%f" in format:

py-polars/polars/functions/__init__.py

@@ -26,6 +26,7 @@
 from polars.functions.business import business_day_count
 from polars.functions.col import col
 from polars.functions.eager import align_frames, concat
+from polars.functions.escape_regex import escape_regex
 from polars.functions.lazy import (
     approx_n_unique,
     arctan2,
@@ -170,4 +171,6 @@
     # polars.functions.whenthen
     "when",
     "sql_expr",
+    # polars.functions.escape_regex
+    "escape_regex",
 ]

py-polars/polars/functions/escape_regex.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import contextlib
+
+with contextlib.suppress(ImportError):  # Module not available when building docs
+    import polars.polars as plr
+import polars._reexport as pl
+
+
+def escape_regex(s: str) -> str:
+    r"""
+    Escapes string regex meta characters.
+
+    Parameters
+    ----------
+    s
+        The string that all of its meta characters will be escaped.
+
+    """
+    if isinstance(s, pl.Expr):
+        msg = "escape_regex function is unsupported for `Expr`, you may want use `Expr.str.escape_regex` instead"
+        raise TypeError(msg)
+    elif not isinstance(s, str):
+        msg = f"escape_regex function supports only `str` type, got `{type(s)}`"
+        raise TypeError(msg)
+
+    return plr.escape_regex(s)

py-polars/src/lib.rs

@@ -275,6 +275,10 @@ fn polars(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(functions::set_random_seed))
         .unwrap();
 
+    // Functions - escape_regex
+    m.add_wrapped(wrap_pyfunction!(functions::escape_regex))
+        .unwrap();
+
     // Exceptions - Errors
     m.add(
         "PolarsError",

py-polars/tests/unit/functions/test_functions.py

@@ -538,3 +538,22 @@ def test_head_tail(fruits_cars: pl.DataFrame) -> None:
     res_expr = fruits_cars.select(pl.tail("A", 2))
     expected = pl.Series("A", [4, 5])
     assert_series_equal(res_expr.to_series(), expected)
+
+
+def test_escape_regex() -> None:
+    result = pl.escape_regex("abc(\\w+)")
+    expected = "abc\\(\\\\w\\+\\)"
+    assert result == expected
+
+    df = pl.DataFrame({"text": ["abc", "def", None, "abc(\\w+)"]})
+    with pytest.raises(
+        TypeError,
+        match="escape_regex function is unsupported for `Expr`, you may want use `Expr.str.escape_regex` instead",
+    ):
+        df.with_columns(escaped=pl.escape_regex(pl.col("text")))  # type: ignore[arg-type]
+
+    with pytest.raises(
+        TypeError,
+        match="escape_regex function supports only `str` type, got `<class 'int'>`",
+    ):
+        pl.escape_regex(3)  # type: ignore[arg-type]

py-polars/tests/unit/operations/namespaces/string/test_string.py

@@ -1793,3 +1793,16 @@ def test_json_decode_struct_schema() -> None:
         ),
         pl.Series([{"a": 1}, {"a": 2}]),
     )
+
+
+def test_escape_regex() -> None:
+    df = pl.DataFrame({"text": ["abc", "def", None, "abc(\\w+)"]})
+    result_df = df.with_columns(pl.col("text").str.escape_regex().alias("escaped"))
+    expected_df = pl.DataFrame(
+        {
+            "text": ["abc", "def", None, "abc(\\w+)"],
+            "escaped": ["abc", "def", None, "abc\\(\\\\w\\+\\)"],
+        }
+    )
+
+    assert_frame_equal(result_df, expected_df)