mirror of
https://github.com/apache/arrow.git
synced 2026-10-29 00:56:31 +08:00
Refactor IPC/File record batch read/write structure to reflect discussion in ARROW-384
Change-Id: Ic119c5647bdb6b1ff1fb31b3712527058c68e40c
This commit is contained in:
1 parent
dbd6ed6d5e
commit
85128f7add
14 files changed
+399
-276
No files matched your search
@@ -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}
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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(¤t_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(¤t_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
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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,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);
|
||||
}
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
Reference in new issue
Block a user