Method TestTransferMemory

be/src/runtime/buffered-tuple-stream-test.cc:786–838

}

void SimpleTupleStreamTest::TestTransferMemory(bool pin_stream, bool read_write) {
  // 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::TestTransferMemory", pin_stream));
  if (read_write) {
    bool got_reservation;
    ASSERT_OK(stream.PrepareForReadWrite(true, &got_reservation));
    ASSERT_TRUE(got_reservation);
  } else {
    bool got_write_reservation;
    ASSERT_OK(stream.PrepareForWrite(&got_write_reservation));
    ASSERT_TRUE(got_write_reservation);
  }
  RowBatch* 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 < batch->num_rows(); ++i) {
      bool ret = stream.AddRow(batch->GetRow(i), &status);
      EXPECT_TRUE(ret);
      ASSERT_OK(status);
    }
  }

  batch->Reset();

  if (read_write) {
    // Read back batch so that we have a read buffer in memory.
    bool eos;
    ASSERT_OK(stream.GetNext(batch, &eos));
    EXPECT_FALSE(eos);
  }
  stream.Close(batch, RowBatch::FlushMode::FLUSH_RESOURCES);
  if (pin_stream) {
    EXPECT_EQ(total_num_pages, batch->num_buffers());
  } else if (read_write) {
    // Read and write buffer should be attached.
    EXPECT_EQ(2, batch->num_buffers());
  } else {
    // Read buffer should be attached.
    EXPECT_EQ(1, batch->num_buffers());
  }
  EXPECT_TRUE(batch->AtCapacity()); // Flush resources flag should have been set.
  batch->Reset();
  EXPECT_EQ(0, batch->num_buffers());
}

/// Test attaching memory to a row batch from a pinned stream.
TEST_F(SimpleTupleStreamTest, TransferMemoryFromPinnedStreamReadWrite) {