ARROW-100: [C++] Computing RowBatch size

cpp/src/arrow/ipc/adapter.cc

@@ -149,10 +149,19 @@ class RowBatchWriter {
   }
 
   // This must be called after invoking AssemblePayload
-  int64_t DataHeaderSize() {
-    // TODO(wesm): In case it is needed, compute the upper bound for the size
-    // of the buffer containing the flatbuffer data header.
-    return 0;
+  Status DataHeaderSize(int64_t* size) {
+    // emulates the behavior of Write without actually writing
+    int64_t offset = 0;
+    for (size_t i = 0; i < buffers_.size(); ++i) {
+      const Buffer* buffer = buffers_[i].get();
+      offset += buffer->size();
+      buffer_meta_.push_back(flatbuf::Buffer(0, 0, 0));
+    }
+    std::shared_ptr<Buffer> data_header;
+    RETURN_NOT_OK(WriteDataHeader(
+        batch_->num_rows(), offset, field_nodes_, buffer_meta_, &data_header));
+    *size = data_header->size();
+    return Status::OK();
   }
 
   // Total footprint of buffers. This must be called after invoking
@@ -179,6 +188,15 @@ Status WriteRowBatch(
   RETURN_NOT_OK(serializer.AssemblePayload());
   return serializer.Write(dst, position, header_offset);
 }
+
+Status GetRowBatchSize(const RowBatch* batch, int64_t* size) {
+  RowBatchWriter serializer(batch);
+  RETURN_NOT_OK(serializer.AssemblePayload());
+  RETURN_NOT_OK(serializer.DataHeaderSize(size));
+  *size += serializer.TotalBytes();
+  return Status::OK();
+}
+
 // ----------------------------------------------------------------------
 // Row batch read path
 

cpp/src/arrow/ipc/adapter.h

@@ -60,7 +60,7 @@ Status WriteRowBatch(
 // Compute the precise number of bytes needed in a contiguous memory segment to
 // write the row batch. This involves generating the complete serialized
 // Flatbuffers metadata.
-int64_t GetRowBatchSize(const RowBatch* batch);
+Status GetRowBatchSize(const RowBatch* batch, int64_t* size);
 
 // ----------------------------------------------------------------------
 // "Read" path; does not copy data if the MemorySource does not

cpp/src/arrow/ipc/ipc-adapter-test.cc

@@ -91,13 +91,21 @@ TEST_F(TestWriteRowBatch, IntegerRoundTrip) {
   int64_t header_location;
   ASSERT_OK(WriteRowBatch(mmap_.get(), &batch, 0, &header_location));
 
+  MockMemorySource mock_source(1 << 16);
+  int64_t mock_header_location;
+  ASSERT_OK(WriteRowBatch(&mock_source, &batch, 0, &mock_header_location));
+
   std::shared_ptr<RowBatchReader> result;
   ASSERT_OK(RowBatchReader::Open(mmap_.get(), header_location, &result));
 
   std::shared_ptr<RowBatch> batch_result;
   ASSERT_OK(result->GetRowBatch(schema, &batch_result));
   EXPECT_EQ(batch.num_rows(), batch_result->num_rows());
 
+  int64_t size;
+  ASSERT_OK(GetRowBatchSize(batch_result.get(), &size));
+  EXPECT_EQ(mock_source.Position(), size);
+
   for (int i = 0; i < batch.num_columns(); ++i) {
     EXPECT_TRUE(batch.column(i)->Equals(batch_result->column(i))) << i
                                                                   << batch.column_name(i);

cpp/src/arrow/ipc/memory.cc

@@ -144,5 +144,29 @@ Status MemoryMappedSource::Write(int64_t position, const uint8_t* data, int64_t
   return Status::OK();
 }
 
+MockMemorySource::MockMemorySource(int64_t size) : size_(size) {}
+
+Status MockMemorySource::Close() {
+  return Status::OK();
+}
+
+Status MockMemorySource::ReadAt(
+    int64_t position, int64_t nbytes, std::shared_ptr<Buffer>* out) {
+  return Status::OK();
+}
+
+Status MockMemorySource::Write(int64_t position, const uint8_t* data, int64_t nbytes) {
+  pos_ = position + nbytes;
+  return Status::OK();
+}
+
+int64_t MockMemorySource::Size() const {
+  return size_;
+}
+
+int64_t MockMemorySource::Position() const {
+  return pos_;
+}
+
 }  // namespace ipc
 }  // namespace arrow

cpp/src/arrow/ipc/memory.h

@@ -121,6 +121,26 @@ class MemoryMappedSource : public MemorySource {
   std::unique_ptr<Impl> impl_;
 };
 
+// A MemorySource that tracks the size of allocations from a memory source
+class MockMemorySource : public MemorySource {
+ public:
+  explicit MockMemorySource(int64_t size);
+
+  Status Close() override;
+
+  Status ReadAt(int64_t position, int64_t nbytes, std::shared_ptr<Buffer>* out) override;
+
+  Status Write(int64_t position, const uint8_t* data, int64_t nbytes) override;
+
+  int64_t Size() const override;
+
+  int64_t Position() const;
+
+ private:
+  int64_t size_;
+  int64_t pos_;
+};
+
 }  // namespace ipc
 }  // namespace arrow