Make record batch writes aligned on word boundaries

Change-Id: Ia0e59d236054db576bcc458231c914a4a865b428

Author: wesm

cpp/src/arrow/ipc/adapter.cc

@@ -142,21 +142,21 @@ class RecordBatchWriter {
 
   Status AssemblePayload() {
     // Perform depth-first traversal of the row-batch
-    for (int i = 0; i < columns_->size(); ++i) {
+    for (size_t i = 0; i < columns_->size(); ++i) {
       const Array* arr = (*columns_)[i].get();
       RETURN_NOT_OK(VisitArray(arr, &field_nodes_, &buffers_, max_recursion_depth_));
     }
     return Status::OK();
   }
 
   Status Write(io::OutputStream* dst, int64_t* data_header_offset) {
-    // Write out all the buffers contiguously and compute the total size of the
-    // memory payload
-    int64_t offset = 0;
-
     // Get the starting position
-    int64_t position;
-    RETURN_NOT_OK(dst->Tell(&position));
+    int64_t start_position;
+    RETURN_NOT_OK(dst->Tell(&start_position));
+
+    // Keep track of the current position so we can determine the size of the
+    // message body
+    int64_t position = start_position;
 
     for (size_t i = 0; i < buffers_.size(); ++i) {
       const Buffer* buffer = buffers_[i].get();
@@ -178,11 +178,11 @@ class RecordBatchWriter {
       // are using from any OS-level shared memory. The thought is that systems
       // may (in the future) associate integer page id's with physical memory
       // pages (according to whatever is the desired shared memory mechanism)
-      buffer_meta_.push_back(flatbuf::Buffer(0, position + offset, size));
+      buffer_meta_.push_back(flatbuf::Buffer(0, position, size));
 
       if (size > 0) {
         RETURN_NOT_OK(dst->Write(buffer->data(), size));
-        offset += size;
+        position += size;
       }
     }
 
@@ -194,13 +194,24 @@ class RecordBatchWriter {
     // determine the data header size then request a buffer such that you can
     // construct the flatbuffer data accessor object (see arrow::ipc::Message)
     std::shared_ptr<Buffer> data_header;
-    RETURN_NOT_OK(
-        WriteDataHeader(num_rows_, offset, field_nodes_, buffer_meta_, &data_header));
+    RETURN_NOT_OK(WriteDataHeader(
+        num_rows_, position - start_position, field_nodes_, buffer_meta_, &data_header));
 
     // Write the data header at the end
     RETURN_NOT_OK(dst->Write(data_header->data(), data_header->size()));
+    *data_header_offset = position;
 
-    *data_header_offset = position + offset;
+    return Align(dst, &position);
+  }
+
+  Status Align(io::OutputStream* dst, int64_t* position) {
+    // Write all buffers here on word boundaries
+    // TODO(wesm): Is there benefit to 64-byte padding in IPC?
+    int64_t remainder = PaddedLength(*position) - *position;
+    if (remainder > 0) {
+      RETURN_NOT_OK(dst->Write(kPaddingBytes, remainder));
+      *position += remainder;
+    }
     return Status::OK();
   }
 

cpp/src/arrow/ipc/file.cc

@@ -22,7 +22,9 @@
 #include <sstream>
 #include <vector>
 
+#include "arrow/ipc/adapter.h"
 #include "arrow/ipc/metadata.h"
+#include "arrow/ipc/util.h"
 #include "arrow/io/interfaces.h"
 #include "arrow/util/buffer.h"
 #include "arrow/util/logging.h"
@@ -33,14 +35,6 @@ namespace ipc {
 
 static constexpr const char* kArrowMagicBytes = "ARROW1";
 
-// Align on 8-byte boundaries
-static constexpr int kArrowAlignment = 8;
-static constexpr uint8_t kPaddingBytes[kArrowAlignment] = {0};
-
-static inline int64_t PaddedLength(int64_t nbytes) {
-  return ((nbytes + kArrowAlignment - 1) / kArrowAlignment) * kArrowAlignment;
-}
-
 // ----------------------------------------------------------------------
 // Writer implementation
 
@@ -87,6 +81,30 @@ Status FileWriter::CheckStarted() {
   return Status::OK();
 }
 
+Status FileWriter::WriteRecordBatch(
+    const std::vector<std::shared_ptr<Array>>& columns, int32_t num_rows) {
+  RETURN_NOT_OK(CheckStarted());
+
+  int64_t offset = position_;
+
+  int64_t header_offset;
+  RETURN_NOT_OK(arrow::ipc::WriteRecordBatch(columns, num_rows, sink_, &header_offset));
+  RETURN_NOT_OK(UpdatePosition());
+
+  DCHECK(position_ % 8 == 0) << "ipc::WriteRecordBatch did not perform aligned writes";
+
+  // We can infer the metadata length and length of the record batch body (the
+  // concatenated buffers) from the heade offset and the new output stream
+  // position
+  int32_t metadata_length = position_ - header_offset;
+  int32_t body_length = position_ - offset - metadata_length;
+
+  // Append metadata, to be written in the footer latera
+  record_batches_.emplace_back(offset, metadata_length, body_length);
+
+  return Status::OK();
+}
+
 Status FileWriter::Close() {
   // Write metadata
   int64_t initial_position = position_;
@@ -105,11 +123,6 @@ Status FileWriter::Close() {
       reinterpret_cast<const uint8_t*>(kArrowMagicBytes), strlen(kArrowMagicBytes));
 }
 
-Status FileWriter::WriteRecordBatch(
-    const std::vector<std::shared_ptr<Array>>& columns, int32_t num_rows) {
-  return Status::OK();
-}
-
 // ----------------------------------------------------------------------
 // Reader implementation
 

cpp/src/arrow/ipc/util.h

@@ -27,6 +27,14 @@
 namespace arrow {
 namespace ipc {
 
+// Align on 8-byte boundaries
+static constexpr int kArrowAlignment = 8;
+static constexpr uint8_t kPaddingBytes[kArrowAlignment] = {0};
+
+static inline int64_t PaddedLength(int64_t nbytes, int64_t alignment = kArrowAlignment) {
+  return ((nbytes + alignment - 1) / alignment) * alignment;
+}
+
 // A helper class to tracks the size of allocations
 class MockOutputStream : public io::OutputStream {
  public: