mirrors/arrow
0
0 Fork 0
mirror of https://github.com/apache/arrow.git synced 2026-10-29 17:15:39 +08:00
Code Activity

Refactor to introduce MessageReader abstract type, use unique_ptr for messages

Browse Source 
instead of shared_ptr. First cut at Message, MessageReader Python API. Add
read_message, C++/Python machinery for message roundtrips to Buffer,
comparison. Add function to read RecordBatch from encapsulated message given
schema.
This commit is contained in:
Wes McKinney committed 2017-07-15 12:23:21 -04:00
1 parent d46b7ea3e8
commit 07f1820ad6
35 files changed
+1134 -642
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

c_glib/arrow-glib/reader.cpp
+2 -2
View File
@@ -173,7 +173,7 @@ garrow_record_batch_reader_get_next_record_batch(GArrowRecordBatchReader *reader
{
auto arrow_reader = garrow_record_batch_reader_get_raw(reader);
std::shared_ptr<arrow::RecordBatch> arrow_record_batch;
auto status = arrow_reader->GetNextRecordBatch(&arrow_record_batch);
auto status = arrow_reader->ReadNextRecordBatch(&arrow_record_batch);
if (garrow_error_check(error,
status,
@@ -410,7 +410,7 @@ garrow_record_batch_file_reader_get_record_batch(GArrowRecordBatchFileReader *re
{
auto arrow_reader = garrow_record_batch_file_reader_get_raw(reader);
std::shared_ptr<arrow::RecordBatch> arrow_record_batch;
auto status = arrow_reader->GetRecordBatch(i, &arrow_record_batch);
auto status = arrow_reader->ReadRecordBatch(i, &arrow_record_batch);
if (garrow_error_check(error,
status,
cpp/src/arrow/buffer.cc
-2
View File
@@ -27,8 +27,6 @@
namespace arrow {
Buffer::~Buffer() {}
Status Buffer::Copy(
int64_t start, int64_t nbytes, MemoryPool* pool, std::shared_ptr<Buffer>* out) const {
// Sanity checks
cpp/src/arrow/buffer.h
+14 -1
View File
@@ -22,6 +22,7 @@
#include <cstdint>
#include <cstring>
#include <memory>
#include <string>
#include "arrow/status.h"
#include "arrow/util/macros.h"
@@ -47,7 +48,8 @@ class ARROW_EXPORT Buffer {
public:
Buffer(const uint8_t* data, int64_t size)
: is_mutable_(false), data_(data), size_(size), capacity_(size) {}
virtual ~Buffer();
virtual ~Buffer() = default;
/// An offset into data that is owned by another buffer, but we want to be
/// able to retain a valid pointer to it even after other shared_ptr's to the
@@ -97,6 +99,17 @@ class ARROW_EXPORT Buffer {
DISALLOW_COPY_AND_ASSIGN(Buffer);
};
/// \brief Create Buffer referencing std::string memory
///
/// Warning: string instance must stay alive
///
/// \param str std::string instance
/// \return std::shared_ptr<Buffer>
static inline std::shared_ptr<Buffer> GetBufferFromString(const std::string& str) {
return std::make_shared<Buffer>(
reinterpret_cast<const uint8_t*>(str.c_str()), static_cast<int64_t>(str.size()));
}
/// Construct a view on passed buffer at the indicated offset and length. This
/// function cannot fail and does not error checking (except in debug builds)
static inline std::shared_ptr<Buffer> SliceBuffer(
cpp/src/arrow/builder.cc
+6 -6
View File
@@ -342,8 +342,8 @@ Status AdaptiveIntBuilder::Append(
sizeof(int64_t) * length);
} else {
#ifdef _MSC_VER
# pragma warning(push)
# pragma warning(disable:4996)
#pragma warning(push)
#pragma warning(disable : 4996)
#endif
// int_size_ may have changed, so we need to recheck
switch (int_size_) {
@@ -366,7 +366,7 @@ Status AdaptiveIntBuilder::Append(
DCHECK(false);
}
#ifdef _MSC_VER
# pragma warning(pop)
#pragma warning(pop)
#endif
}
@@ -497,8 +497,8 @@ Status AdaptiveUIntBuilder::Append(
sizeof(uint64_t) * length);
} else {
#ifdef _MSC_VER
# pragma warning(push)
# pragma warning(disable:4996)
#pragma warning(push)
#pragma warning(disable : 4996)
#endif
// int_size_ may have changed, so we need to recheck
switch (int_size_) {
@@ -521,7 +521,7 @@ Status AdaptiveUIntBuilder::Append(
DCHECK(false);
}
#ifdef _MSC_VER
# pragma warning(pop)
#pragma warning(pop)
#endif
}
cpp/src/arrow/ipc/file-to-stream.cc
+1 -1
View File
@@ -39,7 +39,7 @@ Status ConvertToStream(const char* path) {
RETURN_NOT_OK(RecordBatchStreamWriter::Open(&sink, reader->schema(), &writer));
for (int i = 0; i < reader->num_record_batches(); ++i) {
std::shared_ptr<RecordBatch> chunk;
RETURN_NOT_OK(reader->GetRecordBatch(i, &chunk));
RETURN_NOT_OK(reader->ReadRecordBatch(i, &chunk));
RETURN_NOT_OK(writer->WriteRecordBatch(*chunk));
}
return writer->Close();
cpp/src/arrow/ipc/ipc-json-test.cc
+3 -3
View File
@@ -276,7 +276,7 @@ TEST(TestJsonFileReadWrite, BasicRoundTrip) {
for (int i = 0; i < nbatches; ++i) {
std::shared_ptr<RecordBatch> batch;
ASSERT_OK(reader->GetRecordBatch(i, &batch));
ASSERT_OK(reader->ReadRecordBatch(i, &batch));
ASSERT_TRUE(batch->Equals(*batches[i]));
}
}
@@ -344,7 +344,7 @@ TEST(TestJsonFileReadWrite, MinimalFormatExample) {
ASSERT_EQ(1, reader->num_record_batches());
std::shared_ptr<RecordBatch> batch;
ASSERT_OK(reader->GetRecordBatch(0, &batch));
ASSERT_OK(reader->ReadRecordBatch(0, &batch));
std::vector<bool> foo_valid = {true, false, true, true, true};
std::vector<int32_t> foo_values = {1, 2, 3, 4, 5};
@@ -388,7 +388,7 @@ void CheckRoundtrip(const RecordBatch& batch) {
ASSERT_OK(JsonReader::Open(buffer, &reader));
std::shared_ptr<RecordBatch> result_batch;
ASSERT_OK(reader->GetRecordBatch(0, &result_batch));
ASSERT_OK(reader->ReadRecordBatch(0, &result_batch));
CompareBatch(batch, *result_batch);
}
cpp/src/arrow/ipc/ipc-read-write-test.cc
+43 -27
View File
@@ -51,8 +51,8 @@ class TestSchemaMetadata : public ::testing::Test {
std::shared_ptr<Buffer> buffer;
ASSERT_OK(WriteSchemaMessage(schema, memo, &buffer));
std::shared_ptr<Message> message;
ASSERT_OK(Message::Open(buffer, 0, &message));
std::unique_ptr<Message> message;
ASSERT_OK(Message::Open(buffer, nullptr, &message));
ASSERT_EQ(Message::SCHEMA, message->type());
@@ -65,6 +65,32 @@ class TestSchemaMetadata : public ::testing::Test {
}
};
TEST(TestMessage, Equals) {
std::string metadata = "foo";
std::string body = "bar";
auto b1 = GetBufferFromString(metadata);
auto b2 = GetBufferFromString(metadata);
auto b3 = GetBufferFromString(body);
auto b4 = GetBufferFromString(body);
Message msg1(b1, b3);
Message msg2(b2, b4);
Message msg3(b1, nullptr);
Message msg4(b2, nullptr);
ASSERT_TRUE(msg1.Equals(msg2));
ASSERT_TRUE(msg3.Equals(msg4));
ASSERT_FALSE(msg1.Equals(msg3));
ASSERT_FALSE(msg3.Equals(msg1));
// same metadata as msg1, different body
Message msg5(b2, b1);
ASSERT_FALSE(msg1.Equals(msg5));
ASSERT_FALSE(msg5.Equals(msg1));
}
const std::shared_ptr<DataType> INT32 = std::make_shared<Int32Type>();
TEST_F(TestSchemaMetadata, PrimitiveFields) {
@@ -123,16 +149,12 @@ class IpcTestFixture : public io::MemoryMapFixture {
RETURN_NOT_OK(WriteRecordBatch(
batch, buffer_offset, mmap_.get(), &metadata_length, &body_length, pool_));
std::shared_ptr<Message> message;
std::unique_ptr<Message> message;
RETURN_NOT_OK(ReadMessage(0, metadata_length, mmap_.get(), &message));
// The buffer offsets start at 0, so we must construct a
// RandomAccessFile 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(*message, batch.schema(), &buffer_reader, batch_result);
io::BufferReader buffer_reader(message->body());
return ReadRecordBatch(
*message->metadata(), batch.schema(), &buffer_reader, batch_result);
}
Status DoLargeRoundTrip(
@@ -151,7 +173,7 @@ class IpcTestFixture : public io::MemoryMapFixture {
std::shared_ptr<RecordBatchFileReader> file_reader;
RETURN_NOT_OK(RecordBatchFileReader::Open(mmap_, offset, &file_reader));
return file_reader->GetRecordBatch(0, result);
return file_reader->ReadRecordBatch(0, result);
}
void CheckReadResult(const RecordBatch& result, const RecordBatch& expected) {
@@ -225,7 +247,7 @@ TEST_F(TestIpcRoundTrip, MetadataVersion) {
ASSERT_OK(WriteRecordBatch(
*batch, buffer_offset, mmap_.get(), &metadata_length, &body_length, pool_));
std::shared_ptr<Message> message;
std::unique_ptr<Message> message;
ASSERT_OK(ReadMessage(0, metadata_length, mmap_.get(), &message));
ASSERT_EQ(MetadataVersion::V3, message->metadata_version());
@@ -434,16 +456,13 @@ TEST_F(RecursionLimits, ReadLimit) {
ASSERT_OK(WriteToMmap(
recursion_depth, true, &metadata_length, &body_length, &batch, &schema));
std::shared_ptr<Message> message;
std::unique_ptr<Message> message;
ASSERT_OK(ReadMessage(0, metadata_length, mmap_.get(), &message));
std::shared_ptr<Buffer> payload;
ASSERT_OK(mmap_->ReadAt(metadata_length, body_length, &payload));
io::BufferReader reader(payload);
io::BufferReader reader(message->body());
std::shared_ptr<RecordBatch> result;
ASSERT_RAISES(Invalid, ReadRecordBatch(*message, schema, &reader, &result));
ASSERT_RAISES(Invalid, ReadRecordBatch(*message->metadata(), schema, &reader, &result));
}
TEST_F(RecursionLimits, StressLimit) {
@@ -455,16 +474,13 @@ TEST_F(RecursionLimits, StressLimit) {
ASSERT_OK(WriteToMmap(
recursion_depth, true, &metadata_length, &body_length, &batch, &schema));
std::shared_ptr<Message> message;
std::unique_ptr<Message> message;
ASSERT_OK(ReadMessage(0, metadata_length, mmap_.get(), &message));
std::shared_ptr<Buffer> payload;
ASSERT_OK(mmap_->ReadAt(metadata_length, body_length, &payload));
io::BufferReader reader(payload);
io::BufferReader reader(message->body());
std::shared_ptr<RecordBatch> result;
ASSERT_OK(ReadRecordBatch(*message, schema, recursion_depth + 1, &reader, &result));
ASSERT_OK(ReadRecordBatch(
*message->metadata(), schema, recursion_depth + 1, &reader, &result));
*it_works = result->Equals(*batch);
};
@@ -511,7 +527,7 @@ class TestFileFormat : public ::testing::TestWithParam<MakeRecordBatch*> {
EXPECT_EQ(num_batches, reader->num_record_batches());
for (int i = 0; i < num_batches; ++i) {
std::shared_ptr<RecordBatch> chunk;
RETURN_NOT_OK(reader->GetRecordBatch(i, &chunk));
RETURN_NOT_OK(reader->ReadRecordBatch(i, &chunk));
out_batches->emplace_back(chunk);
}
@@ -571,7 +587,7 @@ class TestStreamFormat : public ::testing::TestWithParam<MakeRecordBatch*> {
std::shared_ptr<RecordBatch> chunk;
while (true) {
RETURN_NOT_OK(reader->GetNextRecordBatch(&chunk));
RETURN_NOT_OK(reader->ReadNextRecordBatch(&chunk));
if (chunk == nullptr) { break; }
out_batches->emplace_back(chunk);
}
cpp/src/arrow/ipc/json-integration-test.cc
+4 -4
View File
@@ -82,7 +82,7 @@ static Status ConvertJsonToArrow(
for (int i = 0; i < reader->num_record_batches(); ++i) {
std::shared_ptr<RecordBatch> batch;
RETURN_NOT_OK(reader->GetRecordBatch(i, &batch));
RETURN_NOT_OK(reader->ReadRecordBatch(i, &batch));
RETURN_NOT_OK(writer->WriteRecordBatch(*batch));
}
return writer->Close();
@@ -109,7 +109,7 @@ static Status ConvertArrowToJson(
for (int i = 0; i < reader->num_record_batches(); ++i) {
std::shared_ptr<RecordBatch> batch;
RETURN_NOT_OK(reader->GetRecordBatch(i, &batch));
RETURN_NOT_OK(reader->ReadRecordBatch(i, &batch));
RETURN_NOT_OK(writer->WriteRecordBatch(*batch));
}
@@ -168,8 +168,8 @@ static Status ValidateArrowVsJson(
std::shared_ptr<RecordBatch> arrow_batch;
std::shared_ptr<RecordBatch> json_batch;
for (int i = 0; i < json_nbatches; ++i) {
RETURN_NOT_OK(json_reader->GetRecordBatch(i, &json_batch));
RETURN_NOT_OK(arrow_reader->GetRecordBatch(i, &arrow_batch));
RETURN_NOT_OK(json_reader->ReadRecordBatch(i, &json_batch));
RETURN_NOT_OK(arrow_reader->ReadRecordBatch(i, &arrow_batch));
if (!json_batch->ApproxEquals(*arrow_batch)) {
std::stringstream ss;
cpp/src/arrow/ipc/json.cc
+3 -3
View File
@@ -115,7 +115,7 @@ class JsonReader::JsonReaderImpl {
return Status::OK();
}
Status GetRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) const {
Status ReadRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) const {
DCHECK_GE(i, 0) << "i out of bounds";
DCHECK_LT(i, static_cast<int>(record_batches_->GetArray().Size()))
<< "i out of bounds";
@@ -164,8 +164,8 @@ int JsonReader::num_record_batches() const {
return impl_->num_record_batches();
}
Status JsonReader::GetRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) const {
return impl_->GetRecordBatch(i, batch);
Status JsonReader::ReadRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) const {
return impl_->ReadRecordBatch(i, batch);
}
} // namespace ipc
cpp/src/arrow/ipc/json.h
+1 -1
View File
@@ -72,7 +72,7 @@ class ARROW_EXPORT JsonReader {
int num_record_batches() const;
// Read a record batch from the file
Status GetRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) const;
Status ReadRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) const;
private:
JsonReader(MemoryPool* pool, const std::shared_ptr<Buffer>& data);
cpp/src/arrow/ipc/metadata.cc
+123 -27
View File
@@ -17,6 +17,7 @@
#include "arrow/ipc/metadata.h"
#include <algorithm>
#include <cstdint>
#include <memory>
#include <sstream>
@@ -834,11 +835,12 @@ Status DictionaryMemo::AddDictionary(
class Message::MessageImpl {
public:
explicit MessageImpl(const std::shared_ptr<Buffer>& buffer, int64_t offset)
: buffer_(buffer), offset_(offset), message_(nullptr) {}
explicit MessageImpl(
const std::shared_ptr<Buffer>& metadata, const std::shared_ptr<Buffer>& body)
: metadata_(metadata), message_(nullptr), body_(body) {}
Status Open() {
message_ = flatbuf::GetMessage(buffer_->data() + offset_);
message_ = flatbuf::GetMessage(metadata_->data());
// Check that the metadata version is supported
if (message_->version() < kMinMetadataVersion) {
@@ -872,7 +874,7 @@ class Message::MessageImpl {
// Arrow 0.2
return MetadataVersion::V2;
case flatbuf::MetadataVersion_V3:
// Arrow 0.3
// Arrow >= 0.3
return MetadataVersion::V3;
// Add cases as other versions become available
default:
@@ -882,28 +884,38 @@ class Message::MessageImpl {
const void* header() const { return message_->header(); }
int64_t body_length() const { return message_->bodyLength(); }
std::shared_ptr<Buffer> body() const { return body_; }
std::shared_ptr<Buffer> metadata() const { return metadata_; }
private:
// Retain reference to memory
std::shared_ptr<Buffer> buffer_;
int64_t offset_;
// The Flatbuffer metadata
std::shared_ptr<Buffer> metadata_;
const flatbuf::Message* message_;
// The message body, if any
std::shared_ptr<Buffer> body_;
};
Message::Message(const std::shared_ptr<Buffer>& buffer, int64_t offset) {
impl_.reset(new MessageImpl(buffer, offset));
Message::Message(
const std::shared_ptr<Buffer>& metadata, const std::shared_ptr<Buffer>& body) {
impl_.reset(new MessageImpl(metadata, body));
}
Status Message::Open(const std::shared_ptr<Buffer>& metadata,
const std::shared_ptr<Buffer>& body, std::unique_ptr<Message>* out) {
out->reset(new Message(metadata, body));
return (*out)->impl_->Open();
}
Message::~Message() {}
Status Message::Open(const std::shared_ptr<Buffer>& buffer, int64_t offset,
std::shared_ptr<Message>* out) {
// ctor is private
std::shared_ptr<Buffer> Message::body() const {
return impl_->body();
}
*out = std::shared_ptr<Message>(new Message(buffer, offset));
return (*out)->impl_->Open();
std::shared_ptr<Buffer> Message::metadata() const {
return impl_->metadata();
}
Message::Type Message::type() const {
@@ -914,14 +926,64 @@ MetadataVersion Message::metadata_version() const {
return impl_->version();
}
int64_t Message::body_length() const {
return impl_->body_length();
}
const void* Message::header() const {
return impl_->header();
}
bool Message::Equals(const Message& other) const {
int64_t metadata_bytes = std::min(metadata()->size(), other.metadata()->size());
if (!metadata()->Equals(*other.metadata(), metadata_bytes)) {
return false;
}
// Compare bodies, if they have them
auto this_body = body();
auto other_body = other.body();
const bool this_has_body = (this_body != nullptr) && (this_body->size() > 0);
const bool other_has_body = (other_body != nullptr) && (other_body->size() > 0);
if (this_has_body && other_has_body) {
return this_body->Equals(*other_body);
} else if (this_has_body ^ other_has_body) {
// One has a body but not the other
return false;
} else {
// Neither has a body
return true;
}
}
Status Message::SerializeTo(io::OutputStream* file, int64_t* output_length) const {
int32_t metadata_length = 0;
RETURN_NOT_OK(WriteMessage(*metadata(), file, &metadata_length));
*output_length = metadata_length;
auto body_buffer = body();
if (body_buffer) {
RETURN_NOT_OK(file->Write(body_buffer->data(), body_buffer->size()));
*output_length += body_buffer->size();
}
return Status::OK();
}
std::string FormatMessageType(Message::Type type) {
switch (type) {
case Message::SCHEMA:
return "schema";
case Message::RECORD_BATCH:
return "record batch";
case Message::DICTIONARY_BATCH:
return "dictionary";
default:
break;
}
return "unknown";
}
// ----------------------------------------------------------------------
static Status VisitField(const flatbuf::Field* field, DictionaryTypeMap* id_to_field) {
@@ -975,10 +1037,11 @@ Status GetSchema(const void* opaque_schema, const DictionaryMemo& dictionary_mem
return Status::OK();
}
Status GetTensorMetadata(const void* opaque_tensor, std::shared_ptr<DataType>* type,
Status GetTensorMetadata(const Buffer& metadata, std::shared_ptr<DataType>* type,
std::vector<int64_t>* shape, std::vector<int64_t>* strides,
std::vector<std::string>* dim_names) {
auto tensor = static_cast<const flatbuf::Tensor*>(opaque_tensor);
auto message = flatbuf::GetMessage(metadata.data());
auto tensor = reinterpret_cast<const flatbuf::Tensor*>(message->header());
int ndim = static_cast<int>(tensor->shape()->size());
@@ -1006,8 +1069,27 @@ Status GetTensorMetadata(const void* opaque_tensor, std::shared_ptr<DataType>* t
// ----------------------------------------------------------------------
// Read and write messages
static Status ReadFullMessage(const std::shared_ptr<Buffer>& metadata,
io::InputStream* stream, std::unique_ptr<Message>* message) {
auto fb_message = flatbuf::GetMessage(metadata->data());
int64_t body_length = fb_message->bodyLength();
std::shared_ptr<Buffer> body;
RETURN_NOT_OK(stream->Read(body_length, &body));
if (body->size() < body_length) {
std::stringstream ss;
ss << "Expected to be able to read " << body_length << " bytes for message body, got "
<< body->size();
return Status::IOError(ss.str());
}
return Message::Open(metadata, body, message);
}
Status ReadMessage(int64_t offset, int32_t metadata_length, io::RandomAccessFile* file,
std::shared_ptr<Message>* message) {
std::unique_ptr<Message>* message) {
std::shared_ptr<Buffer> buffer;
RETURN_NOT_OK(file->ReadAt(offset, metadata_length, &buffer));
@@ -1019,13 +1101,15 @@ Status ReadMessage(int64_t offset, int32_t metadata_length, io::RandomAccessFile
<< ", metadata length: " << metadata_length;
return Status::Invalid(ss.str());
}
return Message::Open(buffer, 4, message);
auto metadata = SliceBuffer(buffer, 4, buffer->size() - 4);
return ReadFullMessage(metadata, file, message);
}
Status ReadMessage(io::InputStream* file, std::shared_ptr<Message>* message) {
Status ReadMessage(io::InputStream* file, std::unique_ptr<Message>* message) {
std::shared_ptr<Buffer> buffer;
RETURN_NOT_OK(file->Read(sizeof(int32_t), &buffer));
RETURN_NOT_OK(file->Read(sizeof(int32_t), &buffer));
if (buffer->size() != sizeof(int32_t)) {
*message = nullptr;
return Status::OK();
@@ -1044,9 +1128,21 @@ Status ReadMessage(io::InputStream* file, std::shared_ptr<Message>* message) {
return Status::IOError("Unexpected end of stream trying to read message");
}
return Message::Open(buffer, 0, message);
return ReadFullMessage(buffer, file, message);
}
// ----------------------------------------------------------------------
// Implement InputStream message reader
Status InputStreamMessageReader::ReadNextMessage(std::unique_ptr<Message>* message) {
return ReadMessage(stream_.get(), message);
}
InputStreamMessageReader::~InputStreamMessageReader() {}
// ----------------------------------------------------------------------
// Implement message writing
Status WriteMessage(
const Buffer& message, io::OutputStream* file, int32_t* message_length) {
// Need to write 4 bytes (message size), the message, plus padding to
cpp/src/arrow/ipc/metadata.h
+70 -28
View File
@@ -59,26 +59,12 @@ static constexpr const char* kArrowMagicBytes = "ARROW1";
constexpr int kMaxNestingDepth = 64;
struct ARROW_EXPORT FieldMetadata {
FieldMetadata() {}
FieldMetadata(int64_t length, int64_t null_count, int64_t offset)
: length(length), null_count(null_count), offset(offset) {}
FieldMetadata(const FieldMetadata& other) {
this->length = other.length;
this->null_count = other.null_count;
this->offset = other.offset;
}
int64_t length;
int64_t null_count;
int64_t offset;
};
struct ARROW_EXPORT BufferMetadata {
BufferMetadata() {}
BufferMetadata(int32_t page, int64_t offset, int64_t length)
: page(page), offset(offset), length(length) {}
/// The shared memory page id where to find this. Set to -1 if unused
int32_t page;
@@ -90,10 +76,6 @@ struct ARROW_EXPORT BufferMetadata {
};
struct FileBlock {
FileBlock() {}
FileBlock(int64_t offset, int32_t metadata_length, int64_t body_length)
: offset(offset), metadata_length(metadata_length), body_length(body_length) {}
int64_t offset;
int32_t metadata_length;
int64_t body_length;
@@ -153,20 +135,46 @@ Status GetDictionaryTypes(const void* opaque_schema, DictionaryTypeMap* id_to_fi
Status ARROW_EXPORT GetSchema(const void* opaque_schema,
const DictionaryMemo& dictionary_memo, std::shared_ptr<Schema>* out);
Status ARROW_EXPORT GetTensorMetadata(const void* opaque_tensor,
Status ARROW_EXPORT GetTensorMetadata(const Buffer& metadata,
std::shared_ptr<DataType>* type, std::vector<int64_t>* shape,
std::vector<int64_t>* strides, std::vector<std::string>* dim_names);
/// \brief An IPC message including metadata and body
class ARROW_EXPORT Message {
public:
enum Type { NONE, SCHEMA, DICTIONARY_BATCH, RECORD_BATCH, TENSOR };
/// \brief Construct message, but do not validate
///
/// Use at your own risk; Message::Open has more metadata validation
Message(const std::shared_ptr<Buffer>& metadata, const std::shared_ptr<Buffer>& body);
~Message();
static Status Open(const std::shared_ptr<Buffer>& buffer, int64_t offset,
std::shared_ptr<Message>* out);
/// \brief Create and validate a Message instance from two buffers
///
/// \param[in] metadata a buffer containing the Flatbuffer metadata
/// \param[in] body a buffer containing the message body, which may be nullptr
/// \param[out] out the created message
static Status Open(const std::shared_ptr<Buffer>& metadata,
const std::shared_ptr<Buffer>& body, std::unique_ptr<Message>* out);
int64_t body_length() const;
/// \brief Write length-prefixed metadata and body to output stream
///
/// \param[in] file output stream to write to
/// \param[out] output_length the number of bytes written
/// \return Status
bool Equals(const Message& other) const;
/// \brief the Message metadata
///
/// \return buffer
std::shared_ptr<Buffer> metadata() const;
/// \brief the Message body, if any
///
/// \return buffer is nullptr if no body
std::shared_ptr<Buffer> body() const;
Type type() const;
@@ -174,9 +182,14 @@ class ARROW_EXPORT Message {
const void* header() const;
private:
Message(const std::shared_ptr<Buffer>& buffer, int64_t offset);
/// \brief Write length-prefixed metadata and body to output stream
///
/// \param[in] file output stream to write to
/// \param[out] output_length the number of bytes written
/// \return Status
Status SerializeTo(io::OutputStream* file, int64_t* output_length) const;
private:
// Hide serialization details from user API
class MessageImpl;
std::unique_ptr<MessageImpl> impl_;
@@ -184,8 +197,34 @@ class ARROW_EXPORT Message {
DISALLOW_COPY_AND_ASSIGN(Message);
};
ARROW_EXPORT std::string FormatMessageType(Message::Type type);
/// \brief Abstract interface for a sequence of messages
class ARROW_EXPORT MessageReader {
public:
virtual ~MessageReader() = default;
virtual Status ReadNextMessage(std::unique_ptr<Message>* message) = 0;
};
class ARROW_EXPORT InputStreamMessageReader : public MessageReader {
public:
explicit InputStreamMessageReader(const std::shared_ptr<io::InputStream>& stream)
: stream_(stream) {}
~InputStreamMessageReader();
Status ReadNextMessage(std::unique_ptr<Message>* message) override;
private:
std::shared_ptr<io::InputStream> stream_;
};
/// \brief Read encapulated RPC message from position in file
///
/// Read a length-prefixed message flatbuffer starting at the indicated file
/// offset
/// offset. If the message has a body with non-zero length, it will also be
/// read
///
/// The metadata_length includes at least the length prefix and the flatbuffer
///
@@ -196,15 +235,18 @@ class ARROW_EXPORT Message {
/// \param[out] message the message read
/// \return Status success or failure
Status ARROW_EXPORT ReadMessage(int64_t offset, int32_t metadata_length,
io::RandomAccessFile* file, std::shared_ptr<Message>* message);
io::RandomAccessFile* file, std::unique_ptr<Message>* message);
/// \brief Read encapulated RPC message (metadata and body) from InputStream
///
/// Read length-prefixed message with as-yet unknown length. Returns nullptr if
/// there are not enough bytes available or the message length is 0 (e.g. EOS
/// in a stream)
Status ARROW_EXPORT ReadMessage(
io::InputStream* stream, std::shared_ptr<Message>* message);
io::InputStream* stream, std::unique_ptr<Message>* message);
/// Write a serialized message with a length-prefix and padding to an 8-byte offset
/// Write a serialized message metadata with a length-prefix and padding to an
/// 8-byte offset
///
/// <message_size: int32><message: const void*><padding>
Status ARROW_EXPORT WriteMessage(
cpp/src/arrow/ipc/reader.cc
+91 -115
View File
@@ -257,11 +257,18 @@ static Status LoadArray(const std::shared_ptr<DataType>& type,
return loader.Load();
}
Status ReadRecordBatch(const Message& metadata, const std::shared_ptr<Schema>& schema,
Status ReadRecordBatch(const Buffer& metadata, const std::shared_ptr<Schema>& schema,
io::RandomAccessFile* file, std::shared_ptr<RecordBatch>* out) {
return ReadRecordBatch(metadata, schema, kMaxNestingDepth, file, out);
}
Status ReadRecordBatch(const Message& message, const std::shared_ptr<Schema>& schema,
std::shared_ptr<RecordBatch>* out) {
io::BufferReader reader(message.body());
DCHECK_EQ(message.type(), Message::RECORD_BATCH);
return ReadRecordBatch(*message.metadata(), schema, kMaxNestingDepth, &reader, out);
}
// ----------------------------------------------------------------------
// Array loading
@@ -294,18 +301,22 @@ static inline Status ReadRecordBatch(const flatbuf::RecordBatch* metadata,
schema, metadata->length(), max_recursion_depth, &source, out);
}
Status ReadRecordBatch(const Message& metadata, const std::shared_ptr<Schema>& schema,
Status ReadRecordBatch(const Buffer& metadata, const std::shared_ptr<Schema>& schema,
int max_recursion_depth, io::RandomAccessFile* file,
std::shared_ptr<RecordBatch>* out) {
DCHECK_EQ(metadata.type(), Message::RECORD_BATCH);
auto batch = reinterpret_cast<const flatbuf::RecordBatch*>(metadata.header());
auto message = flatbuf::GetMessage(metadata.data());
if (message->header_type() != flatbuf::MessageHeader_RecordBatch) {
DCHECK_EQ(message->header_type(), flatbuf::MessageHeader_RecordBatch);
}
auto batch = reinterpret_cast<const flatbuf::RecordBatch*>(message->header());
return ReadRecordBatch(batch, schema, max_recursion_depth, file, out);
}
Status ReadDictionary(const Message& metadata, const DictionaryTypeMap& dictionary_types,
Status ReadDictionary(const Buffer& metadata, const DictionaryTypeMap& dictionary_types,
io::RandomAccessFile* file, int64_t* dictionary_id, std::shared_ptr<Array>* out) {
auto message = flatbuf::GetMessage(metadata.data());
auto dictionary_batch =
reinterpret_cast<const flatbuf::DictionaryBatch*>(metadata.header());
reinterpret_cast<const flatbuf::DictionaryBatch*>(message->header());
int64_t id = *dictionary_id = dictionary_batch->id();
auto it = dictionary_types.find(id);
@@ -335,25 +346,33 @@ Status ReadDictionary(const Message& metadata, const DictionaryTypeMap& dictiona
return Status::OK();
}
static Status ReadMessageAndValidate(MessageReader* reader, Message::Type expected_type,
bool allow_null, std::unique_ptr<Message>* message) {
RETURN_NOT_OK(reader->ReadNextMessage(message));
if (!(*message) && !allow_null) {
std::stringstream ss;
ss << "Expected " << FormatMessageType(expected_type)
<< " message in stream, was null or length 0";
return Status::Invalid(ss.str());
}
if ((*message) == nullptr) { return Status::OK(); }
if ((*message)->type() != expected_type) {
std::stringstream ss;
ss << "Message not expected type: " << FormatMessageType(expected_type)
<< ", was: " << (*message)->type();
return Status::IOError(ss.str());
}
return Status::OK();
}
// ----------------------------------------------------------------------
// RecordBatchStreamReader implementation
static inline FileBlock FileBlockFromFlatbuffer(const flatbuf::Block* block) {
return FileBlock(block->offset(), block->metaDataLength(), block->bodyLength());
}
static inline std::string FormatMessageType(Message::Type type) {
switch (type) {
case Message::SCHEMA:
return "schema";
case Message::RECORD_BATCH:
return "record batch";
case Message::DICTIONARY_BATCH:
return "dictionary";
default:
break;
}
return "unknown";
return FileBlock{block->offset(), block->metaDataLength(), block->bodyLength()};
}
RecordBatchReader::~RecordBatchReader() {}
@@ -363,59 +382,29 @@ class RecordBatchStreamReader::RecordBatchStreamReaderImpl {
RecordBatchStreamReaderImpl() {}
~RecordBatchStreamReaderImpl() {}
Status Open(const std::shared_ptr<io::InputStream>& stream) {
stream_ = stream;
Status Open(std::unique_ptr<MessageReader> message_reader) {
message_reader_ = std::move(message_reader);
return ReadSchema();
}
Status ReadNextMessage(
Message::Type expected_type, bool allow_null, std::shared_ptr<Message>* message) {
RETURN_NOT_OK(ReadMessage(stream_.get(), message));
if (!(*message) && !allow_null) {
std::stringstream ss;
ss << "Expected " << FormatMessageType(expected_type)
<< " message in stream, was null or length 0";
return Status::Invalid(ss.str());
}
if ((*message) == nullptr) { return Status::OK(); }
if ((*message)->type() != expected_type) {
std::stringstream ss;
ss << "Message not expected type: " << FormatMessageType(expected_type)
<< ", was: " << (*message)->type();
return Status::IOError(ss.str());
}
return Status::OK();
}
Status ReadExact(int64_t size, std::shared_ptr<Buffer>* buffer) {
RETURN_NOT_OK(stream_->Read(size, buffer));
if ((*buffer)->size() < size) {
return Status::IOError("Unexpected EOS when reading buffer");
}
return Status::OK();
}
Status ReadNextDictionary() {
std::shared_ptr<Message> message;
RETURN_NOT_OK(ReadNextMessage(Message::DICTIONARY_BATCH, false, &message));
std::unique_ptr<Message> message;
RETURN_NOT_OK(ReadMessageAndValidate(
message_reader_.get(), Message::DICTIONARY_BATCH, false, &message));
std::shared_ptr<Buffer> batch_body;
RETURN_NOT_OK(ReadExact(message->body_length(), &batch_body));
io::BufferReader reader(batch_body);
io::BufferReader reader(message->body());
std::shared_ptr<Array> dictionary;
int64_t id;
RETURN_NOT_OK(ReadDictionary(*message, dictionary_types_, &reader, &id, &dictionary));
RETURN_NOT_OK(ReadDictionary(
*message->metadata(), dictionary_types_, &reader, &id, &dictionary));
return dictionary_memo_.AddDictionary(id, dictionary);
}
Status ReadSchema() {
std::shared_ptr<Message> message;
RETURN_NOT_OK(ReadNextMessage(Message::SCHEMA, false, &message));
std::unique_ptr<Message> message;
RETURN_NOT_OK(
ReadMessageAndValidate(message_reader_.get(), Message::SCHEMA, false, &message));
RETURN_NOT_OK(GetDictionaryTypes(message->header(), &dictionary_types_));
@@ -429,9 +418,10 @@ class RecordBatchStreamReader::RecordBatchStreamReaderImpl {
return GetSchema(message->header(), dictionary_memo_, &schema_);
}
Status GetNextRecordBatch(std::shared_ptr<RecordBatch>* batch) {
std::shared_ptr<Message> message;
RETURN_NOT_OK(ReadNextMessage(Message::RECORD_BATCH, true, &message));
Status ReadNextRecordBatch(std::shared_ptr<RecordBatch>* batch) {
std::unique_ptr<Message> message;
RETURN_NOT_OK(ReadMessageAndValidate(
message_reader_.get(), Message::RECORD_BATCH, true, &message));
if (message == nullptr) {
// End of stream
@@ -439,21 +429,18 @@ class RecordBatchStreamReader::RecordBatchStreamReaderImpl {
return Status::OK();
}
std::shared_ptr<Buffer> batch_body;
RETURN_NOT_OK(ReadExact(message->body_length(), &batch_body));
io::BufferReader reader(batch_body);
return ReadRecordBatch(*message, schema_, &reader, batch);
io::BufferReader reader(message->body());
return ReadRecordBatch(*message->metadata(), schema_, &reader, batch);
}
std::shared_ptr<Schema> schema() const { return schema_; }
private:
std::unique_ptr<MessageReader> message_reader_;
// dictionary_id -> type
DictionaryTypeMap dictionary_types_;
DictionaryMemo dictionary_memo_;
std::shared_ptr<io::InputStream> stream_;
std::shared_ptr<Schema> schema_;
};
@@ -463,19 +450,25 @@ RecordBatchStreamReader::RecordBatchStreamReader() {
RecordBatchStreamReader::~RecordBatchStreamReader() {}
Status RecordBatchStreamReader::Open(const std::shared_ptr<io::InputStream>& stream,
Status RecordBatchStreamReader::Open(std::unique_ptr<MessageReader> message_reader,
std::shared_ptr<RecordBatchStreamReader>* reader) {
// Private ctor
*reader = std::shared_ptr<RecordBatchStreamReader>(new RecordBatchStreamReader());
return (*reader)->impl_->Open(stream);
return (*reader)->impl_->Open(std::move(message_reader));
}
Status RecordBatchStreamReader::Open(const std::shared_ptr<io::InputStream>& stream,
std::shared_ptr<RecordBatchStreamReader>* out) {
std::unique_ptr<MessageReader> message_reader(new InputStreamMessageReader(stream));
return Open(std::move(message_reader), out);
}
std::shared_ptr<Schema> RecordBatchStreamReader::schema() const {
return impl_->schema();
}
Status RecordBatchStreamReader::GetNextRecordBatch(std::shared_ptr<RecordBatch>* batch) {
return impl_->GetNextRecordBatch(batch);
Status RecordBatchStreamReader::ReadNextRecordBatch(std::shared_ptr<RecordBatch>* batch) {
return impl_->ReadNextRecordBatch(batch);
}
// ----------------------------------------------------------------------
@@ -547,22 +540,17 @@ class RecordBatchFileReader::RecordBatchFileReaderImpl {
return FileBlockFromFlatbuffer(footer_->dictionaries()->Get(i));
}
Status GetRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) {
Status ReadRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) {
DCHECK_GE(i, 0);
DCHECK_LT(i, num_record_batches());
FileBlock block = record_batch(i);
std::shared_ptr<Message> message;
std::unique_ptr<Message> message;
RETURN_NOT_OK(
ReadMessage(block.offset, block.metadata_length, file_.get(), &message));
// 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(*message, schema_, &reader, batch);
io::BufferReader reader(message->body());
return ::arrow::ipc::ReadRecordBatch(*message->metadata(), schema_, &reader, batch);
}
Status ReadSchema() {
@@ -571,23 +559,16 @@ class RecordBatchFileReader::RecordBatchFileReaderImpl {
// Read all the dictionaries
for (int i = 0; i < num_dictionaries(); ++i) {
FileBlock block = dictionary(i);
std::shared_ptr<Message> message;
std::unique_ptr<Message> message;
RETURN_NOT_OK(
ReadMessage(block.offset, block.metadata_length, file_.get(), &message));
// TODO(wesm): ARROW-577: This code is a bit duplicated, can be fixed
// with a more invasive refactor
// 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);
io::BufferReader reader(message->body());
std::shared_ptr<Array> dictionary;
int64_t dictionary_id;
RETURN_NOT_OK(ReadDictionary(
*message, dictionary_fields_, &reader, &dictionary_id, &dictionary));
RETURN_NOT_OK(ReadDictionary(*message->metadata(), dictionary_fields_, &reader,
&dictionary_id, &dictionary));
RETURN_NOT_OK(dictionary_memo_->AddDictionary(dictionary_id, dictionary));
}
@@ -653,12 +634,13 @@ MetadataVersion RecordBatchFileReader::version() const {
return impl_->version();
}
Status RecordBatchFileReader::GetRecordBatch(int i, std::shared_ptr<RecordBatch>* batch) {
return impl_->GetRecordBatch(i, batch);
Status RecordBatchFileReader::ReadRecordBatch(
int i, std::shared_ptr<RecordBatch>* batch) {
return impl_->ReadRecordBatch(i, batch);
}
static Status ReadContiguousPayload(int64_t offset, io::RandomAccessFile* file,
std::shared_ptr<Message>* message, std::shared_ptr<Buffer>* payload) {
static Status ReadContiguousPayload(
int64_t offset, io::RandomAccessFile* file, std::unique_ptr<Message>* message) {
std::shared_ptr<Buffer> buffer;
RETURN_NOT_OK(file->Seek(offset));
RETURN_NOT_OK(ReadMessage(file, message));
@@ -666,38 +648,32 @@ static Status ReadContiguousPayload(int64_t offset, io::RandomAccessFile* file,
if (*message == nullptr) {
return Status::Invalid("Unable to read metadata at offset");
}
// TODO(ARROW-388): The buffer offsets start at 0, so we must construct a
// RandomAccessFile according to that frame of reference
RETURN_NOT_OK(file->Read((*message)->body_length(), payload));
return Status::OK();
}
Status ReadRecordBatch(const std::shared_ptr<Schema>& schema, int64_t offset,
io::RandomAccessFile* file, std::shared_ptr<RecordBatch>* out) {
std::shared_ptr<Buffer> payload;
std::shared_ptr<Message> message;
RETURN_NOT_OK(ReadContiguousPayload(offset, file, &message, &payload));
io::BufferReader buffer_reader(payload);
return ReadRecordBatch(*message, schema, kMaxNestingDepth, &buffer_reader, out);
std::unique_ptr<Message> message;
RETURN_NOT_OK(ReadContiguousPayload(offset, file, &message));
io::BufferReader buffer_reader(message->body());
return ReadRecordBatch(
*message->metadata(), schema, kMaxNestingDepth, &buffer_reader, out);
}
Status ReadTensor(
int64_t offset, io::RandomAccessFile* file, std::shared_ptr<Tensor>* out) {
// Respect alignment of Tensor messages (see WriteTensor)
offset = PaddedLength(offset);
std::shared_ptr<Message> message;
std::shared_ptr<Buffer> data;
RETURN_NOT_OK(ReadContiguousPayload(offset, file, &message, &data));
std::unique_ptr<Message> message;
RETURN_NOT_OK(ReadContiguousPayload(offset, file, &message));
std::shared_ptr<DataType> type;
std::vector<int64_t> shape;
std::vector<int64_t> strides;
std::vector<std::string> dim_names;
RETURN_NOT_OK(
GetTensorMetadata(message->header(), &type, &shape, &strides, &dim_names));
*out = std::make_shared<Tensor>(type, data, shape, strides, dim_names);
GetTensorMetadata(*message->metadata(), &type, &shape, &strides, &dim_names));
*out = std::make_shared<Tensor>(type, message->body(), shape, strides, dim_names);
return Status::OK();
}
cpp/src/arrow/ipc/reader.h
+26 -9
View File
@@ -57,7 +57,7 @@ class ARROW_EXPORT RecordBatchReader {
///
/// \param(out) batch the next loaded batch, nullptr at end of stream
/// \return Status
virtual Status GetNextRecordBatch(std::shared_ptr<RecordBatch>* batch) = 0;
virtual Status ReadNextRecordBatch(std::shared_ptr<RecordBatch>* batch) = 0;
};
/// \class RecordBatchStreamReader
@@ -66,16 +66,24 @@ class ARROW_EXPORT RecordBatchStreamReader : public RecordBatchReader {
public:
virtual ~RecordBatchStreamReader();
/// Create batch reader from InputStream
/// Create batch reader from generic MessageReader
///
/// \param(in) message_reader a MessageReader implementation
/// \param(out) out the created RecordBatchStreamReader object
/// \return Status
static Status Open(std::unique_ptr<MessageReader> message_reader,
std::shared_ptr<RecordBatchStreamReader>* out);
/// \Create Record batch stream reader from InputStream
///
/// \param(in) stream an input stream instance
/// \param(out) reader the created reader object
/// \param(out) out the created RecordBatchStreamReader object
/// \return Status
static Status Open(const std::shared_ptr<io::InputStream>& stream,
std::shared_ptr<RecordBatchStreamReader>* reader);
std::shared_ptr<RecordBatchStreamReader>* out);
std::shared_ptr<Schema> schema() const override;
Status GetNextRecordBatch(std::shared_ptr<RecordBatch>* batch) override;
Status ReadNextRecordBatch(std::shared_ptr<RecordBatch>* batch) override;
private:
RecordBatchStreamReader();
@@ -122,7 +130,7 @@ class ARROW_EXPORT RecordBatchFileReader {
/// \param(in) i the index of the record batch to return
/// \param(out) batch the read batch
/// \return Status
Status GetRecordBatch(int i, std::shared_ptr<RecordBatch>* batch);
Status ReadRecordBatch(int i, std::shared_ptr<RecordBatch>* batch);
private:
RecordBatchFileReader();
@@ -133,16 +141,25 @@ class ARROW_EXPORT RecordBatchFileReader {
// Generic read functions; does not copy data if the input supports zero copy reads
/// Read record batch from file given metadata and schema
/// \brief Read record batch from file given metadata and schema
///
/// \param(in) metadata a Message containing the record batch metadata
/// \param(in) schema the record batch schema
/// \param(in) file a random access file
/// \param(out) out the read record batch
Status ARROW_EXPORT ReadRecordBatch(const Message& metadata,
Status ARROW_EXPORT ReadRecordBatch(const Buffer& metadata,
const std::shared_ptr<Schema>& schema, io::RandomAccessFile* file,
std::shared_ptr<RecordBatch>* out);
/// \brief Read record batch from fully encapulated Message
///
/// \param[in] message a message instance containing metadata and body
/// \param[in] schema
/// \param[out] out the resulting RecordBatch
/// \return Status
Status ARROW_EXPORT ReadRecordBatch(const Message& message,
const std::shared_ptr<Schema>& schema, std::shared_ptr<RecordBatch>* out);
/// Read record batch from file given metadata and schema
///
/// \param(in) metadata a Message containing the record batch metadata
@@ -150,7 +167,7 @@ Status ARROW_EXPORT ReadRecordBatch(const Message& metadata,
/// \param(in) file a random access file
/// \param(in) max_recursion_depth the maximum permitted nesting depth
/// \param(out) out the read record batch
Status ARROW_EXPORT ReadRecordBatch(const Message& metadata,
Status ARROW_EXPORT ReadRecordBatch(const Buffer& metadata,
const std::shared_ptr<Schema>& schema, int max_recursion_depth,
io::RandomAccessFile* file, std::shared_ptr<RecordBatch>* out);
cpp/src/arrow/ipc/stream-to-file.cc
+1 -1
View File
@@ -40,7 +40,7 @@ Status ConvertToFile() {
std::shared_ptr<RecordBatch> batch;
while (true) {
RETURN_NOT_OK(reader->GetNextRecordBatch(&batch));
RETURN_NOT_OK(reader->ReadNextRecordBatch(&batch));
if (batch == nullptr) break;
RETURN_NOT_OK(writer->WriteRecordBatch(*batch));
}
cpp/src/arrow/ipc/writer.cc
+2 -2
View File
@@ -110,7 +110,7 @@ class RecordBatchSerializer : public ArrayVisitor {
}
// push back all common elements
field_nodes_.emplace_back(arr.length(), arr.null_count(), 0);
field_nodes_.push_back({arr.length(), arr.null_count(), 0});
if (arr.null_count() > 0) {
std::shared_ptr<Buffer> bitmap;
@@ -680,7 +680,7 @@ class RecordBatchStreamWriter::RecordBatchStreamWriterImpl {
Status WriteRecordBatch(const RecordBatch& batch, bool allow_64bit) {
// Push an empty FileBlock. Can be written in the footer later
record_batches_.emplace_back(0, 0, 0);
record_batches_.push_back({0, 0, 0});
return WriteRecordBatch(
batch, allow_64bit, &record_batches_[record_batches_.size() - 1]);
}
cpp/src/arrow/python/builtin_convert.cc
+2 -3
View File
@@ -477,9 +477,8 @@ class FixedWidthBytesConverter
inline Status AppendItem(const OwnedRef& item) {
PyObject* bytes_obj;
OwnedRef tmp;
Py_ssize_t expected_length =
std::dynamic_pointer_cast<FixedSizeBinaryType>(typed_builder_->type())
->byte_width();
Py_ssize_t expected_length = std::dynamic_pointer_cast<FixedSizeBinaryType>(
typed_builder_->type())->byte_width();
if (item.obj() == Py_None) {
RETURN_NOT_OK(typed_builder_->AppendNull());
return Status::OK();
python/doc/source/api.rst
+8 -4
View File
@@ -137,7 +137,6 @@ Tables and Record Batches
Column
RecordBatch
Table
get_record_batch_size
.. _api.tensor:
@@ -148,9 +147,6 @@ Tensor type and Functions
:toctree: generated/
Tensor
write_tensor
get_tensor_size
read_tensor
.. _api.io:
@@ -177,12 +173,20 @@ Interprocess Communication and Messaging
.. autosummary::
:toctree: generated/
Message
MessageReader
RecordBatchFileReader
RecordBatchFileWriter
RecordBatchStreamReader
RecordBatchStreamWriter
open_file
open_stream
read_message
read_record_batch
get_record_batch_size
read_tensor
write_tensor
get_tensor_size
.. _api.memory_pool:
python/pyarrow/__init__.py
+6 -3
View File
@@ -69,9 +69,8 @@ from pyarrow.lib import (null, bool_,
from pyarrow.lib import (HdfsFile, NativeFile, PythonFile,
Buffer, BufferReader, BufferOutputStream,
OSFile, MemoryMappedFile, memory_map,
frombuffer, read_tensor, write_tensor,
frombuffer,
memory_map, create_memory_map,
get_record_batch_size, get_tensor_size,
have_libhdfs, have_libhdfs3, MockOutputStream)
from pyarrow.lib import (MemoryPool, total_allocated_bytes,
@@ -89,8 +88,12 @@ from pyarrow.lib import (ArrowException,
from pyarrow.filesystem import Filesystem, HdfsClient, LocalFilesystem
from pyarrow.ipc import (RecordBatchFileReader, RecordBatchFileWriter,
from pyarrow.ipc import (Message, MessageReader,
RecordBatchFileReader, RecordBatchFileWriter,
RecordBatchStreamReader, RecordBatchStreamWriter,
read_message, read_record_batch, read_tensor,
write_tensor,
get_record_batch_size, get_tensor_size,
open_stream,
open_file,
serialize_pandas, deserialize_pandas)
python/pyarrow/feather.pxi
+109
View File
@@ -0,0 +1,109 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#----------------------------------------------------------------------
# Implement legacy Feather file format
class FeatherError(Exception):
pass
cdef class FeatherWriter:
cdef:
unique_ptr[CFeatherWriter] writer
cdef public:
int64_t num_rows
def __cinit__(self):
self.num_rows = -1
def open(self, object dest):
cdef shared_ptr[OutputStream] sink
get_writer(dest, &sink)
with nogil:
check_status(CFeatherWriter.Open(sink, &self.writer))
def close(self):
if self.num_rows < 0:
self.num_rows = 0
self.writer.get().SetNumRows(self.num_rows)
check_status(self.writer.get().Finalize())
def write_array(self, object name, object col, object mask=None):
cdef Array arr
if self.num_rows >= 0:
if len(col) != self.num_rows:
raise ValueError('prior column had a different number of rows')
else:
self.num_rows = len(col)
if isinstance(col, Array):
arr = col
else:
arr = Array.from_pandas(col, mask=mask)
cdef c_string c_name = tobytes(name)
with nogil:
check_status(
self.writer.get().Append(c_name, deref(arr.sp_array)))
cdef class FeatherReader:
cdef:
unique_ptr[CFeatherReader] reader
def __cinit__(self):
pass
def open(self, source):
cdef shared_ptr[RandomAccessFile] reader
get_reader(source, &reader)
with nogil:
check_status(CFeatherReader.Open(reader, &self.reader))
property num_rows:
def __get__(self):
return self.reader.get().num_rows()
property num_columns:
def __get__(self):
return self.reader.get().num_columns()
def get_column_name(self, int i):
cdef c_string name = self.reader.get().GetColumnName(i)
return frombytes(name)
def get_column(self, int i):
if i < 0 or i >= self.num_columns:
raise IndexError(i)
cdef shared_ptr[CColumn] sp_column
with nogil:
check_status(self.reader.get()
.GetColumn(i, &sp_column))
cdef Column col = Column()
col.init(sp_column)
return col
python/pyarrow/includes/libarrow.pxd
+38 -28
View File
@@ -546,41 +546,41 @@ cdef extern from "arrow/io/memory.h" namespace "arrow::io" nogil:
int64_t GetExtentBytesWritten()
cdef extern from "arrow/ipc/metadata.h" namespace "arrow::ipc" nogil:
cdef cppclass SchemaMessage:
int num_fields()
CStatus GetField(int i, shared_ptr[CField]* out)
CStatus GetSchema(shared_ptr[CSchema]* out)
cdef cppclass FieldMetadata:
pass
cdef cppclass BufferMetadata:
pass
cdef cppclass RecordBatchMessage:
pass
cdef cppclass DictionaryBatchMessage:
pass
cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil:
enum MessageType" arrow::ipc::Message::Type":
MessageType_SCHEMA" arrow::ipc::Message::SCHEMA"
MessageType_RECORD_BATCH" arrow::ipc::Message::RECORD_BATCH"
MessageType_DICTIONARY_BATCH" arrow::ipc::Message::DICTIONARY_BATCH"
cdef cppclass Message:
CStatus Open(const shared_ptr[CBuffer]& buf,
shared_ptr[Message]* out)
int64_t body_length()
enum MetadataVersion" arrow::ipc::MetadataVersion":
MessageType_V1" arrow::ipc::MetadataVersion::V1"
MessageType_V2" arrow::ipc::MetadataVersion::V2"
MessageType_V3" arrow::ipc::MetadataVersion::V3"
cdef cppclass CMessage" arrow::ipc::Message":
CStatus Open(const shared_ptr[CBuffer]& metadata,
const shared_ptr[CBuffer]& body,
unique_ptr[CMessage]* out)
shared_ptr[CBuffer] body()
c_bool Equals(const CMessage& other)
shared_ptr[CBuffer] metadata()
MetadataVersion metadata_version()
MessageType type()
shared_ptr[SchemaMessage] GetSchema()
shared_ptr[RecordBatchMessage] GetRecordBatch()
shared_ptr[DictionaryBatchMessage] GetDictionaryBatch()
CStatus SerializeTo(OutputStream* stream, int64_t* output_length)
c_string FormatMessageType(MessageType type)
cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil:
cdef cppclass CMessageReader \
" arrow::ipc::MessageReader":
CStatus ReadNextMessage(unique_ptr[CMessage]* out)
cdef cppclass CInputStreamMessageReader \
" arrow::ipc::InputStreamMessageReader":
CInputStreamMessageReader(const shared_ptr[InputStream]& stream)
cdef cppclass CRecordBatchWriter \
" arrow::ipc::RecordBatchWriter":
@@ -590,7 +590,7 @@ cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil:
cdef cppclass CRecordBatchReader \
" arrow::ipc::RecordBatchReader":
shared_ptr[CSchema] schema()
CStatus GetNextRecordBatch(shared_ptr[CRecordBatch]* batch)
CStatus ReadNextRecordBatch(shared_ptr[CRecordBatch]* batch)
cdef cppclass CRecordBatchStreamReader \
" arrow::ipc::RecordBatchStreamReader"(CRecordBatchReader):
@@ -598,6 +598,10 @@ cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil:
CStatus Open(const shared_ptr[InputStream]& stream,
shared_ptr[CRecordBatchStreamReader]* out)
@staticmethod
CStatus Open2" Open"(unique_ptr[CMessageReader] message_reader,
shared_ptr[CRecordBatchStreamReader]* out)
cdef cppclass CRecordBatchStreamWriter \
" arrow::ipc::RecordBatchStreamWriter"(CRecordBatchWriter):
@staticmethod
@@ -625,7 +629,9 @@ cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil:
int num_record_batches()
CStatus GetRecordBatch(int i, shared_ptr[CRecordBatch]* batch)
CStatus ReadRecordBatch(int i, shared_ptr[CRecordBatch]* batch)
CStatus ReadMessage(InputStream* stream, unique_ptr[CMessage]* message)
CStatus GetRecordBatchSize(const CRecordBatch& batch, int64_t* size)
CStatus GetTensorSize(const CTensor& tensor, int64_t* size)
@@ -637,6 +643,10 @@ cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil:
CStatus ReadTensor(int64_t offset, RandomAccessFile* file,
shared_ptr[CTensor]* out)
CStatus ReadRecordBatch(const CMessage& message,
const shared_ptr[CSchema]& schema,
shared_ptr[CRecordBatch]* out)
cdef extern from "arrow/ipc/feather.h" namespace "arrow::ipc::feather" nogil:
python/pyarrow/io.pxi
-353
View File
@@ -944,356 +944,3 @@ cdef class HdfsFile(NativeFile):
def __dealloc__(self):
self.parent = None
# ----------------------------------------------------------------------
# File and stream readers and writers
cdef class _RecordBatchWriter:
cdef:
shared_ptr[CRecordBatchWriter] writer
shared_ptr[OutputStream] sink
bint closed
def __cinit__(self):
self.closed = True
def __dealloc__(self):
if not self.closed:
self.close()
def _open(self, sink, Schema schema):
cdef:
shared_ptr[CRecordBatchStreamWriter] writer
get_writer(sink, &self.sink)
with nogil:
check_status(
CRecordBatchStreamWriter.Open(self.sink.get(),
schema.sp_schema,
&writer))
self.writer = <shared_ptr[CRecordBatchWriter]> writer
self.closed = False
def write_batch(self, RecordBatch batch):
with nogil:
check_status(self.writer.get()
.WriteRecordBatch(deref(batch.batch)))
def close(self):
with nogil:
check_status(self.writer.get().Close())
self.closed = True
cdef class _RecordBatchReader:
cdef:
shared_ptr[CRecordBatchReader] reader
cdef readonly:
Schema schema
def __cinit__(self):
pass
def _open(self, source):
cdef:
shared_ptr[RandomAccessFile] file_handle
shared_ptr[InputStream] in_stream
shared_ptr[CRecordBatchStreamReader] reader
get_reader(source, &file_handle)
in_stream = <shared_ptr[InputStream]> file_handle
with nogil:
check_status(CRecordBatchStreamReader.Open(in_stream, &reader))
self.reader = <shared_ptr[CRecordBatchReader]> reader
self.schema = pyarrow_wrap_schema(self.reader.get().schema())
def get_next_batch(self):
"""
Read next RecordBatch from the stream. Raises StopIteration at end of
stream
"""
cdef shared_ptr[CRecordBatch] batch
with nogil:
check_status(self.reader.get().GetNextRecordBatch(&batch))
if batch.get() == NULL:
raise StopIteration
return pyarrow_wrap_batch(batch)
def read_all(self):
"""
Read all record batches as a pyarrow.Table
"""
cdef:
vector[shared_ptr[CRecordBatch]] batches
shared_ptr[CRecordBatch] batch
shared_ptr[CTable] table
with nogil:
while True:
check_status(self.reader.get().GetNextRecordBatch(&batch))
if batch.get() == NULL:
break
batches.push_back(batch)
check_status(CTable.FromRecordBatches(batches, &table))
return pyarrow_wrap_table(table)
cdef class _RecordBatchFileWriter(_RecordBatchWriter):
def _open(self, sink, Schema schema):
cdef shared_ptr[CRecordBatchFileWriter] writer
get_writer(sink, &self.sink)
with nogil:
check_status(
CRecordBatchFileWriter.Open(self.sink.get(), schema.sp_schema,
&writer))
# Cast to base class, because has same interface
self.writer = <shared_ptr[CRecordBatchWriter]> writer
self.closed = False
cdef class _RecordBatchFileReader:
cdef:
shared_ptr[CRecordBatchFileReader] reader
cdef readonly:
Schema schema
def __cinit__(self):
pass
def _open(self, source, footer_offset=None):
cdef shared_ptr[RandomAccessFile] reader
get_reader(source, &reader)
cdef int64_t offset = 0
if footer_offset is not None:
offset = footer_offset
with nogil:
if offset != 0:
check_status(CRecordBatchFileReader.Open2(
reader, offset, &self.reader))
else:
check_status(CRecordBatchFileReader.Open(reader, &self.reader))
self.schema = pyarrow_wrap_schema(self.reader.get().schema())
property num_record_batches:
def __get__(self):
return self.reader.get().num_record_batches()
def get_batch(self, int i):
cdef shared_ptr[CRecordBatch] batch
if i < 0 or i >= self.num_record_batches:
raise ValueError('Batch number {0} out of range'.format(i))
with nogil:
check_status(self.reader.get().GetRecordBatch(i, &batch))
return pyarrow_wrap_batch(batch)
# TODO(wesm): ARROW-503: Function was renamed. Remove after a period of
# time has passed
get_record_batch = get_batch
def read_all(self):
"""
Read all record batches as a pyarrow.Table
"""
cdef:
vector[shared_ptr[CRecordBatch]] batches
shared_ptr[CTable] table
int i, nbatches
nbatches = self.num_record_batches
batches.resize(nbatches)
with nogil:
for i in range(nbatches):
check_status(self.reader.get().GetRecordBatch(i, &batches[i]))
check_status(CTable.FromRecordBatches(batches, &table))
return pyarrow_wrap_table(table)
#----------------------------------------------------------------------
# Implement legacy Feather file format
class FeatherError(Exception):
pass
cdef class FeatherWriter:
cdef:
unique_ptr[CFeatherWriter] writer
cdef public:
int64_t num_rows
def __cinit__(self):
self.num_rows = -1
def open(self, object dest):
cdef shared_ptr[OutputStream] sink
get_writer(dest, &sink)
with nogil:
check_status(CFeatherWriter.Open(sink, &self.writer))
def close(self):
if self.num_rows < 0:
self.num_rows = 0
self.writer.get().SetNumRows(self.num_rows)
check_status(self.writer.get().Finalize())
def write_array(self, object name, object col, object mask=None):
cdef Array arr
if self.num_rows >= 0:
if len(col) != self.num_rows:
raise ValueError('prior column had a different number of rows')
else:
self.num_rows = len(col)
if isinstance(col, Array):
arr = col
else:
arr = Array.from_pandas(col, mask=mask)
cdef c_string c_name = tobytes(name)
with nogil:
check_status(
self.writer.get().Append(c_name, deref(arr.sp_array)))
cdef class FeatherReader:
cdef:
unique_ptr[CFeatherReader] reader
def __cinit__(self):
pass
def open(self, source):
cdef shared_ptr[RandomAccessFile] reader
get_reader(source, &reader)
with nogil:
check_status(CFeatherReader.Open(reader, &self.reader))
property num_rows:
def __get__(self):
return self.reader.get().num_rows()
property num_columns:
def __get__(self):
return self.reader.get().num_columns()
def get_column_name(self, int i):
cdef c_string name = self.reader.get().GetColumnName(i)
return frombytes(name)
def get_column(self, int i):
if i < 0 or i >= self.num_columns:
raise IndexError(i)
cdef shared_ptr[CColumn] sp_column
with nogil:
check_status(self.reader.get()
.GetColumn(i, &sp_column))
cdef Column col = Column()
col.init(sp_column)
return col
def get_tensor_size(Tensor tensor):
"""
Return total size of serialized Tensor including metadata and padding
"""
cdef int64_t size
with nogil:
check_status(GetTensorSize(deref(tensor.tp), &size))
return size
def get_record_batch_size(RecordBatch batch):
"""
Return total size of serialized RecordBatch including metadata and padding
"""
cdef int64_t size
with nogil:
check_status(GetRecordBatchSize(deref(batch.batch), &size))
return size
def write_tensor(Tensor tensor, NativeFile dest):
"""
Write pyarrow.Tensor to pyarrow.NativeFile object its current position
Parameters
----------
tensor : pyarrow.Tensor
dest : pyarrow.NativeFile
Returns
-------
bytes_written : int
Total number of bytes written to the file
"""
cdef:
int32_t metadata_length
int64_t body_length
dest._assert_writeable()
with nogil:
check_status(
WriteTensor(deref(tensor.tp), dest.wr_file.get(),
&metadata_length, &body_length))
return metadata_length + body_length
def read_tensor(NativeFile source):
"""
Read pyarrow.Tensor from pyarrow.NativeFile object from current
position. If the file source supports zero copy (e.g. a memory map), then
this operation does not allocate any memory
Parameters
----------
source : pyarrow.NativeFile
Returns
-------
tensor : Tensor
"""
cdef:
shared_ptr[CTensor] sp_tensor
source._assert_readable()
cdef int64_t offset = source.tell()
with nogil:
check_status(ReadTensor(offset, source.rd_file.get(), &sp_tensor))
return pyarrow_wrap_tensor(sp_tensor)
python/pyarrow/ipc.pxi
+480
View File
@@ -0,0 +1,480 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
cdef class Message:
cdef:
unique_ptr[CMessage] message
def __cinit__(self):
pass
def __null_check(self):
if self.message.get() == NULL:
raise TypeError('Message improperly initialized (null)')
property type:
def __get__(self):
self.__null_check()
return frombytes(FormatMessageType(self.message.get().type()))
property metadata:
def __get__(self):
self.__null_check()
return pyarrow_wrap_buffer(self.message.get().metadata())
property body:
def __get__(self):
self.__null_check()
cdef shared_ptr[CBuffer] body = self.message.get().body()
if body.get() == NULL:
return None
else:
return pyarrow_wrap_buffer(body)
def equals(self, Message other):
"""
Returns True if the message contents (metadata and body) are identical
Parameters
----------
other : Message
Returns
-------
are_equal : bool
"""
cdef c_bool result
with nogil:
result = self.message.get().Equals(deref(other.message.get()))
return result
def serialize(self, memory_pool=None):
"""
Write message to Buffer with length-prefixed metadata, then body
Parameters
----------
memory_pool : MemoryPool, default None
Uses default memory pool if not specified
Returns
-------
serialized : Buffer
"""
cdef:
BufferOutputStream stream = BufferOutputStream(memory_pool)
int64_t output_length = 0
with nogil:
check_status(self.message.get()
.SerializeTo(stream.wr_file.get(),
&output_length))
return stream.get_result()
def __repr__(self):
metadata_len = self.metadata.size
body = self.body
body_len = 0 if body is None else body.size
return """pyarrow.Message
type: {0}
metadata length: {1}
body length: {2}""".format(self.type, metadata_len, body_len)
cdef class MessageReader:
cdef:
unique_ptr[CMessageReader] reader
def __cinit__(self):
pass
def __null_check(self):
if self.reader.get() == NULL:
raise TypeError('Message improperly initialized (null)')
def __repr__(self):
self.__null_check()
return object.__repr__(self)
@staticmethod
def open_stream(source):
cdef MessageReader result = MessageReader()
cdef shared_ptr[InputStream] in_stream
get_input_stream(source, &in_stream)
with nogil:
result.reader.reset(new CInputStreamMessageReader(in_stream))
return result
def __iter__(self):
while True:
yield self.read_next_message()
def read_next_message(self):
"""
Read next Message from the stream. Raises StopIteration at end of
stream
"""
cdef Message result = Message()
with nogil:
check_status(self.reader.get().ReadNextMessage(&result.message))
if result.message.get() == NULL:
raise StopIteration
return result
# ----------------------------------------------------------------------
# File and stream readers and writers
cdef class _RecordBatchWriter:
cdef:
shared_ptr[CRecordBatchWriter] writer
shared_ptr[OutputStream] sink
bint closed
def __cinit__(self):
self.closed = True
def __dealloc__(self):
if not self.closed:
self.close()
def _open(self, sink, Schema schema):
cdef:
shared_ptr[CRecordBatchStreamWriter] writer
get_writer(sink, &self.sink)
with nogil:
check_status(
CRecordBatchStreamWriter.Open(self.sink.get(),
schema.sp_schema,
&writer))
self.writer = <shared_ptr[CRecordBatchWriter]> writer
self.closed = False
def write_batch(self, RecordBatch batch):
with nogil:
check_status(self.writer.get()
.WriteRecordBatch(deref(batch.batch)))
def close(self):
with nogil:
check_status(self.writer.get().Close())
self.closed = True
cdef get_input_stream(object source, shared_ptr[InputStream]* out):
cdef:
shared_ptr[RandomAccessFile] file_handle
get_reader(source, &file_handle)
out[0] = <shared_ptr[InputStream]> file_handle
cdef class _RecordBatchReader:
cdef:
shared_ptr[CRecordBatchReader] reader
cdef readonly:
Schema schema
def __cinit__(self):
pass
def _open(self, source):
cdef:
shared_ptr[InputStream] in_stream
shared_ptr[CRecordBatchStreamReader] reader
get_input_stream(source, &in_stream)
with nogil:
check_status(CRecordBatchStreamReader.Open(in_stream, &reader))
self.reader = <shared_ptr[CRecordBatchReader]> reader
self.schema = Schema()
self.schema.init_schema(self.reader.get().schema())
def __iter__(self):
while True:
yield self.read_next_batch()
def get_next_batch(self):
import warnings
warnings.warn('Please use read_next_batch instead of '
'get_next_batch', FutureWarning)
return self.read_next_batch()
def read_next_batch(self):
"""
Read next RecordBatch from the stream. Raises StopIteration at end of
stream
"""
cdef shared_ptr[CRecordBatch] batch
with nogil:
check_status(self.reader.get().ReadNextRecordBatch(&batch))
if batch.get() == NULL:
raise StopIteration
return pyarrow_wrap_batch(batch)
def read_all(self):
"""
Read all record batches as a pyarrow.Table
"""
cdef:
vector[shared_ptr[CRecordBatch]] batches
shared_ptr[CRecordBatch] batch
shared_ptr[CTable] table
with nogil:
while True:
check_status(self.reader.get().ReadNextRecordBatch(&batch))
if batch.get() == NULL:
break
batches.push_back(batch)
check_status(CTable.FromRecordBatches(batches, &table))
return pyarrow_wrap_table(table)
cdef class _RecordBatchFileWriter(_RecordBatchWriter):
def _open(self, sink, Schema schema):
cdef shared_ptr[CRecordBatchFileWriter] writer
get_writer(sink, &self.sink)
with nogil:
check_status(
CRecordBatchFileWriter.Open(self.sink.get(), schema.sp_schema,
&writer))
# Cast to base class, because has same interface
self.writer = <shared_ptr[CRecordBatchWriter]> writer
self.closed = False
cdef class _RecordBatchFileReader:
cdef:
shared_ptr[CRecordBatchFileReader] reader
cdef readonly:
Schema schema
def __cinit__(self):
pass
def _open(self, source, footer_offset=None):
cdef shared_ptr[RandomAccessFile] reader
get_reader(source, &reader)
cdef int64_t offset = 0
if footer_offset is not None:
offset = footer_offset
with nogil:
if offset != 0:
check_status(CRecordBatchFileReader.Open2(
reader, offset, &self.reader))
else:
check_status(CRecordBatchFileReader.Open(reader, &self.reader))
self.schema = pyarrow_wrap_schema(self.reader.get().schema())
property num_record_batches:
def __get__(self):
return self.reader.get().num_record_batches()
def get_batch(self, int i):
cdef shared_ptr[CRecordBatch] batch
if i < 0 or i >= self.num_record_batches:
raise ValueError('Batch number {0} out of range'.format(i))
with nogil:
check_status(self.reader.get().ReadRecordBatch(i, &batch))
return pyarrow_wrap_batch(batch)
# TODO(wesm): ARROW-503: Function was renamed. Remove after a period of
# time has passed
get_record_batch = get_batch
def read_all(self):
"""
Read all record batches as a pyarrow.Table
"""
cdef:
vector[shared_ptr[CRecordBatch]] batches
shared_ptr[CTable] table
int i, nbatches
nbatches = self.num_record_batches
batches.resize(nbatches)
with nogil:
for i in range(nbatches):
check_status(self.reader.get().ReadRecordBatch(i, &batches[i]))
check_status(CTable.FromRecordBatches(batches, &table))
return pyarrow_wrap_table(table)
def get_tensor_size(Tensor tensor):
"""
Return total size of serialized Tensor including metadata and padding
"""
cdef int64_t size
with nogil:
check_status(GetTensorSize(deref(tensor.tp), &size))
return size
def get_record_batch_size(RecordBatch batch):
"""
Return total size of serialized RecordBatch including metadata and padding
"""
cdef int64_t size
with nogil:
check_status(GetRecordBatchSize(deref(batch.batch), &size))
return size
def write_tensor(Tensor tensor, NativeFile dest):
"""
Write pyarrow.Tensor to pyarrow.NativeFile object its current position
Parameters
----------
tensor : pyarrow.Tensor
dest : pyarrow.NativeFile
Returns
-------
bytes_written : int
Total number of bytes written to the file
"""
cdef:
int32_t metadata_length
int64_t body_length
dest._assert_writeable()
with nogil:
check_status(
WriteTensor(deref(tensor.tp), dest.wr_file.get(),
&metadata_length, &body_length))
return metadata_length + body_length
def read_tensor(NativeFile source):
"""
Read pyarrow.Tensor from pyarrow.NativeFile object from current
position. If the file source supports zero copy (e.g. a memory map), then
this operation does not allocate any memory
Parameters
----------
source : pyarrow.NativeFile
Returns
-------
tensor : Tensor
"""
cdef:
shared_ptr[CTensor] sp_tensor
source._assert_readable()
cdef int64_t offset = source.tell()
with nogil:
check_status(ReadTensor(offset, source.rd_file.get(), &sp_tensor))
return pyarrow_wrap_tensor(sp_tensor)
def read_message(source):
"""
Read length-prefixed message from file or buffer-like object
Parameters
----------
source : pyarrow.NativeFile, file-like object, or buffer-like object
Returns
-------
message : Message
"""
cdef:
Message result = Message()
NativeFile cpp_file
if not isinstance(source, NativeFile):
if hasattr(source, 'read'):
source = PythonFile(source)
else:
source = BufferReader(source)
if not isinstance(source, NativeFile):
raise ValueError('Unable to read message from object with type: {0}'
.format(type(source)))
source._assert_readable()
cpp_file = source
with nogil:
check_status(ReadMessage(cpp_file.rd_file.get(),
&result.message))
return result
def read_record_batch(Message batch_message, Schema schema):
"""
Read RecordBatch from message, given a known schema
Parameters
----------
batch_message : Message
Such as that obtained from read_message
schema : Schema
Returns
-------
batch : RecordBatch
"""
cdef shared_ptr[CRecordBatch] result
with nogil:
check_status(ReadRecordBatch(deref(batch_message.message.get()),
schema.sp_schema, &result))
return pyarrow_wrap_batch(result)
python/pyarrow/ipc.py
+7 -6
View File
@@ -18,6 +18,11 @@
# Arrow file and stream reader/writer classes, and other messaging tools
import pyarrow as pa
from pyarrow.lib import (Message, MessageReader, # noqa
read_message, read_record_batch,
read_tensor, write_tensor,
get_record_batch_size, get_tensor_size)
import pyarrow.lib as lib
@@ -33,10 +38,6 @@ class RecordBatchStreamReader(lib._RecordBatchReader):
def __init__(self, source):
self._open(source)
def __iter__(self):
while True:
yield self.get_next_batch()
class RecordBatchStreamWriter(lib._RecordBatchWriter):
"""
@@ -136,7 +137,7 @@ def serialize_pandas(df):
"""
batch = pa.RecordBatch.from_pandas(df)
sink = pa.InMemoryOutputStream()
writer = pa.RecordBatchFileWriter(sink, batch.schema)
writer = pa.RecordBatchStreamWriter(sink, batch.schema)
writer.write_batch(batch)
writer.close()
return sink.get_result()
@@ -157,6 +158,6 @@ def deserialize_pandas(buf, nthreads=1):
df : pandas.DataFrame
"""
buffer_reader = pa.BufferReader(buf)
reader = pa.RecordBatchFileReader(buffer_reader)
reader = pa.RecordBatchStreamReader(buffer_reader)
table = reader.read_all()
return table.to_pandas(nthreads=nthreads)
python/pyarrow/lib.pyx
+7 -1
View File
@@ -106,9 +106,15 @@ include "array.pxi"
# Column, Table, Record Batch
include "table.pxi"
# File IO, IPC
# File IO
include "io.pxi"
# IPC / Messaging
include "ipc.pxi"
# Feather format
include "feather.pxi"
#----------------------------------------------------------------------
# Public API
python/pyarrow/pandas_compat.py
+1 -1
View File
@@ -22,7 +22,7 @@ import pandas as pd
import six
import pyarrow as pa
from pyarrow.compat import PY2
from pyarrow.compat import PY2 # noqa
INDEX_LEVEL_NAME_REGEX = re.compile(r'^__index_level_\d+__$')
python/pyarrow/table.pxi
+15 -1
View File
@@ -432,9 +432,23 @@ cdef class RecordBatch:
return self._schema
def __getitem__(self, i):
def column(self, i):
"""
Select single column from record batcha
Returns
-------
column : pyarrow.Array
"""
if not -self.num_columns <= i < self.num_columns:
raise IndexError(
'Record batch column index {:d} is out of range'.format(i)
)
return pyarrow_wrap_array(self.batch.column(i))
def __getitem__(self, i):
return self.column(i)
def slice(self, offset=0, length=None):
"""
Compute zero-copy slice of this RecordBatch
python/pyarrow/tests/conftest.py
+1 -1
View File
@@ -26,7 +26,7 @@ defaults = {
}
try:
import pyarrow.parquet
import pyarrow.parquet # noqa
defaults['parquet'] = True
except ImportError:
pass
python/pyarrow/tests/test_array.py
+1 -1
View File
@@ -225,7 +225,7 @@ def test_simple_type_construction():
(pa.timestamp('us', 'UTC'), 'datetimetz'),
pytest.mark.xfail((pa.time32('s'), None), raises=NotImplementedError),
pytest.mark.xfail((pa.time64('us'), None), raises=NotImplementedError),
]
]
)
def test_logical_type(type, expected):
assert get_logical_type(type) == expected
python/pyarrow/tests/test_convert_builtin.py
+3
View File
@@ -22,6 +22,7 @@ import pyarrow as pa
import datetime
import decimal
class StrangeIterable:
def __init__(self, lst):
self.lst = lst
@@ -29,6 +30,7 @@ class StrangeIterable:
def __iter__(self):
return self.lst.__iter__()
class TestConvertIterable(unittest.TestCase):
def test_iterable_types(self):
@@ -61,6 +63,7 @@ class TestLimitedConvertIterator(unittest.TestCase):
arr2 = pa.array((0, 1, 2))
assert arr1.equals(arr2)
class TestConvertSequence(unittest.TestCase):
def test_sequence_types(self):
python/pyarrow/tests/test_feather.py
+2 -1
View File
@@ -359,7 +359,8 @@ class TestFeatherReader(unittest.TestCase):
expected = df.rename(columns=str)
self._check_pandas_roundtrip(df, expected)
@pytest.mark.skipif(not os.path.supports_unicode_filenames, reason='unicode filenames not supported')
@pytest.mark.skipif(not os.path.supports_unicode_filenames,
reason='unicode filenames not supported')
def test_unicode_filename(self):
# GH #209
name = (b'Besa_Kavaj\xc3\xab.feather').decode('utf-8')
python/pyarrow/tests/test_ipc.py
+54 -5
View File
@@ -38,7 +38,7 @@ class MessagingTest(object):
return io.BytesIO()
def _get_source(self):
return pa.BufferReader(self.sink.getvalue())
return self.sink.getvalue()
def write_batches(self):
nrows = 5
@@ -79,7 +79,7 @@ class TestFile(MessagingTest, unittest.TestCase):
def test_simple_roundtrip(self):
batches = self.write_batches()
file_contents = self._get_source()
file_contents = pa.BufferReader(self._get_source())
reader = pa.open_file(file_contents)
@@ -93,7 +93,7 @@ class TestFile(MessagingTest, unittest.TestCase):
def test_read_all(self):
batches = self.write_batches()
file_contents = self._get_source()
file_contents = pa.BufferReader(self._get_source())
reader = pa.open_file(file_contents)
@@ -114,7 +114,7 @@ class TestStream(MessagingTest, unittest.TestCase):
def test_simple_roundtrip(self):
batches = self.write_batches()
file_contents = self._get_source()
file_contents = pa.BufferReader(self._get_source())
reader = pa.open_stream(file_contents)
assert reader.schema.equals(batches[0].schema)
@@ -131,7 +131,7 @@ class TestStream(MessagingTest, unittest.TestCase):
def test_read_all(self):
batches = self.write_batches()
file_contents = self._get_source()
file_contents = pa.BufferReader(self._get_source())
reader = pa.open_stream(file_contents)
result = reader.read_all()
@@ -139,6 +139,55 @@ class TestStream(MessagingTest, unittest.TestCase):
assert result.equals(expected)
class TestMessageReader(MessagingTest, unittest.TestCase):
def _get_example_messages(self):
batches = self.write_batches()
file_contents = self._get_source()
buf_reader = pa.BufferReader(file_contents)
reader = pa.MessageReader.open_stream(buf_reader)
return batches, list(reader)
def _get_writer(self, sink, schema):
return pa.RecordBatchStreamWriter(sink, schema)
def test_ctors_no_segfault(self):
with pytest.raises(TypeError):
repr(pa.Message())
with pytest.raises(TypeError):
repr(pa.MessageReader())
def test_message_reader(self):
_, messages = self._get_example_messages()
assert len(messages) == 6
assert messages[0].type == 'schema'
for msg in messages[1:]:
assert msg.type == 'record batch'
def test_serialize_read_message(self):
_, messages = self._get_example_messages()
msg = messages[0]
buf = msg.serialize()
restored = pa.read_message(buf)
restored2 = pa.read_message(pa.BufferReader(buf))
restored3 = pa.read_message(buf.to_pybytes())
assert msg.equals(restored)
assert msg.equals(restored2)
assert msg.equals(restored3)
def test_read_record_batch(self):
batches, messages = self._get_example_messages()
for batch, message in zip(batches, messages[1:]):
read_batch = pa.read_record_batch(message, batch.schema)
assert read_batch.equals(batch)
class TestSocket(MessagingTest, unittest.TestCase):
class StreamReaderServer(threading.Thread):
python/pyarrow/tests/test_parquet.py
+5 -2
View File
@@ -453,13 +453,15 @@ def test_date_time_types():
table = pa.Table.from_arrays([a1, a2, a3, a4, a5, a6],
['date32', 'date64', 'timestamp[us]',
'time32[s]', 'time64[us]', 'time32_from64[s]'])
'time32[s]', 'time64[us]',
'time32_from64[s]'])
# date64 as date32
# time32[s] to time32[ms]
expected = pa.Table.from_arrays([a1, a1, a3, a4, a5, ex_a6],
['date32', 'date64', 'timestamp[us]',
'time32[s]', 'time64[us]', 'time32_from64[s]'])
'time32[s]', 'time64[us]',
'time32_from64[s]'])
_check_roundtrip(table, expected=expected, version='2.0')
@@ -848,6 +850,7 @@ def test_read_multiple_files(tmpdir):
with pytest.raises(ValueError):
read_multiple_files(mixed_paths)
@parquet
def test_multiindex_duplicate_values(tmpdir):
num_rows = 3
python/pyarrow/tests/test_table.py
+4
View File
@@ -75,6 +75,10 @@ def test_recordbatch_basics():
('c1', [-10, -5, 0, 5, 10])
])
with pytest.raises(IndexError):
# bounds checking
batch[2]
def test_recordbatch_slice():
data = [
python/pyarrow/tests/test_tensor.py
+1
View File
@@ -115,6 +115,7 @@ def test_tensor_size():
tensor = pa.Tensor.from_numpy(data)
assert pa.get_tensor_size(tensor) > (data.size * 8)
def test_read_tensor(tmpdir):
# Create and write tensor tensor
data = np.random.randn(10, 4)