Finish draft of FileReader::GetRecordBatch. Add body end offset to ipc adapter

Change-Id: Ibf6fd0c4afd0704cdbc7a6501179c9ac47c9d33d

Author: wesm

cpp/src/arrow/ipc/adapter.cc

@@ -149,7 +149,8 @@ class RecordBatchWriter {
     return Status::OK();
   }
 
-  Status Write(io::OutputStream* dst, int64_t* data_header_offset) {
+  Status Write(
+      io::OutputStream* dst, int64_t* body_end_offset, int64_t* header_end_offset) {
     // Get the starting position
     int64_t start_position;
     RETURN_NOT_OK(dst->Tell(&start_position));
@@ -186,6 +187,8 @@ class RecordBatchWriter {
       }
     }
 
+    *body_end_offset = position;
+
     // Now that we have computed the locations of all of the buffers in shared
     // memory, the data header can be converted to a flatbuffer and written out
     //
@@ -199,7 +202,7 @@ class RecordBatchWriter {
 
     // Write the data header at the end
     RETURN_NOT_OK(dst->Write(data_header->data(), data_header->size()));
-    *data_header_offset = position + data_header->size();
+    *header_end_offset = position + data_header->size();
 
     return Align(dst, &position);
   }
@@ -218,9 +221,10 @@ class RecordBatchWriter {
   // This must be called after invoking AssemblePayload
   Status GetTotalSize(int64_t* size) {
     // emulates the behavior of Write without actually writing
+    int64_t body_offset;
     int64_t data_header_offset;
     MockOutputStream dst;
-    RETURN_NOT_OK(Write(&dst, &data_header_offset));
+    RETURN_NOT_OK(Write(&dst, &body_offset, &data_header_offset));
     *size = dst.GetExtentBytesWritten();
     return Status::OK();
   }
@@ -237,12 +241,12 @@ class RecordBatchWriter {
 };
 
 Status WriteRecordBatch(const std::vector<std::shared_ptr<Array>>& columns,
-    int32_t num_rows, io::OutputStream* dst, int64_t* header_offset,
-    int max_recursion_depth) {
+    int32_t num_rows, io::OutputStream* dst, int64_t* body_end_offset,
+    int64_t* header_end_offset, int max_recursion_depth) {
   DCHECK_GT(max_recursion_depth, 0);
   RecordBatchWriter serializer(columns, num_rows, max_recursion_depth);
   RETURN_NOT_OK(serializer.AssemblePayload());
-  return serializer.Write(dst, header_offset);
+  return serializer.Write(dst, body_end_offset, header_end_offset);
 }
 
 Status GetRecordBatchSize(const RecordBatch* batch, int64_t* size) {

cpp/src/arrow/ipc/adapter.h

@@ -61,12 +61,12 @@ constexpr int kMaxIpcRecursionDepth = 64;
 //
 // <int32: metadata size> <uint8*: metadata>
 //
-// Finally, the absolute offset (relative to the start of the output stream) to
-// the end of the metadata / data header (suffixed by the header size) is
-// returned in an out-variable
+// Finally, the absolute offsets (relative to the start of the output stream)
+// to the end of the body and end of the metadata / data header (suffixed by
+// the header size) is returned in out-variables
 ARROW_EXPORT Status WriteRecordBatch(const std::vector<std::shared_ptr<Array>>& columns,
-    int32_t num_rows, io::OutputStream* dst, int64_t* header_offset,
-    int max_recursion_depth = kMaxIpcRecursionDepth);
+    int32_t num_rows, io::OutputStream* dst, int64_t* body_end_offset,
+    int64_t* header_end_offset, int max_recursion_depth = kMaxIpcRecursionDepth);
 
 // int64_t GetRecordBatchMetadata(const RecordBatch* batch);
 

cpp/src/arrow/ipc/file.cc

@@ -87,17 +87,18 @@ Status FileWriter::WriteRecordBatch(
 
   int64_t offset = position_;
 
-  int64_t header_offset;
-  RETURN_NOT_OK(arrow::ipc::WriteRecordBatch(columns, num_rows, sink_, &header_offset));
+  int64_t body_end_offset;
+  int64_t header_end_offset;
+  RETURN_NOT_OK(arrow::ipc::WriteRecordBatch(
+      columns, num_rows, sink_, &body_end_offset, &header_end_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;
+  // There may be padding ever the end of the metadata, so we cannot rely on
+  // position_
+  int32_t metadata_length = body_end_offset - header_end_offset;
+  int32_t body_length = body_end_offset - offset;
 
   // Append metadata, to be written in the footer latera
   record_batches_.emplace_back(offset, metadata_length, body_length);
@@ -177,11 +178,12 @@ Status FileReader::ReadFooter() {
       footer_offset_ - footer_length - file_end_size, footer_length, &buffer));
   RETURN_NOT_OK(FileFooter::Open(buffer, &footer_));
 
-  return Status::OK();
+  // Get the schema
+  return footer_->GetSchema(&schema_);
 }
 
-Status FileReader::GetSchema(std::shared_ptr<Schema>* schema) const {
-  return footer_->GetSchema(schema);
+const std::shared_ptr<Schema>& FileReader::schema() const {
+  return schema_;
 }
 
 int FileReader::num_dictionaries() const {
@@ -199,12 +201,13 @@ MetadataVersion::type FileReader::version() const {
 Status FileReader::GetRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) {
   DCHECK_GT(i, 0);
   DCHECK_LT(i, num_record_batches());
-  // FileBlock block = footer_->record_batch(i);
+  FileBlock block = footer_->record_batch(i);
+  int64_t metadata_end_offset = block.offset + block.body_length + block.metadata_length;
 
-  //
-  std::shared_ptr<Buffer> buffer;
-  // RETURN_NOT_OK(file_->ReadAt
-  return Status::OK();
+  std::shared_ptr<RecordBatchReader> reader;
+  RETURN_NOT_OK(RecordBatchReader::Open(file_.get(), metadata_end_offset, &reader));
+
+  return reader->GetRecordBatch(schema_, batch);
 }
 
 }  // namespace ipc

cpp/src/arrow/ipc/file.h

@@ -106,9 +106,7 @@ class ARROW_EXPORT FileReader {
   static Status Open(const std::shared_ptr<io::ReadableFileInterface>& file,
       int64_t footer_offset, std::shared_ptr<FileReader>* reader);
 
-  // The Arrow schema shared by all of the record batches
-  // @param schema (out): arrow::Schema
-  Status GetSchema(std::shared_ptr<Schema>* schema) const;
+  const std::shared_ptr<Schema>& schema() const;
 
   // Shared dictionaries for dictionary-encoding cross record batches
   // TODO(wesm): Implement dictionary reading when we also have dictionary
@@ -139,6 +137,7 @@ class ARROW_EXPORT FileReader {
   int64_t footer_offset_;
 
   std::unique_ptr<FileFooter> footer_;
+  std::shared_ptr<Schema> schema_;
 };
 
 }  // namespace ipc

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

@@ -60,13 +60,15 @@ class TestWriteRecordBatch : public ::testing::TestWithParam<MakeRecordBatch*>,
       std::shared_ptr<RecordBatch>* batch_result) {
     std::string path = "test-write-row-batch";
     io::MemoryMapFixture::InitMemoryMap(memory_map_size, path, &mmap_);
-    int64_t header_offset;
 
-    RETURN_NOT_OK(
-        WriteRecordBatch(batch.columns(), batch.num_rows(), mmap_.get(), &header_offset));
+    int64_t body_end_offset;
+    int64_t header_end_offset;
+
+    RETURN_NOT_OK(WriteRecordBatch(batch.columns(), batch.num_rows(), mmap_.get(),
+        &body_end_offset, &header_end_offset));
 
     std::shared_ptr<RecordBatchReader> reader;
-    RETURN_NOT_OK(RecordBatchReader::Open(mmap_.get(), header_offset, &reader));
+    RETURN_NOT_OK(RecordBatchReader::Open(mmap_.get(), header_end_offset, &reader));
 
     RETURN_NOT_OK(reader->GetRecordBatch(batch.schema(), batch_result));
     return Status::OK();
@@ -282,9 +284,10 @@ INSTANTIATE_TEST_CASE_P(RoundTripTests, TestWriteRecordBatch,
 void TestGetRecordBatchSize(std::shared_ptr<RecordBatch> batch) {
   ipc::MockOutputStream mock;
   int64_t mock_header_offset = -1;
+  int64_t mock_body_offset = -1;
   int64_t size = -1;
-  ASSERT_OK(
-      WriteRecordBatch(batch->columns(), batch->num_rows(), &mock, &mock_header_offset));
+  ASSERT_OK(WriteRecordBatch(batch->columns(), batch->num_rows(), &mock,
+      &mock_body_offset, &mock_header_offset));
   ASSERT_OK(GetRecordBatchSize(batch.get(), &size));
   ASSERT_EQ(mock.GetExtentBytesWritten(), size);
 }
@@ -335,14 +338,17 @@ class RecursionLimits : public ::testing::Test, public io::MemoryMapFixture {
     std::string path = "test-write-past-max-recursion";
     const int memory_map_size = 1 << 16;
     io::MemoryMapFixture::InitMemoryMap(memory_map_size, path, &mmap_);
+
+    int64_t body_offset;
     int64_t header_offset;
+
     int64_t* header_out_param = header_out == nullptr ? &header_offset : header_out;
     if (override_level) {
       return WriteRecordBatch(batch->columns(), batch->num_rows(), mmap_.get(),
-          header_out_param, recursion_level + 1);
+          &body_offset, header_out_param, recursion_level + 1);
     } else {
-      return WriteRecordBatch(
-          batch->columns(), batch->num_rows(), mmap_.get(), header_out_param);
+      return WriteRecordBatch(batch->columns(), batch->num_rows(), mmap_.get(),
+          &body_offset, header_out_param);
     }
   }