From 168d8e08abd089b7108b79e420b045308f8cba3a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Wed, 2 Sep 2026 14:14:43 +0000 Subject: [PATCH 1/5] [C++][Acero] Fix and expose time normalization --- cpp/src/arrow/acero/time_series_util.cc | 7 ------- cpp/src/arrow/acero/time_series_util.h | 15 +++++++++++++-- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/cpp/src/arrow/acero/time_series_util.cc b/cpp/src/arrow/acero/time_series_util.cc index 60f9044d7bf7..9fbf54666373 100644 --- a/cpp/src/arrow/acero/time_series_util.cc +++ b/cpp/src/arrow/acero/time_series_util.cc @@ -22,13 +22,6 @@ namespace arrow::acero { -template ::value, bool>> -inline uint64_t NormalizeTime(T t) { - uint64_t bias = - std::is_signed::value ? static_cast(1) << (8 * sizeof(T) - 1) : 0; - return t < 0 ? static_cast(t + bias) : static_cast(t); -} - uint64_t GetTime(const RecordBatch* batch, Type::type time_type, int col, uint64_t row) { #define LATEST_VAL_CASE(id, val) \ case Type::id: { \ diff --git a/cpp/src/arrow/acero/time_series_util.h b/cpp/src/arrow/acero/time_series_util.h index 97707f43bf20..852a647002ad 100644 --- a/cpp/src/arrow/acero/time_series_util.h +++ b/cpp/src/arrow/acero/time_series_util.h @@ -17,14 +17,25 @@ #pragma once +#include +#include + #include "arrow/record_batch.h" #include "arrow/type_traits.h" namespace arrow::acero { // normalize the value to unsigned 64-bits while preserving ordering of values -template ::value, bool> = true> -uint64_t NormalizeTime(T t); +template && !std::is_same_v, bool> = true> +uint64_t NormalizeTime(T t) { + using U = std::make_unsigned_t; + U normalized = static_cast(t); + if constexpr (std::is_signed_v) { + normalized ^= U{1} << (std::numeric_limits::digits - 1); + } + return static_cast(normalized); +} uint64_t GetTime(const RecordBatch* batch, Type::type time_type, int col, uint64_t row); From ade8e988a9427fa7bbe50581714134c2aa515137 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Wed, 2 Sep 2026 10:12:58 +0000 Subject: [PATCH 2/5] [C++][Acero] Rework SortedMerge execution and flow control --- cpp/src/arrow/acero/sorted_merge_node.cc | 971 ++++++++++-------- cpp/src/arrow/acero/sorted_merge_node_test.cc | 162 ++- 2 files changed, 725 insertions(+), 408 deletions(-) diff --git a/cpp/src/arrow/acero/sorted_merge_node.cc b/cpp/src/arrow/acero/sorted_merge_node.cc index 43f5b7b930ab..55806fb153d5 100644 --- a/cpp/src/arrow/acero/sorted_merge_node.cc +++ b/cpp/src/arrow/acero/sorted_merge_node.cc @@ -15,36 +15,34 @@ // specific language governing permissions and limitations // under the License. -#include +#include #include +#include +#include +#include #include +#include #include -#include -#include -#include +#include +#include #include -#include "arrow/acero/concurrent_queue_internal.h" +#include "arrow/acero/accumulation_queue.h" #include "arrow/acero/exec_plan.h" #include "arrow/acero/exec_plan_internal.h" #include "arrow/acero/options.h" #include "arrow/acero/query_context.h" #include "arrow/acero/time_series_util.h" -#include "arrow/acero/unmaterialized_table_internal.h" #include "arrow/acero/util.h" #include "arrow/array/builder_base.h" +#include "arrow/array/util.h" #include "arrow/result.h" #include "arrow/type_fwd.h" +#include "arrow/type_traits.h" +#include "arrow/util/checked_cast.h" #include "arrow/util/logging_internal.h" namespace { -template -struct Defer { - Callable callable; - explicit Defer(Callable callable_) : callable(std::move(callable_)) {} - ~Defer() noexcept { callable(); } -}; - std::vector GetInputLabels( const arrow::acero::ExecNode::NodeVector& inputs) { std::vector labels(inputs.size()); @@ -54,205 +52,261 @@ std::vector GetInputLabels( return labels; } -template -inline typename T::const_iterator std_find(const T& container, const V& val) { - return std::find(container.begin(), container.end(), val); -} - -template -inline bool std_has(const T& container, const V& val) { - return container.end() != std_find(container, val); -} - } // namespace namespace arrow::acero { namespace { -// Each slice is associated with a single input source, so we only need 1 record -// batch per slice -using SingleRecordBatchSliceBuilder = arrow::acero::UnmaterializedSliceBuilder<1>; -using SingleRecordBatchCompositeTable = arrow::acero::UnmaterializedCompositeTable<1>; - using row_index_t = uint64_t; using time_unit_t = uint64_t; using col_index_t = int; +using Task = util::SequencingQueue::Task; + +template +Result> ReadTimeValue(const Datum& value, int64_t row) { + using ArrowType = typename TypeIdTraits::Type; + using CType = typename TypeTraits::CType; + using ScalarType = typename TypeTraits::ScalarType; + + if (value.is_scalar()) { + const auto& scalar = + ::arrow::internal::checked_cast(*value.scalar()); + if (!scalar.is_valid) { + return std::nullopt; + } + return NormalizeTime(static_cast(scalar.value)); + } + if (value.is_array()) { + ArraySpan array(*value.array()); + if (array.IsNull(row)) { + return std::nullopt; + } + return NormalizeTime(array.GetValues(1)[row]); + } + return Status::Invalid("SortedMerge sort key must be an array or scalar, but got ", + ::arrow::ToString(value.kind())); +} -constexpr bool kNewTask = true; -constexpr bool kPoisonPill = false; +Result> ReadTimeValue(const Datum& value, int64_t row) { + switch (value.type()->id()) { +#define SORTED_MERGE_TIME_CASE(ID) \ + case Type::ID: \ + return ReadTimeValue(value, row) + SORTED_MERGE_TIME_CASE(INT8); + SORTED_MERGE_TIME_CASE(INT16); + SORTED_MERGE_TIME_CASE(INT32); + SORTED_MERGE_TIME_CASE(INT64); + SORTED_MERGE_TIME_CASE(UINT8); + SORTED_MERGE_TIME_CASE(UINT16); + SORTED_MERGE_TIME_CASE(UINT32); + SORTED_MERGE_TIME_CASE(UINT64); + SORTED_MERGE_TIME_CASE(DATE32); + SORTED_MERGE_TIME_CASE(DATE64); + SORTED_MERGE_TIME_CASE(TIME32); + SORTED_MERGE_TIME_CASE(TIME64); + SORTED_MERGE_TIME_CASE(TIMESTAMP); +#undef SORTED_MERGE_TIME_CASE + default: + return Status::Invalid("Unsupported SortedMerge sort-key type ", + value.type()->ToString()); + } +} -class BackpressureController : public BackpressureControl { - public: - BackpressureController(ExecNode* node, ExecNode* output, - std::atomic& backpressure_counter) - : node_(node), output_(output), backpressure_counter_(backpressure_counter) {} +bool TimeIsLater(const std::optional& left, + const std::optional& right, + compute::NullPlacement null_placement) { + if (left && right) { + return *left > *right; + } + if (!left && !right) { + return false; + } + return null_placement == compute::NullPlacement::AtStart ? left.has_value() + : !left.has_value(); +} - void Pause() override { node_->PauseProducing(output_, ++backpressure_counter_); } - void Resume() override { node_->ResumeProducing(output_, ++backpressure_counter_); } +struct PreparedBatch { + ExecBatch batch; + std::vector> times; +}; - private: - ExecNode* node_; - ExecNode* output_; - std::atomic& backpressure_counter_; +struct SelectionRun { + std::shared_ptr source; + int64_t offset; + int64_t length; }; -/// InputState corresponds to an input. Input record batches are queued up in InputState -/// until processed and turned into output record batches. -class InputState { - public: - InputState(size_t index, BackpressureHandler handler, - const std::shared_ptr& 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()) {} +struct Selection { + std::vector runs; + int64_t length = 0; - template - static arrow::Result Make(size_t index, arrow::acero::ExecNode* input, - arrow::acero::ExecNode* output, - std::atomic& backpressure_counter, - const std::shared_ptr& schema, - const col_index_t time_col_index) { - constexpr size_t low_threshold = 4, high_threshold = 8; - std::unique_ptr backpressure_control = - std::make_unique(input, output, backpressure_counter); - ARROW_ASSIGN_OR_RAISE(auto handler, - BackpressureHandler::Make(low_threshold, high_threshold, - std::move(backpressure_control))); - return PtrType(new InputState(index, std::move(handler), schema, time_col_index)); + void Append(std::shared_ptr source, int64_t offset, int64_t run_length) { + runs.push_back({std::move(source), offset, run_length}); + length += run_length; } - bool IsTimeColumn(col_index_t i) const { - DCHECK_LT(i, schema_->num_fields()); - return (i == time_col_index_); - } + bool empty() const { return runs.empty(); } +}; - // Gets the latest row index, assuming the queue isn't empty - row_index_t GetLatestRow() const { return latest_ref_row_; } +struct FlowAction { + enum class Kind { None, Pause, Resume }; - bool Empty() const { - // cannot be empty if ref row is >0 -- can avoid slow queue lock - // below - if (latest_ref_row_ > 0) { - return false; + void Apply() const { + if (kind == Kind::Pause) { + input->PauseProducing(output, counter); + } else if (kind == Kind::Resume) { + input->ResumeProducing(output, counter); } - return queue_.Empty(); } - size_t index() const { return index_; } + Kind kind = Kind::None; + ExecNode* input = nullptr; + ExecNode* output = nullptr; + int32_t counter = 0; +}; - int total_batches() const { return total_batches_; } +class SortedMergeNode; - // Gets latest batch (precondition: must not be empty) - const std::shared_ptr& GetLatestBatch() const { - return queue_.Front(); - } +/// Sequences and buffers one sorted input. Buffer contents and row position are +/// protected by SortedMergeNode's coordinator mutex; flow-control state uses its own +/// lock because it is touched before and after coordinator work. +class InputState final : public util::SerialSequencingQueue::Processor { + public: + InputState(SortedMergeNode* node, size_t index, ExecNode* input, + col_index_t time_col_index, compute::NullPlacement null_placement); -#define LATEST_VAL_CASE(id, val) \ - case arrow::Type::id: { \ - using T = typename arrow::TypeIdTraits::Type; \ - using CType = typename arrow::TypeTraits::CType; \ - return val(data->GetValues(1)[row]); \ - } + Status InsertBatch(ExecBatch batch); + Status Process(ExecBatch batch) override; - inline time_unit_t GetLatestTime() const { - return GetTime(GetLatestBatch().get(), time_type_id_, time_col_index_, - latest_ref_row_); + // The methods below are called only while the node's coordinator mutex is held. + bool HasData() const { return !batches_.empty(); } + bool AllBatchesReceived() const { + return total_batches_ && received_batches_ == *total_batches_; } + bool Finished() const { return AllBatchesReceived() && !HasData(); } -#undef LATEST_VAL_CASE - - bool Finished() const { return batches_processed_ == total_batches_; } + const std::shared_ptr& GetLatestBatch() const { + DCHECK(HasData()); + return batches_.front(); + } - void Advance(SingleRecordBatchSliceBuilder& builder) { - // Advance the row until a new time is encountered or the record batch - // ends. This will return a range of {-1, -1} and a nullptr if there is - // no input - bool active = - (latest_ref_row_ > 0 /*short circuit the lock on the queue*/) || !queue_.Empty(); + const std::optional& GetLatestTime() const { + return GetLatestBatch()->times[latest_ref_row_]; + } - if (!active) { - return; + bool Advance(const std::optional* upper_bound, int64_t max_length, + Selection* selection) { + DCHECK(HasData()); + DCHECK_GT(max_length, 0); + const row_index_t start = latest_ref_row_; + std::shared_ptr batch = batches_.front(); + const row_index_t rows_in_batch = static_cast(batch->batch.length); + const row_index_t limit = + std::min(rows_in_batch, start + static_cast(max_length)); + + while (latest_ref_row_ < limit && + (upper_bound == nullptr || + !TimeIsLater(GetLatestTime(), *upper_bound, null_placement_))) { + ++latest_ref_row_; + } + DCHECK_GT(latest_ref_row_, start); + selection->Append(batch, static_cast(start), + static_cast(latest_ref_row_ - start)); + if (latest_ref_row_ >= rows_in_batch) { + latest_ref_row_ = 0; + batches_.pop_front(); + return true; } + return false; + } - row_index_t start = latest_ref_row_; - row_index_t end = latest_ref_row_; - time_unit_t startTime = GetLatestTime(); - std::shared_ptr batch = queue_.Front(); - auto rows_in_batch = (row_index_t)batch->num_rows(); - - while (GetLatestTime() == startTime) { - end = ++latest_ref_row_; - if (latest_ref_row_ >= rows_in_batch) { - // hit the end of the batch, need to get the next batch if - // possible. - ++batches_processed_; - latest_ref_row_ = 0; - active &= !queue_.TryPop(); - if (active) { - DCHECK_GT(queue_.Front()->num_rows(), - 0); // empty batches disallowed, sanity check - } - break; - } + Status PushSequenced(std::shared_ptr batch) { + if (total_batches_ && received_batches_ >= *total_batches_) { + return Status::Invalid("SortedMerge input ", index_, + " produced more batches than declared"); + } + ++received_batches_; + if (batch->batch.length > 0) { + batches_.push_back(std::move(batch)); } - builder.AddEntry(batch, start, end); + return Status::OK(); } - arrow::Status Push(const std::shared_ptr& rb) { - if (rb->num_rows() > 0) { - queue_.Push(rb); - } else { - ++batches_processed_; // don't enqueue empty batches, just record - // as processed + Status SetTotal(int total_batches) { + if (total_batches < 0) { + return Status::Invalid("SortedMerge input ", index_, + " reported a negative batch count"); + } + if (total_batches_) { + return *total_batches_ == total_batches + ? Status::OK() + : Status::Invalid("SortedMerge input ", index_, + " changed its total batch count"); + } + if (received_batches_ > total_batches) { + return Status::Invalid("SortedMerge input ", index_, + " declared fewer batches than it produced"); } - return arrow::Status::OK(); + total_batches_ = total_batches; + return Status::OK(); } - const std::shared_ptr& get_schema() const { return schema_; } + void ClearBuffered() { + batches_.clear(); + latest_ref_row_ = 0; + } - void set_total_batches(int n) { total_batches_ = n; } + FlowAction BatchBuffered(); + FlowAction BatchConsumed(); + FlowAction Shutdown(); private: + FlowAction SetUpstreamPausedUnlocked(bool paused); + Result> PrepareBatch(ExecBatch batch); + Status ValidateTime(const std::optional& time); + + static constexpr size_t kLowWatermark = 4; + static constexpr size_t kHighWatermark = 8; + + SortedMergeNode* node_; size_t index_; - // Pending record batches. The latest is the front. Batches cannot be empty. - BackpressureConcurrentQueue> queue_; - // Schema associated with the input - std::shared_ptr schema_; - // Total number of batches (only int because InputFinished uses int) - std::atomic total_batches_{-1}; - // Number of batches processed so far (only int because InputFinished uses - // int) - std::atomic batches_processed_{0}; - // Index of the time col + ExecNode* input_; col_index_t time_col_index_; - // Type id of the time column - arrow::Type::type time_type_id_; - // Index of the latest row reference within; if >0 then queue_ cannot be - // empty Must be < queue_.front()->num_rows() if queue_ is non-empty + compute::NullPlacement null_placement_; + + std::unique_ptr sequencer_; + + std::deque> batches_; + int received_batches_ = 0; + std::optional total_batches_; row_index_t latest_ref_row_ = 0; - // Time of latest row - time_unit_t latest_time_ = std::numeric_limits::lowest(); + + std::optional last_time_; + bool saw_null_ = false; + bool saw_non_null_ = false; + + std::mutex flow_mutex_; + size_t buffered_batches_ = 0; + bool upstream_paused_ = false; + int32_t outgoing_counter_ = 0; + bool shutdown_ = false; }; struct InputStateComparator { - bool operator()(const std::shared_ptr& lhs, - const std::shared_ptr& rhs) const { - // True if lhs is ahead of time of rhs - if (lhs->Finished()) { - return false; - } - if (rhs->Finished()) { - return false; - } - time_unit_t lFirst = lhs->GetLatestTime(); - time_unit_t rFirst = rhs->GetLatestTime(); - return lFirst > rFirst; + explicit InputStateComparator(compute::NullPlacement null_placement) + : null_placement(null_placement) {} + + bool operator()(const InputState* lhs, const InputState* rhs) const { + return TimeIsLater(lhs->GetLatestTime(), rhs->GetLatestTime(), null_placement); } + + compute::NullPlacement null_placement; }; +enum class MergeState { Idle, Running, Terminal }; +enum class OutputGate { Open, Paused, Flushing }; + class SortedMergeNode : public ExecNode { static constexpr int64_t kTargetOutputBatchSize = 1024 * 1024; @@ -262,26 +316,10 @@ class SortedMergeNode : public ExecNode { std::shared_ptr output_schema, arrow::Ordering new_ordering) : ExecNode(plan, inputs, GetInputLabels(inputs), std::move(output_schema)), - ordering_(std::move(new_ordering)), - input_counter(inputs_.size()), - output_counter(inputs_.size()) -#ifdef ARROW_ENABLE_THREADING - , - process_thread() -#endif - { + ordering_(std::move(new_ordering)) { SetLabel("sorted_merge"); } - ~SortedMergeNode() override { - PushTask(kPoisonPill); -#ifdef ARROW_ENABLE_THREADING - if (process_thread.joinable()) { - process_thread.join(); - } -#endif - } - static arrow::Result Make( arrow::acero::ExecPlan* plan, std::vector inputs, const arrow::acero::ExecNodeOptions& options) { @@ -320,18 +358,20 @@ class SortedMergeNode : public ExecNode { const arrow::Ordering& ordering() const override { return ordering_; } arrow::Status Init() override { - ARROW_CHECK(ordering_.sort_keys().size() == 1) << "Only one sort key supported"; + if (ordering_.sort_keys().size() != 1) { + return Status::NotImplemented("SortedMerge supports exactly one sort key"); + } + + 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()) { @@ -340,83 +380,97 @@ class SortedMergeNode : public ExecNode { ARROW_ASSIGN_OR_RAISE(auto match, match_res); ARROW_DCHECK(match.indices().size() == 1); - ARROW_ASSIGN_OR_RAISE(auto input_state, - InputState::Make>( - i, input, this, backpressure_counter, schema, - std::move(match.indices()[0]))); - state.push_back(std::move(input_state)); + state_.push_back(std::make_unique( + this, i, input, std::move(match.indices()[0]), sort_key.null_placement)); } return Status::OK(); } arrow::Status InputReceived(arrow::acero::ExecNode* input, 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 rb, - batch.ToRecordBatch(output_schema_)); - - // Push 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)); - PushTask(kNewTask); - return Status::OK(); - } - - void PushTask(bool ok) { -#ifdef ARROW_ENABLE_THREADING - process_queue.Push(ok); -#else - if (process_task.is_finished()) { - return; + if (terminal_.load()) { + return Status::OK(); } - if (ok == kNewTask) { - PollOnce(); - } else { - EndFromProcessThread(); + auto it = std::find(inputs_.begin(), inputs_.end(), input); + if (it == inputs_.end()) { + return Status::Invalid("SortedMerge received a batch from an unknown input"); } -#endif + const size_t index = static_cast(it - inputs_.begin()); + return state_[index]->InsertBatch(std::move(batch)); } arrow::Status InputFinished(arrow::acero::ExecNode* input, int total_batches) override { - ARROW_DCHECK(std_has(inputs_, input)); + if (terminal_.load()) { + return Status::OK(); + } + auto it = std::find(inputs_.begin(), inputs_.end(), input); + if (it == inputs_.end()) { + return Status::Invalid("SortedMerge received completion from an unknown input"); + } + const size_t index = static_cast(it - inputs_.begin()); + std::optional task; { - std::lock_guard guard(gate); - ARROW_DCHECK(std_has(inputs_, input)); - size_t k = std_find(inputs_, input) - inputs_.begin(); - state.at(k)->set_total_batches(total_batches); + std::lock_guard lock(coordinator_mutex_); + if (merge_state_ == MergeState::Terminal) { + return Status::OK(); + } + ARROW_RETURN_NOT_OK(state_[index]->SetTotal(total_batches)); + task = MaybeStartUnlocked(); } - // Trigger a final process call for stragglers - PushTask(kNewTask); + return task ? std::move(*task)() : Status::OK(); + } + + arrow::Status StartProducing() override { return Status::OK(); } + arrow::Status StopProducingImpl() override { + EnterTerminal(); return Status::OK(); } - arrow::Status StartProducing() override { - ARROW_ASSIGN_OR_RAISE(process_task, plan_->query_context()->BeginExternalTask( - "SortedMergeNode::ProcessThread")); - if (!process_task.is_valid()) { - // Plan has already aborted. Do not start process thread - return Status::OK(); + void PauseProducing(arrow::acero::ExecNode* output, int32_t counter) override { + std::lock_guard lock(coordinator_mutex_); + if (merge_state_ == MergeState::Terminal || counter <= downstream_counter_) { + return; + } + downstream_counter_ = counter; + if (output_gate_ != OutputGate::Flushing) { + output_gate_ = OutputGate::Paused; } -#ifdef ARROW_ENABLE_THREADING - process_thread = std::thread(&SortedMergeNode::StartPoller, this); -#endif - return Status::OK(); } - arrow::Status StopProducingImpl() override { -#ifdef ARROW_ENABLE_THREADING - process_queue.Clear(); -#endif - PushTask(kPoisonPill); - return Status::OK(); + void ResumeProducing(arrow::acero::ExecNode* output, int32_t counter) override { + std::optional task; + { + std::lock_guard lock(coordinator_mutex_); + if (merge_state_ == MergeState::Terminal || counter <= downstream_counter_) { + return; + } + downstream_counter_ = counter; + if (output_gate_ != OutputGate::Flushing) { + output_gate_ = OutputGate::Open; + } + task = MaybeStartUnlocked(); + } + if (task) { + Schedule(std::move(*task), "SortedMergeNode::Resume"); + } + } + + void Schedule(Task task, std::string_view name = "SortedMergeNode::Merge") { + plan()->query_context()->ScheduleTask(std::move(task), name); + } + + Result> OnSequenced(size_t input_index, + std::shared_ptr batch) { + std::lock_guard lock(coordinator_mutex_); + if (merge_state_ == MergeState::Terminal) { + return std::nullopt; + } + ARROW_RETURN_NOT_OK(state_[input_index]->PushSequenced(std::move(batch))); + return MaybeStartUnlocked(); } - // 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 {} + bool IsTerminal() const { return terminal_.load(); } protected: std::string ToStringExtra(int indent) const override { @@ -426,214 +480,319 @@ class SortedMergeNode : public ExecNode { } private: - void EndFromProcessThread(arrow::Status st = arrow::Status::OK()) { - ARROW_CHECK(!cleanup_started); - for (size_t i = 0; i < input_counter.size(); ++i) { - ARROW_CHECK(input_counter[i] == output_counter[i]) - << input_counter[i] << " != " << output_counter[i]; - } - -#ifdef ARROW_ENABLE_THREADING - ARROW_UNUSED( - plan_->query_context()->executor()->Spawn([this, st = std::move(st)]() mutable { - Defer cleanup([this, &st]() { process_task.MarkFinished(st); }); - if (st.ok()) { - st = output_->InputFinished(this, batches_produced); - } - })); -#else - process_task.MarkFinished(st); - if (st.ok()) { - st = output_->InputFinished(this, batches_produced); + void MaybeEnterFlushingUnlocked() { + if (std::all_of(state_.begin(), state_.end(), + [](const auto& input) { return input->AllBatchesReceived(); })) { + output_gate_ = OutputGate::Flushing; } -#endif } - bool CheckEnded() { - bool all_finished = true; - for (const auto& s : state) { - all_finished &= s->Finished(); - } - if (all_finished) { - EndFromProcessThread(); - return false; + bool CanProgressUnlocked() const { + return std::all_of(state_.begin(), state_.end(), [](const auto& input) { + return input->HasData() || input->Finished(); + }); + } + + bool FinishedUnlocked() const { + return std::all_of(state_.begin(), state_.end(), + [](const auto& input) { return input->Finished(); }); + } + + std::optional MaybeStartUnlocked() { + MaybeEnterFlushingUnlocked(); + if (merge_state_ != MergeState::Idle || output_gate_ == OutputGate::Paused || + !CanProgressUnlocked()) { + return std::nullopt; } - return true; + merge_state_ = MergeState::Running; + return Task([this] { return Drain(); }); } - /// Streams the input states in sorted order until we run out of input - arrow::Result> getNextBatch() { - DCHECK(!state.empty()); - for (const auto& s : state) { - if (s->Empty() && !s->Finished()) { - return nullptr; // not enough data, wait + Selection GetNextSelectionUnlocked(std::vector* consumed) { + DCHECK(CanProgressUnlocked()); + Selection selection; + std::vector heap; + heap.reserve(state_.size()); + for (const auto& input : state_) { + if (input->HasData()) { + heap.push_back(input.get()); } } - - std::vector> heap = state; - // filter out finished states - heap.erase(std::remove_if( - heap.begin(), heap.end(), - [](const std::shared_ptr& s) { return s->Finished(); }), - heap.end()); - - // If any are Empty(), then return early since we don't have enough data - if (std::any_of(heap.begin(), heap.end(), - [](const std::shared_ptr& s) { return s->Empty(); })) { - return nullptr; + if (heap.empty()) { + return selection; } - // Currently we only support one sort key - const auto sort_col = *ordering_.sort_keys().at(0).target.name(); - const auto comp = InputStateComparator(); + const auto comp = InputStateComparator(ordering_.sort_keys()[0].null_placement); std::make_heap(heap.begin(), heap.end(), comp); - // Each slice only has one record batch with the same schema as the output - std::unordered_map> output_col_to_src; - for (int i = 0; i < output_schema_->num_fields(); i++) { - output_col_to_src[i] = std::make_pair(0, i); - } - SingleRecordBatchCompositeTable output(output_schema(), 1, - std::move(output_col_to_src), - plan()->query_context()->memory_pool()); - - // Generate rows until we run out of data or we exceed the target output - // size - bool waiting_for_more_data = false; - while (!waiting_for_more_data && !heap.empty() && - output.Size() < kTargetOutputBatchSize) { + // Generate rows until we run out of data or reach the target output size. + while (!heap.empty() && selection.length < kTargetOutputBatchSize) { std::pop_heap(heap.begin(), heap.end(), comp); auto& next_item = heap.back(); - time_unit_t latest_time = std::numeric_limits::min(); - time_unit_t new_time = next_item->GetLatestTime(); - ARROW_CHECK(new_time >= latest_time) - << "Input state " << next_item->index() - << " has out of order data. newTime=" << new_time - << " latestTime=" << latest_time; - - latest_time = new_time; - SingleRecordBatchSliceBuilder builder{&output}; - next_item->Advance(builder); - - if (builder.Size() > 0) { - output_counter[next_item->index()] += builder.Size(); - builder.Finalize(); + // pop_heap leaves the remaining heap's earliest input at the front. The selected + // input can safely contribute every row up to that timestamp; no other input can + // contain an earlier row. + const std::optional* upper_bound = + heap.size() > 1 ? &heap.front()->GetLatestTime() : nullptr; + bool batch_consumed = next_item->Advance( + upper_bound, kTargetOutputBatchSize - selection.length, &selection); + if (batch_consumed) { + consumed->push_back(next_item); } if (next_item->Finished()) { heap.pop_back(); - } else if (next_item->Empty()) { + continue; + } + if (!next_item->HasData()) { // We've run out of data on one of the inputs - waiting_for_more_data = true; - continue; // skip the unnecessary make_heap + break; } - std::make_heap(heap.begin(), heap.end(), comp); + std::push_heap(heap.begin(), heap.end(), comp); } + return selection; + } - // Emit the batch - if (output.Size() == 0) { - return nullptr; - } + Result Materialize(const Selection& selection) { + std::vector values; + values.reserve(output_schema_->num_fields()); + for (int column = 0; column < output_schema_->num_fields(); ++column) { + ARROW_ASSIGN_OR_RAISE(auto builder, + MakeBuilder(output_schema_->field(column)->type(), + plan()->query_context()->memory_pool())); + ARROW_RETURN_NOT_OK(builder->Reserve(selection.length)); + const ArrayData* current_source = nullptr; + std::optional current_span; + for (const SelectionRun& run : selection.runs) { + const Datum& source = run.source->batch.values[column]; + if (source.is_scalar()) { + ARROW_RETURN_NOT_OK(builder->AppendScalar(*source.scalar(), run.length)); + continue; + } + if (!source.is_array()) { + return Status::Invalid("SortedMerge input must be an array or scalar, but got ", + ::arrow::ToString(source.kind()), " in column ", column); + } - ARROW_ASSIGN_OR_RAISE(auto maybe_rb, output.Materialize()); - return maybe_rb.value_or(nullptr); - } - /// Gets a batch. Returns true if there is more data to process, false if we - /// are done or an error occurred - bool PollOnce() { - std::lock_guard guard(gate); - if (!CheckEnded()) { - return false; + if (source.array().get() != current_source) { + current_source = source.array().get(); + current_span.emplace(*source.array()); + } + Status status = builder->AppendArraySlice(*current_span, run.offset, run.length); + if (status.IsNotImplemented()) { + auto source_array = MakeArray(source.array()); + for (int64_t row = run.offset; row < run.offset + run.length; ++row) { + ARROW_ASSIGN_OR_RAISE(auto scalar, source_array->GetScalar(row)); + ARROW_RETURN_NOT_OK(builder->AppendScalar(*scalar)); + } + } else { + ARROW_RETURN_NOT_OK(status); + } + } + ARROW_ASSIGN_OR_RAISE(auto array, builder->Finish()); + values.emplace_back(std::move(array)); } + return ExecBatch(std::move(values), selection.length); + } - // Process batches while we have data + Status Drain() { for (;;) { - Result> result = getNextBatch(); - - if (result.ok()) { - auto out_rb = *result; - if (!out_rb) { - break; + std::vector consumed; + Selection selection; + int32_t output_index = -1; + bool finish = false; + { + std::lock_guard lock(coordinator_mutex_); + if (merge_state_ == MergeState::Terminal) { + return Status::OK(); } - ExecBatch out_b(*out_rb); - out_b.index = batches_produced++; - Status st = output_->InputReceived(this, std::move(out_b)); - if (!st.ok()) { - ARROW_LOG(FATAL) << "Error in output_::InputReceived: " << st.ToString(); - EndFromProcessThread(std::move(st)); + DCHECK_EQ(merge_state_, MergeState::Running); + MaybeEnterFlushingUnlocked(); + if (output_gate_ == OutputGate::Paused || !CanProgressUnlocked()) { + merge_state_ = MergeState::Idle; + return Status::OK(); + } + + selection = GetNextSelectionUnlocked(&consumed); + if (!selection.empty()) { + // The sole merge task owns output numbering, independent of executor mode. + output_index = batches_produced_++; + } else if (FinishedUnlocked()) { + merge_state_ = MergeState::Terminal; + terminal_.store(true); + finish = true; + } else { + merge_state_ = MergeState::Idle; } - } else { - EndFromProcessThread(result.status()); - return false; } - } - // Report to the output the total batch count, if we've already - // finished everything (there are two places where this can happen: - // here and InputFinished) - // - // It may happen here in cases where InputFinished was called before - // we were finished producing results (so we didn't know the output - // size at that time) - if (!CheckEnded()) { - return false; + for (InputState* input : consumed) { + input->BatchConsumed().Apply(); + } + if (finish) { + return FinishNormally(); + } + if (selection.empty()) { + return Status::OK(); + } + + auto materialized = Materialize(selection); + if (!materialized.ok()) { + EnterTerminal(); + return materialized.status(); + } + ExecBatch output = std::move(*materialized); + output.index = output_index; + Status status = output_->InputReceived(this, std::move(output)); + if (!status.ok()) { + EnterTerminal(); + return status; + } } + } - // There is no more we can do now but there is still work remaining - // for later when more data arrives. - return true; + Status FinishNormally() { + for (auto& input : state_) { + input->Shutdown().Apply(); + } + return output_->InputFinished(this, batches_produced_); } -#ifdef ARROW_ENABLE_THREADING - void EmitBatches() { - while (true) { - // Implementation note: If the queue is empty, we will block here - if (process_queue.WaitAndPop() == kPoisonPill) { - EndFromProcessThread(); - } - // Either we're out of data or something went wrong - if (!PollOnce()) { - return; + void EnterTerminal() { + terminal_.store(true); + { + std::lock_guard lock(coordinator_mutex_); + merge_state_ = MergeState::Terminal; + for (auto& input : state_) { + input->ClearBuffered(); } } + for (auto& input : state_) { + input->Shutdown().Apply(); + } } - /// The entry point for processThread - static void StartPoller(SortedMergeNode* node) { node->EmitBatches(); } -#endif - arrow::Ordering ordering_; - - // Each input state corresponds to an input (e.g. a parquet data file) - std::vector> state; - std::vector input_counter; - std::vector output_counter; - std::mutex gate; - - std::atomic cleanup_started{false}; - - // Backpressure counter common to all input states - std::atomic backpressure_counter; - - std::atomic batches_produced{0}; - -#ifdef ARROW_ENABLE_THREADING - // Queue to trigger processing of a given input. False acts as a poison pill - ConcurrentQueue process_queue; - // Once StartProducing is called, we initialize this thread to poll the - // input states and emit batches - std::thread process_thread; -#endif - arrow::Future<> process_task; - - // Map arg index --> completion counter - std::vector counter_; - // Map arg index --> data - std::vector accumulation_queue_; - std::mutex mutex_; - std::atomic total_batches_{0}; + std::vector> state_; + + std::mutex coordinator_mutex_; + MergeState merge_state_ = MergeState::Idle; + OutputGate output_gate_ = OutputGate::Open; + int32_t downstream_counter_ = std::numeric_limits::min(); + int32_t batches_produced_ = 0; + std::atomic terminal_{false}; }; +InputState::InputState(SortedMergeNode* node, size_t index, ExecNode* input, + col_index_t time_col_index, compute::NullPlacement null_placement) + : node_(node), + index_(index), + input_(input), + time_col_index_(time_col_index), + null_placement_(null_placement), + sequencer_(util::SerialSequencingQueue::Make(this)) {} + +Status InputState::InsertBatch(ExecBatch batch) { + if (batch.index == compute::kUnsequencedIndex) { + return Status::Invalid("SortedMerge requires sequenced input"); + } + return sequencer_->InsertBatch(std::move(batch)); +} + +Status InputState::Process(ExecBatch batch) { + if (node_->IsTerminal()) { + return Status::OK(); + } + // Sequence the original ExecBatch. Once its index is current, prepare only the + // sort-key view needed to select rows; payload columns stay unmaterialized until the + // complete output selection is known. + ARROW_ASSIGN_OR_RAISE(auto prepared, PrepareBatch(std::move(batch))); + + // Only sequenced, non-empty batches count toward backpressure. Counting physical + // arrivals can deadlock with a reordering input if later batches reach the high + // watermark while that input still owns the batch which closes the sequencing gap. + FlowAction buffered = prepared->batch.length == 0 ? FlowAction{} : BatchBuffered(); + ARROW_ASSIGN_OR_RAISE(auto task, node_->OnSequenced(index_, std::move(prepared))); + buffered.Apply(); + return task ? std::move(*task)() : Status::OK(); +} + +Result> InputState::PrepareBatch(ExecBatch batch) { + auto prepared = std::make_shared(); + prepared->batch = std::move(batch); + prepared->times.reserve(prepared->batch.length); + const Datum& time_column = prepared->batch.values[time_col_index_]; + for (int64_t row = 0; row < prepared->batch.length; ++row) { + ARROW_ASSIGN_OR_RAISE(auto time, ReadTimeValue(time_column, row)); + ARROW_RETURN_NOT_OK(ValidateTime(time)); + prepared->times.push_back(time); + } + return prepared; +} + +Status InputState::ValidateTime(const std::optional& time) { + if (!time) { + if (null_placement_ == compute::NullPlacement::AtStart && saw_non_null_) { + return Status::Invalid("SortedMerge input ", index_, + " has out-of-order nulls in its sort key"); + } + saw_null_ = true; + return Status::OK(); + } + if ((null_placement_ == compute::NullPlacement::AtEnd && saw_null_) || + (last_time_ && *time < *last_time_)) { + return Status::Invalid("SortedMerge input ", index_, + " has out-of-order values in its sort key"); + } + saw_non_null_ = true; + last_time_ = *time; + return Status::OK(); +} + +FlowAction InputState::SetUpstreamPausedUnlocked(bool paused) { + if (upstream_paused_ == paused || shutdown_) { + return {}; + } + upstream_paused_ = paused; + return {paused ? FlowAction::Kind::Pause : FlowAction::Kind::Resume, input_, node_, + ++outgoing_counter_}; +} + +FlowAction InputState::BatchBuffered() { + std::lock_guard lock(flow_mutex_); + if (shutdown_) { + return {}; + } + ++buffered_batches_; + return buffered_batches_ >= kHighWatermark ? SetUpstreamPausedUnlocked(true) + : FlowAction{}; +} + +FlowAction InputState::BatchConsumed() { + std::lock_guard lock(flow_mutex_); + if (shutdown_) { + return {}; + } + DCHECK_GT(buffered_batches_, 0); + if (buffered_batches_ > 0) { + --buffered_batches_; + } + return buffered_batches_ <= kLowWatermark ? SetUpstreamPausedUnlocked(false) + : FlowAction{}; +} + +FlowAction InputState::Shutdown() { + std::lock_guard lock(flow_mutex_); + if (shutdown_) { + return {}; + } + shutdown_ = true; + upstream_paused_ = false; + buffered_batches_ = 0; + // Always send a final, newer resume so a delayed pause cannot strand an input. + return {FlowAction::Kind::Resume, input_, node_, ++outgoing_counter_}; +} + } // namespace namespace internal { diff --git a/cpp/src/arrow/acero/sorted_merge_node_test.cc b/cpp/src/arrow/acero/sorted_merge_node_test.cc index 82b630420c4a..1a7e5c865f72 100644 --- a/cpp/src/arrow/acero/sorted_merge_node_test.cc +++ b/cpp/src/arrow/acero/sorted_merge_node_test.cc @@ -24,13 +24,16 @@ #include "arrow/array/builder_base.h" #include "arrow/array/concatenate.h" #include "arrow/compute/ordering.h" +#include "arrow/compute/test_util_internal.h" #include "arrow/result.h" #include "arrow/scalar.h" #include "arrow/table.h" +#include "arrow/testing/future_util.h" #include "arrow/testing/generator.h" #include "arrow/testing/gtest_util.h" #include "arrow/type.h" #include "arrow/type_fwd.h" +#include "arrow/util/logging_internal.h" namespace arrow::acero { @@ -65,8 +68,6 @@ TEST(SortedMergeNode, Basic) { auto ops = OrderByNodeOptions(compute::Ordering({compute::SortKey("timestamp")})); Declaration sorted_merge{"sorted_merge", src_decls, ops}; - // We can't use threads for sorted merging since it relies on - // ascending deterministic order of timestamps ASSERT_OK_AND_ASSIGN(auto output, DeclarationToTable(sorted_merge, /*use_threads=*/false)); ASSERT_EQ(output->num_rows(), 18); @@ -83,4 +84,161 @@ TEST(SortedMergeNode, Basic) { AssertArraysEqual(*expected_ts, *output_ts); } +TEST(SortedMergeNode, SignedValuesCrossZero) { + auto table1 = TestTable( + /*start=*/-4, + /*step=*/2, + /*rows_per_batch=*/3, + /*num_batches=*/1); + auto table2 = TestTable( + /*start=*/-3, + /*step=*/2, + /*rows_per_batch=*/4, + /*num_batches=*/1); + std::vector src_decls; + src_decls.emplace_back(Declaration("table_source", TableSourceNodeOptions(table1))); + src_decls.emplace_back(Declaration("table_source", TableSourceNodeOptions(table2))); + + auto options = OrderByNodeOptions(compute::Ordering({compute::SortKey("timestamp")})); + Declaration sorted_merge{"sorted_merge", src_decls, options}; + ASSERT_OK_AND_ASSIGN(auto output, + DeclarationToTable(sorted_merge, /*use_threads=*/false)); + + auto expected = ArrayFromJSON(int32(), "[-4, -3, -2, -1, 0, 1, 3]"); + ASSERT_OK_AND_ASSIGN(auto actual, Concatenate(output->column(0)->chunks())); + AssertArraysEqual(*expected, *actual); +} + +TEST(SortedMergeNode, MergesScalarPayload) { + auto input_schema = schema({field("timestamp", int32()), field("source", utf8())}); + ExecBatch left( + {ArrayFromJSON(int32(), "[0, 2]"), std::make_shared("left")}, 2); + ExecBatch right( + {ArrayFromJSON(int32(), "[1, 3]"), std::make_shared("right")}, 2); + + Declaration merge{ + "sorted_merge", + {Declaration{"exec_batch_source", + ExecBatchSourceNodeOptions(input_schema, {std::move(left)})}, + Declaration{"exec_batch_source", + ExecBatchSourceNodeOptions(input_schema, {std::move(right)})}}, + OrderByNodeOptions(compute::Ordering({compute::SortKey("timestamp")}))}; + ASSERT_OK_AND_ASSIGN(auto output, + DeclarationToTable(std::move(merge), /*use_threads=*/false)); + + ASSERT_OK_AND_ASSIGN(auto timestamps, Concatenate(output->column(0)->chunks())); + AssertArraysEqual(*ArrayFromJSON(int32(), "[0, 1, 2, 3]"), *timestamps); + ASSERT_OK_AND_ASSIGN(auto sources, Concatenate(output->column(1)->chunks())); + AssertArraysEqual(*ArrayFromJSON(utf8(), R"(["left", "right", "left", "right"])"), + *sources); +} + +TEST(SortedMergeNode, ProducesSortedOutputFromJitteredInputsOnBothExecutors) { + // Use enough small batches that delivery and flow control overlap. Correctness must + // depend on logical batch indices rather than physical arrival order. + constexpr int kBatchesPerInput = 64; + RegisterTestNodes(); + auto table0 = TestTable(0, 3, /*rows_per_batch=*/1, kBatchesPerInput); + auto table1 = TestTable(1, 3, /*rows_per_batch=*/1, kBatchesPerInput); + auto table2 = TestTable(2, 3, /*rows_per_batch=*/1, kBatchesPerInput); + + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + std::vector inputs; + inputs.emplace_back(Declaration::Sequence( + {{"table_source", TableSourceNodeOptions(table0)}, + {"jitter", JitterNodeOptions(/*seed=*/42, /*max_jitter_modifier=*/4)}})); + inputs.emplace_back(Declaration::Sequence( + {{"table_source", TableSourceNodeOptions(table1)}, + {"jitter", JitterNodeOptions(/*seed=*/84, /*max_jitter_modifier=*/4)}})); + inputs.emplace_back(Declaration::Sequence( + {{"table_source", TableSourceNodeOptions(table2)}, + {"jitter", JitterNodeOptions(/*seed=*/126, /*max_jitter_modifier=*/4)}})); + + QueryOptions query_options; + query_options.use_threads = use_threads; + Declaration merge{ + "sorted_merge", std::move(inputs), + OrderByNodeOptions(compute::Ordering({compute::SortKey("timestamp")}))}; + ASSERT_OK_AND_ASSIGN(auto output, + DeclarationToTable(std::move(merge), query_options)); + + ASSERT_EQ(output->num_rows(), 3 * kBatchesPerInput); + ASSERT_OK_AND_ASSIGN(auto actual, Concatenate(output->column(0)->chunks())); + ASSERT_OK_AND_ASSIGN(auto expected_builder, + MakeBuilder(int32(), default_memory_pool())); + for (int value = 0; value < 3 * kBatchesPerInput; ++value) { + ASSERT_OK(expected_builder->AppendScalar(*MakeScalar(value))); + } + ASSERT_OK_AND_ASSIGN(auto expected, expected_builder->Finish()); + AssertArraysEqual(*expected, *actual); + } +} + +TEST(SortedMergeNode, DownstreamBackpressureAndStop) { + for (bool stop_while_paused : {false, true}) { + SCOPED_TRACE(stop_while_paused ? "stop" : "resume"); + auto input_schema = schema({field("timestamp", int32())}); + PushGenerator> left_generator; + PushGenerator> right_generator; + AsyncGenerator> sink_generator; + BackpressureMonitor* backpressure_monitor = nullptr; + + Declaration left{ + "source", SourceNodeOptions(input_schema, left_generator, Ordering::Implicit())}; + Declaration right{ + "source", SourceNodeOptions(input_schema, right_generator, Ordering::Implicit())}; + Declaration merge{ + "sorted_merge", + {std::move(left), std::move(right)}, + OrderByNodeOptions(compute::Ordering({compute::SortKey("timestamp")}))}; + Declaration sink{"sink", + {std::move(merge)}, + SinkNodeOptions(&sink_generator, /*schema=*/nullptr, + BackpressureOptions(/*resume_if_below=*/1, + /*pause_if_above=*/1), + &backpressure_monitor, /*sequence_output=*/false)}; + + ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make(*threaded_exec_context())); + ASSERT_OK(sink.AddToPlan(plan.get())); + ASSERT_OK(plan->Validate()); + ASSERT_NE(backpressure_monitor, nullptr); + plan->StartProducing(); + + left_generator.producer().Push(compute::ExecBatchFromJSON({int32()}, "[[0]]")); + right_generator.producer().Push(compute::ExecBatchFromJSON({int32()}, "[[1]]")); + BusyWait(10.0, [&] { return backpressure_monitor->is_paused(); }); + ASSERT_TRUE(backpressure_monitor->is_paused()); + + if (stop_while_paused) { + plan->StopProducing(); + left_generator.producer().Push(IterationEnd>()); + right_generator.producer().Push(IterationEnd>()); + ASSERT_TRUE(plan->finished().Wait(kDefaultAssertFinishesWaitSeconds)); + ASSERT_TRUE(plan->finished().status().IsCancelled()); + continue; + } + + const uint64_t paused_bytes = backpressure_monitor->bytes_in_use(); + left_generator.producer().Push(compute::ExecBatchFromJSON({int32()}, "[[2]]")); + right_generator.producer().Push(compute::ExecBatchFromJSON({int32()}, "[[3]]")); + arrow::internal::GetCpuThreadPool()->WaitForIdle(); + EXPECT_EQ(backpressure_monitor->bytes_in_use(), paused_bytes); + + ASSERT_FINISHES_OK_AND_ASSIGN(auto first_output, sink_generator()); + ASSERT_TRUE(first_output.has_value()); + BusyWait(10.0, [&] { return backpressure_monitor->bytes_in_use() > 0; }); + const bool resumed = backpressure_monitor->bytes_in_use() > 0; + + left_generator.producer().Push(IterationEnd>()); + right_generator.producer().Push(IterationEnd>()); + for (;;) { + ASSERT_FINISHES_OK_AND_ASSIGN(auto output, sink_generator()); + if (!output) break; + } + ASSERT_FINISHES_OK(plan->finished()); + EXPECT_TRUE(resumed); + } +} + } // namespace arrow::acero From 01eb9dd00d88d3d3ac7c97af9369532fb25fb525 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Wed, 2 Sep 2026 13:35:43 +0000 Subject: [PATCH 3/5] [C++][Acero] Simplify SortedMerge shutdown flow control --- cpp/src/arrow/acero/sorted_merge_node.cc | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/cpp/src/arrow/acero/sorted_merge_node.cc b/cpp/src/arrow/acero/sorted_merge_node.cc index 55806fb153d5..9c11743e7c6f 100644 --- a/cpp/src/arrow/acero/sorted_merge_node.cc +++ b/cpp/src/arrow/acero/sorted_merge_node.cc @@ -259,7 +259,7 @@ class InputState final : public util::SerialSequencingQueue::Processor { FlowAction BatchBuffered(); FlowAction BatchConsumed(); - FlowAction Shutdown(); + void Shutdown(); private: FlowAction SetUpstreamPausedUnlocked(bool paused); @@ -653,7 +653,7 @@ class SortedMergeNode : public ExecNode { Status FinishNormally() { for (auto& input : state_) { - input->Shutdown().Apply(); + input->Shutdown(); } return output_->InputFinished(this, batches_produced_); } @@ -668,7 +668,7 @@ class SortedMergeNode : public ExecNode { } } for (auto& input : state_) { - input->Shutdown().Apply(); + input->Shutdown(); } } @@ -781,16 +781,13 @@ FlowAction InputState::BatchConsumed() { : FlowAction{}; } -FlowAction InputState::Shutdown() { +void InputState::Shutdown() { std::lock_guard lock(flow_mutex_); if (shutdown_) { - return {}; + return; } shutdown_ = true; - upstream_paused_ = false; buffered_batches_ = 0; - // Always send a final, newer resume so a delayed pause cannot strand an input. - return {FlowAction::Kind::Resume, input_, node_, ++outgoing_counter_}; } } // namespace From 8470111ac3e4c8dcc9e2f06929471d27752d3dff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Wed, 2 Sep 2026 13:58:00 +0000 Subject: [PATCH 4/5] [C++][Acero] Preserve scheduler lifetime across SortedMerge resume --- cpp/src/arrow/acero/sorted_merge_node.cc | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/cpp/src/arrow/acero/sorted_merge_node.cc b/cpp/src/arrow/acero/sorted_merge_node.cc index 9c11743e7c6f..743a6643020a 100644 --- a/cpp/src/arrow/acero/sorted_merge_node.cc +++ b/cpp/src/arrow/acero/sorted_merge_node.cc @@ -440,6 +440,7 @@ class SortedMergeNode : public ExecNode { void ResumeProducing(arrow::acero::ExecNode* output, int32_t counter) override { std::optional task; + Future<> handoff; { std::lock_guard lock(coordinator_mutex_); if (merge_state_ == MergeState::Terminal || counter <= downstream_counter_) { @@ -450,9 +451,22 @@ class SortedMergeNode : public ExecNode { output_gate_ = OutputGate::Open; } task = MaybeStartUnlocked(); + if (task) { + // Resume may be called outside the scheduler by a consuming sink. Reserve a + // scheduler task before releasing the coordinator lock so a concurrent final + // InputFinished cannot end the scheduler before the drain is registered. + auto maybe_handoff = + plan()->query_context()->BeginExternalTask("SortedMergeNode::ResumeHandoff"); + if (!maybe_handoff.ok() || !maybe_handoff->is_valid()) { + merge_state_ = MergeState::Idle; + return; + } + handoff = std::move(*maybe_handoff); + } } if (task) { Schedule(std::move(*task), "SortedMergeNode::Resume"); + handoff.MarkFinished(); } } From 334e55cc2e6b544e05ea534efdbb2c72071f6c1d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Thu, 3 Sep 2026 13:48:00 +0000 Subject: [PATCH 5/5] [C++][Acero] Limit SortedMerge output batch size --- cpp/src/arrow/acero/sorted_merge_node.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/acero/sorted_merge_node.cc b/cpp/src/arrow/acero/sorted_merge_node.cc index 743a6643020a..0c80b40494c5 100644 --- a/cpp/src/arrow/acero/sorted_merge_node.cc +++ b/cpp/src/arrow/acero/sorted_merge_node.cc @@ -308,7 +308,7 @@ enum class MergeState { Idle, Running, Terminal }; enum class OutputGate { Open, Paused, Flushing }; class SortedMergeNode : public ExecNode { - static constexpr int64_t kTargetOutputBatchSize = 1024 * 1024; + static constexpr int64_t kTargetOutputBatchSize = ExecPlan::kMaxBatchSize; public: SortedMergeNode(arrow::acero::ExecPlan* plan,