diff --git a/cpp/src/arrow/compute/function.cc b/cpp/src/arrow/compute/function.cc index b0b12a690f86..b799f7d08b49 100644 --- a/cpp/src/arrow/compute/function.cc +++ b/cpp/src/arrow/compute/function.cc @@ -21,7 +21,10 @@ #include #include +#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& args) { + return function.is_pure() && !function.arity().is_varargs && + function.arity().num_args == 1 && args.size() == 1 && + args[0].type()->id() == Type::DICTIONARY; +} + +Result ExecuteDictionaryArray(const ScalarFunction& function, const Datum& arg, + const FunctionOptions* options, ExecContext* ctx) { + auto input_array = arg.make_array(); + const auto& input = checked_cast(*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 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(*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 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(*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 Function::DispatchExact( const std::vector& values) const { if (kind_ == Function::META) { @@ -355,6 +423,51 @@ Result Function::Execute(const ExecBatch& batch, const FunctionOptions* o return ExecuteInternal(*this, batch.values, batch.length, options, ctx); } +Result ScalarFunction::Execute(const std::vector& 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 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) { diff --git a/cpp/src/arrow/compute/function.h b/cpp/src/arrow/compute/function.h index 399081e2a737..b28e5045627a 100644 --- a/cpp/src/arrow/compute/function.h +++ b/cpp/src/arrow/compute/function.h @@ -304,6 +304,12 @@ class ARROW_EXPORT ScalarFunction : public detail::FunctionImpl { std::move(doc), default_options), is_pure_(is_pure) {} + Result Execute(const std::vector& args, const FunctionOptions* options, + ExecContext* ctx) const override; + + Result 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). diff --git a/cpp/src/arrow/compute/function_test.cc b/cpp/src/arrow/compute/function_test.cc index cdbb21862aaf..0cebabc9b55d 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,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(*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)); + + 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(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(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 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()); diff --git a/cpp/src/arrow/compute/kernels/scalar_string_test.cc b/cpp/src/arrow/compute/kernels/scalar_string_test.cc index b279f991f6ea..67ce265bf832 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.: