Add support for running str.match in arrow compute. (#865)

Co-authored-by: Todd A. Anderson <[email protected]>

Author: DrTodd13

bodo/pandas/physical/expression.cpp

@@ -166,7 +166,8 @@ std::shared_ptr<array_info> do_arrow_compute_binary(
 }
 
 std::shared_ptr<array_info> do_arrow_compute_unary(
-    std::shared_ptr<ExprResult> left_res, const std::string& comparator) {
+    std::shared_ptr<ExprResult> left_res, const std::string& comparator,
+    const arrow::compute::FunctionOptions* func_options) {
     // Try to convert the results of our children into array
     // or scalar results to see which one they are.
     std::shared_ptr<ArrayExprResult> left_as_array =
@@ -187,7 +188,7 @@ std::shared_ptr<array_info> do_arrow_compute_unary(
     }
 
     arrow::Result<arrow::Datum> cmp_res =
-        arrow::compute::CallFunction(comparator, {src1});
+        arrow::compute::CallFunction(comparator, {src1}, func_options);
     if (!cmp_res.ok()) [[unlikely]] {
         throw std::runtime_error(
             "do_array_compute_unary: Error in Arrow compute: " +
@@ -665,6 +666,31 @@ std::shared_ptr<ExprResult> PhysicalArrowExpression::ProcessBatch(
         // which returns a struct. To match the output dtype of Pandas, we Cast
         // to Date32 instead.
         result = do_arrow_compute_cast(res, duckdb::LogicalType::DATE);
+    } else if (scalar_func_data.arrow_func_name == "match_substring_regex") {
+        if (!PyTuple_Check(scalar_func_data.args) ||
+            PyTuple_Size(scalar_func_data.args) != 1) {
+            throw std::runtime_error(
+                "match_substring_regex args not a 1-element tuple.");
+        }
+
+        // Get the first element (borrowed reference)
+        PyObject* py_str = PyTuple_GetItem(scalar_func_data.args, 0);
+
+        if (!PyUnicode_Check(py_str)) {
+            throw std::runtime_error(
+                "match_substring_regex args element is not a Python string.");
+        }
+
+        // Convert to UTF‑8 C string
+        const char* c_str = PyUnicode_AsUTF8(py_str);
+        if (!c_str) {
+            throw std::runtime_error(
+                "match_substring_regex error extracting Python string.");
+        }
+
+        arrow::compute::MatchSubstringOptions opts(c_str);
+        result = do_arrow_compute_unary(res, scalar_func_data.arrow_func_name,
+                                        &opts);
     } else {
         result = do_arrow_compute_unary(res, scalar_func_data.arrow_func_name);
     }

bodo/pandas/physical/expression.h

@@ -137,7 +137,8 @@ extern std::function<bool(int)> less_equal_test;
  *
  */
 std::shared_ptr<array_info> do_arrow_compute_unary(
-    std::shared_ptr<ExprResult> left_res, const std::string &comparator);
+    std::shared_ptr<ExprResult> left_res, const std::string &comparator,
+    const arrow::compute::FunctionOptions *func_options = nullptr);
 
 /**
  * @brief Convert two ExprResults to arrow and run compute operation on them.

bodo/pandas/plan.py

@@ -739,6 +739,11 @@ def function_name(self):
         """Return the function name."""
         return self.args[2]
 
+    @property
+    def function_args(self):
+        """Return the function args."""
+        return self.args[3]
+
     def update_func_expr_source(self, new_source_plan: LazyPlan, col_index_offset: int):
         """Update the source and column index of the function expression."""
         if self.source != new_source_plan:
@@ -763,6 +768,7 @@ def update_func_expr_source(self, new_source_plan: LazyPlan, col_index_offset: i
                 new_source_plan,
                 (in_col_ind + col_index_offset,) + index_cols,
                 self.function_name,
+                self.function_args,
             )
             expr.is_series = self.is_series
             return expr

bodo/pandas/plan_optimizer.pyx

@@ -676,11 +676,12 @@ cdef class ArrowScalarFuncExpression(Expression):
         object out_schema,
         LogicalOperator source,
         vector[int] input_column_indices,
-        str function_name):
+        str function_name,
+        object args):
 
         self.out_schema = out_schema
         self.c_expression = make_scalar_func_expr(
-            source.c_logical_operator, out_schema, None, input_column_indices, False, False, function_name.encode())
+            source.c_logical_operator, out_schema, args, input_column_indices, False, False, function_name.encode())
 
     def __str__(self):
         return f"ArrowScalarFuncExpression({self.function_name})"

bodo/pandas/series.py

@@ -2554,10 +2554,7 @@ def make_expr(expr, plan, first, schema, index_cols, side="right"):
         idx = get_new_idx(idx, first, side)
         empty_data = arrow_to_empty_df(pa.schema([expr.pa_schema[0]]))
         return ArrowScalarFuncExpression(
-            empty_data,
-            plan,
-            (idx,) + tuple(index_cols),
-            expr.function_name,
+            empty_data, plan, (idx,) + tuple(index_cols), expr.function_name, ()
         )
     elif is_arith_expr(expr):
         # TODO: recursively traverse arithmetic expr tree to update col idx.
@@ -2729,6 +2726,7 @@ def _get_series_func_plan(
         "str.swapcase",
         "str.title",
         "str.reverse",
+        "str.match",
     )
 
     def get_arrow_func(name):
@@ -2740,20 +2738,23 @@ def get_arrow_func(name):
         if name.startswith("str.is"):
             body = name.split(".")[1]
             return "utf8_" + body[:2] + "_" + body[2:]
+        if name == "str.match":
+            return "match_substring_regex"
         if name.startswith("str."):
             return "utf8_" + name.split(".")[1]
         return name.split(".")[1]
 
-    if func in arrow_compute_list:
+    if func in arrow_compute_list and len(kwargs) == 0:
         func_name = get_arrow_func(func)
-        func_args = ()  # TODO: expand this to enable arrow compute calls with args
+        func_args = tuple(args)
         is_cfunc = False
         has_state = False
         expr = ArrowScalarFuncExpression(
             empty_data,
             source_data,
             (col_index,) + tuple(index_cols),
             func_name,
+            func_args,
         )
     else:
         # Empty func_name separates Python calls from Arrow calls.

bodo/tests/test_df_lib/test_end_to_end.py

@@ -3763,3 +3763,18 @@ def test_join_non_equi_key_not_in_output():
         reset_index=True,
         sort_output=True,
     )
+
+
+def test_series_str_match():
+    s = pd.Series(["abc", "a1c", "zzz", None], dtype="string")
+    bs = bd.Series(s)
+
+    # Match strings that start with 'a' and end with 'c'
+    pmask = s.str.match(r"a.*c")
+    bmask = bs.str.match(r"a.*c")
+
+    _test_equal(
+        bmask,
+        pmask,
+        check_pandas_types=False,
+    )