apply suggested changes

Author: arashandishgar

cpp/src/arrow/compute/api_scalar.h

@@ -272,8 +272,6 @@ class ARROW_EXPORT ExtractRegexSpanOptions : public FunctionOptions {
 
   /// Regular expression with named capture fields
   std::string pattern;
-
-  /// Shows the matched string
 };
 /// Options for IsIn and IndexIn functions
 class ARROW_EXPORT SetLookupOptions : public FunctionOptions {

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

@@ -43,7 +43,6 @@ namespace compute {
 namespace internal {
 
 namespace {
-
 // ----------------------------------------------------------------------
 // re2 utilities
 
@@ -2185,9 +2184,34 @@ void AddAsciiStringReplaceSubstring(FunctionRegistry* registry) {
 
 using ExtractRegexState = OptionsWrapper<ExtractRegexOptions>;
 
+struct BaseExtractRegexData {
+  Status Init() {
+    RETURN_NOT_OK(RegexStatus(*regex));
+    const int group_count = regex->NumberOfCapturingGroups();
+    const auto& name_map = regex->CapturingGroupNames();
+    group_names.reserve(group_count);
+
+    for (int i = 0; i < group_count; i++) {
+      auto item = name_map.find(i + 1);  // re2 starts counting from 1
+      if (item == name_map.end()) {
+        // XXX should we instead just create fields with an empty name?
+        return Status::Invalid("Regular expression contains unnamed groups");
+      }
+      group_names.emplace_back(item->second);
+    }
+    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;
+
+ protected:
+  explicit BaseExtractRegexData(const std::string& pattern, bool is_utf8 = true)
+      : regex(new RE2(pattern, MakeRE2Options(is_utf8))) {}
+};
+
 // TODO cache this once per ExtractRegexOptions
-class ExtractRegexData {
- public:
+struct ExtractRegexData : public BaseExtractRegexData {
   static Result<ExtractRegexData> Make(const ExtractRegexOptions& options,
                                        bool is_utf8 = true) {
     ExtractRegexData data(options.pattern, is_utf8);
@@ -2197,50 +2221,24 @@ class ExtractRegexData {
 
   Result<TypeHolder> ResolveOutputType(const std::vector<TypeHolder>& types) const {
     const DataType* input_type = types[0].type;
-    // as mentioned here
-    // https://arrow.apache.org/docs/developers/cpp/development.html#code-style-linting-and-ci
-    // nullptr should not be used
-    if (input_type == NULLPTR) {
+    if (input_type == nullptr) {
       // No input type specified
-      return NULLPTR;
+      return nullptr;
     }
     // Input type is either [Large]Binary or [Large]String and is also the type
     // of each field in the output struct type.
     DCHECK(is_base_binary_like(input_type->id()));
     FieldVector fields;
-    fields.reserve(group_names_.size());
+    fields.reserve(num_groups());
     std::shared_ptr<DataType> owned_type = input_type->GetSharedPtr();
-    std::transform(group_names_.begin(), group_names_.end(), std::back_inserter(fields),
+    std::transform(group_names.begin(), group_names.end(), std::back_inserter(fields),
                    [&](const std::string& name) { return field(name, owned_type); });
-    return struct_(fields);
+    return struct_(std::move(fields));
   }
-  int64_t num_group() const { return group_names_.size(); }
-  std::shared_ptr<RE2> regex() const { return regex_; }
 
- protected:
+ private:
   explicit ExtractRegexData(const std::string& pattern, bool is_utf8 = true)
-      : regex_(new RE2(pattern, MakeRE2Options(is_utf8))) {}
-
-  Status Init() {
-    RETURN_NOT_OK(RegexStatus(*regex_));
-
-    const int group_count = regex_->NumberOfCapturingGroups();
-    const auto& name_map = regex_->CapturingGroupNames();
-    group_names_.reserve(group_count);
-
-    for (int i = 0; i < group_count; i++) {
-      auto item = name_map.find(i + 1);  // re2 starts counting from 1
-      if (item == name_map.end()) {
-        // XXX should we instead just create fields with an empty name?
-        return Status::Invalid("Regular expression contains unnamed groups");
-      }
-      group_names_.emplace_back(item->second);
-    }
-    return Status::OK();
-  }
-
-  std::shared_ptr<RE2> regex_;
-  std::vector<std::string> group_names_;
+      : BaseExtractRegexData(pattern, is_utf8) {}
 };
 
 Result<TypeHolder> ResolveExtractRegexOutput(KernelContext* ctx,
@@ -2251,17 +2249,17 @@ Result<TypeHolder> ResolveExtractRegexOutput(KernelContext* ctx,
 }
 
 struct ExtractRegexBase {
-  const ExtractRegexData& data;
+  const BaseExtractRegexData& data;
   const int group_count;
   std::vector<re2::StringPiece> found_values;
   std::vector<RE2::Arg> args;
   std::vector<const RE2::Arg*> args_pointers;
   const RE2::Arg** args_pointers_start;
   const RE2::Arg* null_arg = nullptr;
 
-  explicit ExtractRegexBase(const ExtractRegexData& data)
+  explicit ExtractRegexBase(const BaseExtractRegexData& data)
       : data(data),
-        group_count(static_cast<int>(data.num_group())),
+        group_count(static_cast<int>(data.num_groups())),
         found_values(group_count) {
     args.reserve(group_count);
     args_pointers.reserve(group_count);
@@ -2276,7 +2274,7 @@ struct ExtractRegexBase {
   }
 
   bool Match(std::string_view s) {
-    return RE2::PartialMatchN(ToStringPiece(s), *data.regex(), args_pointers_start,
+    return RE2::PartialMatchN(ToStringPiece(s), *data.regex, args_pointers_start,
                               group_count);
   }
 };
@@ -2291,18 +2289,16 @@ struct ExtractRegex : public ExtractRegexBase {
   static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
     ExtractRegexOptions options = ExtractRegexState::Get(ctx);
     ARROW_ASSIGN_OR_RAISE(auto data, ExtractRegexData::Make(options, Type::is_utf8));
-    return ExtractRegex{data}.Extract(ctx, batch, out);
+    return ExtractRegex(data).Extract(ctx, batch, out);
   }
 
   Status Extract(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
-    ExtractRegexOptions options = ExtractRegexState::Get(ctx);
-    DCHECK_NE(out->array_data(), NULLPTR);
+    DCHECK_NE(out->array_data(), nullptr);
     std::shared_ptr<DataType> type = out->array_data()->type;
-    DCHECK_NE(type, NULLPTR);
-
-    std::unique_ptr<ArrayBuilder> array_builder;
-    RETURN_NOT_OK(MakeBuilder(ctx->memory_pool(), type, &array_builder));
-    StructBuilder* struct_builder = checked_cast<StructBuilder*>(array_builder.get());
+    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));
 
     std::vector<BuilderType*> field_builders;
     field_builders.reserve(group_count);
@@ -2357,82 +2353,83 @@ void AddAsciiStringExtractRegex(FunctionRegistry* registry) {
   }
   DCHECK_OK(registry->AddFunction(std::move(func)));
 }
-class ExtractRegexSpanData : public ExtractRegexData {
- public:
-  static Result<ExtractRegexSpanData> Make(const std::string& pattern) {
-    auto data = ExtractRegexSpanData(pattern, true);
+struct ExtractRegexSpanData : public BaseExtractRegexData {
+  static Result<ExtractRegexSpanData> Make(const std::string& pattern,
+                                           bool is_utf8 = true) {
+    auto data = ExtractRegexSpanData(pattern, is_utf8);
     ARROW_RETURN_NOT_OK(data.Init());
     return data;
   }
 
   Result<TypeHolder> ResolveOutputType(const std::vector<TypeHolder>& types) const {
     const DataType* input_type = types[0].type;
-    if (input_type == NULLPTR) {
-      return NULLPTR;
+    if (input_type == nullptr) {
+      return nullptr;
     }
     DCHECK(is_base_binary_like(input_type->id()));
-    const size_t field_count = group_names_.size();
+    const size_t field_count = num_groups();
     FieldVector fields;
     fields.reserve(field_count);
-    const auto owned_type = input_type->GetSharedPtr();
-    for (const auto& group_name : group_names_) {
-      auto type = is_binary_like(owned_type->id()) ? int32() : int64();
+    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
-      fields.push_back(field(group_name + "_span", fixed_size_list(type, 2)));
+      fields.push_back(field(group_name, fixed_size_list(index_type, 2)));
     }
-    return struct_(fields);
+    return struct_(std::move(fields));
   }
 
  private:
   ExtractRegexSpanData(const std::string& pattern, const bool is_utf8)
-      : ExtractRegexData(pattern, is_utf8) {}
+      : BaseExtractRegexData(pattern, is_utf8) {}
 };
 
 template <typename Type>
 struct ExtractRegexSpan : ExtractRegexBase {
   using ArrayType = typename TypeTraits<Type>::ArrayType;
   using BuilderType = typename TypeTraits<Type>::BuilderType;
+  using offset_type = typename Type::offset_type;
+  using OffsetBuilderType =
+      typename TypeTraits<typename CTypeTraits<offset_type>::ArrowType>::BuilderType;
+  using OffsetCType =
+      typename TypeTraits<typename CTypeTraits<offset_type>::ArrowType>::CType;
+
   using ExtractRegexBase::ExtractRegexBase;
 
   static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
     auto options = OptionsWrapper<ExtractRegexSpanOptions>::Get(ctx);
-    ARROW_ASSIGN_OR_RAISE(auto data, ExtractRegexSpanData::Make(options.pattern));
+    ARROW_ASSIGN_OR_RAISE(auto data,
+                          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);
+    DCHECK_NE(out->array_data(), nullptr);
     std::shared_ptr<DataType> out_type = out->array_data()->type;
-    DCHECK_NE(out_type, NULLPTR);
-    std::unique_ptr<ArrayBuilder> out_builder;
-    ARROW_RETURN_NOT_OK(
-        MakeBuilder(ctx->memory_pool(), out->type()->GetSharedPtr(), &out_builder));
+    ARROW_ASSIGN_OR_RAISE(auto out_builder, MakeBuilder(out_type, ctx->memory_pool()));
     auto struct_builder = checked_pointer_cast<StructBuilder>(std::move(out_builder));
+    ARROW_RETURN_NOT_OK(struct_builder->Reserve(batch[0].array.length));
     std::vector<FixedSizeListBuilder*> span_builders;
-    std::vector<ArrayBuilder*> array_builders;
+    std::vector<OffsetBuilderType*> array_builders;
     span_builders.reserve(group_count);
     array_builders.reserve(group_count);
     for (int i = 0; i < group_count; i++) {
       span_builders.push_back(
           checked_cast<FixedSizeListBuilder*>(struct_builder->field_builder(i)));
-      array_builders.push_back(span_builders[i]->value_builder());
+      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));
     }
+
     auto visit_null = [&]() { return struct_builder->AppendNull(); };
     auto visit_value = [&](std::string_view element) -> Status {
       if (Match(element)) {
         for (int i = 0; i < group_count; i++) {
           // https://github.com/google/re2/issues/24#issuecomment-97653183
-          if (found_values[i].data() != NULLPTR) {
+          if (found_values[i].data() != nullptr) {
             int64_t begin = found_values[i].data() - element.data();
             int64_t size = found_values[i].size();
-            if (is_binary_like(batch.GetTypes()[0].id())) {
-              ARROW_RETURN_NOT_OK(checked_cast<Int32Builder*>(array_builders[i])
-                                      ->AppendValues({static_cast<int32_t>(begin),
-                                                      static_cast<int32_t>(size)}));
-            } else {
-              ARROW_RETURN_NOT_OK(checked_cast<Int64Builder*>(array_builders[i])
-                                      ->AppendValues({begin, size}));
-            }
-
+            array_builders[i]->UnsafeAppend(static_cast<OffsetCType>(begin));
+            array_builders[i]->UnsafeAppend(static_cast<OffsetCType>(size));
             ARROW_RETURN_NOT_OK(span_builders[i]->Append());
           } else {
             ARROW_RETURN_NOT_OK(span_builders[i]->AppendNull());
@@ -2448,25 +2445,33 @@ struct ExtractRegexSpan : ExtractRegexBase {
         VisitArraySpanInline<Type>(batch[0].array, visit_value, visit_null));
 
     ARROW_ASSIGN_OR_RAISE(auto out_array, struct_builder->Finish());
-    out->value = out_array->data();
+    out->value = std::move(out_array->data());
     return Status::OK();
   }
 };
 
-const FunctionDoc extract_regex_doc_span(
-    "likes extract_regex; however, it contains the position and length of results", "",
+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);
 
-Result<TypeHolder> resolver(KernelContext* ctx, const std::vector<TypeHolder>& types) {
+Result<TypeHolder> ResolveExtractRegexSpanOutputType(
+    KernelContext* ctx, const std::vector<TypeHolder>& types) {
   auto options = OptionsWrapper<ExtractRegexSpanOptions>::Get(*ctx->state());
   ARROW_ASSIGN_OR_RAISE(auto span, ExtractRegexSpanData::Make(options.pattern));
   return span.ResolveOutputType(types);
 }
 
 void AddAsciiStringExtractRegexSpan(FunctionRegistry* registry) {
   auto func = std::make_shared<ScalarFunction>("extract_regex_span", Arity::Unary(),
-                                               extract_regex_doc_span);
-  OutputType output_type(resolver);
+                                               extract_regex_span_doc);
+  OutputType output_type(ResolveExtractRegexSpanOutputType);
   for (const auto& type : BaseBinaryTypes()) {
     ScalarKernel kernel({type}, output_type,
                         GenerateVarBinaryToVarBinary<ExtractRegexSpan>(type),

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

@@ -1960,28 +1960,59 @@ TYPED_TEST(TestBaseBinaryKernels, ExtractRegex) {
                    R"([{"letter": "a", "digit": "1"}, {"letter": "b", "digit": "3"}])",
                    &options);
 }
-TYPED_TEST(TestBaseBinaryKernels, ExtractRegexSapn) {
-  ExtractRegexSpanOptions options{"(?P<letter>[ab])(?P<digit>\\d)"};
+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();
-  auto out_type = struct_({field("letter_span", fixed_size_list(type_fixe_size_list, 2)),
-                           field("digit_span", fixed_size_list(type_fixe_size_list, 2))});
+  auto out_type = struct_({field("letter", fixed_size_list(type_fixe_size_list, 2)),
+                           field("digit", fixed_size_list(type_fixe_size_list, 2))});
   this->CheckUnary("extract_regex_span", R"([])", out_type, R"([])", &options);
+  this->CheckUnary("extract_regex_span", R"([ null,"123ab","cd123ab","cd123abef"])",
+                   out_type, R"([null,null,null,null])", &options);
   this->CheckUnary(
-      "extract_regex_span", R"(["a1", "b2", "c3", null])", out_type,
-      R"([{"letter_span":[0,1], "digit_span":[1,1]}, {"letter_span":[0,1], "digit_span":[1,1]}, null, null])",
-      &options);
-  this->CheckUnary(
-      "extract_regex_span", R"(["a1", "c3", null, "b2"])", out_type,
-      R"([{"letter_span":[0,1], "digit_span": [1,1]}, null, null, {"letter_span":[0,1], "digit_span":[1,1]}])",
-      &options);
-  this->CheckUnary(
-      "extract_regex_span", R"(["a1", "b2"])", out_type,
-      R"([{"letter_span": [0,1], "digit_span": [1,1]}, {"letter_span": [0,1], "digit_span": [1,1]}])",
-      &options);
-  this->CheckUnary(
-      "extract_regex_span", R"(["a1", "zb3z"])", out_type,
-      R"([{"letter_span": [0,1], "digit_span": [1,1]}, {"letter_span": [1,1], "digit_span": [2,1]}])",
+      "extract_regex_span",
+      R"(["a1", "b2", "c3", null,"123ab","abb12","abc13","cedbb15","cedaabb125efg"])",
+      out_type,
+      R"([{"letter":[0,1], "digit":[1,1]},
+                    {"letter":[0,1], "digit":[1,1]},
+                    null,
+                    null,
+                    null,
+                    {"letter":[0,3], "digit":[3,2]},
+                    null,
+                    {"letter":[3,2], "digit":[5,2]},
+                    {"letter":[3,4], "digit":[7,3]}])",
       &options);
+  this->CheckUnary("extract_regex_span", R"([ "a3","b2","cdaa123","cdab123ef"])",
+                   out_type,
+                   R"([{"letter":[0,1], "digit":[1,1]},
+                                  {"letter":[0,1], "digit":[1,1]},
+                                  {"letter":[2,2], "digit":[4,3]},
+                                  {"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();
+  auto out_type = struct_({field("foo", fixed_size_list(type_fixe_size_list, 2)),
+                           field("digit", fixed_size_list(type_fixe_size_list, 2))});
+  this->CheckUnary("extract_regex_span", R"([])", out_type, R"([])", &options);
+  this->CheckUnary("extract_regex_span", R"(["foo","foo123","abcfoo123","abc",null])",
+                   out_type,
+                   R"([{"foo":[0,3],"digit":null},
+                                  {"foo":[0,3],"digit":[3,3]},
+                                  {"foo":null,"digit":null},
+                                  {"foo":null,"digit":null},
+                                  null])",
+                   &options);
+  options = ExtractRegexSpanOptions{"(?P<foo>foo)(?P<digit>\\d+)?"};
+  this->CheckUnary("extract_regex_span", R"(["foo123","foo","123","abc","abcfoo"])",
+                   out_type,
+                   R"([{"foo":[0,3],"digit":[3,3]},
+                                  {"foo":[0,3],"digit":null},
+                                  null,
+                                  null,
+                                  {"foo":[3,3],"digit":null}])",
+                   &options);
 }
 
 TYPED_TEST(TestBaseBinaryKernels, ExtractRegexNoCapture) {