GH-50615: [C++] Reduce code generation for string kernels

cpp/src/arrow/compute/kernels/codegen_internal.h

@@ -1307,10 +1307,14 @@ KernelType GenerateTypeAgnosticPrimitive(detail::GetTypeId get_id) {
   }
 }
 
-// similar to GenerateTypeAgnosticPrimitive, but for base variable binary types
-template <template <typename...> class Generator, typename KernelType = ArrayKernelExec,
-          typename... Args>
-KernelType GenerateTypeAgnosticVarBinaryBase(detail::GetTypeId get_id) {
+// Similar to GenerateTypeAgnosticPrimitive, but for base variable binary types
+//
+// Note that we don't offer to generate separate code for String types, because
+// the utf8-ness of a type can be retrieved and handled efficiently at runtime.
+// This helps cut down on code generation (see GH-50615).
+template <template <typename...> class Generator, typename... Args>
+auto GenerateTypeAgnosticVarBinaryBase(detail::GetTypeId get_id) {
+  using KernelType = decltype(&Generator<BinaryType, Args...>::Exec);
   switch (get_id.id) {
     case Type::BINARY:
     case Type::STRING:
@@ -1324,24 +1328,6 @@ KernelType GenerateTypeAgnosticVarBinaryBase(detail::GetTypeId get_id) {
   }
 }
 
-// Generate a kernel given a templated functor for binary and string types
-template <template <typename...> class Generator, typename... Args>
-ArrayKernelExec GenerateVarBinaryToVarBinary(detail::GetTypeId get_id) {
-  switch (get_id.id) {
-    case Type::BINARY:
-      return Generator<BinaryType, Args...>::Exec;
-    case Type::STRING:
-      return Generator<StringType, Args...>::Exec;
-    case Type::LARGE_BINARY:
-      return Generator<LargeBinaryType, Args...>::Exec;
-    case Type::LARGE_STRING:
-      return Generator<LargeStringType, Args...>::Exec;
-    default:
-      ARROW_DCHECK(false);
-      return nullptr;
-  }
-}
-
 // Generate a kernel given a templated functor for base binary types. Generates
 // a single kernel for binary/string and large binary/large string. If your kernel
 // implementation needs access to the specific type at compile time, please use
@@ -1363,24 +1349,6 @@ ArrayKernelExec GenerateVarBinaryBase(detail::GetTypeId get_id) {
   }
 }
 
-// See BaseBinary documentation
-template <template <typename...> class Generator, typename Type0, typename... Args>
-ArrayKernelExec GenerateVarBinary(detail::GetTypeId get_id) {
-  switch (get_id.id) {
-    case Type::BINARY:
-      return Generator<Type0, BinaryType, Args...>::Exec;
-    case Type::STRING:
-      return Generator<Type0, StringType, Args...>::Exec;
-    case Type::LARGE_BINARY:
-      return Generator<Type0, LargeBinaryType, Args...>::Exec;
-    case Type::LARGE_STRING:
-      return Generator<Type0, LargeStringType, Args...>::Exec;
-    default:
-      ARROW_DCHECK(false);
-      return nullptr;
-  }
-}
-
 // Generate a kernel given a templated functor for binary-view types. Generates a
 // single kernel for binary/string-view.
 //

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

@@ -825,8 +825,7 @@ std::shared_ptr<ScalarFunction> MakeScalarMinMax(std::string name, FunctionDoc d
     DCHECK_OK(func->AddKernel(std::move(kernel)));
   }
   for (const auto& ty : BaseBinaryTypes()) {
-    auto exec =
-        GenerateTypeAgnosticVarBinaryBase<BinaryScalarMinMax, ArrayKernelExec, Op>(ty);
+    auto exec = GenerateTypeAgnosticVarBinaryBase<BinaryScalarMinMax, Op>(ty);
     ScalarKernel kernel{KernelSignature::Make({ty}, ty, /*is_varargs=*/true), exec,
                         MinMaxState::Init};
     kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE;

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

@@ -1327,7 +1327,7 @@ void AddBinaryIfElseKernels(const std::shared_ptr<IfElseFunction>& scalar_functi
                             const std::vector<std::shared_ptr<DataType>>& types) {
   for (auto&& type : types) {
     auto exec =
-        internal::GenerateTypeAgnosticVarBinaryBase<ResolveIfElseExec, ArrayKernelExec,
+        internal::GenerateTypeAgnosticVarBinaryBase<ResolveIfElseExec,
                                                     /*AllocateMem=*/std::true_type>(
             *type);
     // cond array needs to be boolean always

cpp/src/arrow/compute/kernels/scalar_string_internal.h

@@ -21,6 +21,7 @@
 
 #include "arrow/compute/api_scalar.h"
 #include "arrow/compute/kernels/common_internal.h"
+#include "arrow/type_traits.h"
 
 namespace arrow {
 namespace compute {
@@ -68,10 +69,13 @@ static int64_t GetVarBinaryValuesLength(const ArraySpan& span) {
 ///
 /// and returns the number of codeunits of the `output` sequence or a negative
 /// value if an invalid input sequence is detected.
-template <typename Type, typename StringTransform>
+template <typename PhysicalType, typename StringTransform>
 struct StringTransformExecBase {
-  using offset_type = typename Type::offset_type;
-  using ArrayType = typename TypeTraits<Type>::ArrayType;
+  using offset_type = typename PhysicalType::offset_type;
+  using ArrayType = typename TypeTraits<PhysicalType>::ArrayType;
+
+  static_assert(!is_string_or_string_view(PhysicalType::type_id),
+                "should only codegen on physical types");
 
   static Status Execute(KernelContext* ctx, StringTransform* transform,
                         const ExecSpan& batch, ExecResult* out) {
@@ -121,9 +125,11 @@ struct StringTransformExecBase {
   }
 };
 
-template <typename Type, typename StringTransform>
-struct StringTransformExec : public StringTransformExecBase<Type, StringTransform> {
-  using StringTransformExecBase<Type, StringTransform>::Execute;
+template <typename Type, typename StringTransform,
+          typename PhysicalType = typename Type::PhysicalType>
+struct StringTransformExec
+    : public StringTransformExecBase<PhysicalType, StringTransform> {
+  using StringTransformExecBase<PhysicalType, StringTransform>::Execute;
 
   static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
     StringTransform transform;
@@ -132,11 +138,12 @@ struct StringTransformExec : public StringTransformExecBase<Type, StringTransfor
   }
 };
 
-template <typename Type, typename StringTransform>
+template <typename Type, typename StringTransform,
+          typename PhysicalType = typename Type::PhysicalType>
 struct StringTransformExecWithState
-    : public StringTransformExecBase<Type, StringTransform> {
+    : public StringTransformExecBase<PhysicalType, StringTransform> {
   using State = typename StringTransform::State;
-  using StringTransformExecBase<Type, StringTransform>::Execute;
+  using StringTransformExecBase<PhysicalType, StringTransform>::Execute;
 
   static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
     StringTransform transform(State::Get(ctx));
@@ -151,7 +158,7 @@ void MakeUnaryStringBatchKernel(
     MemAllocation::type mem_allocation = MemAllocation::PREALLOCATE) {
   auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));
   for (const auto& ty : StringTypes()) {
-    auto exec = GenerateVarBinaryToVarBinary<ExecFunctor>(ty);
+    auto exec = GenerateTypeAgnosticVarBinaryBase<ExecFunctor>(ty);
     ScalarKernel kernel{{ty}, ty, std::move(exec)};
     kernel.mem_allocation = mem_allocation;
     ARROW_DCHECK_OK(func->AddKernel(std::move(kernel)));
@@ -216,6 +223,9 @@ static inline FunctionDoc StringClassifyDoc(std::string class_summary,
 
 template <typename Type, typename Predicate>
 struct StringPredicateFunctor {
+  static_assert(!is_string_or_string_view(Type::type_id),
+                "should only codegen on physical types");
+
   static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
     Status st = Status::OK();
     EnsureUtf8LookupTablesFilled();
@@ -237,7 +247,7 @@ void AddUnaryStringPredicate(std::string name, FunctionRegistry* registry,
                              FunctionDoc doc) {
   auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));
   for (const auto& ty : StringTypes()) {
-    auto exec = GenerateVarBinaryToVarBinary<StringPredicateFunctor, Predicate>(ty);
+    auto exec = GenerateTypeAgnosticVarBinaryBase<StringPredicateFunctor, Predicate>(ty);
     ARROW_DCHECK_OK(func->AddKernel({ty}, boolean(), std::move(exec)));
   }
   ARROW_DCHECK_OK(registry->AddFunction(std::move(func)));
@@ -281,7 +291,7 @@ struct ReplaceStringSliceTransformBase : public StringTransformBase {
 template <typename Options>
 struct StringSplitFinderBase {
   virtual ~StringSplitFinderBase() = default;
-  virtual Status PreExec(const Options& options) { return Status::OK(); }
+  virtual Status PreExec(const Options& options, bool is_utf8) { return Status::OK(); }
 
   // Derived classes should also define these methods:
   //   static bool Find(const uint8_t* begin, const uint8_t* end,
@@ -319,8 +329,9 @@ struct StringSplitExec {
   }
 
   Status Execute(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
+    const bool is_utf8 = is_string_or_string_view(batch[0].type()->id());
     SplitFinder finder;
-    RETURN_NOT_OK(finder.PreExec(options));
+    RETURN_NOT_OK(finder.PreExec(options, is_utf8));
     // TODO(wesm): refactor to not require creating ArrayData
     const ArrayType input(batch[0].array.ToArrayData());
 
@@ -347,9 +358,11 @@ struct StringSplitExec {
       *list_offsets++ = static_cast<list_offset_type>(builder.length());
     }
     // Assign string array to list child data
-    std::shared_ptr<Array> string_array;
-    RETURN_NOT_OK(builder.Finish(&string_array));
-    output_list->child_data.push_back(string_array->data());
+    ARROW_ASSIGN_OR_RAISE(auto physical_array, builder.Finish());
+    // We got the physical type (e.g. binary instead of utf8), need to patch it
+    auto child_data = physical_array->data()->Copy();
+    child_data->type = batch[0].type()->GetSharedPtr();
+    output_list->child_data.push_back(std::move(child_data));
     return Status::OK();
   }
 

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

@@ -336,6 +336,44 @@ TYPED_TEST(TestBinaryKernels, NonUtf8Regex) {
                          {"letter": [1, 1], "digit": [2, 1]}])",
                      &options);
   }
+
+  // On non-UTF8 input types, case insensitive matches should only account for
+  // ASCII letters.
+  // (this stresses that the "is_utf8" option is passed down correctly to RE2)
+  auto letters_array = this->MakeArray({"1ab2AB3", "1àb2ÀB3"});
+  {
+    MatchSubstringOptions options("(?i)(A|À)B");
+    // 5 is the byte index of "ÀB" in "1àb2ÀB3"
+    this->CheckUnary("find_substring_regex", letters_array, this->offset_type(), "[1, 5]",
+                     &options);
+    this->CheckUnary("count_substring_regex", letters_array, this->offset_type(),
+                     "[2, 1]", &options);
+    this->CheckUnary("match_substring_regex", letters_array, boolean(), "[true, true]",
+                     &options);
+  }
+  {
+    SplitPatternOptions options("(?i)(A|À)B");
+    this->CheckUnary("split_pattern_regex", letters_array, list(this->type()),
+                     R"([["1", "2", "3"], ["1àb2", "3"]])", &options);
+  }
+  {
+    ReplaceSubstringOptions options("(?i)(A|À)B", "XY");
+    this->CheckUnary("replace_substring_regex", letters_array,
+                     this->MakeArray({"1XY2XY3", "1àb2XY3"}), &options);
+  }
+  {
+    ExtractRegexOptions options("(?P<letters>(?i:(?:A|À)B))");
+    auto out_type = struct_({{"letters", this->type()}});
+    this->CheckUnary("extract_regex", letters_array, out_type,
+                     R"([{"letters": "ab"}, {"letters": "ÀB"}])", &options);
+  }
+  {
+    ExtractRegexSpanOptions options("(?P<letters>(?i:(?:A|À)B))");
+    auto out_type = struct_({{"letters", fixed_size_list(this->offset_type(), 2)}});
+    // 5 is the byte index of "ÀB" in "1àb2ÀB3", 3 is its byte length
+    this->CheckUnary("extract_regex_span", letters_array, out_type,
+                     R"([{"letters": [1, 2]}, {"letters": [5, 3]}])", &options);
+  }
 }
 
 TYPED_TEST(TestBinaryKernels, NonUtf8WithNullRegex) {
@@ -1041,6 +1079,48 @@ TEST_F(TestFixedSizeBinaryKernels, FindSubstringIgnoreCase) {
 }
 #endif
 
+#ifdef ARROW_WITH_RE2
+TYPED_TEST(TestStringKernels, Utf8Regex) {
+  // On UTF8 input types, case insensitive matches should account for all letters.
+  // (this stresses that the "is_utf8" option is passed down correctly to RE2)
+  auto letters_array = this->MakeArray({"🙂ab2AB3", "🙂àb2ÀB3"});
+  {
+    MatchSubstringOptions options("(?i)(A|À)B");
+    // 4 is the byte index of "ab" in "🙂ab2AB3"
+    this->CheckUnary("find_substring_regex", letters_array, this->offset_type(), "[4, 4]",
+                     &options);
+    this->CheckUnary("count_substring_regex", letters_array, this->offset_type(),
+                     "[2, 2]", &options);
+    this->CheckUnary("match_substring_regex", letters_array, boolean(), "[true, true]",
+                     &options);
+  }
+  {
+    SplitPatternOptions options("(?i)(A|À)B");
+    this->CheckUnary("split_pattern_regex", letters_array, list(this->type()),
+                     R"([["🙂", "2", "3"], ["🙂", "2", "3"]])", &options);
+  }
+  {
+    ReplaceSubstringOptions options("(?i)(A|À)B", "XY");
+    this->CheckUnary("replace_substring_regex", letters_array,
+                     this->MakeArray({"🙂XY2XY3", "🙂XY2XY3"}), &options);
+  }
+  {
+    ExtractRegexOptions options("(?P<letters>(?i:(?:A|À)B))");
+    auto out_type = struct_({{"letters", this->type()}});
+    this->CheckUnary("extract_regex", letters_array, out_type,
+                     R"([{"letters": "ab"}, {"letters": "àb"}])", &options);
+  }
+  {
+    ExtractRegexSpanOptions options("(?P<letters>(?i:(?:A|À)B))");
+    auto out_type = struct_({{"letters", fixed_size_list(this->offset_type(), 2)}});
+    // 4 is the byte index of "ab" in "🙂ab2AB3"
+    // 3 is the byte length of "àb"
+    this->CheckUnary("extract_regex_span", letters_array, out_type,
+                     R"([{"letters": [4, 2]}, {"letters": [4, 3]}])", &options);
+  }
+}
+#endif
+
 TYPED_TEST(TestStringKernels, AsciiUpper) {
   this->CheckUnary("ascii_upper", "[]", this->type(), "[]");
   this->CheckUnary("ascii_upper", "[\"aAazZæÆ&\", null, \"\", \"bbb\"]", this->type(),

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

@@ -45,7 +45,7 @@ void MakeUnaryStringUTF8TransformKernel(std::string name, FunctionRegistry* regi
                                         FunctionDoc doc) {
   auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));
   for (const auto& ty : StringTypes()) {
-    auto exec = GenerateVarBinaryToVarBinary<Transformer>(ty);
+    auto exec = GenerateTypeAgnosticVarBinaryBase<Transformer>(ty);
     DCHECK_OK(func->AddKernel({ty}, ty, std::move(exec)));
   }
   DCHECK_OK(registry->AddFunction(std::move(func)));
@@ -1144,7 +1144,7 @@ void AddUtf8StringReplaceSlice(FunctionRegistry* registry) {
                                                utf8_replace_slice_doc);
 
   for (const auto& ty : StringTypes()) {
-    auto exec = GenerateVarBinaryToVarBinary<Utf8ReplaceSlice>(ty);
+    auto exec = GenerateTypeAgnosticVarBinaryBase<Utf8ReplaceSlice>(ty);
     DCHECK_OK(func->AddKernel({ty}, ty, std::move(exec),
                               ReplaceStringSliceTransformBase::State::Init));
   }
@@ -1340,7 +1340,7 @@ void AddUtf8StringSlice(FunctionRegistry* registry) {
   auto func = std::make_shared<ScalarFunction>("utf8_slice_codeunits", Arity::Unary(),
                                                utf8_slice_codeunits_doc);
   for (const auto& ty : StringTypes()) {
-    auto exec = GenerateVarBinaryToVarBinary<SliceCodeunits>(ty);
+    auto exec = GenerateTypeAgnosticVarBinaryBase<SliceCodeunits>(ty);
     DCHECK_OK(
         func->AddKernel({ty}, ty, std::move(exec), SliceCodeunitsTransform::State::Init));
   }
@@ -1355,7 +1355,8 @@ void AddUtf8StringSlice(FunctionRegistry* registry) {
 struct SplitWhitespaceUtf8Finder : public StringSplitFinderBase<SplitOptions> {
   using Options = SplitOptions;
 
-  Status PreExec(const SplitOptions& options) override {
+  Status PreExec(const SplitOptions& options, bool is_utf8) override {
+    DCHECK(is_utf8);
     EnsureUtf8LookupTablesFilled();
     return Status::OK();
   }
@@ -1426,7 +1427,7 @@ void AddUtf8StringSplitWhitespace(FunctionRegistry* registry) {
       std::make_shared<ScalarFunction>("utf8_split_whitespace", Arity::Unary(),
                                        utf8_split_whitespace_doc, &default_options);
   for (const auto& ty : StringTypes()) {
-    auto exec = GenerateVarBinaryToVarBinary<SplitWhitespaceUtf8Exec, ListType>(ty);
+    auto exec = GenerateTypeAgnosticVarBinaryBase<SplitWhitespaceUtf8Exec, ListType>(ty);
     DCHECK_OK(func->AddKernel({ty}, {list(ty)}, std::move(exec), StringSplitState::Init));
   }
   DCHECK_OK(registry->AddFunction(std::move(func)));

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

@@ -586,9 +586,8 @@ void AddArraySortingKernels(VectorKernel base, VectorFunction* func) {
     DCHECK_OK(func->AddKernel(base));
   }
   for (const auto& ty : BaseBinaryTypes()) {
-    auto physical_type = GetPhysicalType(ty);
     base.signature = KernelSignature::Make({ty}, uint64());
-    base.exec = GenerateVarBinaryBase<ExecTemplate, UInt64Type>(*physical_type);
+    base.exec = GenerateVarBinaryBase<ExecTemplate, UInt64Type>(ty);
     DCHECK_OK(func->AddKernel(base));
   }
   base.signature = KernelSignature::Make({Type::FIXED_SIZE_BINARY}, uint64());

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

@@ -857,11 +857,10 @@ void RegisterVectorFunction(FunctionRegistry* registry,
             Functor<FixedSizeBinaryType>::Exec, ChunkedFunctor<FixedSizeBinaryType>::Exec,
             registry, func.get());
   for (const auto& ty : BaseBinaryTypes()) {
-    AddKernel(
-        ty->id(), Functor<FixedSizeBinaryType>::GetSignature(ty->id()),
-        GenerateTypeAgnosticVarBinaryBase<Functor, ArrayKernelExec>(*ty),
-        GenerateTypeAgnosticVarBinaryBase<ChunkedFunctor, VectorKernel::ChunkedExec>(*ty),
-        registry, func.get());
+    AddKernel(ty->id(), Functor<FixedSizeBinaryType>::GetSignature(ty->id()),
+              GenerateTypeAgnosticVarBinaryBase<Functor>(*ty),
+              GenerateTypeAgnosticVarBinaryBase<ChunkedFunctor>(*ty), registry,
+              func.get());
   }
   // TODO: list types
   DCHECK_OK(registry->AddFunction(std::move(func)));