Function RunComputeRegister

cpp/examples/arrow/acero_register_example.cc:116–161

    "ExampleFunctionOptions"};

arrow::Status RunComputeRegister(int argc, char** argv) {
  const std::string name = "compute_register_example";
  auto func = std::make_shared<cp::ScalarFunction>(name, cp::Arity::Unary(), func_doc);
  cp::ScalarKernel kernel({arrow::int64()}, arrow::int64(), ExampleFunctionImpl);
  kernel.mem_allocation = cp::MemAllocation::NO_PREALLOCATE;
  ARROW_RETURN_NOT_OK(func->AddKernel(std::move(kernel)));

  auto registry = cp::GetFunctionRegistry();
  ARROW_RETURN_NOT_OK(registry->AddFunction(std::move(func)));

  arrow::Int64Builder builder(arrow::default_memory_pool());
  std::shared_ptr<arrow::Array> arr;
  ARROW_RETURN_NOT_OK(builder.Append(42));
  ARROW_RETURN_NOT_OK(builder.Finish(&arr));
  auto options = std::make_shared<ExampleFunctionOptions>();
  auto maybe_result = cp::CallFunction(name, {arr}, options.get());
  ARROW_RETURN_NOT_OK(maybe_result.status());

  std::cout << maybe_result->make_array()->ToString() << std::endl;

  // Expression serialization will raise NotImplemented if an expression includes
  // FunctionOptions for which serialization is not supported.
  auto expr = cp::call(name, {}, options);
  auto maybe_serialized = cp::Serialize(expr);
  std::cerr << maybe_serialized.status().ToString() << std::endl;

  auto exec_registry = ac::default_exec_factory_registry();
  ARROW_RETURN_NOT_OK(
      exec_registry->AddFactory("compute_register_example", ExampleExecNodeFactory));

  auto maybe_plan = ac::ExecPlan::Make();
  ARROW_RETURN_NOT_OK(maybe_plan.status());
  ARROW_ASSIGN_OR_RAISE(auto plan, maybe_plan);

  arrow::AsyncGenerator<std::optional<cp::ExecBatch>> source_gen, sink_gen;
  ARROW_RETURN_NOT_OK(
      ac::Declaration::Sequence(
          {
              {"source", ac::SourceNodeOptions{arrow::schema({}), source_gen}},
              {"compute_register_example", ExampleNodeOptions{}},
              {"sink", ac::SinkNodeOptions{&sink_gen}},
          })
          .AddToPlan(plan.get())
          .status());
  return arrow::Status::OK();
}

int main(int argc, char** argv) {
  auto status = RunComputeRegister(argc, argv);