ARROW-9115: [C++] Implementation of ascii_lower/ascii_upper by processing input data buffers in batch

cpp/src/arrow/compute/kernels/CMakeLists.txt

@@ -29,6 +29,7 @@ add_arrow_compute_test(scalar_test
                        test_util.cc)
 
 add_arrow_benchmark(scalar_compare_benchmark PREFIX "arrow-compute")
+add_arrow_benchmark(scalar_string_benchmark PREFIX "arrow-compute")
 
 # ----------------------------------------------------------------------
 # Vector kernels

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

@@ -28,6 +28,8 @@ namespace arrow {
 namespace compute {
 namespace internal {
 
+namespace {
+
 // TODO: optional ascii validation
 
 struct AsciiLength {
@@ -37,26 +39,108 @@ struct AsciiLength {
   }
 };
 
-struct AsciiUpper {
-  // XXX: the Scalar codegen path passes template arguments that are unused
-  template <typename... Ignored>
-  static std::string Call(KernelContext*, const util::string_view& val) {
-    std::string result = val.to_string();
-    std::transform(result.begin(), result.end(), result.begin(),
-                   [](unsigned char c) { return std::toupper(c); });
-    return result;
+using TransformFunc = std::function<void(const uint8_t*, int64_t, uint8_t*)>;
+
+void StringDataTransform(KernelContext* ctx, const ExecBatch& batch,
+                         TransformFunc transform, Datum* out) {
+  if (batch[0].kind() == Datum::ARRAY) {
+    const ArrayData& input = *batch[0].array();
+    ArrayData* out_arr = out->mutable_array();
+    // Reuse offsets from input
+    out_arr->buffers[1] = input.buffers[1];
+    int64_t data_nbytes = input.buffers[2]->size();
+    KERNEL_RETURN_IF_ERROR(ctx, ctx->Allocate(data_nbytes).Value(&out_arr->buffers[2]));
+    transform(input.buffers[2]->data(), data_nbytes, out_arr->buffers[2]->mutable_data());
+  } else {
+    const auto& input = checked_cast<const BaseBinaryScalar&>(*batch[0].scalar());
+    auto result = checked_pointer_cast<BaseBinaryScalar>(MakeNullScalar(out->type()));
+    if (input.is_valid) {
+      result->is_valid = true;
+      int64_t data_nbytes = input.value->size();
+      KERNEL_RETURN_IF_ERROR(ctx, ctx->Allocate(data_nbytes).Value(&result->value));
+      transform(input.value->data(), data_nbytes, result->value->mutable_data());
+    }
+    out->value = result;
   }
-};
+}
 
-struct AsciiLower {
-  template <typename... Ignored>
-  static std::string Call(KernelContext*, const util::string_view& val) {
-    std::string result = val.to_string();
-    std::transform(result.begin(), result.end(), result.begin(),
-                   [](unsigned char c) { return std::tolower(c); });
-    return result;
+// Generated with
+//
+// print("static constexpr uint8_t kAsciiUpperTable[] = {")
+// for i in range(256):
+//     if i > 0: print(', ', end='')
+//     if i >= ord('a') and i <= ord('z'):
+//         print(i - 32, end='')
+//     else:
+//         print(i, end='')
+// print("};")
+
+static constexpr uint8_t kAsciiUpperTable[] = {
+    0,   1,   2,   3,   4,   5,   6,   7,   8,   9,   10,  11,  12,  13,  14,  15,
+    16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,  28,  29,  30,  31,
+    32,  33,  34,  35,  36,  37,  38,  39,  40,  41,  42,  43,  44,  45,  46,  47,
+    48,  49,  50,  51,  52,  53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,
+    64,  65,  66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,  78,  79,
+    80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,
+    96,  65,  66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,  78,  79,
+    80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,  123, 124, 125, 126, 127,
+    128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143,
+    144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159,
+    160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175,
+    176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191,
+    192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207,
+    208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223,
+    224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239,
+    240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255};
+
+void TransformAsciiUpper(const uint8_t* input, int64_t length, uint8_t* output) {
+  for (int64_t i = 0; i < length; ++i) {
+    *output++ = kAsciiUpperTable[*input++];
   }
-};
+}
+
+void AsciiUpperExec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+  StringDataTransform(ctx, batch, TransformAsciiUpper, out);
+}
+
+// Generated with
+//
+// print("static constexpr uint8_t kAsciiLowerTable[] = {")
+// for i in range(256):
+//     if i > 0: print(', ', end='')
+//     if i >= ord('A') and i <= ord('Z'):
+//         print(i + 32, end='')
+//     else:
+//         print(i, end='')
+// print("};")
+
+static constexpr uint8_t kAsciiLowerTable[] = {
+    0,   1,   2,   3,   4,   5,   6,   7,   8,   9,   10,  11,  12,  13,  14,  15,
+    16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,  28,  29,  30,  31,
+    32,  33,  34,  35,  36,  37,  38,  39,  40,  41,  42,  43,  44,  45,  46,  47,
+    48,  49,  50,  51,  52,  53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,
+    64,  97,  98,  99,  100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
+    112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 91,  92,  93,  94,  95,
+    96,  97,  98,  99,  100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
+    112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127,
+    128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143,
+    144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159,
+    160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175,
+    176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191,
+    192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207,
+    208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223,
+    224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239,
+    240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255};
+
+void TransformAsciiLower(const uint8_t* input, int64_t length, uint8_t* output) {
+  for (int64_t i = 0; i < length; ++i) {
+    *output++ = kAsciiLowerTable[*input++];
+  }
+}
+
+void AsciiLowerExec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+  StringDataTransform(ctx, batch, TransformAsciiLower, out);
+}
 
 void AddAsciiLength(FunctionRegistry* registry) {
   auto func = std::make_shared<ScalarFunction>("ascii_length", Arity::Unary());
@@ -114,11 +198,19 @@ void AddStrptime(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunction(std::move(func)));
 }
 
-// ----------------------------------------------------------------------
+void MakeUnaryStringBatchKernel(std::string name, ArrayKernelExec exec,
+                                FunctionRegistry* registry) {
+  auto func = std::make_shared<ScalarFunction>(name, Arity::Unary());
+  DCHECK_OK(func->AddKernel({utf8()}, utf8(), exec));
+  DCHECK_OK(func->AddKernel({large_utf8()}, large_utf8(), exec));
+  DCHECK_OK(registry->AddFunction(std::move(func)));
+}
+
+}  // namespace
 
 void RegisterScalarStringAscii(FunctionRegistry* registry) {
-  MakeUnaryStringToString<AsciiUpper>("ascii_upper", registry);
-  MakeUnaryStringToString<AsciiLower>("ascii_lower", registry);
+  MakeUnaryStringBatchKernel("ascii_upper", AsciiUpperExec, registry);
+  MakeUnaryStringBatchKernel("ascii_lower", AsciiLowerExec, registry);
   AddAsciiLength(registry);
   AddStrptime(registry);
 }

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

@@ -0,0 +1,59 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/api_scalar.h"
+#include "arrow/compute/benchmark_util.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/random.h"
+
+namespace arrow {
+namespace compute {
+
+constexpr auto kSeed = 0x94378165;
+
+static void UnaryStringBenchmark(benchmark::State& state, const std::string& func_name) {
+  const int64_t array_length = 1 << 20;
+  const int64_t value_min_size = 0;
+  const int64_t value_max_size = 32;
+  const double null_probability = 0.01;
+  random::RandomArrayGenerator rng(kSeed);
+
+  auto values =
+      rng.String(array_length, value_min_size, value_max_size, null_probability);
+  for (auto _ : state) {
+    ABORT_NOT_OK(CallFunction(func_name, {values}));
+  }
+  state.SetItemsProcessed(state.iterations() * array_length);
+  state.SetBytesProcessed(state.iterations() * values->data()->buffers[2]->size());
+}
+
+static void AsciiLower(benchmark::State& state) {
+  UnaryStringBenchmark(state, "ascii_lower");
+}
+
+static void AsciiUpper(benchmark::State& state) {
+  UnaryStringBenchmark(state, "ascii_upper");
+}
+
+BENCHMARK(AsciiLower);
+BENCHMARK(AsciiUpper);
+
+}  // namespace compute
+}  // namespace arrow