feat: expose regex_count function (#1066)

* Added wrapper for regex_count function

* fix comment

---------

Co-authored-by: Nirnay Roy <[email protected]>

Author: nirnayroy

python/datafusion/functions.py

@@ -217,6 +217,7 @@
     "random",
     "range",
     "rank",
+    "regexp_count",
     "regexp_like",
     "regexp_match",
     "regexp_replace",
@@ -779,6 +780,23 @@ def regexp_replace(
     return Expr(f.regexp_replace(string.expr, pattern.expr, replacement.expr, flags))
 
 
+def regexp_count(
+    string: Expr, pattern: Expr, start: Expr, flags: Expr | None = None
+) -> Expr:
+    """Returns the number of matches in a string.
+
+    Optional start position (the first position is 1) to search for the regular
+    expression.
+    """
+    if flags is not None:
+        flags = flags.expr
+    if start is not None:
+        start = start.expr
+    else:
+        start = Expr.expr
+    return Expr(f.regexp_count(string.expr, pattern.expr, start, flags))
+
+
 def repeat(string: Expr, n: Expr) -> Expr:
     """Repeats the ``string`` to ``n`` times."""
     return Expr(f.repeat(string.expr, n.expr))

python/tests/test_functions.py

@@ -740,6 +740,10 @@ def test_array_function_obj_tests(stmt, py_expr):
             f.regexp_replace(column("a"), literal("(ell|orl)"), literal("-")),
             pa.array(["H-o", "W-d", "!"]),
         ),
+        (
+            f.regexp_count(column("a"), literal("(ell|orl)"), literal(1)),
+            pa.array([1, 1, 0], type=pa.int64()),
+        ),
     ],
 )
 def test_string_functions(df, function, expected_result):

src/functions.rs

@@ -173,6 +173,25 @@ fn regexp_replace(
     )
     .into())
 }
+
+#[pyfunction]
+#[pyo3(signature = (string, pattern, start, flags=None))]
+/// Returns the number of matches found in the string.
+fn regexp_count(
+    string: PyExpr,
+    pattern: PyExpr,
+    start: Option<PyExpr>,
+    flags: Option<PyExpr>,
+) -> PyResult<PyExpr> {
+    Ok(functions::expr_fn::regexp_count(
+        string.expr,
+        pattern.expr,
+        start.map(|x| x.expr),
+        flags.map(|x| x.expr),
+    )
+    .into())
+}
+
 /// Creates a new Sort Expr
 #[pyfunction]
 fn order_by(expr: PyExpr, asc: bool, nulls_first: bool) -> PyResult<PySortExpr> {
@@ -943,6 +962,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(power))?;
     m.add_wrapped(wrap_pyfunction!(radians))?;
     m.add_wrapped(wrap_pyfunction!(random))?;
+    m.add_wrapped(wrap_pyfunction!(regexp_count))?;
     m.add_wrapped(wrap_pyfunction!(regexp_like))?;
     m.add_wrapped(wrap_pyfunction!(regexp_match))?;
     m.add_wrapped(wrap_pyfunction!(regexp_replace))?;