| 301 | } |
| 302 | |
| 303 | Status Finalize(KernelContext* ctx, Datum* out) override { |
| 304 | // Exclude the last column which is the group id |
| 305 | const int num_args = input_schema->num_fields() - 1; |
| 306 | |
| 307 | ARROW_ASSIGN_OR_RAISE(auto groups_buffer, groups.Finish()); |
| 308 | ARROW_ASSIGN_OR_RAISE(auto groupings, |
| 309 | Grouper::MakeGroupings(UInt32Array(num_values, groups_buffer), |
| 310 | static_cast<uint32_t>(num_groups))); |
| 311 | |
| 312 | ARROW_ASSIGN_OR_RAISE(auto table, |
| 313 | arrow::Table::FromRecordBatches(input_schema, values)); |
| 314 | ARROW_ASSIGN_OR_RAISE(auto rb, table->CombineChunksToBatch(ctx->memory_pool())); |
| 315 | UdfContext udf_context{ctx->memory_pool(), table->num_rows()}; |
| 316 | |
| 317 | if (rb->num_rows() == 0) { |
| 318 | *out = Datum(); |
| 319 | return Status::OK(); |
| 320 | } |
| 321 | |
| 322 | ARROW_ASSIGN_OR_RAISE(RecordBatchVector rbs, ApplyGroupings(*groupings, rb)); |
| 323 | |
| 324 | return SafeCallIntoPython([&] { |
| 325 | ARROW_ASSIGN_OR_RAISE(std::unique_ptr<ArrayBuilder> builder, |
| 326 | MakeBuilder(output_type, ctx->memory_pool())); |
| 327 | for (auto& group_rb : rbs) { |
| 328 | std::unique_ptr<OwnedRef> result; |
| 329 | OwnedRef arg_tuple(PyTuple_New(num_args)); |
| 330 | RETURN_NOT_OK(CheckPyError()); |
| 331 | |
| 332 | for (int arg_id = 0; arg_id < num_args; arg_id++) { |
| 333 | // Since we combined chunks there is only one chunk |
| 334 | std::shared_ptr<Array> c_data = group_rb->column(arg_id); |
| 335 | PyObject* data = wrap_array(c_data); |
| 336 | PyTuple_SetItem(arg_tuple.obj(), arg_id, data); |
| 337 | } |
| 338 | |
| 339 | result = |
| 340 | std::make_unique<OwnedRef>(cb(function->obj(), udf_context, arg_tuple.obj())); |
| 341 | RETURN_NOT_OK(CheckPyError()); |
| 342 | |
| 343 | // unwrapping the output for expected output type |
| 344 | if (is_scalar(result->obj())) { |
| 345 | ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Scalar> val, |
| 346 | unwrap_scalar(result->obj())); |
| 347 | if (*output_type != *val->type) { |
| 348 | return Status::TypeError("Expected output datatype ", output_type->ToString(), |
| 349 | ", but function returned datatype ", |
| 350 | val->type->ToString()); |
| 351 | } |
| 352 | ARROW_RETURN_NOT_OK(builder->AppendScalar(std::move(*val))); |
| 353 | } else { |
| 354 | return Status::TypeError("Unexpected output type: ", |
| 355 | Py_TYPE(result->obj())->tp_name, " (expected Scalar)"); |
| 356 | } |
| 357 | } |
| 358 | ARROW_ASSIGN_OR_RAISE(auto result, builder->Finish()); |
| 359 | out->value = std::move(result->data()); |
| 360 | return Status::OK(); |
no test coverage detected