Method getNextBatch

cpp/src/arrow/acero/sorted_merge_node.cc:465–540

Streams the input states in sorted order until we run out of input

  /// Streams the input states in sorted order until we run out of input
  arrow::Result<std::shared_ptr<arrow::RecordBatch>> getNextBatch() {
    DCHECK(!state.empty());
    for (const auto& s : state) {
      if (s->Empty() && !s->Finished()) {
        return nullptr;  // not enough data, wait
      }
    }

    std::vector<std::shared_ptr<InputState>> heap = state;
    // filter out finished states
    heap.erase(std::remove_if(
                   heap.begin(), heap.end(),
                   [](const std::shared_ptr<InputState>& 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<InputState>& s) { return s->Empty(); })) {
      return nullptr;
    }

    // Currently we only support one sort key
    const auto sort_col = *ordering_.sort_keys().at(0).target.name();
    const auto comp = InputStateComparator();
    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<int, std::pair<int, int>> 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) {
      std::pop_heap(heap.begin(), heap.end(), comp);

      auto& next_item = heap.back();
      time_unit_t latest_time = std::numeric_limits<time_unit_t>::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();
      }