Skip to content

Commit 498b645

Browse files
westonpacewestonpace
authored
ARROW-15732: [C++] Do not use any CPU threads in execution plan when use_threads is false (apache#15104)
This PR gets rid of the "executor==nullptr" style of execution from Acero. Now a CPU executor must always be defined. This also updates the DeclarationToXyz methods to support 5 "standard" modes of operation: * sync / threaded - All CPU work is done on the CPU thread pool. The calling thread sleeps until the plan is done. * sync / serial - All CPU work is done on the calling thread. The calling thread becomes a serial executor until plan is done. * async / threaded - All CPU work is done on the CPU thread pool. The calling thread returns immediately. * async / serial - All CPU work is done on a plan-specific 1-thread pool. The calling thread returns immediately. * async / custom exec context - All CPU work is done on the provided executor which cannot be nullptr and the caller is responsible for keeping it alive. For backwards compatibility the "old style" of serial execution (which is still used by R/python) will still be accepted. Under the hood a 1-thread pool will be created and its lifetime will be bound to the exec plan (this is similar to async / serial when using DeclarationToXyz). Since DeclarationToXyz doesn't support "ordered sinks" we cannot yet migrate the bindings and so they will continue to rely on the deprecated API. All unit tests are updated to use the new API. Some unit tests were not actually running in parallel (even though they said they were) and now that they are running in parallel a few changes had to be made to accept / handle non-deterministic result order. As part of this process some unit tests were migrated away from StartAndCollect in favor of DeclarationToXyz. Hopefully we can eventually migrate almost all of the unit tests to DeclarationToXyz but that is a longer term effort. Authored-by: Weston Pace <weston.pace@gmail.com> Signed-off-by: Weston Pace <weston.pace@gmail.com>
1 parent 237bc30 commit 498b645

40 files changed

Lines changed: 1002 additions & 1169 deletions

‎cpp/examples/arrow/execution_plan_documentation_examples.cc‎

Lines changed: 155 additions & 283 deletions
Large diffs are not rendered by default.

‎cpp/examples/arrow/join_example.cc‎

Lines changed: 6 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -82,18 +82,8 @@ arrow::Result<std::shared_ptr<arrow::dataset::Dataset>> CreateDataSetFromCSVData
8282
}
8383

8484
arrow::Status DoHashJoin() {
85-
cp::ExecContext exec_context;
86-
8785
arrow::dataset::internal::Initialize();
8886

89-
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<cp::ExecPlan> plan,
90-
cp::ExecPlan::Make(&exec_context));
91-
92-
arrow::AsyncGenerator<std::optional<cp::ExecBatch>> sink_gen;
93-
94-
cp::ExecNode* left_source;
95-
cp::ExecNode* right_source;
96-
9787
ARROW_ASSIGN_OR_RAISE(auto l_dataset, CreateDataSetFromCSVData(true));
9888
ARROW_ASSIGN_OR_RAISE(auto r_dataset, CreateDataSetFromCSVData(false));
9989

@@ -111,10 +101,8 @@ arrow::Status DoHashJoin() {
111101
auto l_scan_node_options = arrow::dataset::ScanNodeOptions{l_dataset, l_options};
112102
auto r_scan_node_options = arrow::dataset::ScanNodeOptions{r_dataset, r_options};
113103

114-
ARROW_ASSIGN_OR_RAISE(left_source,
115-
cp::MakeExecNode("scan", plan.get(), {}, l_scan_node_options));
116-
ARROW_ASSIGN_OR_RAISE(right_source,
117-
cp::MakeExecNode("scan", plan.get(), {}, r_scan_node_options));
104+
arrow::compute::Declaration left{"scan", std::move(l_scan_node_options)};
105+
arrow::compute::Declaration right{"scan", std::move(r_scan_node_options)};
118106

119107
arrow::compute::HashJoinNodeOptions join_opts{arrow::compute::JoinType::INNER,
120108
/*in_left_keys=*/{"lkey"},
@@ -123,26 +111,12 @@ arrow::Status DoHashJoin() {
123111
/*output_suffix_for_left*/ "_l",
124112
/*output_suffix_for_right*/ "_r"};
125113

126-
ARROW_ASSIGN_OR_RAISE(
127-
auto hashjoin,
128-
cp::MakeExecNode("hashjoin", plan.get(), {left_source, right_source}, join_opts));
114+
arrow::compute::Declaration hashjoin{
115+
"hashjoin", {std::move(left), std::move(right)}, join_opts};
129116

130-
ARROW_ASSIGN_OR_RAISE(std::ignore, cp::MakeExecNode("sink", plan.get(), {hashjoin},
131-
cp::SinkNodeOptions{&sink_gen}));
132117
// expected columns l_a, l_b
133-
std::shared_ptr<arrow::RecordBatchReader> sink_reader = cp::MakeGeneratorReader(
134-
hashjoin->output_schema(), std::move(sink_gen), exec_context.memory_pool());
135-
136-
// validate the ExecPlan
137-
ARROW_RETURN_NOT_OK(plan->Validate());
138-
// start the ExecPlan
139-
ARROW_RETURN_NOT_OK(plan->StartProducing());
140-
141-
// collect sink_reader into a Table
142-
std::shared_ptr<arrow::Table> response_table;
143-
144-
ARROW_ASSIGN_OR_RAISE(response_table,
145-
arrow::Table::FromRecordBatchReader(sink_reader.get()));
118+
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Table> response_table,
119+
arrow::compute::DeclarationToTable(std::move(hashjoin)));
146120

147121
std::cout << "Results : " << response_table->ToString() << std::endl;
148122

‎cpp/src/arrow/compute/exec.cc‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
#include "arrow/util/checked_cast.h"
4949
#include "arrow/util/cpu_info.h"
5050
#include "arrow/util/logging.h"
51+
#include "arrow/util/thread_pool.h"
5152
#include "arrow/util/vector.h"
5253

5354
namespace arrow {
@@ -56,6 +57,7 @@ using internal::BitmapAnd;
5657
using internal::checked_cast;
5758
using internal::CopyBitmap;
5859
using internal::CpuInfo;
60+
using internal::GetCpuThreadPool;
5961

6062
namespace compute {
6163

@@ -64,6 +66,11 @@ ExecContext* default_exec_context() {
6466
return &default_ctx;
6567
}
6668

69+
ExecContext* threaded_exec_context() {
70+
static ExecContext threaded_ctx(default_memory_pool(), GetCpuThreadPool());
71+
return &threaded_ctx;
72+
}
73+
6774
ExecBatch::ExecBatch(const RecordBatch& batch)
6875
: values(batch.num_columns()), length(batch.num_rows()) {
6976
auto columns = batch.column_data();

‎cpp/src/arrow/compute/exec/asof_join_benchmark.cc‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,6 @@ static void TableJoinOverhead(benchmark::State& state,
5454
TableGenerationProperties right_table_properties,
5555
int batch_size, int num_right_tables,
5656
std::string factory_name, ExecNodeOptions& options) {
57-
ExecContext ctx(default_memory_pool(), nullptr);
5857
left_table_properties.column_prefix = "lt";
5958
left_table_properties.seed = 0;
6059
ASSERT_OK_AND_ASSIGN(TableStats left_table_stats, MakeTable(left_table_properties));
@@ -75,23 +74,18 @@ static void TableJoinOverhead(benchmark::State& state,
7574

7675
for (auto _ : state) {
7776
state.PauseTiming();
78-
ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::compute::ExecPlan> plan,
79-
ExecPlan::Make(&ctx));
80-
std::vector<ExecNode*> input_nodes = {*arrow::compute::MakeExecNode(
81-
"table_source", plan.get(), {},
77+
std::vector<Declaration::Input> input_nodes = {Declaration(
78+
"table_source",
8279
arrow::compute::TableSourceNodeOptions(left_table_stats.table, batch_size))};
8380
input_nodes.reserve(right_input_tables.size() + 1);
8481
for (TableStats table_stats : right_input_tables) {
85-
input_nodes.push_back(*arrow::compute::MakeExecNode(
86-
"table_source", plan.get(), {},
82+
input_nodes.push_back(Declaration(
83+
"table_source",
8784
arrow::compute::TableSourceNodeOptions(table_stats.table, batch_size)));
8885
}
89-
ASSERT_OK_AND_ASSIGN(arrow::compute::ExecNode * join_node,
90-
MakeExecNode(factory_name, plan.get(), input_nodes, options));
91-
AsyncGenerator<std::optional<ExecBatch>> sink_gen;
92-
ASSERT_OK(MakeExecNode("sink", plan.get(), {join_node}, SinkNodeOptions{&sink_gen}));
86+
Declaration join_node{factory_name, {input_nodes}, options};
9387
state.ResumeTiming();
94-
ASSERT_FINISHES_OK(StartAndCollect(plan.get(), sink_gen));
88+
ASSERT_OK(DeclarationToStatus(std::move(join_node)));
9589
}
9690

9791
state.counters["input_rows_per_second"] = benchmark::Counter(
@@ -104,7 +98,7 @@ static void TableJoinOverhead(benchmark::State& state,
10498
benchmark::Counter::kIsRate);
10599

106100
state.counters["maximum_peak_memory"] =
107-
benchmark::Counter(static_cast<double>(ctx.memory_pool()->max_memory()));
101+
benchmark::Counter(static_cast<double>(default_memory_pool()->max_memory()));
108102
}
109103

110104
static void AsOfJoinOverhead(benchmark::State& state) {

‎cpp/src/arrow/compute/exec/asof_join_node.cc‎

Lines changed: 43 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -807,10 +807,36 @@ class AsofJoinNode : public ExecNode {
807807
}
808808
}
809809

810-
void Process() {
810+
template <typename Callable>
811+
struct Defer {
812+
Callable callable;
813+
explicit Defer(Callable callable) : callable(std::move(callable)) {}
814+
~Defer() noexcept { callable(); }
815+
};
816+
817+
void EndFromProcessThread() {
818+
// We must spawn a new task to transfer off the process thread when
819+
// marking this finished. Otherwise there is a chance that doing so could
820+
// mark the plan finished which may destroy the plan which will destroy this
821+
// node which will cause us to join on ourselves.
822+
ErrorIfNotOk(plan_->query_context()->executor()->Spawn([this] {
823+
Defer cleanup([this]() { finished_.MarkFinished(); });
824+
outputs_[0]->InputFinished(this, batches_produced_);
825+
}));
826+
}
827+
828+
bool CheckEnded() {
829+
if (state_.at(0)->Finished()) {
830+
EndFromProcessThread();
831+
return false;
832+
}
833+
return true;
834+
}
835+
836+
bool Process() {
811837
std::lock_guard<std::mutex> guard(gate_);
812-
if (finished_.is_finished()) {
813-
return;
838+
if (!CheckEnded()) {
839+
return false;
814840
}
815841

816842
// Process batches while we have data
@@ -825,7 +851,8 @@ class AsofJoinNode : public ExecNode {
825851
outputs_[0]->InputReceived(this, std::move(out_b));
826852
} else {
827853
ErrorIfNotOk(result.status());
828-
return;
854+
EndFromProcessThread();
855+
return false;
829856
}
830857
}
831858

@@ -834,18 +861,24 @@ class AsofJoinNode : public ExecNode {
834861
//
835862
// It may happen here in cases where InputFinished was called before we were finished
836863
// producing results (so we didn't know the output size at that time)
837-
if (state_.at(0)->Finished()) {
838-
outputs_[0]->InputFinished(this, batches_produced_);
839-
finished_.MarkFinished();
864+
if (!CheckEnded()) {
865+
return false;
840866
}
867+
868+
// There is no more we can do now but there is still work remaining for later when
869+
// more data arrives.
870+
return true;
841871
}
842872

843873
void ProcessThread() {
844874
for (;;) {
845875
if (!process_.Pop()) {
876+
EndFromProcessThread();
877+
return;
878+
}
879+
if (!Process()) {
846880
return;
847881
}
848-
Process();
849882
}
850883
}
851884

@@ -1120,10 +1153,7 @@ class AsofJoinNode : public ExecNode {
11201153
// finished.
11211154
process_.Push(true);
11221155
}
1123-
Status StartProducing() override {
1124-
finished_ = arrow::Future<>::Make();
1125-
return Status::OK();
1126-
}
1156+
Status StartProducing() override { return Status::OK(); }
11271157
void PauseProducing(ExecNode* output, int32_t counter) override {}
11281158
void ResumeProducing(ExecNode* output, int32_t counter) override {}
11291159
void StopProducing(ExecNode* output) override {
@@ -1137,7 +1167,6 @@ class AsofJoinNode : public ExecNode {
11371167
arrow::Future<> finished() override { return finished_; }
11381168

11391169
private:
1140-
arrow::Future<> finished_;
11411170
std::vector<col_index_t> indices_of_on_key_;
11421171
std::vector<std::vector<col_index_t>> indices_of_by_key_;
11431172
std::vector<std::unique_ptr<KeyHasher>> key_hashers_;
@@ -1176,9 +1205,7 @@ AsofJoinNode::AsofJoinNode(ExecPlan* plan, NodeVector inputs,
11761205
may_rehash_(may_rehash),
11771206
tolerance_(tolerance),
11781207
process_(),
1179-
process_thread_(&AsofJoinNode::ProcessThreadWrapper, this) {
1180-
finished_ = arrow::Future<>::MakeFinished();
1181-
}
1208+
process_thread_(&AsofJoinNode::ProcessThreadWrapper, this) {}
11821209

11831210
namespace internal {
11841211
void RegisterAsofJoinNode(ExecFactoryRegistry* registry) {

‎cpp/src/arrow/compute/exec/asof_join_node_test.cc‎

Lines changed: 3 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -223,9 +223,6 @@ void CheckRunOutput(const BatchesWithSchema& l_batches,
223223
const BatchesWithSchema& r1_batches,
224224
const BatchesWithSchema& exp_batches,
225225
const AsofJoinNodeOptions join_options) {
226-
auto exec_ctx = std::make_unique<ExecContext>(default_memory_pool(), nullptr);
227-
ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make(exec_ctx.get()));
228-
229226
Declaration join{"asofjoin", join_options};
230227

231228
join.inputs.emplace_back(Declaration{
@@ -235,21 +232,12 @@ void CheckRunOutput(const BatchesWithSchema& l_batches,
235232
join.inputs.emplace_back(Declaration{
236233
"source", SourceNodeOptions{r1_batches.schema, r1_batches.gen(false, false)}});
237234

238-
AsyncGenerator<std::optional<ExecBatch>> sink_gen;
239-
240-
ASSERT_OK(Declaration::Sequence({join, {"sink", SinkNodeOptions{&sink_gen}}})
241-
.AddToPlan(plan.get()));
242-
243-
ASSERT_FINISHES_OK_AND_ASSIGN(auto res, StartAndCollect(plan.get(), sink_gen));
244-
for (auto batch : res) {
245-
ASSERT_EQ(exp_batches.schema->num_fields(), batch.values.size());
246-
}
235+
ASSERT_OK_AND_ASSIGN(auto res_table,
236+
DeclarationToTable(std::move(join), /*use_threads=*/false));
247237

248238
ASSERT_OK_AND_ASSIGN(auto exp_table,
249239
TableFromExecBatches(exp_batches.schema, exp_batches.batches));
250240

251-
ASSERT_OK_AND_ASSIGN(auto res_table, TableFromExecBatches(exp_batches.schema, res));
252-
253241
AssertTablesEqual(*exp_table, *res_table,
254242
/*same_chunk_layout=*/true, /*flatten=*/true);
255243
}
@@ -270,8 +258,7 @@ void DoInvalidPlanTest(const BatchesWithSchema& l_batches,
270258
const AsofJoinNodeOptions& join_options,
271259
const std::string& expected_error_str,
272260
bool fail_on_plan_creation = false) {
273-
ExecContext exec_ctx;
274-
ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make(&exec_ctx));
261+
ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make(*threaded_exec_context()));
275262

276263
Declaration join{"asofjoin", join_options};
277264
join.inputs.emplace_back(Declaration{

‎cpp/src/arrow/compute/exec/benchmark_util.cc‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,6 @@ namespace compute {
3535
// calling InputFinished and InputReceived.
3636

3737
Status BenchmarkIsolatedNodeOverhead(benchmark::State& state,
38-
arrow::compute::ExecContext ctx,
3938
arrow::compute::Expression expr, int32_t num_batches,
4039
int32_t batch_size,
4140
arrow::compute::BatchesWithSchema data,
@@ -46,7 +45,7 @@ Status BenchmarkIsolatedNodeOverhead(benchmark::State& state,
4645
AsyncGenerator<std::optional<arrow::compute::ExecBatch>> sink_gen;
4746

4847
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::compute::ExecPlan> plan,
49-
arrow::compute::ExecPlan::Make(&ctx));
48+
arrow::compute::ExecPlan::Make());
5049
// Source and sink nodes have no effect on the benchmark.
5150
// Used for dummy purposes as they are referenced in InputReceived and InputFinished.
5251
ARROW_ASSIGN_OR_RAISE(arrow::compute::ExecNode * source_node,
@@ -113,13 +112,13 @@ Status BenchmarkIsolatedNodeOverhead(benchmark::State& state,
113112
// a source -> node_declarations -> sink sequence.
114113

115114
Status BenchmarkNodeOverhead(
116-
benchmark::State& state, arrow::compute::ExecContext ctx, int32_t num_batches,
117-
int32_t batch_size, arrow::compute::BatchesWithSchema data,
115+
benchmark::State& state, int32_t num_batches, int32_t batch_size,
116+
arrow::compute::BatchesWithSchema data,
118117
std::vector<arrow::compute::Declaration>& node_declarations) {
119118
for (auto _ : state) {
120119
state.PauseTiming();
121120
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::compute::ExecPlan> plan,
122-
arrow::compute::ExecPlan::Make(&ctx));
121+
arrow::compute::ExecPlan::Make());
123122
AsyncGenerator<std::optional<arrow::compute::ExecBatch>> sink_gen;
124123
arrow::compute::Declaration source = arrow::compute::Declaration(
125124
{"source",

‎cpp/src/arrow/compute/exec/benchmark_util.h‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,11 @@ namespace arrow {
2929

3030
namespace compute {
3131

32-
Status BenchmarkNodeOverhead(benchmark::State& state, arrow::compute::ExecContext ctx,
33-
int32_t num_batches, int32_t batch_size,
34-
arrow::compute::BatchesWithSchema data,
32+
Status BenchmarkNodeOverhead(benchmark::State& state, int32_t num_batches,
33+
int32_t batch_size, arrow::compute::BatchesWithSchema data,
3534
std::vector<arrow::compute::Declaration>& node_declarations);
3635

3736
Status BenchmarkIsolatedNodeOverhead(benchmark::State& state,
38-
arrow::compute::ExecContext ctx,
3937
arrow::compute::Expression expr, int32_t num_batches,
4038
int32_t batch_size,
4139
arrow::compute::BatchesWithSchema data,

0 commit comments

Comments
 (0)