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"
@@ -295,6 +298,71 @@ struct FunctionExecutorImpl : public FunctionExecutor {
 
 }  // namespace detail
 
+namespace {
+
+bool CanExecuteDictionaryValues(const ScalarFunction& function,
+                                const std::vector<Datum>& args) {
+  return function.is_pure() && !function.arity().is_varargs &&
+         function.arity().num_args == 1 && args.size() == 1 &&
+         args[0].type()->id() == Type::DICTIONARY;
+}
+
+Result<Datum> ExecuteDictionaryArray(const ScalarFunction& function, const Datum& arg,
+                                     const FunctionOptions* options, ExecContext* ctx) {
+  auto input_array = arg.make_array();
+  const auto& input = checked_cast<const DictionaryArray&>(*input_array);
+  ARROW_ASSIGN_OR_RAISE(Datum transformed_values,
+                        function.Execute({input.dictionary()}, options, ctx));
+  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, input.indices(), TakeOptions::Defaults(), ctx);
+}
+
+Result<Datum> ExecuteDictionaryValues(const ScalarFunction& function, const Datum& arg,
+                                      const FunctionOptions* options, ExecContext* ctx) {
+  switch (arg.kind()) {
+    case Datum::ARRAY:
+      return ExecuteDictionaryArray(function, arg, options, ctx);
+    case Datum::SCALAR: {
+      const auto& input = checked_cast<const DictionaryScalar&>(*arg.scalar());
+      ARROW_ASSIGN_OR_RAISE(auto value, input.GetEncodedValue());
+      return function.Execute({std::move(value)}, options, ctx);
+    }
+    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(function, Datum(chunk), options, ctx));
+        DCHECK(output.is_array());
+        output_type = output.type();
+        output_chunks.push_back(output.make_array());
+      }
+
+      if (output_type == nullptr) {
+        const auto& input_type = checked_cast<const DictionaryType&>(*arg.type());
+        ARROW_ASSIGN_OR_RAISE(auto empty_values, MakeEmptyArray(input_type.value_type()));
+        ARROW_ASSIGN_OR_RAISE(Datum output,
+                              function.Execute({std::move(empty_values)}, options, ctx));
+        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");
+  }
+}
+
+}  // namespace
+
 Result<const Kernel*> Function::DispatchExact(
     const std::vector<TypeHolder>& values) const {
   if (kind_ == Function::META) {
@@ -355,6 +423,51 @@ Result<Datum> Function::Execute(const ExecBatch& batch, const FunctionOptions* o
   return ExecuteInternal(*this, batch.values, batch.length, options, ctx);
 }
 
+Result<Datum> ScalarFunction::Execute(const std::vector<Datum>& args,
+                                      const FunctionOptions* options,
+                                      ExecContext* ctx) const {
+  if (!CanExecuteDictionaryValues(*this, args)) {
+    return Function::Execute(args, options, ctx);
+  }
+
+  ARROW_ASSIGN_OR_RAISE(auto types, internal::GetFunctionArgumentTypes(args));
+  auto direct_kernel = DispatchBest(&types);
+  if (direct_kernel.ok()) {
+    return Function::Execute(args, options, ctx);
+  }
+  if (!direct_kernel.status().IsNotImplemented()) {
+    return direct_kernel.status();
+  }
+  return ExecuteDictionaryValues(*this, args[0], options, ctx);
+}
+
+Result<Datum> ScalarFunction::Execute(const ExecBatch& batch,
+                                      const FunctionOptions* options,
+                                      ExecContext* ctx) const {
+  if (!CanExecuteDictionaryValues(*this, batch.values)) {
+    return Function::Execute(batch, options, ctx);
+  }
+
+  ARROW_ASSIGN_OR_RAISE(auto types, internal::GetFunctionArgumentTypes(batch.values));
+  auto direct_kernel = DispatchBest(&types);
+  if (direct_kernel.ok()) {
+    return Function::Execute(batch, options, ctx);
+  }
+  if (!direct_kernel.status().IsNotImplemented()) {
+    return direct_kernel.status();
+  }
+  if (batch.length != -1) {
+    ARROW_ASSIGN_OR_RAISE(auto inferred_length, ExecBatch::InferLength(batch.values));
+    if (batch.length != inferred_length) {
+      return Status::Invalid(
+          "Passed batch length for execution did not match actual"
+          " length of values for execution of scalar function '",
+          name(), "'");
+    }
+  }
+  return ExecuteDictionaryValues(*this, batch.values[0], options, ctx);
+}
+
 namespace {
 
 Status ValidateFunctionSummary(const std::string& s) {

cpp/src/arrow/compute/function.h

@@ -304,6 +304,12 @@ class ARROW_EXPORT ScalarFunction : public detail::FunctionImpl<ScalarKernel> {
                                            std::move(doc), default_options),
         is_pure_(is_pure) {}
 
+  Result<Datum> Execute(const std::vector<Datum>& args, const FunctionOptions* options,
+                        ExecContext* ctx) const override;
+
+  Result<Datum> Execute(const ExecBatch& batch, const FunctionOptions* options,
+                        ExecContext* ctx) const override;
+
   /// \brief Add a kernel with given input/output types, no required state
   /// initialization, preallocation for fixed-width types, and default null
   /// handling (intersect validity bitmaps of inputs).

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,80 @@ 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));
+
+  auto expected = ArrayFromJSON(int64(), "[10, 20, 10, null, 20]");
+  ASSERT_TRUE(result.is_array());
+  AssertArraysEqual(*expected, *result.make_array());
+  ASSERT_EQ(counter->values_processed, 3);
+
+  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, 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.: