GH-35815: [C++][Compute] Support dictionary inputs in UTF-8 trim

cpp/src/arrow/compute/function.cc

@@ -21,7 +21,10 @@
 #include <memory>
 #include <sstream>
 
+#include "arrow/array/array_dict.h"
+#include "arrow/array/util.h"
 #include "arrow/compute/api_scalar.h"
+#include "arrow/compute/api_vector.h"
 #include "arrow/compute/cast.h"
 #include "arrow/compute/exec.h"
 #include "arrow/compute/exec_internal.h"
@@ -293,8 +296,142 @@ struct FunctionExecutorImpl : public FunctionExecutor {
   bool inited;
 };
 
+class DictionaryUnaryFunctionExecutor : public FunctionExecutor {
+ public:
+  DictionaryUnaryFunctionExecutor(TypeHolder input_type,
+                                  std::shared_ptr<FunctionExecutor> value_executor,
+                                  const ScalarFunction& function)
+      : input_type_(std::move(input_type)),
+        value_type_(checked_cast<const DictionaryType&>(*input_type_.GetSharedPtr())
+                        .value_type()),
+        value_executor_(std::move(value_executor)),
+        function_(function) {}
+
+  Status Init(const FunctionOptions* options, ExecContext* exec_ctx) override {
+    if (exec_ctx == NULLPTR) {
+      exec_ctx = default_exec_context();
+    }
+    RETURN_NOT_OK(value_executor_->Init(options, exec_ctx));
+    exec_ctx_ = exec_ctx;
+    inited_ = true;
+    return Status::OK();
+  }
+
+  Result<Datum> Execute(const std::vector<Datum>& args, int64_t passed_length) override {
+    if (args.size() != 1) {
+      return Status::Invalid("Execution of '", function_.name(),
+                             "' expected 1 argument but got ", args.size());
+    }
+    if (!inited_) {
+      ARROW_RETURN_NOT_OK(Init(NULLPTR, default_exec_context()));
+    }
+
+    Datum arg = args[0];
+    if (input_type_ != arg.type()) {
+      ARROW_ASSIGN_OR_RAISE(arg, Cast(arg, CastOptions::Safe(input_type_), exec_ctx_));
+    }
+    if (arg.type()->id() != Type::DICTIONARY) {
+      return Status::Invalid("Dictionary-unary executor for '", function_.name(),
+                             "' received input type ", arg.type()->ToString());
+    }
+    if (passed_length != -1) {
+      ARROW_ASSIGN_OR_RAISE(auto inferred_length,
+                            ExecBatch::InferLength(std::vector<Datum>{arg}));
+      if (passed_length != inferred_length) {
+        return Status::Invalid(
+            "Passed batch length for execution did not match actual"
+            " length of values for execution of scalar function '",
+            function_.name(), "'");
+      }
+    }
+    return ExecuteDictionary(arg);
+  }
+
+ private:
+  Result<Datum> ExecuteDictionaryArray(const Datum& arg) {
+    auto input_array = arg.make_array();
+    const auto& input = checked_cast<const DictionaryArray&>(*input_array);
+
+    ARROW_ASSIGN_OR_RAISE(
+        Datum encoded_indices,
+        DictionaryEncode(input.indices(), DictionaryEncodeOptions::Defaults(),
+                         exec_ctx_));
+    auto encoded_indices_array = encoded_indices.make_array();
+    const auto& referenced_indices =
+        checked_cast<const DictionaryArray&>(*encoded_indices_array);
+
+    ARROW_ASSIGN_OR_RAISE(Datum referenced_values,
+                          Take(input.dictionary(), referenced_indices.dictionary(),
+                               TakeOptions::Defaults(), exec_ctx_));
+    ARROW_ASSIGN_OR_RAISE(Datum transformed_values,
+                          value_executor_->Execute({std::move(referenced_values)}));
+    if (!transformed_values.is_array()) {
+      return Status::Invalid("Unary scalar function '", function_.name(),
+                             "' returned a non-array result for dictionary values");
+    }
+
+    return Take(transformed_values, referenced_indices.indices(), TakeOptions::Defaults(),
+                exec_ctx_);
+  }
+
+  Result<Datum> ExecuteDictionary(const Datum& arg) {
+    switch (arg.kind()) {
+      case Datum::ARRAY:
+        return ExecuteDictionaryArray(arg);
+      case Datum::SCALAR: {
+        const auto& input = checked_cast<const DictionaryScalar&>(*arg.scalar());
+        ARROW_ASSIGN_OR_RAISE(auto value, input.GetEncodedValue());
+        return value_executor_->Execute({std::move(value)});
+      }
+      case Datum::CHUNKED_ARRAY: {
+        ArrayVector output_chunks;
+        output_chunks.reserve(arg.chunked_array()->num_chunks());
+        std::shared_ptr<DataType> output_type;
+        for (const auto& chunk : arg.chunked_array()->chunks()) {
+          ARROW_ASSIGN_OR_RAISE(Datum output, ExecuteDictionaryArray(Datum(chunk)));
+          DCHECK(output.is_array());
+          output_type = output.type();
+          output_chunks.push_back(output.make_array());
+        }
+
+        if (output_type == nullptr) {
+          ARROW_ASSIGN_OR_RAISE(auto empty_values, MakeEmptyArray(value_type_));
+          ARROW_ASSIGN_OR_RAISE(Datum output,
+                                value_executor_->Execute({std::move(empty_values)}));
+          if (!output.is_array()) {
+            return Status::Invalid("Unary scalar function '", function_.name(),
+                                   "' returned a non-array result for dictionary values");
+          }
+          output_type = output.type();
+        }
+        return ChunkedArray::Make(std::move(output_chunks), std::move(output_type));
+      }
+      default:
+        return Status::Invalid("Unsupported dictionary datum kind");
+    }
+  }
+
+  TypeHolder input_type_;
+  std::shared_ptr<DataType> value_type_;
+  std::shared_ptr<FunctionExecutor> value_executor_;
+  const ScalarFunction& function_;
+  ExecContext* exec_ctx_ = NULLPTR;
+  bool inited_ = false;
+};
+
 }  // namespace detail
 
+namespace {
+
+bool CanExecuteDictionaryValues(const ScalarFunction& function,
+                                const std::vector<TypeHolder>& inputs) {
+  return function.is_pure() && !function.arity().is_varargs &&
+         function.arity().num_args == 1 && inputs.size() == 1 &&
+         inputs[0].id() == Type::DICTIONARY;
+}
+
+}  // namespace
+
 Result<const Kernel*> Function::DispatchExact(
     const std::vector<TypeHolder>& values) const {
   if (kind_ == Function::META) {
@@ -326,7 +463,26 @@ Result<std::shared_ptr<FunctionExecutor>> Function::GetBestExecutor(
     return Status::NotImplemented("Direct execution of HASH_AGGREGATE functions");
   }
 
-  ARROW_ASSIGN_OR_RAISE(const Kernel* kernel, DispatchBest(&inputs));
+  auto dispatched = DispatchBest(&inputs);
+  if (!dispatched.ok()) {
+    if (!dispatched.status().IsNotImplemented() || kind() != Function::SCALAR) {
+      return dispatched.status();
+    }
+
+    const auto& scalar_function = checked_cast<const ScalarFunction&>(*this);
+    if (!CanExecuteDictionaryValues(scalar_function, inputs)) {
+      return dispatched.status();
+    }
+
+    const auto& dictionary_type =
+        checked_cast<const DictionaryType&>(*inputs[0].GetSharedPtr());
+    std::vector<TypeHolder> value_inputs = {dictionary_type.value_type()};
+    ARROW_ASSIGN_OR_RAISE(auto value_executor,
+                          Function::GetBestExecutor(std::move(value_inputs)));
+    return std::make_shared<detail::DictionaryUnaryFunctionExecutor>(
+        std::move(inputs[0]), std::move(value_executor), scalar_function);
+  }
+  const Kernel* kernel = *dispatched;
 
   return std::make_shared<detail::FunctionExecutorImpl>(std::move(inputs), kernel,
                                                         std::move(executor), *this);

cpp/src/arrow/compute/function_test.cc

@@ -23,7 +23,9 @@
 #include <string>
 #include <vector>
 
+#include "arrow/array/array_dict.h"
 #include "arrow/array/builder_primitive.h"
+#include "arrow/array/util.h"
 #include "arrow/compute/api_aggregate.h"
 #include "arrow/compute/api_scalar.h"
 #include "arrow/compute/api_vector.h"
@@ -296,6 +298,108 @@ TEST(ScalarVectorFunction, DispatchExact) {
   CheckAddDispatch(&func2, ExecNYI);
 }
 
+namespace {
+
+struct DictionaryValuesCounter : KernelState {
+  int64_t values_processed = 0;
+};
+
+Status CountAndCastDictionaryValues(KernelContext* ctx, const ExecSpan& args,
+                                    ExecResult* out) {
+  auto& counter = checked_cast<DictionaryValuesCounter&>(*ctx->kernel()->data);
+  counter.values_processed += args.length;
+  ARROW_ASSIGN_OR_RAISE(Datum result, Cast(args[0].array.ToArrayData(), int64(),
+                                           CastOptions::Safe(), ctx->exec_context()));
+  out->value = result.array();
+  return Status::OK();
+}
+
+Status ExecuteDirectDictionaryKernel(KernelContext* ctx, const ExecSpan& args,
+                                     ExecResult* out) {
+  ARROW_ASSIGN_OR_RAISE(
+      auto result, MakeArrayFromScalar(Int64Scalar(42), args.length, ctx->memory_pool()));
+  out->value = result->data();
+  return Status::OK();
+}
+
+}  // namespace
+
+TEST(ScalarFunction, DictionaryUnaryAppliesToReferencedDictionaryValues) {
+  ScalarFunction func("dictionary_unary_test", Arity::Unary(), FunctionDoc::Empty());
+  auto counter = std::make_shared<DictionaryValuesCounter>();
+  ScalarKernel kernel({int32()}, int64(), CountAndCastDictionaryValues);
+  kernel.data = counter;
+  kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE;
+  kernel.mem_allocation = MemAllocation::NO_PREALLOCATE;
+  ASSERT_OK(func.AddKernel(std::move(kernel)));
+
+  ASSERT_OK_AND_ASSIGN(
+      auto input, DictionaryArray::FromArrays(ArrayFromJSON(int8(), "[0, 1, 0, null, 1]"),
+                                              ArrayFromJSON(int32(), "[10, 20, 999]")));
+  ASSERT_OK_AND_ASSIGN(Datum result, func.Execute({input}, nullptr, nullptr));
+
+  auto expected = ArrayFromJSON(int64(), "[10, 20, 10, null, 20]");
+  ASSERT_TRUE(result.is_array());
+  AssertArraysEqual(*expected, *result.make_array());
+  ASSERT_EQ(counter->values_processed, 2);
+
+  ASSERT_OK_AND_ASSIGN(auto scalar_input, input->GetScalar(1));
+  ASSERT_OK_AND_ASSIGN(Datum scalar_result,
+                       func.Execute({scalar_input}, nullptr, nullptr));
+  ASSERT_TRUE(scalar_result.is_scalar());
+  AssertScalarsEqual(Int64Scalar(20), *scalar_result.scalar());
+
+  auto chunked_input =
+      std::make_shared<ChunkedArray>(ArrayVector{input->Slice(0, 2), input->Slice(2)});
+  ASSERT_OK_AND_ASSIGN(Datum chunked_result,
+                       func.Execute({chunked_input}, nullptr, nullptr));
+  ASSERT_TRUE(chunked_result.is_chunked_array());
+  AssertChunkedEqual(*chunked_result.chunked_array(),
+                     ArrayVector{expected->Slice(0, 2), expected->Slice(2)});
+
+  auto empty_chunked_input = std::make_shared<ChunkedArray>(ArrayVector{}, input->type());
+  ASSERT_OK_AND_ASSIGN(Datum empty_chunked_result,
+                       func.Execute({empty_chunked_input}, nullptr, nullptr));
+  ASSERT_TRUE(empty_chunked_result.is_chunked_array());
+  ASSERT_EQ(empty_chunked_result.length(), 0);
+  ASSERT_TRUE(empty_chunked_result.type()->Equals(*int64()));
+
+  ASSERT_RAISES(Invalid,
+                func.Execute(ExecBatch({Datum(scalar_input)}, 2), nullptr, nullptr));
+
+  std::vector<TypeHolder> dictionary_types = {dictionary(int8(), int32())};
+  ASSERT_RAISES(NotImplemented, func.DispatchBest(&dictionary_types));
+  ASSERT_RAISES(NotImplemented, func.DispatchExact({dictionary(int8(), utf8())}));
+}
+
+TEST(ScalarFunction, DictionaryUnaryPrefersDirectDictionaryKernel) {
+  ScalarFunction func("dictionary_unary_direct_test", Arity::Unary(),
+                      FunctionDoc::Empty());
+  ASSERT_OK(func.AddKernel({int32()}, int64(), ExecNYI));
+
+  ScalarKernel direct_kernel({dictionary(int8(), int32())}, int64(),
+                             ExecuteDirectDictionaryKernel);
+  direct_kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE;
+  direct_kernel.mem_allocation = MemAllocation::NO_PREALLOCATE;
+  ASSERT_OK(func.AddKernel(std::move(direct_kernel)));
+
+  ASSERT_OK_AND_ASSIGN(auto input,
+                       DictionaryArray::FromArrays(ArrayFromJSON(int8(), "[0, 1, 0]"),
+                                                   ArrayFromJSON(int32(), "[10, 20]")));
+  ASSERT_OK_AND_ASSIGN(Datum result, func.Execute({input}, nullptr, nullptr));
+
+  ASSERT_TRUE(result.is_array());
+  AssertArraysEqual(*ArrayFromJSON(int64(), "[42, 42, 42]"), *result.make_array());
+}
+
+TEST(ScalarFunction, ImpureUnaryRejectsDictionaryInput) {
+  ScalarFunction func("impure_unary_test", Arity::Unary(), FunctionDoc::Empty(),
+                      /*default_options=*/nullptr, /*is_pure=*/false);
+  ASSERT_OK(func.AddKernel({int32()}, int32(), ExecNYI));
+
+  ASSERT_RAISES(NotImplemented, func.DispatchExact({dictionary(int8(), int32())}));
+}
+
 TEST(ArrayFunction, VarArgs) {
   ScalarFunction va_func("va_test", Arity::VarArgs(1), /*doc=*/FunctionDoc::Empty());
 

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

@@ -23,6 +23,7 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
+#include "arrow/array/array_dict.h"
 #include "arrow/compute/api_scalar.h"
 #include "arrow/compute/exec.h"
 #include "arrow/compute/kernels/codegen_internal.h"
@@ -2136,6 +2137,16 @@ TYPED_TEST(TestStringKernels, Strptime) {
       Strptime(ArrayFromJSON(this->type(), input1), options));
 }
 
+TYPED_TEST(TestStringKernels, StrptimeDictionaryIgnoresUnreferencedValues) {
+  auto values = ArrayFromJSON(this->type(), R"(["5/1/2020", "not-a-date"])");
+  ASSERT_OK_AND_ASSIGN(auto input, DictionaryArray::FromArrays(
+                                       ArrayFromJSON(int8(), "[0, 0, null]"), values));
+  StrptimeOptions options("%m/%d/%Y", TimeUnit::MICRO, /*error_is_null=*/false);
+
+  this->CheckUnary("strptime", input, timestamp(TimeUnit::MICRO),
+                   R"(["2020-05-01", "2020-05-01", null])", &options);
+}
+
 TYPED_TEST(TestStringKernels, StrptimeZoneOffset) {
 #ifdef __EMSCRIPTEN__
   GTEST_SKIP()
@@ -2331,6 +2342,17 @@ TYPED_TEST(TestStringKernels, TrimUTF8) {
   EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, testing::HasSubstr("Invalid UTF8"),
                                   CallFunction("utf8_trim", {input}, &options_invalid));
 }
+
+TYPED_TEST(TestStringKernels, TrimUTF8Dictionary) {
+  auto input =
+      ArrayFromJSON(dictionary(int64(), this->type()), R"(["bcabc", "b", "a", null])");
+  auto options = TrimOptions{"bc"};
+  this->CheckUnary("utf8_trim", input, this->type(), R"(["a", "", "a", null])", &options);
+  this->CheckUnary("utf8_ltrim", input, this->type(), R"(["abc", "", "a", null])",
+                   &options);
+  this->CheckUnary("utf8_rtrim", input, this->type(), R"(["bca", "", "a", null])",
+                   &options);
+}
 #endif
 
 // produce test data with e.g.: