Nits + Python bindings

Author: pitrou

cpp/src/arrow/compute/api_scalar.cc

@@ -30,6 +30,7 @@
 #include "arrow/type.h"
 #include "arrow/util/checked_cast.h"
 #include "arrow/util/logging.h"
+
 namespace arrow {
 
 namespace internal {

cpp/src/arrow/compute/api_scalar.h

@@ -264,6 +264,7 @@ class ARROW_EXPORT ExtractRegexOptions : public FunctionOptions {
   /// Regular expression with named capture fields
   std::string pattern;
 };
+
 class ARROW_EXPORT ExtractRegexSpanOptions : public FunctionOptions {
  public:
   explicit ExtractRegexSpanOptions(std::string pattern);
@@ -273,6 +274,7 @@ class ARROW_EXPORT ExtractRegexSpanOptions : public FunctionOptions {
   /// Regular expression with named capture fields
   std::string pattern;
 };
+
 /// Options for IsIn and IndexIn functions
 class ARROW_EXPORT SetLookupOptions : public FunctionOptions {
  public:

cpp/src/arrow/compute/kernels/scalar_string_ascii.cc

@@ -43,6 +43,7 @@ namespace compute {
 namespace internal {
 
 namespace {
+
 // ----------------------------------------------------------------------
 // re2 utilities
 
@@ -2201,7 +2202,9 @@ struct BaseExtractRegexData {
     }
     return Status::OK();
   }
+
   int64_t num_groups() const { return static_cast<int64_t>(group_names.size()); }
+
   std::unique_ptr<RE2> regex;
   std::vector<std::string> group_names;
 
@@ -2297,14 +2300,15 @@ struct ExtractRegex : public ExtractRegexBase {
     std::shared_ptr<DataType> type = out->array_data()->type;
     ARROW_ASSIGN_OR_RAISE(std::unique_ptr<ArrayBuilder> array_builder,
                           MakeBuilder(type, ctx->memory_pool()));
-    auto struct_builder = checked_pointer_cast<StructBuilder>(std::move(array_builder));
-    ARROW_RETURN_NOT_OK(struct_builder->Reserve(batch[0].array.length));
+    StructBuilder* struct_builder = checked_cast<StructBuilder*>(array_builder.get());
+    ARROW_RETURN_NOT_OK(struct_builder->Reserve(batch[0].length()));
 
     std::vector<BuilderType*> field_builders;
     field_builders.reserve(group_count);
     for (int i = 0; i < group_count; i++) {
       field_builders.push_back(
           checked_cast<BuilderType*>(struct_builder->field_builder(i)));
+      RETURN_NOT_OK(field_builders.back()->Reserve(batch[0].length()));
     }
 
     auto visit_null = [&]() { return struct_builder->AppendNull(); };
@@ -2353,6 +2357,7 @@ void AddAsciiStringExtractRegex(FunctionRegistry* registry) {
   }
   DCHECK_OK(registry->AddFunction(std::move(func)));
 }
+
 struct ExtractRegexSpanData : public BaseExtractRegexData {
   static Result<ExtractRegexSpanData> Make(const std::string& pattern,
                                            bool is_utf8 = true) {
@@ -2367,12 +2372,11 @@ struct ExtractRegexSpanData : public BaseExtractRegexData {
       return nullptr;
     }
     DCHECK(is_base_binary_like(input_type->id()));
-    const size_t field_count = num_groups();
     FieldVector fields;
-    fields.reserve(field_count);
+    fields.reserve(num_groups());
     auto index_type = is_binary_like(input_type->id()) ? int32() : int64();
     for (const auto& group_name : group_names) {
-      // size list is 2 as every span contains position and length
+      // list size is 2 as every span contains position and length
       fields.push_back(field(group_name, fixed_size_list(index_type, 2)));
     }
     return struct_(std::move(fields));
@@ -2401,12 +2405,14 @@ struct ExtractRegexSpan : ExtractRegexBase {
                           ExtractRegexSpanData::Make(options.pattern, Type::is_utf8));
     return ExtractRegexSpan{data}.Extract(ctx, batch, out);
   }
+
   Status Extract(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
     DCHECK_NE(out->array_data(), nullptr);
     std::shared_ptr<DataType> out_type = out->array_data()->type;
     ARROW_ASSIGN_OR_RAISE(auto out_builder, MakeBuilder(out_type, ctx->memory_pool()));
-    auto struct_builder = checked_pointer_cast<StructBuilder>(std::move(out_builder));
+    StructBuilder* struct_builder = checked_cast<StructBuilder*>(out_builder.get());
     ARROW_RETURN_NOT_OK(struct_builder->Reserve(batch[0].array.length));
+
     std::vector<FixedSizeListBuilder*> span_builders;
     std::vector<OffsetBuilderType*> array_builders;
     span_builders.reserve(group_count);
@@ -2416,8 +2422,8 @@ struct ExtractRegexSpan : ExtractRegexBase {
           checked_cast<FixedSizeListBuilder*>(struct_builder->field_builder(i)));
       array_builders.push_back(
           checked_cast<OffsetBuilderType*>(span_builders.back()->value_builder()));
-      RETURN_NOT_OK(span_builders.back()->Reserve(batch[0].array.length));
-      RETURN_NOT_OK(array_builders.back()->Reserve(2 * batch[0].array.length));
+      RETURN_NOT_OK(span_builders.back()->Reserve(batch[0].length()));
+      RETURN_NOT_OK(array_builders.back()->Reserve(2 * batch[0].length()));
     }
 
     auto visit_null = [&]() { return struct_builder->AppendNull(); };
@@ -2451,15 +2457,20 @@ struct ExtractRegexSpan : ExtractRegexBase {
 };
 
 const FunctionDoc extract_regex_span_doc(
-    "Extract substrings captured by a regex pattern and Save the result in the form of "
-    "(offset,length)",
-    "For each string in strings, match the regular expression and, if\n"
-    "successful, emit a struct with field names and values coming from the\n"
-    "regular expression's named capture groups, which are stored in a form of a\n "
-    "fixed_size_list(offset, length). If the input is null or the regular \n"
-    "expression Fails matching, a null output value is emitted.\n"
-    "Regular expression matching is done using the Google RE2 library.",
-    {"strings"}, "ExtractRegexSpanOptions", true);
+    "Extract string spans captured by a regex pattern",
+    ("For each string in strings, match the regular expression and, if\n"
+     "successful, emit a struct with field names and values coming from the\n"
+     "regular expression's named capture groups. Each struct field value\n"
+     "will be a fixed_size_list(offset_type, 2) where offset_type is int32\n"
+     "or int64, depending on the input string type. The two elements in\n"
+     "each fixed-size list are the index and the length of the substring\n"
+     "matched by the corresponding named capture group.\n"
+     "\n"
+     "If the input is null or the regular expression fails matching,\n"
+     "a null output value is emitted.\n"
+     "\n"
+     "Regular expression matching is done using the Google RE2 library."),
+    {"strings"}, "ExtractRegexSpanOptions", /*options_required=*/true);
 
 Result<TypeHolder> ResolveExtractRegexSpanOutputType(
     KernelContext* ctx, const std::vector<TypeHolder>& types) {

cpp/src/arrow/compute/kernels/scalar_string_test.cc

@@ -314,7 +314,7 @@ TYPED_TEST(TestBinaryKernels, NonUtf8Regex) {
                      this->MakeArray({"\xfc\x40", "this \xfc\x40 that \xfc\x40"}),
                      this->MakeArray({"bazz", "this bazz that \xfc\x40"}), &options);
   }
-  // TODO the following test is broken
+  // TODO the following test is broken (GH-45735)
   {
     ExtractRegexOptions options("(?P<letter>[\\xfc])(?P<digit>\\d)");
     auto null_bitmap = std::make_shared<Buffer>("0");
@@ -371,7 +371,7 @@ TYPED_TEST(TestBinaryKernels, NonUtf8WithNullRegex) {
                      this->template MakeArray<std::string>({{"\x00\x40", 2}}),
                      this->type(), R"(["bazz"])", &options);
   }
-  // TODO the following test is broken
+  // TODO the following test is broken (GH-45735)
   {
     ExtractRegexOptions options("(?P<null>[\\x00])(?P<digit>\\d)");
     auto null_bitmap = std::make_shared<Buffer>("0");
@@ -1960,6 +1960,7 @@ TYPED_TEST(TestBaseBinaryKernels, ExtractRegex) {
                    R"([{"letter": "a", "digit": "1"}, {"letter": "b", "digit": "3"}])",
                    &options);
 }
+
 TYPED_TEST(TestBaseBinaryKernels, ExtractRegexSpan) {
   ExtractRegexSpanOptions options{"(?P<letter>[ab]+)(?P<digit>\\d+)"};
   auto type_fixe_size_list = is_binary_like(this->type()->id()) ? int32() : int64();
@@ -1990,6 +1991,7 @@ TYPED_TEST(TestBaseBinaryKernels, ExtractRegexSpan) {
                                   {"letter":[2,2], "digit":[4,3]}])",
                    &options);
 }
+
 TYPED_TEST(TestBaseBinaryKernels, ExtractRegexSpanCaptureOption) {
   ExtractRegexSpanOptions options{"(?P<foo>foo)?(?P<digit>\\d+)?"};
   auto type_fixe_size_list = is_binary_like(this->type()->id()) ? int32() : int64();
@@ -2022,13 +2024,15 @@ TYPED_TEST(TestBaseBinaryKernels, ExtractRegexNoCapture) {
   this->CheckUnary("extract_regex", R"(["oofoo", "bar", null])", type,
                    R"([{}, null, null])", &options);
 }
+
 TYPED_TEST(TestBaseBinaryKernels, ExtractRegexSpanNoCapture) {
   // XXX Should we accept this or is it a user error?
   ExtractRegexSpanOptions options{"foo"};
   auto type = struct_({});
   this->CheckUnary("extract_regex_span", R"(["oofoo", "bar", null])", type,
                    R"([{}, null, null])", &options);
 }
+
 TYPED_TEST(TestBaseBinaryKernels, ExtractRegexNoOptions) {
   Datum input = ArrayFromJSON(this->type(), "[]");
   ASSERT_RAISES(Invalid, CallFunction("extract_regex", {input}));
@@ -2051,6 +2055,7 @@ TYPED_TEST(TestBaseBinaryKernels, ExtractRegexInvalid) {
       Invalid, ::testing::HasSubstr("Regular expression contains unnamed groups"),
       CallFunction("extract_regex", {input}, &options));
 }
+
 TYPED_TEST(TestBaseBinaryKernels, ExtractRegexSpanInvalid) {
   Datum input = ArrayFromJSON(this->type(), "[]");
   ExtractRegexSpanOptions options{"invalid["};

docs/source/cpp/compute.rst

@@ -1140,7 +1140,8 @@ String component extraction
   library.  The output struct field names refer to the named capture groups,
   e.g. 'letter' and 'digit' for the regular expression
   ``(?P<letter>[ab])(?P<digit>\\d)``.
-* \(2) Extract  the offset and length of substrings defined by a regular expression
+
+* \(2) Extract the offset and length of substrings defined by a regular expression
   using the Google RE2 library.  The output struct field names refer to the named
   capture groups, e.g. 'letter' and 'digit' for the regular expression
   ``(?P<letter>[ab])(?P<digit>\\d)``. Each output struct field is a fixed size list

python/pyarrow/_compute.pyx

@@ -1222,6 +1222,25 @@ class ExtractRegexOptions(_ExtractRegexOptions):
         self._set_options(pattern)
 
 
+cdef class _ExtractRegexSpanOptions(FunctionOptions):
+    def _set_options(self, pattern):
+        self.wrapped.reset(new CExtractRegexSpanOptions(tobytes(pattern)))
+
+
+class ExtractRegexSpanOptions(_ExtractRegexSpanOptions):
+    """
+    Options for the `extract_regex_span` function.
+
+    Parameters
+    ----------
+    pattern : str
+        Regular expression with named capture fields.
+    """
+
+    def __init__(self, pattern):
+        self._set_options(pattern)
+
+
 cdef class _SliceOptions(FunctionOptions):
     def _set_options(self, start, stop, step):
         self.wrapped.reset(new CSliceOptions(start, stop, step))

python/pyarrow/compute.py

@@ -40,6 +40,7 @@
     RunEndEncodeOptions,
     ElementWiseAggregateOptions,
     ExtractRegexOptions,
+    ExtractRegexSpanOptions,
     FilterOptions,
     IndexOptions,
     JoinOptions,

python/pyarrow/includes/libarrow.pxd

@@ -2500,6 +2500,11 @@ cdef extern from "arrow/compute/api.h" namespace "arrow::compute" nogil:
         CExtractRegexOptions(c_string pattern)
         c_string pattern
 
+    cdef cppclass CExtractRegexSpanOptions \
+            "arrow::compute::ExtractRegexSpanOptions"(CFunctionOptions):
+        CExtractRegexSpanOptions(c_string pattern)
+        c_string pattern
+
     cdef cppclass CCastOptions" arrow::compute::CastOptions"(CFunctionOptions):
         CCastOptions()
         CCastOptions(c_bool safe)

python/pyarrow/tests/test_compute.py

@@ -152,6 +152,7 @@ def test_option_class_equality(request):
         pc.RunEndEncodeOptions(),
         pc.ElementWiseAggregateOptions(skip_nulls=True),
         pc.ExtractRegexOptions("pattern"),
+        pc.ExtractRegexSpanOptions("pattern"),
         pc.FilterOptions(),
         pc.IndexOptions(pa.scalar(1)),
         pc.JoinOptions(),
@@ -1092,6 +1093,16 @@ def test_extract_regex():
     assert struct.tolist() == expected
 
 
+def test_extract_regex_span():
+    ar = pa.array(['a1', 'zb234z'])
+    expected = [{'letter': [0, 1], 'digit': [1, 1]},
+                {'letter': [1, 1], 'digit': [2, 3]}]
+    struct = pc.extract_regex_span(ar, pattern=r'(?P<letter>[ab])(?P<digit>\d+)')
+    assert struct.tolist() == expected
+    struct = pc.extract_regex_span(ar, r'(?P<letter>[ab])(?P<digit>\d+)')
+    assert struct.tolist() == expected
+
+
 def test_binary_join():
     ar_list = pa.array([['foo', 'bar'], None, []])
     expected = pa.array(['foo-bar', None, ''])