Function TestPauseThenStop

cpp/src/arrow/acero/source_node_test.cc:113–155

// backpressure is applied to the source node.
template <typename ThisNode>
void TestPauseThenStop() {
  ASSERT_OK(ThisNode::Register());

  // number of batches, number of batches to pass before pausing, batch size
  constexpr int num_batches = 10, num_pass = 2, batch_size = 1;
  // the above constants can be changed subject to the following restrictions
  // to ensure that the test works
  ASSERT_GE(num_pass, 1);            // must pass at least one batch before pausing
  ASSERT_GT(num_batches, num_pass);  // must have more batches after pausing
  auto t_schema = schema({field("time", int32()), field("value", int32())});
  ASSERT_OK_AND_ASSIGN(auto t_batches,
                       MakeIntegerBatches({[](int row) -> int64_t { return row; },
                                           [](int row) -> int64_t { return row + 1; }},
                                          t_schema, num_batches, batch_size));

  Declaration t_src = {
      "source", SourceNodeOptions(t_batches.schema,
                                  MakeDelayedGen(t_batches, "t_src", /*delay_sec=*/0.5,
                                                 /*noisy=*/false))};
  Declaration ctrl = {
      ThisNode::kFactoryName, {t_src}, PauseThenStopNodeOptions(num_pass)};

  ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatchReader> batch_reader,
                       DeclarationToReader(ctrl, /*use_threads=*/false));

  if (!ThisNode::kPlanExitsEarly) {
    int64_t total_length = 0;
    for (;;) {
      ASSERT_OK_AND_ASSIGN(auto batch, batch_reader->Next());
      if (!batch) {
        break;
      }
      total_length += batch->num_rows();
    }
    ASSERT_EQ(static_cast<int64_t>(num_pass * batch_size), total_length);
  } else {
    for (int i = 0; i < num_pass - 1; i++) {
      ASSERT_OK_AND_ASSIGN(auto batch, batch_reader->Next());
      ASSERT_TRUE(batch);
    }
    ASSERT_RAISES(Cancelled, batch_reader->Next());
  }
}

TEST(SourceNode, PauseThenStopInput) { TestPauseThenStop<PauseThenStopInputNode>(); }