Start drafting FileReader IPC implementation. Change record batch data header to write metadata size int32_t as suffix rather than prefix

Change-Id: I769a8ad47747f6a885701cab08d53c5a0c73c492

Author: wesm

cpp/src/arrow/ipc/adapter.cc

@@ -199,7 +199,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_offset = position + data_header->size();
 
     return Align(dst, &position);
   }
@@ -256,8 +256,6 @@ Status GetRecordBatchSize(const RecordBatch* batch, int64_t* size) {
 // ----------------------------------------------------------------------
 // Record batch read path
 
-static constexpr int64_t INIT_METADATA_SIZE = 4096;
-
 class RecordBatchReader::RecordBatchReaderImpl {
  public:
   RecordBatchReaderImpl(io::ReadableFileInterface* file,
@@ -377,29 +375,31 @@ class RecordBatchReader::RecordBatchReaderImpl {
   int num_flattened_fields_;
 };
 
-Status RecordBatchReader::Open(io::ReadableFileInterface* file, int64_t position,
+Status RecordBatchReader::Open(io::ReadableFileInterface* file, int64_t offset,
     std::shared_ptr<RecordBatchReader>* out) {
-  return Open(file, position, kMaxIpcRecursionDepth, out);
+  return Open(file, offset, kMaxIpcRecursionDepth, out);
 }
 
-Status RecordBatchReader::Open(io::ReadableFileInterface* file, int64_t position,
+Status RecordBatchReader::Open(io::ReadableFileInterface* file, int64_t offset,
     int max_recursion_depth, std::shared_ptr<RecordBatchReader>* out) {
-  std::shared_ptr<Buffer> metadata;
-  RETURN_NOT_OK(file->ReadAt(position, INIT_METADATA_SIZE, &metadata));
+  std::shared_ptr<Buffer> buffer;
+  RETURN_NOT_OK(file->ReadAt(offset - sizeof(int32_t), sizeof(int32_t), &buffer));
 
-  int32_t metadata_size = *reinterpret_cast<const int32_t*>(metadata->data());
+  int32_t metadata_size = *reinterpret_cast<const int32_t*>(buffer->data());
 
-  // We may not need to call ReadAt again
-  if (metadata_size > static_cast<int>(INIT_METADATA_SIZE - sizeof(int32_t))) {
-    // We don't have enough data, read the indicated metadata size.
-    RETURN_NOT_OK(file->ReadAt(position + sizeof(int32_t), metadata_size, &metadata));
+  if (metadata_size + static_cast<int>(sizeof(int32_t)) > offset) {
+    return Status::Invalid("metadata size invalid");
   }
 
+  // Read the metadata
+  RETURN_NOT_OK(
+      file->ReadAt(offset - metadata_size - sizeof(int32_t), metadata_size, &buffer));
+
   // TODO(wesm): buffer slicing here would be better in case ReadAt returns
   // allocated memory
 
   std::shared_ptr<Message> message;
-  RETURN_NOT_OK(Message::Open(metadata, &message));
+  RETURN_NOT_OK(Message::Open(buffer, &message));
 
   if (message->type() != Message::RECORD_BATCH) {
     return Status::Invalid("Metadata message is not a record batch");

cpp/src/arrow/ipc/adapter.h

@@ -62,7 +62,8 @@ constexpr int kMaxIpcRecursionDepth = 64;
 // <int32: metadata size> <uint8*: metadata>
 //
 // Finally, the absolute offset (relative to the start of the output stream) to
-// the start of the metadata / data header is returned in an out-variable
+// the end of the metadata / data header (suffixed by the header size) is
+// returned in an out-variable
 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);
@@ -79,10 +80,12 @@ ARROW_EXPORT Status GetRecordBatchSize(const RecordBatch* batch, int64_t* size);
 
 class ARROW_EXPORT RecordBatchReader {
  public:
-  static Status Open(io::ReadableFileInterface* file, int64_t position,
+  // The offset is the absolute position to the *end* of the record batch data
+  // header
+  static Status Open(io::ReadableFileInterface* file, int64_t offset,
       std::shared_ptr<RecordBatchReader>* out);
 
-  static Status Open(io::ReadableFileInterface* file, int64_t position,
+  static Status Open(io::ReadableFileInterface* file, int64_t offset,
       int max_recursion_depth, std::shared_ptr<RecordBatchReader>* out);
 
   virtual ~RecordBatchReader();

cpp/src/arrow/ipc/file.cc

@@ -126,5 +126,86 @@ Status FileWriter::Close() {
 // ----------------------------------------------------------------------
 // Reader implementation
 
+FileReader::FileReader(
+    const std::shared_ptr<io::ReadableFileInterface>& file, int64_t footer_offset)
+    : file_(file), footer_offset_(footer_offset) {}
+
+FileReader::~FileReader() {}
+
+Status FileReader::Open(const std::shared_ptr<io::ReadableFileInterface>& file,
+    std::shared_ptr<FileReader>* reader) {
+  int64_t footer_offset;
+  RETURN_NOT_OK(file->GetSize(&footer_offset));
+  return Open(file, footer_offset, reader);
+}
+
+Status FileReader::Open(const std::shared_ptr<io::ReadableFileInterface>& file,
+    int64_t footer_offset, std::shared_ptr<FileReader>* reader) {
+  *reader = std::shared_ptr<FileReader>(new FileReader(file, footer_offset));
+  return Status::OK();
+}
+
+Status FileReader::ReadFooter() {
+  int magic_size = static_cast<int>(strlen(kArrowMagicBytes));
+
+  if (footer_offset_ <= magic_size * 2 + 4) {
+    std::stringstream ss;
+    ss << "File is too small: " << footer_offset_;
+    return Status::Invalid(ss.str());
+  }
+
+  RETURN_NOT_OK(file_->Seek(footer_offset_));
+
+  std::shared_ptr<Buffer> buffer;
+
+  int file_end_size = magic_size + sizeof(int32_t);
+
+  RETURN_NOT_OK(file_->ReadAt(footer_offset_ - file_end_size, file_end_size, &buffer));
+
+  if (memcmp(buffer->data(), kArrowMagicBytes, magic_size)) {
+    return Status::Invalid("Not an Arrow file");
+  }
+
+  int32_t footer_length = *reinterpret_cast<const int32_t*>(buffer->data());
+
+  if (footer_length <= 0 || footer_length + magic_size * 2 + 4 > footer_offset_) {
+    return Status::Invalid("File is smaller than indicated metadata size");
+  }
+
+  // Now read the footer
+  RETURN_NOT_OK(file_->ReadAt(
+      footer_offset_ - footer_length - file_end_size, footer_length, &buffer));
+  RETURN_NOT_OK(FileFooter::Open(buffer, &footer_));
+
+  return Status::OK();
+}
+
+Status FileReader::GetSchema(std::shared_ptr<Schema>* schema) const {
+  return footer_->GetSchema(schema);
+}
+
+int FileReader::num_dictionaries() const {
+  return footer_->num_dictionaries();
+}
+
+int FileReader::num_record_batches() const {
+  return footer_->num_record_batches();
+}
+
+MetadataVersion::type FileReader::version() const {
+  return footer_->version();
+}
+
+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);
+
+  //
+  std::shared_ptr<Buffer> buffer;
+  // RETURN_NOT_OK(file_->ReadAt
+  return Status::OK();
+}
+
 }  // namespace ipc
 }  // namespace arrow

cpp/src/arrow/ipc/file.h

@@ -86,12 +86,15 @@ class ARROW_EXPORT FileWriter {
 
 class ARROW_EXPORT FileReader {
  public:
+  ~FileReader();
+
   // Open a file-like object that is assumed to be self-contained; i.e., the
   // end of the file interface is the end of the Arrow file. Note that there
   // can be any amount of data preceding the Arrow-formatted data, because we
   // need only locate the end of the Arrow file stream to discover the metadata
   // and then proceed to read the data into memory.
-  static Status Open(const std::shared_ptr<io::ReadableFileInterface>& file);
+  static Status Open(const std::shared_ptr<io::ReadableFileInterface>& file,
+      std::shared_ptr<FileReader>* reader);
 
   // If the file is embedded within some larger file or memory region, you can
   // pass the absolute memory offset to the end of the file (which contains the
@@ -100,12 +103,12 @@ class ARROW_EXPORT FileReader {
   //
   // @param file: the data source
   // @param footer_offset: the position of the end of the Arrow "file"
-  static Status Open(
-      const std::shared_ptr<io::ReadableFileInterface>& file, int64_t footer_offset);
+  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);
+  Status GetSchema(std::shared_ptr<Schema>* schema) const;
 
   // Shared dictionaries for dictionary-encoding cross record batches
   // TODO(wesm): Implement dictionary reading when we also have dictionary
@@ -114,6 +117,8 @@ class ARROW_EXPORT FileReader {
 
   int num_record_batches() const;
 
+  MetadataVersion::type version() const;
+
   // Read a record batch from the file. Does not copy memory if the input
   // source supports zero-copy.
   //
@@ -122,13 +127,18 @@ class ARROW_EXPORT FileReader {
   Status GetRecordBatch(int i, std::shared_ptr<RecordBatch>* batch);
 
  private:
-  FileReader(const std::shared_ptr<io::ReadableFileInterface>& file);
+  FileReader(
+      const std::shared_ptr<io::ReadableFileInterface>& file, int64_t footer_offset);
+
+  Status ReadFooter();
 
   std::shared_ptr<io::ReadableFileInterface> file_;
 
   // The location where the Arrow file layout ends. May be the end of the file
   // or some other location if embedded in a larger file.
   int64_t footer_offset_;
+
+  std::unique_ptr<FileFooter> footer_;
 };
 
 }  // namespace ipc

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

@@ -60,13 +60,13 @@ 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_location;
+    int64_t header_offset;
 
-    RETURN_NOT_OK(WriteRecordBatch(
-        batch.columns(), batch.num_rows(), mmap_.get(), &header_location));
+    RETURN_NOT_OK(
+        WriteRecordBatch(batch.columns(), batch.num_rows(), mmap_.get(), &header_offset));
 
     std::shared_ptr<RecordBatchReader> reader;
-    RETURN_NOT_OK(RecordBatchReader::Open(mmap_.get(), header_location, &reader));
+    RETURN_NOT_OK(RecordBatchReader::Open(mmap_.get(), header_offset, &reader));
 
     RETURN_NOT_OK(reader->GetRecordBatch(batch.schema(), batch_result));
     return Status::OK();
@@ -281,10 +281,10 @@ INSTANTIATE_TEST_CASE_P(RoundTripTests, TestWriteRecordBatch,
 
 void TestGetRecordBatchSize(std::shared_ptr<RecordBatch> batch) {
   ipc::MockOutputStream mock;
-  int64_t mock_header_location = -1;
+  int64_t mock_header_offset = -1;
   int64_t size = -1;
-  ASSERT_OK(WriteRecordBatch(
-      batch->columns(), batch->num_rows(), &mock, &mock_header_location));
+  ASSERT_OK(
+      WriteRecordBatch(batch->columns(), batch->num_rows(), &mock, &mock_header_offset));
   ASSERT_OK(GetRecordBatchSize(batch.get(), &size));
   ASSERT_EQ(mock.GetExtentBytesWritten(), size);
 }
@@ -335,8 +335,8 @@ 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 header_location;
-    int64_t* header_out_param = header_out == nullptr ? &header_location : header_out;
+    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);
@@ -356,12 +356,12 @@ TEST_F(RecursionLimits, WriteLimit) {
 }
 
 TEST_F(RecursionLimits, ReadLimit) {
-  int64_t header_location = -1;
+  int64_t header_offset = -1;
   std::shared_ptr<Schema> schema;
-  ASSERT_OK(WriteToMmap(64, true, &header_location, &schema));
+  ASSERT_OK(WriteToMmap(64, true, &header_offset, &schema));
 
   std::shared_ptr<RecordBatchReader> reader;
-  ASSERT_OK(RecordBatchReader::Open(mmap_.get(), header_location, &reader));
+  ASSERT_OK(RecordBatchReader::Open(mmap_.get(), header_offset, &reader));
   std::shared_ptr<RecordBatch> batch_result;
   ASSERT_RAISES(Invalid, reader->GetRecordBatch(schema, &batch_result));
 }

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

@@ -108,7 +108,7 @@ class TestFileFooter : public ::testing::Test {
 
     ASSERT_OK(WriteFileFooter(schema, dictionaries, record_batches, &stream));
 
-    std::shared_ptr<FileFooter> footer;
+    std::unique_ptr<FileFooter> footer;
     ASSERT_OK(FileFooter::Open(buffer, &footer));
 
     ASSERT_EQ(MetadataVersion::V1_SNAPSHOT, footer->version());

cpp/src/arrow/ipc/metadata-internal.cc

@@ -304,17 +304,17 @@ Status MessageBuilder::Finish() {
 }
 
 Status MessageBuilder::GetBuffer(std::shared_ptr<Buffer>* out) {
-  // The message buffer is prefixed by the size of the complete flatbuffer as
+  // The message buffer is suffixed by the size of the complete flatbuffer as
   // int32_t
-  // <int32_t: flatbuffer size><uint8_t*: flatbuffer data>
+  // <uint8_t*: flatbuffer data><int32_t: flatbuffer size>
   int32_t size = fbb_.GetSize();
 
   auto result = std::make_shared<PoolBuffer>();
   RETURN_NOT_OK(result->Resize(size + sizeof(int32_t)));
 
   uint8_t* dst = result->mutable_data();
-  memcpy(dst, reinterpret_cast<int32_t*>(&size), sizeof(int32_t));
-  memcpy(dst + sizeof(int32_t), fbb_.GetBufferPointer(), size);
+  memcpy(dst, fbb_.GetBufferPointer(), size);
+  memcpy(dst + size, reinterpret_cast<int32_t*>(&size), sizeof(int32_t));
 
   *out = result;
   return Status::OK();

cpp/src/arrow/ipc/metadata.cc

@@ -84,9 +84,7 @@ Status Message::Open(
     const std::shared_ptr<Buffer>& buffer, std::shared_ptr<Message>* out) {
   std::shared_ptr<Message> result(new Message());
 
-  // The buffer is prefixed by its size as int32_t
-  const uint8_t* fb_head = buffer->data() + sizeof(int32_t);
-  const flatbuf::Message* message = flatbuf::GetMessage(fb_head);
+  const flatbuf::Message* message = flatbuf::GetMessage(buffer->data());
 
   // TODO(wesm): verify message
   result->impl_.reset(new MessageImpl(buffer, message));
@@ -300,11 +298,13 @@ class FileFooter::FileFooterImpl {
 
 FileFooter::FileFooter() {}
 
+FileFooter::~FileFooter() {}
+
 Status FileFooter::Open(
-    const std::shared_ptr<Buffer>& buffer, std::shared_ptr<FileFooter>* out) {
+    const std::shared_ptr<Buffer>& buffer, std::unique_ptr<FileFooter>* out) {
   const flatbuf::Footer* footer = flatbuf::GetFooter(buffer->data());
 
-  *out = std::shared_ptr<FileFooter>(new FileFooter());
+  *out = std::unique_ptr<FileFooter>(new FileFooter());
 
   // TODO(wesm): Verify the footer
   (*out)->impl_.reset(new FileFooterImpl(buffer, footer));
@@ -318,7 +318,7 @@ int FileFooter::num_dictionaries() const {
 
 int FileFooter::num_record_batches() const {
   return impl_->num_record_batches();
-};
+}
 
 MetadataVersion::type FileFooter::version() const {
   return impl_->version();

cpp/src/arrow/ipc/metadata.h

@@ -163,8 +163,10 @@ Status WriteFileFooter(const Schema* schema, const std::vector<FileBlock>& dicti
 
 class ARROW_EXPORT FileFooter {
  public:
+  ~FileFooter();
+
   static Status Open(
-      const std::shared_ptr<Buffer>& buffer, std::shared_ptr<FileFooter>* out);
+      const std::shared_ptr<Buffer>& buffer, std::unique_ptr<FileFooter>* out);
 
   int num_dictionaries() const;
   int num_record_batches() const;