ARROW-9160: [C++] Implement contains for exact matches

cpp/src/arrow/compute/api_scalar.h

@@ -259,6 +259,15 @@ Result<Datum> IsValid(const Datum& values, ExecContext* ctx = NULLPTR);
 ARROW_EXPORT
 Result<Datum> IsNull(const Datum& values, ExecContext* ctx = NULLPTR);
 
+// ----------------------------------------------------------------------
+// String functions
+
+struct ARROW_EXPORT BinaryContainsExactOptions : public FunctionOptions {
+  explicit BinaryContainsExactOptions(std::string pattern) : pattern(pattern) {}
+
+  std::string pattern;
+};
+
 // ----------------------------------------------------------------------
 // Temporal functions
 

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

@@ -297,6 +297,118 @@ void AddAsciiLength(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunction(std::move(func)));
 }
 
+// ----------------------------------------------------------------------
+// exact pattern detection
+
+using StrToBoolTransformFunc =
+    std::function<void(const void*, const uint8_t*, int64_t, int64_t, uint8_t*)>;
+
+// Apply `transform` to input character data- this function cannot change the
+// length
+template <typename Type>
+void StringBoolTransform(KernelContext* ctx, const ExecBatch& batch,
+                         StrToBoolTransformFunc transform, Datum* out) {
+  using offset_type = typename Type::offset_type;
+
+  if (batch[0].kind() == Datum::ARRAY) {
+    const ArrayData& input = *batch[0].array();
+    ArrayData* out_arr = out->mutable_array();
+    if (input.length > 0) {
+      transform(
+          reinterpret_cast<const offset_type*>(input.buffers[1]->data()) + input.offset,
+          input.buffers[2]->data(), input.length, out_arr->offset,
+          out_arr->buffers[1]->mutable_data());
+    }
+  } else {
+    const auto& input = checked_cast<const BaseBinaryScalar&>(*batch[0].scalar());
+    if (input.is_valid) {
+      auto result = checked_pointer_cast<BooleanScalar>(MakeNullScalar(out->type()));
+      uint8_t result_value = 0;
+      result->is_valid = true;
+      std::array<offset_type, 2> offsets{0,
+                                         static_cast<offset_type>(input.value->size())};
+      transform(offsets.data(), input.value->data(), 1, /*output_offset=*/0,
+                &result_value);
+      out->value = std::make_shared<BooleanScalar>(result_value > 0);
+    }
+  }
+}
+
+template <typename offset_type>
+void TransformBinaryContainsExact(const uint8_t* pattern, int64_t pattern_length,
+                                  const offset_type* offsets, const uint8_t* data,
+                                  int64_t length, int64_t output_offset,
+                                  uint8_t* output) {
+  // This is an implementation of the Knuth-Morris-Pratt algorithm
+
+  // Phase 1: Build the prefix table
+  std::vector<offset_type> prefix_table(pattern_length + 1);
+  offset_type prefix_length = -1;
+  prefix_table[0] = -1;
+  for (offset_type pos = 0; pos < pattern_length; ++pos) {
+    // The prefix cannot be expanded, reset.
+    if (prefix_length >= 0 && pattern[pos] != pattern[prefix_length]) {
+      prefix_length = prefix_table[prefix_length];
+    }
+    prefix_length++;
+    prefix_table[pos + 1] = prefix_length;
+  }
+
+  // Phase 2: Find the prefix in the data
+  FirstTimeBitmapWriter bitmap_writer(output, output_offset, length);
+  for (int64_t i = 0; i < length; ++i) {
+    const uint8_t* current_data = data + offsets[i];
+    int64_t current_length = offsets[i + 1] - offsets[i];
+
+    int64_t pattern_pos = 0;
+    for (int64_t k = 0; k < current_length; k++) {
+      if (pattern[pattern_pos] == current_data[k]) {
+        pattern_pos++;
+        if (pattern_pos == pattern_length) {
+          bitmap_writer.Set();
+          break;
+        }
+      } else {
+        pattern_pos = std::max<offset_type>(0, prefix_table[pattern_pos]);
+      }
+    }
+    bitmap_writer.Next();
+  }
+  bitmap_writer.Finish();
+}
+
+using BinaryContainsExactState = OptionsWrapper<BinaryContainsExactOptions>;
+
+template <typename Type>
+struct BinaryContainsExact {
+  using offset_type = typename Type::offset_type;
+  static void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+    BinaryContainsExactOptions arg = BinaryContainsExactState::Get(ctx);
+    const uint8_t* pat = reinterpret_cast<const uint8_t*>(arg.pattern.c_str());
+    const int64_t pat_size = arg.pattern.length();
+    StringBoolTransform<Type>(
+        ctx, batch,
+        [pat, pat_size](const void* offsets, const uint8_t* data, int64_t length,
+                        int64_t output_offset, uint8_t* output) {
+          TransformBinaryContainsExact<offset_type>(
+              pat, pat_size, reinterpret_cast<const offset_type*>(offsets), data, length,
+              output_offset, output);
+        },
+        out);
+  }
+};
+
+void AddBinaryContainsExact(FunctionRegistry* registry) {
+  auto func = std::make_shared<ScalarFunction>("binary_contains_exact", Arity::Unary());
+  auto exec_32 = BinaryContainsExact<StringType>::Exec;
+  auto exec_64 = BinaryContainsExact<LargeStringType>::Exec;
+  DCHECK_OK(
+      func->AddKernel({utf8()}, boolean(), exec_32, BinaryContainsExactState::Init));
+  DCHECK_OK(func->AddKernel({large_utf8()}, boolean(), exec_64,
+                            BinaryContainsExactState::Init));
+  DCHECK_OK(registry->AddFunction(std::move(func)));
+}
+
 // ----------------------------------------------------------------------
 // strptime string parsing
 
@@ -377,6 +489,7 @@ void RegisterScalarStringAscii(FunctionRegistry* registry) {
   MakeUnaryStringUtf8TransformKernel<Utf8Lower>("utf8_lower", registry);
 #endif
   AddAsciiLength(registry);
+  AddBinaryContainsExact(registry);
   AddStrptime(registry);
 }
 

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

@@ -28,7 +28,8 @@ namespace compute {
 
 constexpr auto kSeed = 0x94378165;
 
-static void UnaryStringBenchmark(benchmark::State& state, const std::string& func_name) {
+static void UnaryStringBenchmark(benchmark::State& state, const std::string& func_name,
+                                 const FunctionOptions* options = nullptr) {
   const int64_t array_length = 1 << 20;
   const int64_t value_min_size = 0;
   const int64_t value_max_size = 32;
@@ -39,10 +40,10 @@ static void UnaryStringBenchmark(benchmark::State& state, const std::string& fun
   auto values =
       rng.String(array_length, value_min_size, value_max_size, null_probability);
   // Make sure lookup tables are initialized before measuring
-  ABORT_NOT_OK(CallFunction(func_name, {values}));
+  ABORT_NOT_OK(CallFunction(func_name, {values}, options));
 
   for (auto _ : state) {
-    ABORT_NOT_OK(CallFunction(func_name, {values}));
+    ABORT_NOT_OK(CallFunction(func_name, {values}, options));
   }
   state.SetItemsProcessed(state.iterations() * array_length);
   state.SetBytesProcessed(state.iterations() * values->data()->buffers[2]->size());
@@ -56,6 +57,11 @@ static void AsciiUpper(benchmark::State& state) {
   UnaryStringBenchmark(state, "ascii_upper");
 }
 
+static void BinaryContainsExact(benchmark::State& state) {
+  BinaryContainsExactOptions options("abac");
+  UnaryStringBenchmark(state, "binary_contains_exact", &options);
+}
+
 #ifdef ARROW_WITH_UTF8PROC
 static void Utf8Upper(benchmark::State& state) {
   UnaryStringBenchmark(state, "utf8_upper");
@@ -68,6 +74,7 @@ static void Utf8Lower(benchmark::State& state) {
 
 BENCHMARK(AsciiLower);
 BENCHMARK(AsciiUpper);
+BENCHMARK(BinaryContainsExact);
 #ifdef ARROW_WITH_UTF8PROC
 BENCHMARK(Utf8Lower);
 BENCHMARK(Utf8Upper);

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

@@ -147,6 +147,17 @@ TYPED_TEST(TestStringKernels, Utf8Lower) {
 
 #endif  // ARROW_WITH_UTF8PROC
 
+TYPED_TEST(TestStringKernels, BinaryContainsExact) {
+  BinaryContainsExactOptions options{"ab"};
+  this->CheckUnary("binary_contains_exact", "[]", boolean(), "[]", &options);
+  this->CheckUnary("binary_contains_exact", R"(["abc", "acb", "cab", null, "bac"])",
+                   boolean(), "[true, false, true, null, false]", &options);
+
+  BinaryContainsExactOptions options_repeated{"abab"};
+  this->CheckUnary("binary_contains_exact", R"(["abab", "ab", "cababc", null, "bac"])",
+                   boolean(), "[true, false, true, null, false]", &options_repeated);
+}
+
 TYPED_TEST(TestStringKernels, Strptime) {
   std::string input1 = R"(["5/1/2020", null, "12/11/1900"])";
   std::string output1 = R"(["2020-05-01", null, "1900-12-11"])";

python/pyarrow/_compute.pyx

@@ -398,6 +398,18 @@ cdef class CastOptions(FunctionOptions):
         self.options.allow_invalid_utf8 = flag
 
 
+cdef class BinaryContainsExactOptions(FunctionOptions):
+    cdef:
+        unique_ptr[CBinaryContainsExactOptions] binary_contains_exact_options
+
+    def __init__(self, pattern):
+        self.binary_contains_exact_options.reset(
+            new CBinaryContainsExactOptions(tobytes(pattern)))
+
+    cdef const CFunctionOptions* get_options(self) except NULL:
+        return self.binary_contains_exact_options.get()
+
+
 cdef class FilterOptions(FunctionOptions):
     cdef:
         CFilterOptions filter_options

python/pyarrow/compute.py

@@ -113,6 +113,24 @@ def func(left, right):
 multiply = _simple_binary_function('multiply')
 
 
+def binary_contains_exact(array, pattern):
+    """
+    Test if pattern is contained within a value of a binary array.
+
+    Parameters
+    ----------
+    array : pyarrow.Array or pyarrow.ChunkedArray
+    pattern : str
+        pattern to search for exact matches
+
+    Returns
+    -------
+    result : pyarrow.Array or pyarrow.ChunkedArray
+    """
+    return call_function("binary_contains_exact", [array],
+                         _pc.BinaryContainsExactOptions(pattern))
+
+
 def sum(array):
     """
     Sum the values in a numerical (chunked) array.

python/pyarrow/includes/libarrow.pxd

@@ -1566,6 +1566,11 @@ cdef extern from "arrow/compute/api.h" namespace "arrow::compute" nogil:
 
     CFunctionRegistry* GetFunctionRegistry()
 
+    cdef cppclass CBinaryContainsExactOptions \
+            "arrow::compute::BinaryContainsExactOptions"(CFunctionOptions):
+        CBinaryContainsExactOptions(c_string pattern)
+        c_string pattern
+
     cdef cppclass CCastOptions" arrow::compute::CastOptions"(CFunctionOptions):
         CCastOptions()
         CCastOptions(c_bool safe)

python/pyarrow/tests/test_compute.py

@@ -90,6 +90,13 @@ def test_sum_chunked_array(arrow_type):
     assert pc.sum(arr) == None  # noqa: E711
 
 
+def test_binary_contains_exact():
+    arr = pa.array(["ab", "abc", "ba", None])
+    result = pc.binary_contains_exact(arr, "ab")
+    expected = pa.array([True, True, False, None])
+    assert expected.equals(result)
+
+
 @pytest.mark.parametrize(('ty', 'values'), all_array_types)
 def test_take(ty, values):
     arr = pa.array(values, type=ty)