pitrou commented on code in PR #36932:
URL: https://github.com/apache/arrow/pull/36932#discussion_r1288517725
##########
cpp/src/arrow/compute/kernels/vector_cumulative_ops.cc:
##########
@@ -217,53 +253,102 @@ const FunctionDoc cumulative_min_doc{
"start as the new minimum)."),
{"values"},
"CumulativeOptions"};
-} // namespace
-template <typename Op, typename OptionsType>
-void MakeVectorCumulativeFunction(FunctionRegistry* registry, const
std::string func_name,
- const FunctionDoc doc) {
+const FunctionDoc cumulative_mean_doc{
+ "Compute the cumulative mean over a numeric input",
+ ("`values` must be numeric. Return an array/chunked array which is the\n"
+ "cumulative mean computed over `values`. CumulativeOptions::start_value
is \n"
+ "ignored."),
+ {"values"},
+ "CumulativeOptions"};
+
+// Kernel factory for complex stateful computations.
+template <template <typename ArgType> typename State, typename OptionsType>
+struct CumulativeStatefulKernelFactory {
+ VectorKernel kernel;
+
+ CumulativeStatefulKernelFactory() {
+ kernel.can_execute_chunkwise = false;
+ kernel.null_handling = NullHandling::type::COMPUTED_NO_PREALLOCATE;
+ kernel.mem_allocation = MemAllocation::type::NO_PREALLOCATE;
+ kernel.init = CumulativeOptionsWrapper<OptionsType>::Init;
+ }
+
+ template <typename Type>
+ enable_if_number<Type, Status> Visit(const Type& type) {
+ kernel.signature = KernelSignature::Make(
+ {type.GetSharedPtr()},
+ OutputType(TypeTraits<typename
State<Type>::OutType>::type_singleton()));
+ kernel.exec = CumulativeKernel<Type, State<Type>, OptionsType>::Exec;
+ kernel.exec_chunked = CumulativeKernelChunked<Type, State<Type>,
OptionsType>::Exec;
+ return arrow::Status::OK();
+ }
+
+ Status Visit(const DataType& type) {
+ return Status::NotImplemented("Cumulative kernel not implemented for type
",
+ type.ToString());
+ }
+
+ Result<VectorKernel> Make(const DataType& type) {
+ RETURN_NOT_OK(VisitTypeInline(type, this));
+ return kernel;
+ }
+};
+
+template <template <typename ArgType> typename State, typename OptionsType>
+void MakeVectorCumulativeStatefulFunction(FunctionRegistry* registry,
+ const std::string func_name,
+ const FunctionDoc doc) {
static const OptionsType kDefaultOptions = OptionsType::Defaults();
auto func =
std::make_shared<VectorFunction>(func_name, Arity::Unary(), doc,
&kDefaultOptions);
std::vector<std::shared_ptr<DataType>> types;
types.insert(types.end(), NumericTypes().begin(), NumericTypes().end());
+ CumulativeStatefulKernelFactory<State, OptionsType> kernel_factory;
for (const auto& ty : types) {
- VectorKernel kernel;
- kernel.can_execute_chunkwise = false;
- kernel.null_handling = NullHandling::type::COMPUTED_NO_PREALLOCATE;
- kernel.mem_allocation = MemAllocation::type::NO_PREALLOCATE;
- kernel.signature = KernelSignature::Make({ty}, OutputType(ty));
- kernel.exec =
- ArithmeticExecFromOp<CumulativeKernel, Op, ArrayKernelExec,
OptionsType>(ty);
- kernel.exec_chunked =
- ArithmeticExecFromOp<CumulativeKernelChunked, Op,
VectorKernel::ChunkedExec,
- OptionsType>(ty);
- kernel.init = CumulativeOptionsWrapper<OptionsType>::Init;
+ auto kernel = kernel_factory.Make(*ty).ValueOrDie();
DCHECK_OK(func->AddKernel(std::move(kernel)));
}
DCHECK_OK(registry->AddFunction(std::move(func)));
}
+// A kernel factory that forwards to CumulativeBinaryOp<Op, ...> for the given
type.
+// Need to use a struct because template-using declarations cannot appear in
+// function scope.
+template <typename Op, typename OptionsType>
+struct MakeVectorCumulativeBinaryOpFunction {
+ template <typename ArgType>
+ using State = CumulativeBinaryOp<Op, ArgType>;
+
+ static void Call(FunctionRegistry* registry, std::string func_name,
FunctionDoc doc) {
+ MakeVectorCumulativeStatefulFunction<State, OptionsType>(
+ registry, std::move(func_name), std::move(doc));
+ }
+};
Review Comment:
@bkietz Can you think of a more straightforward way to do this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
