| 36 | // calling InputFinished and InputReceived. |
| 37 | |
| 38 | Status BenchmarkIsolatedNodeOverhead( |
| 39 | benchmark::State& state, arrow::compute::Expression expr, int32_t num_batches, |
| 40 | int32_t batch_size, arrow::acero::BatchesWithSchema data, std::string factory_name, |
| 41 | arrow::acero::ExecNodeOptions& options, arrow::MemoryPool* pool) { |
| 42 | for (auto _ : state) { |
| 43 | state.PauseTiming(); |
| 44 | AsyncGenerator<std::optional<arrow::compute::ExecBatch>> sink_gen; |
| 45 | |
| 46 | ExecContext ctx(pool, ::arrow::internal::GetCpuThreadPool()); |
| 47 | |
| 48 | ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::acero::ExecPlan> plan, |
| 49 | arrow::acero::ExecPlan::Make(ctx)); |
| 50 | // Source and sink nodes have no effect on the benchmark. |
| 51 | // Used for dummy purposes as they are referenced in InputReceived and InputFinished. |
| 52 | ARROW_ASSIGN_OR_RAISE(arrow::acero::ExecNode * source_node, |
| 53 | MakeExecNode("source", plan.get(), {}, |
| 54 | arrow::acero::SourceNodeOptions{ |
| 55 | data.schema, data.gen(/*parallel=*/true, |
| 56 | /*slow=*/false)})); |
| 57 | |
| 58 | ARROW_ASSIGN_OR_RAISE(arrow::acero::ExecNode * node, |
| 59 | MakeExecNode(factory_name, plan.get(), {source_node}, options)); |
| 60 | ARROW_RETURN_NOT_OK(MakeExecNode("sink", plan.get(), {node}, |
| 61 | arrow::acero::SinkNodeOptions{&sink_gen})); |
| 62 | |
| 63 | std::unique_ptr<arrow::acero::TaskScheduler> scheduler = |
| 64 | arrow::acero::TaskScheduler::Make(); |
| 65 | std::condition_variable all_tasks_finished_cv; |
| 66 | std::mutex mutex; |
| 67 | |
| 68 | int task_group_id = scheduler->RegisterTaskGroup( |
| 69 | [&](size_t thread_id, int64_t task_id) { |
| 70 | return node->InputReceived(source_node, data.batches[task_id]); |
| 71 | }, |
| 72 | [&](size_t thread_id) { |
| 73 | RETURN_NOT_OK( |
| 74 | node->InputFinished(source_node, static_cast<int>(data.batches.size()))); |
| 75 | std::unique_lock<std::mutex> lk(mutex); |
| 76 | all_tasks_finished_cv.notify_one(); |
| 77 | return Status::OK(); |
| 78 | }); |
| 79 | scheduler->RegisterEnd(); |
| 80 | |
| 81 | arrow::acero::ThreadIndexer thread_indexer; |
| 82 | |
| 83 | state.ResumeTiming(); |
| 84 | arrow::internal::ThreadPool* thread_pool = arrow::internal::GetCpuThreadPool(); |
| 85 | ARROW_RETURN_NOT_OK(scheduler->StartScheduling( |
| 86 | thread_indexer(), |
| 87 | [&](std::function<Status(size_t)> task) -> Status { |
| 88 | return thread_pool->Spawn([&, task]() { |
| 89 | size_t tid = thread_indexer(); |
| 90 | ARROW_DCHECK_OK(task(tid)); |
| 91 | }); |
| 92 | }, |
| 93 | thread_pool->GetCapacity(), |
| 94 | /*use_sync_execution=*/false)); |
| 95 | std::unique_lock<std::mutex> lk(mutex); |
no test coverage detected