ARROW-15641: [C++][Python] UDF Aggregate Function Implementation

cpp/examples/arrow/udf_example.cc

@@ -83,15 +83,154 @@ arrow::Status Execute() {
 
   ARROW_ASSIGN_OR_RAISE(auto res, cp::CallFunction(name, {x, y, z}));
   auto res_array = res.make_array();
-  std::cout << "Result" << std::endl;
+  std::cout << "Scalar UDF Result" << std::endl;
   std::cout << res_array->ToString() << std::endl;
   return arrow::Status::OK();
 }
 
+// User-defined Scalar Aggregate Function Example
+struct ScalarUdfAggregator : public cp::KernelState {
+  virtual arrow::Status Consume(cp::KernelContext* ctx, const cp::ExecSpan& batch) = 0;
+  virtual arrow::Status MergeFrom(cp::KernelContext* ctx, cp::KernelState&& src) = 0;
+  virtual arrow::Status Finalize(cp::KernelContext* ctx, arrow::Datum* out) = 0;
+};
+
+class SimpleCountFunctionOptionsType : public cp::FunctionOptionsType {
+  const char* type_name() const override { return "SimpleCountFunctionOptionsType"; }
+  std::string Stringify(const cp::FunctionOptions&) const override {
+    return "SimpleCountFunctionOptionsType";
+  }
+  bool Compare(const cp::FunctionOptions&, const cp::FunctionOptions&) const override {
+    return true;
+  }
+  std::unique_ptr<cp::FunctionOptions> Copy(const cp::FunctionOptions&) const override;
+};
+
+cp::FunctionOptionsType* GetSimpleCountFunctionOptionsType() {
+  static SimpleCountFunctionOptionsType options_type;
+  return &options_type;
+}
+
+class SimpleCountOptions : public cp::FunctionOptions {
+ public:
+  SimpleCountOptions() : cp::FunctionOptions(GetSimpleCountFunctionOptionsType()) {}
+  static constexpr char const kTypeName[] = "SimpleCountOptions";
+  static SimpleCountOptions Defaults() { return SimpleCountOptions{}; }
+};
+
+std::unique_ptr<cp::FunctionOptions> SimpleCountFunctionOptionsType::Copy(
+    const cp::FunctionOptions&) const {
+  return std::make_unique<SimpleCountOptions>();
+}
+
+const cp::FunctionDoc simple_count_doc{
+    "SimpleCount the number of null / non-null values",
+    ("By default, only non-null values are counted.\n"
+     "This can be changed through SimpleCountOptions."),
+    {"array"},
+    "SimpleCountOptions"};
+
+// Need Python interface for this Class
+struct SimpleCountImpl : public ScalarUdfAggregator {
+  explicit SimpleCountImpl(SimpleCountOptions options) : options(std::move(options)) {}
+
+  arrow::Status Consume(cp::KernelContext*, const cp::ExecSpan& batch) override {
+    if (batch[0].is_array()) {
+      const arrow::ArraySpan& input = batch[0].array;
+      const int64_t nulls = input.GetNullCount();
+      this->non_nulls += input.length - nulls;
+    } else {
+      const arrow::Scalar& input = *batch[0].scalar;
+      this->non_nulls += input.is_valid * batch.length;
+    }
+    return arrow::Status::OK();
+  }
+
+  arrow::Status MergeFrom(cp::KernelContext*, cp::KernelState&& src) override {
+    const auto& other_state = arrow::internal::checked_cast<const SimpleCountImpl&>(src);
+    std::cout << "This non_nulls: " << this->non_nulls << std::endl;
+    this->non_nulls += other_state.non_nulls;
+    return arrow::Status::OK();
+  }
+
+  arrow::Status Finalize(cp::KernelContext* ctx, arrow::Datum* out) override {
+    const auto& state =
+        arrow::internal::checked_cast<const SimpleCountImpl&>(*ctx->state());
+    *out = arrow::Datum(state.non_nulls);
+    return arrow::Status::OK();
+  }
+
+  SimpleCountOptions options;
+  int64_t non_nulls = 0;
+};
+
+// TODO: need a Python interface for this function
+arrow::Result<std::unique_ptr<cp::KernelState>> SimpleCountInit(
+    cp::KernelContext*, const cp::KernelInitArgs& args) {
+  return std::make_unique<SimpleCountImpl>(
+      static_cast<const SimpleCountOptions&>(*args.options));
+}
+
+arrow::Status AggregateUdfConsume(cp::KernelContext* ctx, const cp::ExecSpan& batch) {
+  return arrow::internal::checked_cast<ScalarUdfAggregator*>(ctx->state())
+      ->Consume(ctx, batch);
+}
+
+arrow::Status AggregateUdfMerge(cp::KernelContext* ctx, cp::KernelState&& src,
+                                cp::KernelState* dst) {
+  return arrow::internal::checked_cast<ScalarUdfAggregator*>(dst)->MergeFrom(
+      ctx, std::move(src));
+}
+
+arrow::Status AggregateUdfFinalize(cp::KernelContext* ctx, arrow::Datum* out) {
+  return arrow::internal::checked_cast<ScalarUdfAggregator*>(ctx->state())
+      ->Finalize(ctx, out);
+}
+
+arrow::Status AddAggKernel(std::shared_ptr<cp::KernelSignature> sig, cp::KernelInit init,
+                           cp::ScalarAggregateFunction* func) {
+  cp::ScalarAggregateKernel kernel(std::move(sig), std::move(init), AggregateUdfConsume,
+                                   AggregateUdfMerge, AggregateUdfFinalize);
+  ARROW_RETURN_NOT_OK(func->AddKernel(std::move(kernel)));
+  return arrow::Status::OK();
+}
+
+arrow::Status ExecuteAggregate() {
+  auto registry = cp::GetFunctionRegistry();
+  static auto default_scalar_aggregate_options = cp::ScalarAggregateOptions::Defaults();
+  static auto default_count_options = SimpleCountOptions::Defaults();
+  const std::string name = "simple_count";
+  auto func = std::make_shared<cp::ScalarAggregateFunction>(
+      name, cp::Arity::Unary(), simple_count_doc, &default_count_options);
+
+  // Takes any input, outputs int64 scalar
+  ARROW_RETURN_NOT_OK(
+      AddAggKernel(cp::KernelSignature::Make({arrow::int64()}, arrow::int64()),
+                   SimpleCountInit, func.get()));
+  ARROW_RETURN_NOT_OK(registry->AddFunction(std::move(func)));
+
+  ARROW_ASSIGN_OR_RAISE(auto x, GetArrayDataSample<arrow::Int64Type>({1, 2, 3, 4, 5, 6}));
+
+  ARROW_ASSIGN_OR_RAISE(auto res, cp::CallFunction(name, {x}));
+  auto res_scalar = res.scalar();
+  std::cout << "Aggregate UDF Result" << std::endl;
+  std::cout << res_scalar->ToString() << std::endl;
+
+  return arrow::Status::OK();
+}
+
 int main(int argc, char** argv) {
-  auto status = Execute();
-  if (!status.ok()) {
-    std::cerr << "Error occurred : " << status.message() << std::endl;
+  std::cout << "Sample Scalar UDF Execution" << std::endl;
+  auto s1 = Execute();
+  if (!s1.ok()) {
+    std::cerr << "Error occurred : " << s1.message() << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  std::cout << "Sample Aggregate UDF Execution" << std::endl;
+  auto s2 = ExecuteAggregate();
+  if (!s2.ok()) {
+    std::cerr << "Error occurred : " << s2.message() << std::endl;
     return EXIT_FAILURE;
   }
   return EXIT_SUCCESS;

python/pyarrow/_compute.pxd

@@ -27,6 +27,12 @@ cdef class ScalarUdfContext(_Weakrefable):
 
     cdef void init(self, const CScalarUdfContext& c_context)
 
+cdef class ScalarAggregateUdfContext(_Weakrefable):
+    cdef:
+        CScalarAggregateUdfContext c_context
+
+    cdef void init(self, const CScalarAggregateUdfContext& c_context)
+
 cdef class FunctionOptions(_Weakrefable):
     cdef:
         shared_ptr[CFunctionOptions] wrapped