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

cpp/src/arrow/compute/function.cc

@@ -21,7 +21,9 @@
 #include <memory>
 #include <sstream>
 
+#include "arrow/array/array_dict.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"
@@ -295,6 +297,118 @@ struct FunctionExecutorImpl : public FunctionExecutor {
 
 }  // namespace detail
 
+namespace {
+
+struct DictionaryUnaryKernelState : KernelState {
+  DictionaryUnaryKernelState(const ScalarFunction* function,
+                             const FunctionOptions* options)
+      : function(function), options(options) {}
+
+  const ScalarFunction* function;
+  const FunctionOptions* options;
+};
+
+class DictionaryUnaryTypeMatcher : public TypeMatcher {
+ public:
+  explicit DictionaryUnaryTypeMatcher(const ScalarFunction* function)
+      : function_(function) {}
+
+  bool Matches(const DataType& type) const override {
+    if (type.id() != Type::DICTIONARY) {
+      return false;
+    }
+    const auto& dictionary_type = checked_cast<const DictionaryType&>(type);
+    std::vector<TypeHolder> value_types = {dictionary_type.value_type()};
+    return function_->DispatchBest(&value_types).ok();
+  }
+
+  std::string ToString() const override {
+    return "dictionary with values accepted by " + function_->name();
+  }
+
+  bool Equals(const TypeMatcher& other) const override {
+    const auto* other_matcher = dynamic_cast<const DictionaryUnaryTypeMatcher*>(&other);
+    return other_matcher != nullptr && function_ == other_matcher->function_;
+  }
+
+ private:
+  const ScalarFunction* function_;
+};
+
+Result<TypeHolder> ResolveDictionaryUnaryOutput(const ScalarFunction& function,
+                                                KernelContext* ctx,
+                                                const std::vector<TypeHolder>& types) {
+  DCHECK_EQ(types.size(), 1);
+  DCHECK_EQ(types[0].id(), Type::DICTIONARY);
+
+  const auto& dictionary_type =
+      checked_cast<const DictionaryType&>(*types[0].GetSharedPtr());
+  std::vector<TypeHolder> value_types = {dictionary_type.value_type()};
+  ARROW_ASSIGN_OR_RAISE(const Kernel* kernel, function.DispatchBest(&value_types));
+
+  const FunctionOptions* options = function.default_options();
+  if (ctx->state() != nullptr) {
+    options = checked_cast<DictionaryUnaryKernelState*>(ctx->state())->options;
+  }
+
+  KernelContext value_ctx(ctx->exec_context(), kernel);
+  std::unique_ptr<KernelState> value_state;
+  if (kernel->init) {
+    ARROW_ASSIGN_OR_RAISE(value_state,
+                          kernel->init(&value_ctx, {kernel, value_types, options}));
+    value_ctx.SetState(value_state.get());
+  }
+  return kernel->signature->out_type().Resolve(&value_ctx, value_types);
+}
+
+Status ExecuteDictionaryUnary(KernelContext* ctx, const ExecSpan& batch,
+                              ExecResult* out) {
+  const auto& state = checked_cast<const DictionaryUnaryKernelState&>(*ctx->state());
+  DictionaryArray input(batch[0].array.ToArrayData());
+
+  ARROW_ASSIGN_OR_RAISE(
+      Datum transformed_values,
+      state.function->Execute({input.dictionary()}, state.options, ctx->exec_context()));
+  if (!transformed_values.is_array()) {
+    return Status::Invalid("Unary scalar function '", state.function->name(),
+                           "' returned a non-array result for dictionary values");
+  }
+
+  ARROW_ASSIGN_OR_RAISE(Datum result, Take(transformed_values, input.indices(),
+                                           TakeOptions::Defaults(), ctx->exec_context()));
+  DCHECK(result.is_array());
+  out->value = result.array();
+  return Status::OK();
+}
+
+ScalarKernel MakeDictionaryUnaryKernel(const ScalarFunction* function) {
+  ScalarKernel kernel(
+      {InputType(std::make_shared<DictionaryUnaryTypeMatcher>(function))},
+      OutputType([function](KernelContext* ctx, const std::vector<TypeHolder>& types) {
+        return ResolveDictionaryUnaryOutput(*function, ctx, types);
+      }),
+      ExecuteDictionaryUnary,
+      [function](KernelContext*,
+                 const KernelInitArgs& args) -> Result<std::unique_ptr<KernelState>> {
+        return std::make_unique<DictionaryUnaryKernelState>(function, args.options);
+      });
+  kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE;
+  kernel.mem_allocation = MemAllocation::NO_PREALLOCATE;
+  return kernel;
+}
+
+}  // namespace
+
+ScalarFunction::ScalarFunction(std::string name, const Arity& arity, FunctionDoc doc,
+                               const FunctionOptions* default_options, bool is_pure)
+    : detail::FunctionImpl<ScalarKernel>(std::move(name), Function::SCALAR, arity,
+                                         std::move(doc), default_options),
+      is_pure_(is_pure) {
+  if (is_pure && !arity.is_varargs && arity.num_args == 1) {
+    DCHECK_OK(AddKernel(MakeDictionaryUnaryKernel(this)));
+  }
+}
+
 Result<const Kernel*> Function::DispatchExact(
     const std::vector<TypeHolder>& values) const {
   if (kind_ == Function::META) {

cpp/src/arrow/compute/function.h

@@ -299,10 +299,7 @@ class ARROW_EXPORT ScalarFunction : public detail::FunctionImpl<ScalarKernel> {
   using KernelType = ScalarKernel;
 
   ScalarFunction(std::string name, const Arity& arity, FunctionDoc doc,
-                 const FunctionOptions* default_options = NULLPTR, bool is_pure = true)
-      : detail::FunctionImpl<ScalarKernel>(std::move(name), Function::SCALAR, arity,
-                                           std::move(doc), default_options),
-        is_pure_(is_pure) {}
+                 const FunctionOptions* default_options = NULLPTR, bool is_pure = true);
 
   /// \brief Add a kernel with given input/output types, no required state
   /// initialization, preallocation for fixed-width types, and default null

cpp/src/arrow/compute/function_test.cc

@@ -23,6 +23,7 @@
 #include <string>
 #include <vector>
 
+#include "arrow/array/array_dict.h"
 #include "arrow/array/builder_primitive.h"
 #include "arrow/compute/api_aggregate.h"
 #include "arrow/compute/api_scalar.h"
@@ -296,6 +297,53 @@ 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();
+}
+
+}  // namespace
+
+TEST(ScalarFunction, DictionaryUnaryAppliesToDictionaryValues) {
+  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));
+
+  ASSERT_TRUE(result.is_array());
+  AssertArraysEqual(*ArrayFromJSON(int64(), "[10, 20, 10, null, 20]"),
+                    *result.make_array());
+  ASSERT_EQ(counter->values_processed, 3);
+  ASSERT_RAISES(NotImplemented, func.DispatchExact({dictionary(int8(), utf8())}));
+}
+
+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

@@ -2331,6 +2331,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.: