Refactor IPC/File record batch read/write structure to reflect discussion in ARROW-384

Change-Id: Ic119c5647bdb6b1ff1fb31b3712527058c68e40c
This commit is contained in:
Wes McKinney committed 2016-11-28 13:13:12 -05:00
1 parent dbd6ed6d5e
commit 85128f7add
14 files changed
+399 -276

No files matched your search

-1
View File
@@ -528,7 +528,6 @@ if(ARROW_BUILD_TESTS)
ExternalProject_Add(gflags_ep
GIT_REPOSITORY https://github.com/gflags/gflags.git
GIT_TAG cce68f0c9c5d054017425e6e6fd54f696d36e8ee
# URL "https://github.com/gflags/gflags/archive/v${GFLAGS_VERSION}.tar.gz"
BUILD_IN_SOURCE 1
CMAKE_ARGS -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
-DCMAKE_INSTALL_PREFIX=${GFLAGS_PREFIX}
+17 -8
View File
@@ -258,8 +258,11 @@ Status BufferOutputStream::Reserve(int64_t nbytes) {
// ----------------------------------------------------------------------
// In-memory buffer reader
BufferReader::BufferReader(const uint8_t* buffer, int buffer_size)
: buffer_(buffer), buffer_size_(buffer_size), position_(0) {}
BufferReader::BufferReader(const std::shared_ptr<Buffer>& buffer)
: buffer_(buffer), data_(buffer->data()), size_(buffer->size()), position_(0) {}
BufferReader::BufferReader(const uint8_t* data, int64_t size)
: buffer_(nullptr), data_(data), size_(size), position_(0) {}
BufferReader::~BufferReader() {}
@@ -278,26 +281,32 @@ bool BufferReader::supports_zero_copy() const {
}
Status BufferReader::Read(int64_t nbytes, int64_t* bytes_read, uint8_t* buffer) {
memcpy(buffer, buffer_ + position_, nbytes);
*bytes_read = std::min(nbytes, buffer_size_ - position_);
memcpy(buffer, data_ + position_, nbytes);
*bytes_read = std::min(nbytes, size_ - position_);
position_ += *bytes_read;
return Status::OK();
}
Status BufferReader::Read(int64_t nbytes, std::shared_ptr<Buffer>* out) {
int64_t size = std::min(nbytes, buffer_size_ - position_);
*out = std::make_shared<Buffer>(buffer_ + position_, size);
int64_t size = std::min(nbytes, size_ - position_);
if (buffer_ != nullptr) {
*out = SliceBuffer(buffer_, position_, size);
} else {
*out = std::make_shared<Buffer>(data_ + position_, size);
}
position_ += nbytes;
return Status::OK();
}
Status BufferReader::GetSize(int64_t* size) {
*size = buffer_size_;
*size = size_;
return Status::OK();
}
Status BufferReader::Seek(int64_t position) {
if (position < 0 || position >= buffer_size_) {
if (position < 0 || position >= size_) {
return Status::IOError("position out of bounds");
}
+5 -3
View File
@@ -99,7 +99,8 @@ class ARROW_EXPORT MemoryMappedFile : public ReadWriteFileInterface {
class ARROW_EXPORT BufferReader : public ReadableFileInterface {
public:
BufferReader(const uint8_t* buffer, int buffer_size);
BufferReader(const std::shared_ptr<Buffer>& buffer);
BufferReader(const uint8_t* data, int64_t size);
~BufferReader();
Status Close() override;
@@ -116,8 +117,9 @@ class ARROW_EXPORT BufferReader : public ReadableFileInterface {
bool supports_zero_copy() const override;
private:
const uint8_t* buffer_;
int buffer_size_;
std::shared_ptr<Buffer> buffer_;
const uint8_t* data_;
int64_t size_;
int64_t position_;
};
+128 -99
View File
@@ -124,30 +124,30 @@ Status VisitArray(const Array* arr, std::vector<flatbuf::FieldNode>* field_nodes
class RecordBatchWriter {
public:
RecordBatchWriter(const std::vector<std::shared_ptr<Array>>& columns, int32_t num_rows,
int max_recursion_depth)
int64_t buffer_start_offset, int max_recursion_depth)
: columns_(&columns),
num_rows_(num_rows),
buffer_start_offset_(buffer_start_offset),
max_recursion_depth_(max_recursion_depth) {}
Status AssemblePayload() {
Status AssemblePayload(int64_t* body_length) {
if (field_nodes_.size() > 0) {
field_nodes_.clear();
buffer_meta_.clear();
buffers_.clear();
}
// Perform depth-first traversal of the row-batch
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* body_end_offset, int64_t* header_end_offset) {
// Get the starting 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;
// The position for the start of a buffer relative to the passed frame of
// reference. May be 0 or some other position in an address space
int64_t offset = buffer_start_offset_;
// Construct the buffer metadata for the record batch header
for (size_t i = 0; i < buffers_.size(); ++i) {
const Buffer* buffer = buffers_[i].get();
int64_t size = 0;
@@ -161,65 +161,103 @@ class RecordBatchWriter {
// TODO(wesm): We currently have no notion of shared memory page id's,
// but we've included it in the metadata IDL for when we have it in the
// future. Use page=0 for now
// future. Use page = -1 for now
//
// Note that page ids are a bespoke notion for Arrow and not a feature we
// 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, size + padding));
if (size > 0) {
RETURN_NOT_OK(dst->Write(buffer->data(), size));
position += size;
}
if (padding > 0) {
RETURN_NOT_OK(dst->Write(kPaddingBytes, padding));
position += padding;
}
buffer_meta_.push_back(flatbuf::Buffer(-1, offset, size + padding));
offset += size + padding;
}
*body_end_offset = position;
*body_length = offset - buffer_start_offset_;
DCHECK(BitUtil::IsMultipleOf64(*body_length));
return Status::OK();
}
Status WriteMetadata(
int64_t body_length, io::OutputStream* dst, int32_t* metadata_length) {
// 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
//
// Note: The memory written here is prefixed by the size of the flatbuffer
// itself as an int32_t. On reading from a input, you will have to
// 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;
// itself as an int32_t.
std::shared_ptr<Buffer> metadata_fb;
RETURN_NOT_OK(WriteRecordBatchMetadata(
num_rows_, position - start_position, field_nodes_, buffer_meta_, &data_header));
num_rows_, body_length, field_nodes_, buffer_meta_, &metadata_fb));
// Write the data header at the end
RETURN_NOT_OK(dst->Write(data_header->data(), data_header->size()));
// Need to write 4 bytes (metadata size), the metadata, plus padding to
// fall on a 64-byte offset
int64_t padded_metadata_length =
BitUtil::RoundUpToMultipleOf64(metadata_fb->size() + 4);
position += data_header->size();
*header_end_offset = position;
// The returned metadata size includes the length prefix, the flatbuffer,
// plus padding
*metadata_length = padded_metadata_length;
return Align(dst, &position);
// Write the flatbuffer size prefix
int32_t flatbuffer_size = metadata_fb->size();
RETURN_NOT_OK(
dst->Write(reinterpret_cast<const uint8_t*>(&flatbuffer_size), sizeof(int32_t)));
// Write the flatbuffer
RETURN_NOT_OK(dst->Write(metadata_fb->data(), metadata_fb->size()));
// Write any padding
int64_t padding = padded_metadata_length - metadata_fb->size() - 4;
if (padding > 0) { RETURN_NOT_OK(dst->Write(kPaddingBytes, padding)); }
return Status::OK();
}
Status Write(io::OutputStream* dst, int32_t* metadata_length, int64_t* body_length) {
RETURN_NOT_OK(AssemblePayload(body_length));
#ifndef NDEBUG
int64_t start_position, current_position;
RETURN_NOT_OK(dst->Tell(&start_position));
#endif
RETURN_NOT_OK(WriteMetadata(*body_length, dst, metadata_length));
#ifndef NDEBUG
RETURN_NOT_OK(dst->Tell(&current_position));
DCHECK(BitUtil::IsMultipleOf64(current_position));
#endif
// Now write the buffers
for (size_t i = 0; i < buffers_.size(); ++i) {
const Buffer* buffer = buffers_[i].get();
int64_t size = 0;
int64_t padding = 0;
// The buffer might be null if we are handling zero row lengths.
if (buffer) {
size = buffer->size();
padding = BitUtil::RoundUpToMultipleOf64(size) - size;
}
if (size > 0) { RETURN_NOT_OK(dst->Write(buffer->data(), size)); }
if (padding > 0) { RETURN_NOT_OK(dst->Write(kPaddingBytes, padding)); }
}
#ifndef NDEBUG
RETURN_NOT_OK(dst->Tell(&current_position));
DCHECK(BitUtil::IsMultipleOf64(current_position));
#endif
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();
}
// 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;
int32_t metadata_length;
int64_t body_length;
MockOutputStream dst;
RETURN_NOT_OK(Write(&dst, &body_offset, &data_header_offset));
RETURN_NOT_OK(Write(&dst, &metadata_length, &body_length));
*size = dst.GetExtentBytesWritten();
return Status::OK();
}
@@ -228,6 +266,7 @@ class RecordBatchWriter {
// Do not copy this vector. Ownership must be retained elsewhere
const std::vector<std::shared_ptr<Array>>* columns_;
int32_t num_rows_;
int64_t buffer_start_offset_;
std::vector<flatbuf::FieldNode> field_nodes_;
std::vector<flatbuf::Buffer> buffer_meta_;
@@ -236,18 +275,17 @@ class RecordBatchWriter {
};
Status WriteRecordBatch(const std::vector<std::shared_ptr<Array>>& columns,
int32_t num_rows, io::OutputStream* dst, int64_t* body_end_offset,
int64_t* header_end_offset, int max_recursion_depth) {
int32_t num_rows, int64_t buffer_start_offset, io::OutputStream* dst,
int32_t* metadata_length, int64_t* body_length, 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, body_end_offset, header_end_offset);
RecordBatchWriter serializer(
columns, num_rows, buffer_start_offset, max_recursion_depth);
return serializer.Write(dst, metadata_length, body_length);
}
Status GetRecordBatchSize(const RecordBatch* batch, int64_t* size) {
RecordBatchWriter serializer(
batch->columns(), batch->num_rows(), kMaxIpcRecursionDepth);
RETURN_NOT_OK(serializer.AssemblePayload());
batch->columns(), batch->num_rows(), 0, kMaxIpcRecursionDepth);
RETURN_NOT_OK(serializer.GetTotalSize(size));
return Status::OK();
}
@@ -255,30 +293,33 @@ Status GetRecordBatchSize(const RecordBatch* batch, int64_t* size) {
// ----------------------------------------------------------------------
// Record batch read path
class RecordBatchReader::RecordBatchReaderImpl {
class RecordBatchReader {
public:
RecordBatchReaderImpl(io::ReadableFileInterface* file,
const std::shared_ptr<RecordBatchMessage>& metadata, int max_recursion_depth)
: file_(file), metadata_(metadata), max_recursion_depth_(max_recursion_depth) {
RecordBatchReader(const std::shared_ptr<RecordBatchMetadata>& metadata,
const std::shared_ptr<Schema>& schema, int max_recursion_depth,
io::ReadableFileInterface* file)
: metadata_(metadata),
schema_(schema),
max_recursion_depth_(max_recursion_depth),
file_(file) {
num_buffers_ = metadata->num_buffers();
num_flattened_fields_ = metadata->num_fields();
}
Status AssembleBatch(
const std::shared_ptr<Schema>& schema, std::shared_ptr<RecordBatch>* out) {
std::vector<std::shared_ptr<Array>> arrays(schema->num_fields());
Status Read(std::shared_ptr<RecordBatch>* out) {
std::vector<std::shared_ptr<Array>> arrays(schema_->num_fields());
// The field_index and buffer_index are incremented in NextArray based on
// how much of the batch is "consumed" (through nested data reconstruction,
// for example)
field_index_ = 0;
buffer_index_ = 0;
for (int i = 0; i < schema->num_fields(); ++i) {
const Field* field = schema->field(i).get();
for (int i = 0; i < schema_->num_fields(); ++i) {
const Field* field = schema_->field(i).get();
RETURN_NOT_OK(NextArray(field, max_recursion_depth_, &arrays[i]));
}
*out = std::make_shared<RecordBatch>(schema, metadata_->length(), arrays);
*out = std::make_shared<RecordBatch>(schema_, metadata_->length(), arrays);
return Status::OK();
}
@@ -375,65 +416,53 @@ class RecordBatchReader::RecordBatchReaderImpl {
}
private:
std::shared_ptr<RecordBatchMetadata> metadata_;
std::shared_ptr<Schema> schema_;
int max_recursion_depth_;
io::ReadableFileInterface* file_;
std::shared_ptr<RecordBatchMessage> metadata_;
int field_index_;
int buffer_index_;
int max_recursion_depth_;
int num_buffers_;
int num_flattened_fields_;
};
Status RecordBatchReader::Open(io::ReadableFileInterface* file, int64_t offset,
std::shared_ptr<RecordBatchReader>* out) {
return Open(file, offset, kMaxIpcRecursionDepth, out);
}
Status RecordBatchReader::Open(io::ReadableFileInterface* file, int64_t offset,
int max_recursion_depth, std::shared_ptr<RecordBatchReader>* out) {
Status ReadRecordBatchMetadata(int64_t offset, int32_t metadata_length,
io::ReadableFileInterface* file, std::shared_ptr<RecordBatchMetadata>* metadata) {
std::shared_ptr<Buffer> buffer;
RETURN_NOT_OK(file->ReadAt(offset - sizeof(int32_t), sizeof(int32_t), &buffer));
RETURN_NOT_OK(file->ReadAt(offset, metadata_length, &buffer));
int32_t metadata_size = *reinterpret_cast<const int32_t*>(buffer->data());
int32_t flatbuffer_size = *reinterpret_cast<const int32_t*>(buffer->data());
if (metadata_size + static_cast<int>(sizeof(int32_t)) > offset) {
if (flatbuffer_size + static_cast<int>(sizeof(int32_t)) > metadata_length) {
std::stringstream ss;
ss << "metadata size " << metadata_size
<< " invalid. File offset: " << offset;
ss << "flatbuffer size " << metadata_length << " invalid. File offset: " << offset
<< ", metadata length: " << metadata_length;
return Status::Invalid(ss.str());
}
// 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(buffer, &message));
RETURN_NOT_OK(Message::Open(buffer, sizeof(int32_t), &message));
if (message->type() != Message::RECORD_BATCH) {
return Status::Invalid("Metadata message is not a record batch");
}
std::shared_ptr<RecordBatchMessage> batch_meta = message->GetRecordBatch();
std::shared_ptr<RecordBatchReader> result(new RecordBatchReader());
result->impl_.reset(new RecordBatchReaderImpl(file, batch_meta, max_recursion_depth));
*out = result;
*metadata = std::make_shared<RecordBatchMetadata>(message);
return Status::OK();
}
// Here the explicit destructor is required for compilers to be aware of
// the complete information of RecordBatchReader::RecordBatchReaderImpl class
RecordBatchReader::~RecordBatchReader() {}
Status ReadRecordBatch(const std::shared_ptr<RecordBatchMetadata>& metadata,
const std::shared_ptr<Schema>& schema, io::ReadableFileInterface* file,
std::shared_ptr<RecordBatch>* out) {
return ReadRecordBatch(metadata, schema, kMaxIpcRecursionDepth, file, out);
}
Status RecordBatchReader::GetRecordBatch(
const std::shared_ptr<Schema>& schema, std::shared_ptr<RecordBatch>* out) {
return impl_->AssembleBatch(schema, out);
Status ReadRecordBatch(const std::shared_ptr<RecordBatchMetadata>& metadata,
const std::shared_ptr<Schema>& schema, int max_recursion_depth,
io::ReadableFileInterface* file, std::shared_ptr<RecordBatch>* out) {
RecordBatchReader reader(metadata, schema, max_recursion_depth, file);
return reader.Read(out);
}
} // namespace ipc
+31 -30
View File
@@ -43,7 +43,7 @@ class OutputStream;
namespace ipc {
class RecordBatchMessage;
class RecordBatchMetadata;
// ----------------------------------------------------------------------
// Write path
@@ -51,22 +51,30 @@ class RecordBatchMessage;
// TODO(emkornfield) investigate this more
constexpr int kMaxIpcRecursionDepth = 64;
// Write the RecordBatch (collection of equal-length Arrow arrays) to the output
// stream
// Write the RecordBatch (collection of equal-length Arrow arrays) to the
// output stream in a contiguous block. The record batch metadata is written as
// a flatbuffer (see format/Message.fbs -- the RecordBatch message type)
// prefixed by its size, followed by each of the memory buffers in the batch
// written end to end (with appropriate alignment and padding):
//
// First, each of the memory buffers are written out end-to-end
//
// Then, this function writes the batch metadata as a flatbuffer (see
// format/Message.fbs -- the RecordBatch message type) like so:
//
// <int32: metadata size> <uint8*: metadata>
// <int32: metadata size> <uint8*: metadata> <buffers>
//
// 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
//
// @param(in) buffer_start_offset: the start offset to use in the buffer metadata,
// default should be 0
//
// @param(out) metadata_length: the size of the length-prefixed flatbuffer
// including padding to a 64-byte boundary
//
// @param(out) body_length: the size of the contiguous buffer block plus
// padding bytes
ARROW_EXPORT Status WriteRecordBatch(const std::vector<std::shared_ptr<Array>>& columns,
int32_t num_rows, io::OutputStream* dst, int64_t* body_end_offset,
int64_t* header_end_offset, int max_recursion_depth = kMaxIpcRecursionDepth);
int32_t num_rows, int64_t buffer_start_offset, io::OutputStream* dst,
int32_t* metadata_length, int64_t* body_length,
int max_recursion_depth = kMaxIpcRecursionDepth);
// int64_t GetRecordBatchMetadata(const RecordBatch* batch);
@@ -78,27 +86,20 @@ ARROW_EXPORT Status GetRecordBatchSize(const RecordBatch* batch, int64_t* size);
// ----------------------------------------------------------------------
// "Read" path; does not copy data if the input supports zero copy reads
class ARROW_EXPORT RecordBatchReader {
public:
// 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);
// Read the record batch flatbuffer metadata starting at the indicated file offset
//
// The flatbuffer is expected to be length-prefixed, so the metadata_length
// includes at least the length prefix and the flatbuffer
Status ARROW_EXPORT ReadRecordBatchMetadata(int64_t offset, int32_t metadata_length,
io::ReadableFileInterface* file, std::shared_ptr<RecordBatchMetadata>* metadata);
static Status Open(io::ReadableFileInterface* file, int64_t offset,
int max_recursion_depth, std::shared_ptr<RecordBatchReader>* out);
Status ARROW_EXPORT ReadRecordBatch(const std::shared_ptr<RecordBatchMetadata>& metadata,
const std::shared_ptr<Schema>& schema, io::ReadableFileInterface* file,
std::shared_ptr<RecordBatch>* out);
virtual ~RecordBatchReader();
// Reassemble the record batch. A Schema is required to be able to construct
// the right array containers
Status GetRecordBatch(
const std::shared_ptr<Schema>& schema, std::shared_ptr<RecordBatch>* out);
private:
class RecordBatchReaderImpl;
std::unique_ptr<RecordBatchReaderImpl> impl_;
};
Status ARROW_EXPORT ReadRecordBatch(const std::shared_ptr<RecordBatchMetadata>& metadata,
const std::shared_ptr<Schema>& schema, int max_recursion_depth,
io::ReadableFileInterface* file, std::shared_ptr<RecordBatch>* out);
} // namespace ipc
} // namespace arrow
+19 -12
View File
@@ -23,6 +23,7 @@
#include <vector>
#include "arrow/io/interfaces.h"
#include "arrow/io/memory.h"
#include "arrow/ipc/adapter.h"
#include "arrow/ipc/metadata.h"
#include "arrow/ipc/util.h"
@@ -87,19 +88,19 @@ Status FileWriter::WriteRecordBatch(
int64_t offset = position_;
int64_t body_end_offset;
int64_t header_end_offset;
// There may be padding ever the end of the metadata, so we cannot rely on
// position_
int32_t metadata_length;
int64_t body_length;
// Frame of reference in file format is 0, see ARROW-384
const int64_t buffer_start_offset = 0;
RETURN_NOT_OK(arrow::ipc::WriteRecordBatch(
columns, num_rows, sink_, &body_end_offset, &header_end_offset));
columns, num_rows, buffer_start_offset, sink_, &metadata_length, &body_length));
RETURN_NOT_OK(UpdatePosition());
DCHECK(position_ % 8 == 0) << "ipc::WriteRecordBatch did not perform aligned writes";
// There may be padding ever the end of the metadata, so we cannot rely on
// position_
int32_t metadata_length = header_end_offset - body_end_offset;
int32_t body_length = body_end_offset - offset;
// Append metadata, to be written in the footer later
record_batches_.emplace_back(offset, metadata_length, body_length);
@@ -198,12 +199,18 @@ Status FileReader::GetRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) {
DCHECK_GE(i, 0);
DCHECK_LT(i, num_record_batches());
FileBlock block = footer_->record_batch(i);
int64_t metadata_end_offset = block.offset + block.body_length + block.metadata_length;
std::shared_ptr<RecordBatchReader> reader;
RETURN_NOT_OK(RecordBatchReader::Open(file_.get(), metadata_end_offset, &reader));
std::shared_ptr<RecordBatchMetadata> metadata;
RETURN_NOT_OK(ReadRecordBatchMetadata(block.offset, block.metadata_length, file_.get(),
&metadata));
return reader->GetRecordBatch(schema_, batch);
// TODO(wesm): ARROW-388 -- the buffer frame of reference is 0 (see
// ARROW-384).
std::shared_ptr<Buffer> buffer_block;
RETURN_NOT_OK(file_->Read(block.body_length, &buffer_block));
io::BufferReader reader(buffer_block);
return ReadRecordBatch(metadata, schema_, &reader, batch);
}
} // namespace ipc
+50 -35
View File
@@ -54,17 +54,24 @@ class TestWriteRecordBatch : public ::testing::TestWithParam<MakeRecordBatch*>,
std::string path = "test-write-row-batch";
io::MemoryMapFixture::InitMemoryMap(memory_map_size, path, &mmap_);
int64_t body_end_offset;
int64_t header_end_offset;
int32_t metadata_length;
int64_t body_length;
RETURN_NOT_OK(WriteRecordBatch(batch.columns(), batch.num_rows(), mmap_.get(),
&body_end_offset, &header_end_offset));
const int64_t buffer_offset = 0;
std::shared_ptr<RecordBatchReader> reader;
RETURN_NOT_OK(RecordBatchReader::Open(mmap_.get(), header_end_offset, &reader));
RETURN_NOT_OK(WriteRecordBatch(batch.columns(), batch.num_rows(), buffer_offset,
mmap_.get(), &metadata_length, &body_length));
RETURN_NOT_OK(reader->GetRecordBatch(batch.schema(), batch_result));
return Status::OK();
std::shared_ptr<RecordBatchMetadata> metadata;
RETURN_NOT_OK(ReadRecordBatchMetadata(0, metadata_length, mmap_.get(), &metadata));
// The buffer offsets start at 0, so we must construct a
// ReadableFileInterface according to that frame of reference
std::shared_ptr<Buffer> buffer_payload;
RETURN_NOT_OK(mmap_->ReadAt(metadata_length, body_length, &buffer_payload));
io::BufferReader buffer_reader(buffer_payload);
return ReadRecordBatch(metadata, batch.schema(), &buffer_reader, batch_result);
}
protected:
@@ -96,11 +103,11 @@ 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;
int32_t mock_metadata_length = -1;
int64_t mock_body_length = -1;
int64_t size = -1;
ASSERT_OK(WriteRecordBatch(batch->columns(), batch->num_rows(), &mock,
&mock_body_offset, &mock_header_offset));
ASSERT_OK(WriteRecordBatch(batch->columns(), batch->num_rows(), 0, &mock,
&mock_metadata_length, &mock_body_length));
ASSERT_OK(GetRecordBatchSize(batch.get(), &size));
ASSERT_EQ(mock.GetExtentBytesWritten(), size);
}
@@ -129,39 +136,36 @@ class RecursionLimits : public ::testing::Test, public io::MemoryMapFixture {
void SetUp() { pool_ = default_memory_pool(); }
void TearDown() { io::MemoryMapFixture::TearDown(); }
Status WriteToMmap(int recursion_level, bool override_level,
int64_t* header_out = nullptr, std::shared_ptr<Schema>* schema_out = nullptr) {
Status WriteToMmap(int recursion_level, bool override_level, int32_t* metadata_length,
int64_t* body_length, std::shared_ptr<Schema>* schema) {
const int batch_length = 5;
TypePtr type = kInt32;
TypePtr type = int32();
ArrayPtr array;
const bool include_nulls = true;
RETURN_NOT_OK(MakeRandomInt32Array(1000, include_nulls, pool_, &array));
for (int i = 0; i < recursion_level; ++i) {
type = std::static_pointer_cast<DataType>(std::make_shared<ListType>(type));
type = list(type);
RETURN_NOT_OK(
MakeRandomListArray(array, batch_length, include_nulls, pool_, &array));
}
auto f0 = std::make_shared<Field>("f0", type);
std::shared_ptr<Schema> schema(new Schema({f0}));
if (schema_out != nullptr) { *schema_out = schema; }
auto f0 = field("f0", type);
*schema = std::shared_ptr<Schema>(new Schema({f0}));
std::vector<ArrayPtr> arrays = {array};
auto batch = std::make_shared<RecordBatch>(schema, batch_length, arrays);
auto batch = std::make_shared<RecordBatch>(*schema, batch_length, arrays);
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(),
&body_offset, header_out_param, recursion_level + 1);
return WriteRecordBatch(batch->columns(), batch->num_rows(), 0, mmap_.get(),
metadata_length, body_length, recursion_level + 1);
} else {
return WriteRecordBatch(batch->columns(), batch->num_rows(), mmap_.get(),
&body_offset, header_out_param);
return WriteRecordBatch(batch->columns(), batch->num_rows(), 0, mmap_.get(),
metadata_length, body_length);
}
}
@@ -171,18 +175,29 @@ class RecursionLimits : public ::testing::Test, public io::MemoryMapFixture {
};
TEST_F(RecursionLimits, WriteLimit) {
ASSERT_RAISES(Invalid, WriteToMmap((1 << 8) + 1, false));
int32_t metadata_length = -1;
int64_t body_length = -1;
std::shared_ptr<Schema> schema;
ASSERT_RAISES(
Invalid, WriteToMmap((1 << 8) + 1, false, &metadata_length, &body_length, &schema));
}
TEST_F(RecursionLimits, ReadLimit) {
int64_t header_offset = -1;
int32_t metadata_length = -1;
int64_t body_length = -1;
std::shared_ptr<Schema> schema;
ASSERT_OK(WriteToMmap(64, true, &header_offset, &schema));
ASSERT_OK(WriteToMmap(64, true, &metadata_length, &body_length, &schema));
std::shared_ptr<RecordBatchReader> reader;
ASSERT_OK(RecordBatchReader::Open(mmap_.get(), header_offset, &reader));
std::shared_ptr<RecordBatch> batch_result;
ASSERT_RAISES(Invalid, reader->GetRecordBatch(schema, &batch_result));
std::shared_ptr<RecordBatchMetadata> metadata;
ASSERT_OK(ReadRecordBatchMetadata(0, metadata_length, mmap_.get(), &metadata));
std::shared_ptr<Buffer> payload;
ASSERT_OK(mmap_->ReadAt(metadata_length, body_length, &payload));
io::BufferReader reader(payload);
std::shared_ptr<RecordBatch> batch;
ASSERT_RAISES(Invalid, ReadRecordBatch(metadata, schema, &reader, &batch));
}
} // namespace ipc
+1 -1
View File
@@ -68,7 +68,7 @@ class TestFileFormat : public ::testing::TestWithParam<MakeRecordBatch*> {
RETURN_NOT_OK(sink_->Tell(&footer_offset));
// Open the file
auto reader = std::make_shared<io::BufferReader>(buffer_->data(), buffer_->size());
auto reader = std::make_shared<io::BufferReader>(buffer_);
RETURN_NOT_OK(FileReader::Open(reader, footer_offset, &file_reader_));
EXPECT_EQ(num_batches, file_reader_->num_record_batches());
+5 -5
View File
@@ -43,7 +43,7 @@ static inline void assert_schema_equal(const Schema* lhs, const Schema* rhs) {
}
}
class TestSchemaMessage : public ::testing::Test {
class TestSchemaMetadata : public ::testing::Test {
public:
void SetUp() {}
@@ -52,11 +52,11 @@ class TestSchemaMessage : public ::testing::Test {
ASSERT_OK(WriteSchema(schema, &buffer));
std::shared_ptr<Message> message;
ASSERT_OK(Message::Open(buffer, &message));
ASSERT_OK(Message::Open(buffer, 0, &message));
ASSERT_EQ(Message::SCHEMA, message->type());
std::shared_ptr<SchemaMessage> schema_msg = message->GetSchema();
auto schema_msg = std::make_shared<SchemaMetadata>(message);
ASSERT_EQ(schema->num_fields(), schema_msg->num_fields());
std::shared_ptr<Schema> schema2;
@@ -68,7 +68,7 @@ class TestSchemaMessage : public ::testing::Test {
const std::shared_ptr<DataType> INT32 = std::make_shared<Int32Type>();
TEST_F(TestSchemaMessage, PrimitiveFields) {
TEST_F(TestSchemaMetadata, PrimitiveFields) {
auto f0 = std::make_shared<Field>("f0", std::make_shared<Int8Type>());
auto f1 = std::make_shared<Field>("f1", std::make_shared<Int16Type>());
auto f2 = std::make_shared<Field>("f2", std::make_shared<Int32Type>());
@@ -85,7 +85,7 @@ TEST_F(TestSchemaMessage, PrimitiveFields) {
CheckRoundtrip(&schema);
}
TEST_F(TestSchemaMessage, NestedFields) {
TEST_F(TestSchemaMetadata, NestedFields) {
auto type = std::make_shared<ListType>(std::make_shared<Int32Type>());
auto f0 = std::make_shared<Field>("f0", type);
+50 -48
View File
@@ -50,9 +50,17 @@ Status WriteSchema(const Schema* schema, std::shared_ptr<Buffer>* out) {
class Message::MessageImpl {
public:
explicit MessageImpl(
const std::shared_ptr<Buffer>& buffer, const flatbuf::Message* message)
: buffer_(buffer), message_(message) {}
explicit MessageImpl(const std::shared_ptr<Buffer>& buffer, int64_t offset)
: buffer_(buffer),
offset_(offset),
message_(nullptr) {}
Status Open() {
message_ = flatbuf::GetMessage(buffer_->data() + offset_);
// TODO(wesm): verify the message
return Status::OK();
}
Message::Type type() const {
switch (message_->header_type()) {
@@ -72,25 +80,23 @@ class Message::MessageImpl {
int64_t body_length() const { return message_->bodyLength(); }
private:
// Owns the memory this message accesses
// Retain reference to memory
std::shared_ptr<Buffer> buffer_;
int64_t offset_;
const flatbuf::Message* message_;
};
Message::Message() {}
Message::Message(const std::shared_ptr<Buffer>& buffer, int64_t offset) {
impl_.reset(new MessageImpl(buffer, offset));
}
Status Message::Open(
const std::shared_ptr<Buffer>& buffer, std::shared_ptr<Message>* out) {
std::shared_ptr<Message> result(new Message());
Status Message::Open(const std::shared_ptr<Buffer>& buffer, int64_t offset,
std::shared_ptr<Message>* out) {
// ctor is private
const flatbuf::Message* message = flatbuf::GetMessage(buffer->data());
// TODO(wesm): verify message
result->impl_.reset(new MessageImpl(buffer, message));
*out = result;
return Status::OK();
*out = std::shared_ptr<Message>(new Message(buffer, offset));
return (*out)->impl_->Open();
}
Message::Type Message::type() const {
@@ -101,20 +107,12 @@ int64_t Message::body_length() const {
return impl_->body_length();
}
std::shared_ptr<Message> Message::get_shared_ptr() {
return this->shared_from_this();
}
std::shared_ptr<SchemaMessage> Message::GetSchema() {
return std::make_shared<SchemaMessage>(this->shared_from_this(), impl_->header());
}
// ----------------------------------------------------------------------
// SchemaMessage
// SchemaMetadata
class SchemaMessage::SchemaMessageImpl {
class SchemaMetadata::SchemaMetadataImpl {
public:
explicit SchemaMessageImpl(const void* schema)
explicit SchemaMetadataImpl(const void* schema)
: schema_(static_cast<const flatbuf::Schema*>(schema)) {}
const flatbuf::Field* field(int i) const { return schema_->fields()->Get(i); }
@@ -125,22 +123,29 @@ class SchemaMessage::SchemaMessageImpl {
const flatbuf::Schema* schema_;
};
SchemaMessage::SchemaMessage(
const std::shared_ptr<Message>& message, const void* schema) {
SchemaMetadata::SchemaMetadata(
const std::shared_ptr<Message>& message, const void* flatbuf) {
message_ = message;
impl_.reset(new SchemaMessageImpl(schema));
impl_.reset(new SchemaMetadataImpl(flatbuf));
}
int SchemaMessage::num_fields() const {
SchemaMetadata::SchemaMetadata(const std::shared_ptr<Message>& message) {
message_ = message;
impl_.reset(new SchemaMetadataImpl(message->impl_->header()));
}
SchemaMetadata::~SchemaMetadata() {}
int SchemaMetadata::num_fields() const {
return impl_->num_fields();
}
Status SchemaMessage::GetField(int i, std::shared_ptr<Field>* out) const {
Status SchemaMetadata::GetField(int i, std::shared_ptr<Field>* out) const {
const flatbuf::Field* field = impl_->field(i);
return FieldFromFlatbuffer(field, out);
}
Status SchemaMessage::GetSchema(std::shared_ptr<Schema>* out) const {
Status SchemaMetadata::GetSchema(std::shared_ptr<Schema>* out) const {
std::vector<std::shared_ptr<Field>> fields(num_fields());
for (int i = 0; i < this->num_fields(); ++i) {
RETURN_NOT_OK(GetField(i, &fields[i]));
@@ -150,11 +155,11 @@ Status SchemaMessage::GetSchema(std::shared_ptr<Schema>* out) const {
}
// ----------------------------------------------------------------------
// RecordBatchMessage
// RecordBatchMetadata
class RecordBatchMessage::RecordBatchMessageImpl {
class RecordBatchMetadata::RecordBatchMetadataImpl {
public:
explicit RecordBatchMessageImpl(const void* batch)
explicit RecordBatchMetadataImpl(const void* batch)
: batch_(static_cast<const flatbuf::RecordBatch*>(batch)) {
nodes_ = batch_->nodes();
buffers_ = batch_->buffers();
@@ -176,19 +181,16 @@ class RecordBatchMessage::RecordBatchMessageImpl {
const flatbuffers::Vector<const flatbuf::Buffer*>* buffers_;
};
std::shared_ptr<RecordBatchMessage> Message::GetRecordBatch() {
return std::make_shared<RecordBatchMessage>(this->shared_from_this(), impl_->header());
RecordBatchMetadata::RecordBatchMetadata(const std::shared_ptr<Message>& message) {
message_ = message;
impl_.reset(new RecordBatchMetadataImpl(message->impl_->header()));
}
RecordBatchMessage::RecordBatchMessage(
const std::shared_ptr<Message>& message, const void* batch) {
message_ = message;
impl_.reset(new RecordBatchMessageImpl(batch));
}
RecordBatchMetadata::~RecordBatchMetadata() {}
// TODO(wesm): Copying the flatbuffer data isn't great, but this will do for
// now
FieldMetadata RecordBatchMessage::field(int i) const {
FieldMetadata RecordBatchMetadata::field(int i) const {
const flatbuf::FieldNode* node = impl_->field(i);
FieldMetadata result;
@@ -197,7 +199,7 @@ FieldMetadata RecordBatchMessage::field(int i) const {
return result;
}
BufferMetadata RecordBatchMessage::buffer(int i) const {
BufferMetadata RecordBatchMetadata::buffer(int i) const {
const flatbuf::Buffer* buffer = impl_->buffer(i);
BufferMetadata result;
@@ -207,15 +209,15 @@ BufferMetadata RecordBatchMessage::buffer(int i) const {
return result;
}
int32_t RecordBatchMessage::length() const {
int32_t RecordBatchMetadata::length() const {
return impl_->length();
}
int RecordBatchMessage::num_buffers() const {
int RecordBatchMetadata::num_buffers() const {
return impl_->num_buffers();
}
int RecordBatchMessage::num_fields() const {
int RecordBatchMetadata::num_fields() const {
return impl_->num_fields();
}
@@ -285,7 +287,7 @@ class FileFooter::FileFooterImpl {
}
Status GetSchema(std::shared_ptr<Schema>* out) const {
auto schema_msg = std::make_shared<SchemaMessage>(nullptr, footer_->schema());
auto schema_msg = std::make_shared<SchemaMetadata>(nullptr, footer_->schema());
return schema_msg->GetSchema(out);
}
+22 -22
View File
@@ -58,10 +58,14 @@ Status WriteSchema(const Schema* schema, std::shared_ptr<Buffer>* out);
class Message;
// Container for serialized Schema metadata contained in an IPC message
class ARROW_EXPORT SchemaMessage {
class ARROW_EXPORT SchemaMetadata {
public:
SchemaMetadata(const std::shared_ptr<Message>& message);
// Accepts an opaque flatbuffer pointer
SchemaMessage(const std::shared_ptr<Message>& message, const void* schema);
SchemaMetadata(const std::shared_ptr<Message>& message, const void* schema);
~SchemaMetadata();
int num_fields() const;
@@ -76,8 +80,8 @@ class ARROW_EXPORT SchemaMessage {
// Parent, owns the flatbuffer data
std::shared_ptr<Message> message_;
class SchemaMessageImpl;
std::unique_ptr<SchemaMessageImpl> impl_;
class SchemaMetadataImpl;
std::unique_ptr<SchemaMetadataImpl> impl_;
};
// Field metadata
@@ -93,10 +97,10 @@ struct BufferMetadata {
};
// Container for serialized record batch metadata contained in an IPC message
class ARROW_EXPORT RecordBatchMessage {
class ARROW_EXPORT RecordBatchMetadata {
public:
// Accepts an opaque flatbuffer pointer
RecordBatchMessage(const std::shared_ptr<Message>& message, const void* batch_meta);
RecordBatchMetadata(const std::shared_ptr<Message>& message);
~RecordBatchMetadata();
FieldMetadata field(int i) const;
BufferMetadata buffer(int i) const;
@@ -109,36 +113,32 @@ class ARROW_EXPORT RecordBatchMessage {
// Parent, owns the flatbuffer data
std::shared_ptr<Message> message_;
class RecordBatchMessageImpl;
std::unique_ptr<RecordBatchMessageImpl> impl_;
class RecordBatchMetadataImpl;
std::unique_ptr<RecordBatchMetadataImpl> impl_;
};
class ARROW_EXPORT DictionaryBatchMessage {
class ARROW_EXPORT DictionaryBatchMetadata {
public:
int64_t id() const;
std::unique_ptr<RecordBatchMessage> data() const;
std::unique_ptr<RecordBatchMetadata> data() const;
};
class ARROW_EXPORT Message : public std::enable_shared_from_this<Message> {
class ARROW_EXPORT Message {
public:
enum Type { NONE, SCHEMA, DICTIONARY_BATCH, RECORD_BATCH };
static Status Open(
const std::shared_ptr<Buffer>& buffer, std::shared_ptr<Message>* out);
std::shared_ptr<Message> get_shared_ptr();
static Status Open(const std::shared_ptr<Buffer>& buffer, int64_t offset,
std::shared_ptr<Message>* out);
int64_t body_length() const;
Type type() const;
// These methods only to be invoked if you have checked the message type
std::shared_ptr<SchemaMessage> GetSchema();
std::shared_ptr<RecordBatchMessage> GetRecordBatch();
std::shared_ptr<DictionaryBatchMessage> GetDictionaryBatch();
private:
Message();
Message(const std::shared_ptr<Buffer>& buffer, int64_t offset);
friend class RecordBatchMetadata;
friend class SchemaMetadata;
// Hide serialization details from user API
class MessageImpl;
+7 -8
View File
@@ -39,8 +39,7 @@
namespace arrow {
namespace ipc {
const auto kInt32 = std::make_shared<Int32Type>();
const auto kListInt32 = list(kInt32);
const auto kListInt32 = list(int32());
const auto kListListInt32 = list(kListInt32);
Status MakeRandomInt32Array(
@@ -99,8 +98,8 @@ Status MakeIntRecordBatch(std::shared_ptr<RecordBatch>* out) {
const int length = 1000;
// Make the schema
auto f0 = std::make_shared<Field>("f0", kInt32);
auto f1 = std::make_shared<Field>("f1", kInt32);
auto f0 = std::make_shared<Field>("f0", int32());
auto f1 = std::make_shared<Field>("f1", int32());
std::shared_ptr<Schema> schema(new Schema({f0, f1}));
// Example data
@@ -161,7 +160,7 @@ Status MakeListRecordBatch(std::shared_ptr<RecordBatch>* out) {
// Make the schema
auto f0 = std::make_shared<Field>("f0", kListInt32);
auto f1 = std::make_shared<Field>("f1", kListListInt32);
auto f2 = std::make_shared<Field>("f2", kInt32);
auto f2 = std::make_shared<Field>("f2", int32());
std::shared_ptr<Schema> schema(new Schema({f0, f1, f2}));
// Example data
@@ -184,7 +183,7 @@ Status MakeZeroLengthRecordBatch(std::shared_ptr<RecordBatch>* out) {
// Make the schema
auto f0 = std::make_shared<Field>("f0", kListInt32);
auto f1 = std::make_shared<Field>("f1", kListListInt32);
auto f2 = std::make_shared<Field>("f2", kInt32);
auto f2 = std::make_shared<Field>("f2", int32());
std::shared_ptr<Schema> schema(new Schema({f0, f1, f2}));
// Example data
@@ -205,7 +204,7 @@ Status MakeNonNullRecordBatch(std::shared_ptr<RecordBatch>* out) {
// Make the schema
auto f0 = std::make_shared<Field>("f0", kListInt32);
auto f1 = std::make_shared<Field>("f1", kListListInt32);
auto f2 = std::make_shared<Field>("f2", kInt32);
auto f2 = std::make_shared<Field>("f2", int32());
std::shared_ptr<Schema> schema(new Schema({f0, f1, f2}));
// Example data
@@ -226,7 +225,7 @@ Status MakeNonNullRecordBatch(std::shared_ptr<RecordBatch>* out) {
Status MakeDeeplyNestedList(std::shared_ptr<RecordBatch>* out) {
const int batch_length = 5;
TypePtr type = kInt32;
TypePtr type = int32();
MemoryPool* pool = default_memory_pool();
ArrayPtr array;
+2 -4
View File
@@ -28,12 +28,10 @@ namespace arrow {
namespace ipc {
// Align on 8-byte boundaries
static constexpr int kArrowAlignment = 8;
// Buffers are padded to 64-byte boundaries (for SIMD)
static constexpr int kArrowBufferAlignment = 64;
static constexpr int kArrowAlignment = 64;
static constexpr uint8_t kPaddingBytes[kArrowBufferAlignment] = {0};
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;
+62
View File
@@ -15,3 +15,65 @@
# Interprocess messaging / communication (IPC)
## File format
We define a self-contained "file format" containing an Arrow schema along with
one or more record batches defining a dataset. See [format/File.fbs][1] for the
precise details of the file metadata.
In general, the file looks like:
```
<magic number "ARROW1">
<empty padding bytes [to 8 byte boundary]>
<DICTIONARY 0>
...
<DICTIONARY k - 1>
<RECORD BATCH 0>
...
<RECORD BATCH n - 1>
<METADATA org.apache.arrow.flatbuf.Footer>
<metadata_size: int32>
<magic number "ARROW1">
```
See the File.fbs document for details about the Flatbuffers metadata. The
record batches have a particular structure, defined next.
### Record batches
The record batch metadata is written as
a flatbuffer (see format/Message.fbs -- the RecordBatch message type)
prefixed by its size, followed by each of the memory buffers in the batch
written end to end (with appropriate alignment and padding):
```
<int32: metadata flatbuffer size>
<metadata: org.apache.arrow.flatbuf.RecordBatch>
<padding bytes [to 64-byte boundary]>
<body: buffers end to end>
```
The location of a record batch and the size of the metadata block as well as
the body of buffers is stored in the file footer:
```
struct Block {
offset: long;
metaDataLength: int;
bodyLength: long;
}
```
Some notes about this
* The `Block` offset indicates the starting byte of the record batch.
* The metadata length includes the flatbuffer size, the record batch metadata
flatbuffer, and any padding bytes
### Dictionary batches
Dictionary batches have not yet been implemented, while they are provided for
in the metadata. For the time being, the `DICTIONARY` segments shown above in
the file do not appear in any of the file implementations.
[1]: https://github.com/apache/arrow/blob/master/format/File.fbs