diff --git a/cpp/src/arrow/compute/function.cc b/cpp/src/arrow/compute/function.cc index b0b12a690f8..a98e34cd7a8 100644 --- a/cpp/src/arrow/compute/function.cc +++ b/cpp/src/arrow/compute/function.cc @@ -21,7 +21,9 @@ #include #include +#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(type); + std::vector 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(&other); + return other_matcher != nullptr && function_ == other_matcher->function_; + } + + private: + const ScalarFunction* function_; +}; + +Result ResolveDictionaryUnaryOutput(const ScalarFunction& function, + KernelContext* ctx, + const std::vector& types) { + DCHECK_EQ(types.size(), 1); + DCHECK_EQ(types[0].id(), Type::DICTIONARY); + + const auto& dictionary_type = + checked_cast(*types[0].GetSharedPtr()); + std::vector 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(ctx->state())->options; + } + + KernelContext value_ctx(ctx->exec_context(), kernel); + std::unique_ptr 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(*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(function))}, + OutputType([function](KernelContext* ctx, const std::vector& types) { + return ResolveDictionaryUnaryOutput(*function, ctx, types); + }), + ExecuteDictionaryUnary, + [function](KernelContext*, + const KernelInitArgs& args) -> Result> { + return std::make_unique(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(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 Function::DispatchExact( const std::vector& values) const { if (kind_ == Function::META) { diff --git a/cpp/src/arrow/compute/function.h b/cpp/src/arrow/compute/function.h index 399081e2a73..380ba333a35 100644 --- a/cpp/src/arrow/compute/function.h +++ b/cpp/src/arrow/compute/function.h @@ -299,10 +299,7 @@ class ARROW_EXPORT ScalarFunction : public detail::FunctionImpl { using KernelType = ScalarKernel; ScalarFunction(std::string name, const Arity& arity, FunctionDoc doc, - const FunctionOptions* default_options = NULLPTR, bool is_pure = true) - : detail::FunctionImpl(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 diff --git a/cpp/src/arrow/compute/function_test.cc b/cpp/src/arrow/compute/function_test.cc index cdbb21862aa..f7d5ec761a2 100644 --- a/cpp/src/arrow/compute/function_test.cc +++ b/cpp/src/arrow/compute/function_test.cc @@ -23,6 +23,7 @@ #include #include +#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(*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(); + 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()); diff --git a/cpp/src/arrow/compute/kernels/scalar_string_test.cc b/cpp/src/arrow/compute/kernels/scalar_string_test.cc index b279f991f6e..67ce265bf83 100644 --- a/cpp/src/arrow/compute/kernels/scalar_string_test.cc +++ b/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.: