Skip to content

GH-47393: [C++][Acero] Support for multi threaded input: SortedMergeNode#47394

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
X-Lemon-X wants to merge 5 commits into
base:main
Choose a base branch
from
Open

GH-47393: [C++][Acero] Support for multi threaded input: SortedMergeNode #47394

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
54f614a
[C++][sorted_merge][sequencer][threaded] Added threading capabilitie…
Aug 21, 2025
0f234fb
Merge branch 'main' of github.com:X-Lemon-X/arrow into sorted_merge_t…
Aug 21, 2025
6eb8c1c
[C++][sorted_merge][sequencer][threaded][tests] added test for Sorted…
Aug 21, 2025
8df5569
[C++][sorted_merge][sequencer][threaded] Fix for SortedMergeNode when…
Aug 21, 2025
ae70393
[C++][sorted_merge][sequencer][threaded][tests] Disabled test that re…
Aug 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 2 additions & 53 deletions cpp/src/arrow/acero/asof_join_node_test.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -1373,56 +1373,6 @@ TRACED_TEST(AsofJoinTest, TestUnorderedOnKey, {
schema({field("time", int64()), field("key", int32()), field("r0_v0", float64())}));
})

struct BackpressureCounters {
std::atomic<int32_t> pause_count = 0;
std::atomic<int32_t> resume_count = 0;
};

struct BackpressureCountingNodeOptions : public ExecNodeOptions {
BackpressureCountingNodeOptions(BackpressureCounters* counters) : counters(counters) {}

BackpressureCounters* counters;
};

struct BackpressureCountingNode : public MapNode {
static constexpr const char* kKindName = "BackpressureCountingNode";
static constexpr const char* kFactoryName = "backpressure_count";

static void Register() {
auto exec_reg = default_exec_factory_registry();
if (!exec_reg->GetFactory(kFactoryName).ok()) {
ASSERT_OK(exec_reg->AddFactory(kFactoryName, BackpressureCountingNode::Make));
}
}

BackpressureCountingNode(ExecPlan* plan, std::vector<ExecNode*> inputs,
std::shared_ptr<Schema> output_schema,
const BackpressureCountingNodeOptions& options)
: MapNode(plan, inputs, output_schema), counters(options.counters) {}

static Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
const ExecNodeOptions& options) {
RETURN_NOT_OK(ValidateExecNodeInputs(plan, inputs, 1, kKindName));
auto bp_options = static_cast<const BackpressureCountingNodeOptions&>(options);
return plan->EmplaceNode<BackpressureCountingNode>(
plan, inputs, inputs[0]->output_schema(), bp_options);
}

const char* kind_name() const override { return kKindName; }
Result<ExecBatch> ProcessBatch(ExecBatch batch) override { return batch; }

void PauseProducing(ExecNode* output, int32_t counter) override {
++counters->pause_count;
inputs()[0]->PauseProducing(this, counter);
}
void ResumeProducing(ExecNode* output, int32_t counter) override {
++counters->resume_count;
inputs()[0]->ResumeProducing(this, counter);
}

BackpressureCounters* counters;
};

AsyncGenerator<std::optional<ExecBatch>> GetGen(
AsyncGenerator<std::optional<ExecBatch>> gen) {
return gen;
Expand DownExpand Up@@ -1453,8 +1403,7 @@ void TestBackpressure(BatchesMaker maker, int batch_size, int num_l_batches,
ASSERT_OK_AND_ASSIGN(auto r0_batches, make_shift(num_r0_batches, r0_schema, 1));
ASSERT_OK_AND_ASSIGN(auto r1_batches, make_shift(num_r1_batches, r1_schema, 2));

BackpressureCountingNode::Register();
RegisterTestNodes(); // for GatedNode
RegisterTestNodes(); // for GatedNode and BackpressureCountingNode

struct BackpressureSourceConfig {
std::string name_prefix;
Expand DownExpand Up@@ -1499,7 +1448,7 @@ void TestBackpressure(BatchesMaker maker, int batch_size, int num_l_batches,
std::make_shared<BackpressureCountingNodeOptions>(&bp_counters[i]));
std::shared_ptr<ExecNodeOptions> options = bp_options.back();
std::vector<Declaration::Input> bp_in = {src_decls.back()};
Declaration bp_decl = {BackpressureCountingNode::kFactoryName, bp_in,
Declaration bp_decl = {BackpressureCountingNodeOptions::kName, bp_in,
std::move(options)};
if (config.is_gated) {
bp_decl = {std::string{GatedNodeOptions::kName}, {bp_decl}, gate_options};
Expand Down
110 changes: 94 additions & 16 deletions cpp/src/arrow/acero/sorted_merge_node.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -17,13 +17,16 @@

#include <any>
#include <atomic>
#include <memory>
#include <mutex>
#include <sstream>
#include <thread>
#include <tuple>
#include <unordered_map>
#include <utility>
#include <vector>

#include "arrow/acero/accumulation_queue.h"
#include "arrow/acero/concurrent_queue_internal.h"
#include "arrow/acero/exec_plan.h"
#include "arrow/acero/exec_plan_internal.h"
Expand All@@ -34,6 +37,7 @@
#include "arrow/acero/util.h"
#include "arrow/array/builder_base.h"
#include "arrow/result.h"
#include "arrow/testing/process.h"
#include "arrow/type_fwd.h"
#include "arrow/util/logging_internal.h"

Expand DownExpand Up@@ -99,15 +103,17 @@ class BackpressureController : public BackpressureControl {

/// InputState corresponds to an input. Input record batches are queued up in InputState
/// until processed and turned into output record batches.
class InputState {
class InputState : public util::SerialSequencingQueue::Processor {
public:
InputState(size_t index, BackpressureHandler handler,
const std::shared_ptr<arrow::Schema>& schema, const int time_col_index)
: index_(index),
queue_(std::move(handler)),
schema_(schema),
time_col_index_(time_col_index),
time_type_id_(schema_->fields()[time_col_index_]->type()->id()) {}
time_type_id_(schema_->fields()[time_col_index_]->type()->id()) {
sequencer_ = util::SerialSequencingQueue::Make(this);
}

template <typename PtrType>
static arrow::Result<PtrType> Make(size_t index, arrow::acero::ExecNode* input,
Expand DownExpand Up@@ -211,6 +217,15 @@ class InputState {
return arrow::Status::OK();
}

Status InsertBatch(ExecBatch batch) {
return sequencer_->InsertBatch(std::move(batch));
}

arrow::Status Process(arrow::ExecBatch batch) final {
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<RecordBatch> rb, batch.ToRecordBatch(schema_));
return Push(rb);
}

const std::shared_ptr<arrow::Schema>& get_schema() const { return schema_; }

void set_total_batches(int n) { total_batches_ = n; }
Expand All@@ -219,6 +234,7 @@ class InputState {
size_t index_;
// Pending record batches. The latest is the front. Batches cannot be empty.
BackpressureConcurrentQueue<std::shared_ptr<arrow::RecordBatch>> queue_;
std::unique_ptr<util::SerialSequencingQueue> sequencer_;
// Schema associated with the input
std::shared_ptr<arrow::Schema> schema_;
// Total number of batches (only int because InputFinished uses int)
Expand DownExpand Up@@ -301,6 +317,9 @@ class SortedMergeNode : public ExecNode {
"was: ",
schema->ToString(), " got schema: ", input->output_schema()->ToString());
}
if (input->ordering().is_unordered()) {
return Status::Invalid("Input have to be ordered");
}
}

const auto& order_options =
Expand All@@ -322,16 +341,16 @@ class SortedMergeNode : public ExecNode {
arrow::Status Init() override {
ARROW_CHECK(ordering_.sort_keys().size() == 1) << "Only one sort key supported";

const auto& sort_key = ordering_.sort_keys()[0];
if (sort_key.order != arrow::compute::SortOrder::Ascending) {
return Status::NotImplemented("Only ascending sort order is supported");
}

auto inputs = this->inputs();
for (size_t i = 0; i < inputs.size(); i++) {
ExecNode* input = inputs[i];
const auto& schema = input->output_schema();

const auto& sort_key = ordering_.sort_keys()[0];
if (sort_key.order != arrow::compute::SortOrder::Ascending) {
return Status::NotImplemented("Only ascending sort order is supported");
}

const FieldRef& ref = sort_key.target;
auto match_res = ref.FindOne(*schema);
if (!match_res.ok()) {
Expand All@@ -353,13 +372,12 @@ class SortedMergeNode : public ExecNode {
arrow::ExecBatch batch) override {
ARROW_DCHECK(std_has(inputs_, input));
const size_t index = std_find(inputs_, input) - inputs_.begin();
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<RecordBatch> rb,
batch.ToRecordBatch(output_schema_));

// Push into the queue. Note that we don't need to lock since
// Sequencer menages incoming batches order. then pushes it into
// the queue. Note that we don't need to lock since
// InputState's ConcurrentQueue manages locking
input_counter[index] += rb->num_rows();
ARROW_RETURN_NOT_OK(state[index]->Push(rb));
input_counter[index] += batch.length;
ARROW_RETURN_NOT_OK(state[index]->InsertBatch(std::move(batch)));
PushTask(kNewTask);
return Status::OK();
}
Expand DownExpand Up@@ -406,6 +424,24 @@ class SortedMergeNode : public ExecNode {
return Status::OK();
}

arrow::Status StopProducing() override {
#ifdef ARROW_ENABLE_THREADING
Future<> to_finish;
{
std::lock_guard<std::mutex> lg(backpressure_mutex_);
if (!backpressure_future_.is_finished()) {
to_finish = backpressure_future_;
backpressure_future_ = Future<>::MakeFinished();
}
}
if (to_finish.is_valid()) {
to_finish.MarkFinished();
}
#endif

return ExecNode::StopProducing();
}

arrow::Status StopProducingImpl() override {
#ifdef ARROW_ENABLE_THREADING
process_queue.Clear();
Expand All@@ -415,8 +451,39 @@ class SortedMergeNode : public ExecNode {
}

// handled by the backpressure controller
void PauseProducing(arrow::acero::ExecNode* output, int32_t counter) override {}
void ResumeProducing(arrow::acero::ExecNode* output, int32_t counter) override {}
void PauseProducing(ExecNode* output, int32_t counter) override {
#ifdef ARROW_ENABLE_THREADING
std::lock_guard<std::mutex> lg(backpressure_mutex_);
if (counter <= last_backpressure_counter_) {
return;
}
last_backpressure_counter_ = counter;
if (!backpressure_future_.is_finished()) {
// Could happen if we get something like Pause(1) Pause(3) Resume(2)
return;
}
backpressure_future_ = Future<>::Make();
#endif
}

void ResumeProducing(ExecNode* output, int32_t counter) override {
#ifdef ARROW_ENABLE_THREADING
Future<> to_finish;
{
std::lock_guard<std::mutex> lg(backpressure_mutex_);
if (counter <= last_backpressure_counter_) {
return;
}
last_backpressure_counter_ = counter;
if (backpressure_future_.is_finished()) {
return;
}
to_finish = backpressure_future_;
backpressure_future_ = Future<>::MakeFinished();
}
to_finish.MarkFinished();
#endif
}

protected:
std::string ToStringExtra(int indent) const override {
Expand DownExpand Up@@ -587,6 +654,13 @@ class SortedMergeNode : public ExecNode {
#ifdef ARROW_ENABLE_THREADING
void EmitBatches() {
while (true) {
Future<> to_wait;
{
std::lock_guard<std::mutex> lg(backpressure_mutex_);
to_wait = backpressure_future_;
}
to_wait.Wait();

// Implementation note: If the queue is empty, we will block here
if (process_queue.WaitAndPop() == kPoisonPill) {
EndFromProcessThread();
Expand All@@ -613,7 +687,7 @@ class SortedMergeNode : public ExecNode {
std::atomic<bool> cleanup_started{false};

// Backpressure counter common to all input states
std::atomic<int32_t> backpressure_counter;
std::atomic<int32_t> backpressure_counter{0};

std::atomic<int32_t> batches_produced{0};

Expand All@@ -623,6 +697,10 @@ class SortedMergeNode : public ExecNode {
// Once StartProducing is called, we initialize this thread to poll the
// input states and emit batches
std::thread process_thread;

std::mutex backpressure_mutex_;
std::atomic<int32_t> last_backpressure_counter_{0};
Future<> backpressure_future_ = Future<>::MakeFinished();
#endif
arrow::Future<> process_task;

Expand All@@ -642,4 +720,4 @@ void RegisterSortedMergeNode(ExecFactoryRegistry* registry) {
}
} // namespace internal

} // namespace arrow::acero
} // namespace arrow::acero
Loading