Method TestAttachMemory

be/src/runtime/buffered-tuple-stream-test.cc:859–932

Test iteration over a stream with and without attaching memory.

/// Test iteration over a stream with and without attaching memory.
void SimpleTupleStreamTest::TestAttachMemory(bool pin_stream, bool attach_on_read) {
  // Use smaller buffers so that the explicit FLUSH_RESOURCES flag is required to
  // make the batch at capacity.
  int buffer_size = 4 * 1024;
  Init(100 * buffer_size);

  BufferedTupleStream stream(
      runtime_state_, int_desc_, &client_, buffer_size, buffer_size);
  ASSERT_OK(stream.Init("SimpleTupleStreamTest::TestAttachMemory", pin_stream));
  bool got_write_reservation;
  ASSERT_OK(stream.PrepareForWrite(&got_write_reservation));
  ASSERT_TRUE(got_write_reservation);
  RowBatch* in_batch = CreateIntBatch(0, 1024, false);

  // Construct a stream with 4 pages.
  const int total_num_pages = 4;
  while (stream.byte_size() < total_num_pages * buffer_size) {
    Status status;
    for (int i = 0; i < in_batch->num_rows(); ++i) {
      bool ret = stream.AddRow(in_batch->GetRow(i), &status);
      EXPECT_TRUE(ret);
      ASSERT_OK(status);
    }
  }

  RowBatch* out_batch = pool_.Add(new RowBatch(int_desc_, 100, &tracker_));
  int num_buffers_attached = 0;
  int num_flushes = 0;
  int64_t num_rows_returned = 0;
  bool got_read_reservation;
  ASSERT_OK(stream.PrepareForRead(attach_on_read, &got_read_reservation));
  ASSERT_TRUE(got_read_reservation);
  bool eos;
  do {
    ASSERT_EQ(0, out_batch->num_buffers());
    ASSERT_OK(stream.GetNext(out_batch, &eos));
    EXPECT_LE(out_batch->num_buffers(), 1) << "Should only attach one buffer at a time";
    if (out_batch->num_buffers() > 0) {
      EXPECT_TRUE(out_batch->AtCapacity()) << "Flush resources flag should have been set";
    }
    num_buffers_attached += out_batch->num_buffers();
    for (int i = 0; i < out_batch->num_rows(); ++i) {
      int slot_offset = int_desc_->tuple_descriptors()[0]->slots()[0]->tuple_offset();
      TupleRow* in_row = in_batch->GetRow(num_rows_returned % in_batch->num_rows());
      EXPECT_EQ(*in_row->GetTuple(0)->GetIntSlot(slot_offset),
          *out_batch->GetRow(i)->GetTuple(0)->GetIntSlot(slot_offset));
      ++num_rows_returned;
    }
    num_flushes += out_batch->flush_mode() == RowBatch::FlushMode::FLUSH_RESOURCES;
    out_batch->Reset();
  } while (!eos);

  if (attach_on_read) {
    EXPECT_EQ(4, num_buffers_attached) << "All buffers attached during iteration.";
  } else {
    EXPECT_EQ(0, num_buffers_attached) << "No buffers attached during iteration.";
  }
  if (attach_on_read || !pin_stream) {