| 588 | } |
| 589 | |
| 590 | Status RegisterHashAggregateFunction(PyObject* function, UdfWrapperCallback cb, |
| 591 | const UdfOptions& options, |
| 592 | compute::FunctionRegistry* registry) { |
| 593 | if (!PyCallable_Check(function)) { |
| 594 | return Status::TypeError("Expected a callable Python object."); |
| 595 | } |
| 596 | |
| 597 | if (registry == NULLPTR) { |
| 598 | registry = compute::GetFunctionRegistry(); |
| 599 | } |
| 600 | |
| 601 | UdfOptions hash_options = AdjustForHashAggregate(options); |
| 602 | |
| 603 | std::vector<compute::InputType> input_types; |
| 604 | for (const auto& in_dtype : hash_options.input_types) { |
| 605 | input_types.emplace_back(in_dtype); |
| 606 | } |
| 607 | compute::OutputType output_type(hash_options.output_type); |
| 608 | |
| 609 | static auto default_hash_aggregate_options = |
| 610 | compute::ScalarAggregateOptions::Defaults(); |
| 611 | auto hash_aggregate_func = std::make_shared<compute::HashAggregateFunction>( |
| 612 | hash_options.func_name, hash_options.arity, hash_options.func_doc, |
| 613 | &default_hash_aggregate_options); |
| 614 | |
| 615 | // Take reference before wrapping with OwnedRefNoGIL |
| 616 | Py_INCREF(function); |
| 617 | auto function_ref = std::make_shared<OwnedRefNoGIL>(function); |
| 618 | compute::KernelInit init = [function_ref, cb, hash_options]( |
| 619 | compute::KernelContext* ctx, |
| 620 | const compute::KernelInitArgs& args) |
| 621 | -> Result<std::unique_ptr<compute::KernelState>> { |
| 622 | return std::make_unique<PythonUdfHashAggregatorImpl>( |
| 623 | function_ref, cb, hash_options.input_types, hash_options.output_type); |
| 624 | }; |
| 625 | |
| 626 | auto sig = compute::KernelSignature::Make( |
| 627 | std::move(input_types), std::move(output_type), hash_options.arity.is_varargs); |
| 628 | |
| 629 | compute::HashAggregateKernel kernel( |
| 630 | std::move(sig), std::move(init), HashAggregateUdfResize, HashAggregateUdfConsume, |
| 631 | HashAggregateUdfMerge, HashAggregateUdfFinalize, /*ordered=*/false); |
| 632 | RETURN_NOT_OK(hash_aggregate_func->AddKernel(std::move(kernel))); |
| 633 | RETURN_NOT_OK(registry->AddFunction(std::move(hash_aggregate_func))); |
| 634 | return Status::OK(); |
| 635 | } |
| 636 | |
| 637 | Status RegisterAggregateFunction(PyObject* function, UdfWrapperCallback cb, |
| 638 | const UdfOptions& options, |
no test coverage detected