| 110 | // a source -> node_declarations -> sink sequence. |
| 111 | |
| 112 | Status BenchmarkNodeOverhead(benchmark::State& state, int32_t num_batches, |
| 113 | int32_t batch_size, arrow::acero::BatchesWithSchema data, |
| 114 | std::vector<arrow::acero::Declaration>& node_declarations, |
| 115 | MemoryPool* pool) { |
| 116 | ExecContext ctx(pool, ::arrow::internal::GetCpuThreadPool()); |
| 117 | for (auto _ : state) { |
| 118 | state.PauseTiming(); |
| 119 | ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::acero::ExecPlan> plan, |
| 120 | arrow::acero::ExecPlan::Make(ctx)); |
| 121 | AsyncGenerator<std::optional<arrow::compute::ExecBatch>> sink_gen; |
| 122 | arrow::acero::Declaration source = arrow::acero::Declaration( |
| 123 | {"source", |
| 124 | arrow::acero::SourceNodeOptions{data.schema, |
| 125 | data.gen(/*parallel=*/true, /*slow=*/false)}, |
| 126 | "custom_source_label"}); |
| 127 | arrow::acero::Declaration sink = arrow::acero::Declaration( |
| 128 | {"sink", arrow::acero::SinkNodeOptions{&sink_gen}, "custom_sink_label"}); |
| 129 | std::vector<arrow::acero::Declaration> sequence = {source}; |
| 130 | sequence.insert(sequence.end(), node_declarations.begin(), node_declarations.end()); |
| 131 | sequence.push_back(sink); |
| 132 | ARROW_RETURN_NOT_OK( |
| 133 | arrow::acero::Declaration::Sequence(sequence).AddToPlan(plan.get())); |
| 134 | state.ResumeTiming(); |
| 135 | ARROW_RETURN_NOT_OK(StartAndCollect(plan.get(), sink_gen).status()); |
| 136 | } |
| 137 | |
| 138 | state.counters["rows_per_second"] = benchmark::Counter( |
| 139 | static_cast<double>(state.iterations() * num_batches * batch_size), |
| 140 | benchmark::Counter::kIsRate); |
| 141 | |
| 142 | state.counters["batches_per_second"] = benchmark::Counter( |
| 143 | static_cast<double>(state.iterations() * num_batches), benchmark::Counter::kIsRate); |
| 144 | return Status::OK(); |
| 145 | } |
| 146 | |
| 147 | } // namespace acero |
| 148 | } // namespace arrow |
no test coverage detected