From 7dc8fd2df1548c4832d0ded72539bc403d8cc49a Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 10 May 2026 10:03:14 +0000 Subject: [PATCH] GH-45847: [C++][Acero] Two-pass merge in GroupByNode to amortise kernel resizes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The hash-aggregate `Merge()` step was a single pass that, for every contributing thread, fed its uniques into state0->grouper and then iterated every kernel — calling `kernel->resize(state0->grouper->num_groups())` once per (thread, kernel) pair. With N threads and K kernels that is N*K resize calls, while only K are needed (the grouper grows monotonically so the only resize that does real work is the final one). This commit implements the two-pass algorithm noted in the existing in-code TODO at lines 285-289: Pass 1 Compute every contributing thread's transposition into state0->grouper. After this pass the grouper is fully grown. Pass 2 For each kernel, resize state0->agg_states[k] once to the final group count, then merge every contributing thread's kernel state in turn. Semantically equivalent to the previous code: the grouper grows monotonically across Consume() calls, so each thread's transposition remains a valid index into state0's final kernel state. Per-row output (count, sum, etc.) is unchanged. Verified with a join-then-aggregate reproducer of GH-45847 that exercises the multi-threaded merge path on a worst-case unique-key inner join (build x probe = 64 x 64 batches of 32 768 rows = ~2M groups across 4 worker threads). out_rows and per-group counts match the unmodified Merge() exactly. The catastrophic 9000x slowdown reported in GH-45847 itself is largely addressed by the JoinResultMaterialize::Flush task group that landed in #45918; this change is the complementary aggregator-side fix called out by the in-code TODO and reduces resize() call count in Merge() from O(num_threads * num_kernels) to O(num_kernels). Reproducer: cpp/src/arrow/acero/repro_45847.cc Standalone program that wires source(probe) + source(build) -> hashjoin -> aggregate(hash_count) and times the run. Validates the output against the analytical expected count so it can be used to check correctness of further changes. https://claude.ai/code/session_01BX3n7pizaHw9yVpSetFmXo --- cpp/src/arrow/acero/groupby_aggregate_node.cc | 60 +++--- cpp/src/arrow/acero/repro_45847.cc | 198 ++++++++++++++++++ 2 files changed, 234 insertions(+), 24 deletions(-) create mode 100644 cpp/src/arrow/acero/repro_45847.cc diff --git a/cpp/src/arrow/acero/groupby_aggregate_node.cc b/cpp/src/arrow/acero/groupby_aggregate_node.cc index 5c62ddf15a25..a056ad8c3bbb 100644 --- a/cpp/src/arrow/acero/groupby_aggregate_node.cc +++ b/cpp/src/arrow/acero/groupby_aggregate_node.cc @@ -257,40 +257,52 @@ Status GroupByNode::Merge() { START_COMPUTE_SPAN(span, "Merge", {{"group_by", ToStringExtra(0)}, {"node.label", label()}}); ThreadLocalState* state0 = &local_states_[0]; + + // Pass 1: gather every contributing thread's transposition into state0's + // grouper. Each Consume() call may grow state0->grouper; we finish growing + // it before touching any kernel state. + std::vector contributing_threads; + contributing_threads.reserve(local_states_.size()); + std::vector transpositions(local_states_.size()); for (size_t i = 1; i < local_states_.size(); ++i) { ThreadLocalState* state = &local_states_[i]; if (!state->grouper) { continue; } - ARROW_ASSIGN_OR_RAISE(ExecBatch other_keys, state->grouper->GetUniques()); - ARROW_ASSIGN_OR_RAISE(Datum transposition, + ARROW_ASSIGN_OR_RAISE(transpositions[i], state0->grouper->Consume(ExecSpan(other_keys))); state->grouper.reset(); + contributing_threads.push_back(i); + } + + // Pass 2: resize each kernel state exactly once at the final group count, + // then merge every contributing thread's state into state0's. This + // replaces O(num_threads * num_kernels) resize() calls with O(num_kernels), + // which avoids quadratic behaviour when many threads each grow num_groups + // by a small amount. + const int64_t final_num_groups = state0->grouper->num_groups(); + for (size_t span_i = 0; span_i < agg_kernels_.size(); ++span_i) { + arrow::util::tracing::Span span_item; + START_COMPUTE_SPAN( + span_item, aggs_[span_i].function, + {{"function.name", aggs_[span_i].function}, + {"function.options", + aggs_[span_i].options ? aggs_[span_i].options->ToString() : ""}, + {"function.kind", std::string(kind_name()) + "::Merge"}}); + + auto ctx = plan_->query_context()->exec_context(); + KernelContext batch_ctx{ctx}; + DCHECK(state0->agg_states[span_i]); + batch_ctx.SetState(state0->agg_states[span_i].get()); + + RETURN_NOT_OK(agg_kernels_[span_i]->resize(&batch_ctx, final_num_groups)); - for (size_t span_i = 0; span_i < agg_kernels_.size(); ++span_i) { - arrow::util::tracing::Span span_item; - START_COMPUTE_SPAN( - span_item, aggs_[span_i].function, - {{"function.name", aggs_[span_i].function}, - {"function.options", - aggs_[span_i].options ? aggs_[span_i].options->ToString() : ""}, - {"function.kind", std::string(kind_name()) + "::Merge"}}); - - auto ctx = plan_->query_context()->exec_context(); - KernelContext batch_ctx{ctx}; - DCHECK(state0->agg_states[span_i]); - batch_ctx.SetState(state0->agg_states[span_i].get()); - - // XXX this resizes each KernelState (state0->agg_states[span_i]) multiple times. - // An alternative would be a two-pass algorithm: - // 1. Compute all transpositions (one per local state) and the final number of - // groups. - // 2. Process all agg kernels, resizing each KernelState only once. - RETURN_NOT_OK( - agg_kernels_[span_i]->resize(&batch_ctx, state0->grouper->num_groups())); + for (size_t i : contributing_threads) { + ThreadLocalState* state = &local_states_[i]; RETURN_NOT_OK(agg_kernels_[span_i]->merge( - &batch_ctx, std::move(*state->agg_states[span_i]), *transposition.array())); + &batch_ctx, std::move(*state->agg_states[span_i]), + *transpositions[i].array())); state->agg_states[span_i].reset(); } } diff --git a/cpp/src/arrow/acero/repro_45847.cc b/cpp/src/arrow/acero/repro_45847.cc new file mode 100644 index 000000000000..dca2b3d3bab6 --- /dev/null +++ b/cpp/src/arrow/acero/repro_45847.cc @@ -0,0 +1,198 @@ +// Reproducer for apache/arrow#45847. +// +// Reproduces the topology reported by uchenily: +// +// aggregate +// | +// hashjoin +// / \ +// source_0 source_1 +// (probe) (build) +// +// The reporter sees ~9000x slowdown and single-core utilization at large +// build_batch counts. This program runs the same topology with controllable +// dimensions and prints wall-clock timings and effective parallelism. +// +// Usage: +// repro_45847 [build_batches] [probe_batches] [batch_size] [threads] [runs] +// +// Defaults match the issue: batch_size=1<<15, threads=hw, runs=3. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/acero/exec_plan.h" +#include "arrow/acero/options.h" +#include "arrow/api.h" +#include "arrow/compute/exec.h" +#include "arrow/compute/initialize.h" +#include "arrow/util/thread_pool.h" + +using arrow::Field; +using arrow::FieldRef; +using arrow::Int64Builder; +using arrow::Schema; +using arrow::Status; +using arrow::Table; +using arrow::TableBatchReader; +using arrow::compute::Aggregate; +using arrow::acero::AggregateNodeOptions; +using arrow::acero::Declaration; +using arrow::acero::DeclarationToTable; +using arrow::acero::HashJoinNodeOptions; +using arrow::acero::JoinKeyCmp; +using arrow::acero::JoinType; +using arrow::acero::QueryOptions; +using arrow::acero::TableSourceNodeOptions; + +#define ABORT_NOT_OK(expr) \ + do { \ + auto _st = (expr); \ + if (!_st.ok()) { \ + std::cerr << "FATAL: " << _st.ToString() << "\n"; \ + std::abort(); \ + } \ + } while (0) + +namespace { + +// Build a table of `num_batches * batch_size` rows with two int64 columns: +// key: start + i (each row gets a distinct key within [start, start+rows)) +// payload: i +// The returned Table is single-chunk; TableSourceNodeOptions will slice it +// into per-batch_size morsels. +arrow::Result> MakeKeyedTable(int64_t num_batches, + int64_t batch_size, + int64_t key_start, + const std::string& key_name, + const std::string& payload_name) { + const int64_t total_rows = num_batches * batch_size; + Int64Builder key_b, pay_b; + ARROW_RETURN_NOT_OK(key_b.Resize(total_rows)); + ARROW_RETURN_NOT_OK(pay_b.Resize(total_rows)); + for (int64_t i = 0; i < total_rows; ++i) { + key_b.UnsafeAppend(key_start + i); + pay_b.UnsafeAppend(i); + } + std::shared_ptr key_arr, pay_arr; + ARROW_RETURN_NOT_OK(key_b.Finish(&key_arr)); + ARROW_RETURN_NOT_OK(pay_b.Finish(&pay_arr)); + auto schema = arrow::schema({arrow::field(key_name, arrow::int64()), + arrow::field(payload_name, arrow::int64())}); + return Table::Make(schema, {key_arr, pay_arr}, total_rows); +} + +struct Args { + int64_t build_batches = 32; + int64_t probe_batches = 1; + int64_t batch_size = 1 << 15; // 32 768 + int threads = 0; // 0 -> hardware + int runs = 3; + int num_aggs = 1; // number of hash_count aggregations stacked on the same key +}; + +Args ParseArgs(int argc, char** argv) { + Args a; + if (argc > 1) a.build_batches = std::atoll(argv[1]); + if (argc > 2) a.probe_batches = std::atoll(argv[2]); + if (argc > 3) a.batch_size = std::atoll(argv[3]); + if (argc > 4) a.threads = std::atoi(argv[4]); + if (argc > 5) a.runs = std::atoi(argv[5]); + if (argc > 6) a.num_aggs = std::atoi(argv[6]); + return a; +} + +} // namespace + +int main(int argc, char** argv) { + ABORT_NOT_OK(arrow::compute::Initialize()); + Args args = ParseArgs(argc, argv); + if (args.threads > 0) { + ABORT_NOT_OK(arrow::SetCpuThreadPoolCapacity(args.threads)); + } + const int reported_threads = + arrow::internal::GetCpuThreadPool()->GetCapacity(); + + std::cout << "build_batches=" << args.build_batches + << " probe_batches=" << args.probe_batches + << " batch_size=" << args.batch_size + << " threads=" << reported_threads + << " runs=" << args.runs << "\n"; + + // Both sides start at key 0, so probe rows match build rows 1:1 within the + // overlap [0, min(build_rows, probe_rows)). This is intentionally cheap on + // the join side so the cost we measure is the aggregator's. + auto build_table = MakeKeyedTable(args.build_batches, args.batch_size, + /*key_start=*/0, "rk", "rp") + .ValueOrDie(); + auto probe_table = MakeKeyedTable(args.probe_batches, args.batch_size, + /*key_start=*/0, "lk", "lp") + .ValueOrDie(); + + std::cout << "build_rows=" << build_table->num_rows() + << " probe_rows=" << probe_table->num_rows() << "\n"; + + for (int run = 0; run < args.runs; ++run) { + Declaration probe_src{ + "table_source", TableSourceNodeOptions{probe_table, args.batch_size}}; + Declaration build_src{ + "table_source", TableSourceNodeOptions{build_table, args.batch_size}}; + + HashJoinNodeOptions join_opts{JoinType::INNER, + /*left_keys=*/{FieldRef("lk")}, + /*right_keys=*/{FieldRef("rk")}}; + Declaration join{"hashjoin", + {std::move(probe_src), std::move(build_src)}, + std::move(join_opts)}; + + std::vector aggs; + for (int i = 0; i < args.num_aggs; ++i) { + aggs.push_back(Aggregate{"hash_count", nullptr, FieldRef("lk"), + "cnt" + std::to_string(i)}); + } + AggregateNodeOptions agg_opts{std::move(aggs), /*keys=*/{FieldRef("lk")}}; + Declaration agg{"aggregate", {std::move(join)}, std::move(agg_opts)}; + + auto t0 = std::chrono::steady_clock::now(); + auto table_or = DeclarationToTable(std::move(agg), /*use_threads=*/true); + auto t1 = std::chrono::steady_clock::now(); + ABORT_NOT_OK(table_or.status()); + double secs = std::chrono::duration(t1 - t0).count(); + auto table = table_or.ValueOrDie(); + + // Validate: each row's cnt0 must equal 1 (unique keys, 1:1 inner join). + // The aggregator runs Merge across all threads on every run; if + // transposition or kernel-state ordering breaks, this will fail. + int64_t sum_cnt = 0; + auto cnt_chunked = table->GetColumnByName("cnt0"); + for (int c = 0; c < cnt_chunked->num_chunks(); ++c) { + auto arr = + std::static_pointer_cast(cnt_chunked->chunk(c)); + for (int64_t i = 0; i < arr->length(); ++i) { + if (arr->Value(i) != 1) { + std::cerr << "FATAL: chunk=" << c << " row=" << i + << " cnt=" << arr->Value(i) << " (expected 1)\n"; + std::abort(); + } + sum_cnt += arr->Value(i); + } + } + const int64_t expected = + std::min(build_table->num_rows(), probe_table->num_rows()); + if (sum_cnt != expected) { + std::cerr << "FATAL: sum_cnt=" << sum_cnt << " expected=" << expected + << "\n"; + std::abort(); + } + std::cout << "run=" << run << " seconds=" << secs + << " out_rows=" << table->num_rows() << " sum_cnt=" << sum_cnt + << "\n"; + } + return 0; +}