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

Refactoring sweep and cleanup of public IPC API. Move non-public APIs from metadata.h to metadata-internal.h and create message.h, dictionary.h

Browse Source 
Change-Id: I3fbacc6a0f5d02b734a04b9717d426cb1b89ace2
This commit is contained in:
Wes McKinney committed 2017-08-24 15:02:28 -04:00
1 parent b646f965ee
commit 122a759156
25 files changed
+953 -628
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cpp/CMakeLists.txt
+5 -3
View File
@@ -217,12 +217,12 @@ include(BuildUtils)
# Compiler flags
############################################################
include(SetupCxxFlags)
if (ARROW_NO_DEPRECATED_API)
add_definitions(-DARROW_NO_DEPRECATED_API)
endif()
include(SetupCxxFlags)
############################################################
# Dependencies
############################################################
@@ -780,10 +780,12 @@ endif()
if (ARROW_IPC)
set(ARROW_SRCS ${ARROW_SRCS}
src/arrow/ipc/dictionary.cc
src/arrow/ipc/feather.cc
src/arrow/ipc/json.cc
src/arrow/ipc/json-internal.cc
src/arrow/ipc/metadata.cc
src/arrow/ipc/message.cc
src/arrow/ipc/metadata-internal.cc
src/arrow/ipc/reader.cc
src/arrow/ipc/writer.cc
)
cpp/src/arrow/ipc/CMakeLists.txt
+2 -1
View File
@@ -83,9 +83,10 @@ add_custom_target(metadata_fbs DEPENDS ${FBS_OUTPUT_FILES})
# Headers: top level
install(FILES
api.h
dictionary.h
feather.h
json.h
metadata.h
message.h
reader.h
writer.h
DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/arrow/ipc")
cpp/src/arrow/ipc/api.h
+2 -1
View File
@@ -18,9 +18,10 @@
#ifndef ARROW_IPC_API_H
#define ARROW_IPC_API_H
#include "arrow/ipc/dictionary.h"
#include "arrow/ipc/feather.h"
#include "arrow/ipc/json.h"
#include "arrow/ipc/metadata.h"
#include "arrow/ipc/message.h"
#include "arrow/ipc/reader.h"
#include "arrow/ipc/writer.h"
cpp/src/arrow/ipc/dictionary.cc
+85
View File
@@ -0,0 +1,85 @@
// 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.
#include "arrow/ipc/dictionary.h"
#include <cstdint>
#include <memory>
#include <sstream>
#include "arrow/array.h"
#include "arrow/status.h"
#include "arrow/type.h"
namespace arrow {
namespace ipc {
DictionaryMemo::DictionaryMemo() {}
// Returns KeyError if dictionary not found
Status DictionaryMemo::GetDictionary(int64_t id,
std::shared_ptr<Array>* dictionary) const {
auto it = id_to_dictionary_.find(id);
if (it == id_to_dictionary_.end()) {
std::stringstream ss;
ss << "Dictionary with id " << id << " not found";
return Status::KeyError(ss.str());
}
*dictionary = it->second;
return Status::OK();
}
int64_t DictionaryMemo::GetId(const std::shared_ptr<Array>& dictionary) {
intptr_t address = reinterpret_cast<intptr_t>(dictionary.get());
auto it = dictionary_to_id_.find(address);
if (it != dictionary_to_id_.end()) {
// Dictionary already observed, return the id
return it->second;
} else {
int64_t new_id = static_cast<int64_t>(dictionary_to_id_.size());
dictionary_to_id_[address] = new_id;
id_to_dictionary_[new_id] = dictionary;
return new_id;
}
}
bool DictionaryMemo::HasDictionary(const std::shared_ptr<Array>& dictionary) const {
intptr_t address = reinterpret_cast<intptr_t>(dictionary.get());
auto it = dictionary_to_id_.find(address);
return it != dictionary_to_id_.end();
}
bool DictionaryMemo::HasDictionaryId(int64_t id) const {
auto it = id_to_dictionary_.find(id);
return it != id_to_dictionary_.end();
}
Status DictionaryMemo::AddDictionary(int64_t id,
const std::shared_ptr<Array>& dictionary) {
if (HasDictionaryId(id)) {
std::stringstream ss;
ss << "Dictionary with id " << id << " already exists";
return Status::KeyError(ss.str());
}
intptr_t address = reinterpret_cast<intptr_t>(dictionary.get());
id_to_dictionary_[id] = dictionary;
dictionary_to_id_[address] = id;
return Status::OK();
}
} // namespace ipc
} // namespace arrow
cpp/src/arrow/ipc/dictionary.h
+88
View File
@@ -0,0 +1,88 @@
// 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.
// Tools for dictionaries in IPC context
#ifndef ARROW_IPC_DICTIONARY_H
#define ARROW_IPC_DICTIONARY_H
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "arrow/status.h"
#include "arrow/util/macros.h"
#include "arrow/util/visibility.h"
namespace arrow {
class Array;
class Buffer;
class Field;
namespace io {
class InputStream;
class OutputStream;
class RandomAccessFile;
} // namespace io
namespace ipc {
using DictionaryMap = std::unordered_map<int64_t, std::shared_ptr<Array>>;
using DictionaryTypeMap = std::unordered_map<int64_t, std::shared_ptr<Field>>;
// Memoization data structure for handling shared dictionaries
class ARROW_EXPORT DictionaryMemo {
public:
DictionaryMemo();
// Returns KeyError if dictionary not found
Status GetDictionary(int64_t id, std::shared_ptr<Array>* dictionary) const;
/// Return id for dictionary, computing new id if necessary
int64_t GetId(const std::shared_ptr<Array>& dictionary);
bool HasDictionary(const std::shared_ptr<Array>& dictionary) const;
bool HasDictionaryId(int64_t id) const;
// Add a dictionary to the memo with a particular id. Returns KeyError if
// that dictionary already exists
Status AddDictionary(int64_t id, const std::shared_ptr<Array>& dictionary);
const DictionaryMap& id_to_dictionary() const { return id_to_dictionary_; }
int size() const { return static_cast<int>(id_to_dictionary_.size()); }
private:
// Dictionary memory addresses, to track whether a dictionary has been seen
// before
std::unordered_map<intptr_t, int64_t> dictionary_to_id_;
// Map of dictionary id to dictionary array
DictionaryMap id_to_dictionary_;
DISALLOW_COPY_AND_ASSIGN(DictionaryMemo);
};
} // namespace ipc
} // namespace arrow
#endif // ARROW_IPC_DICTIONARY_H
cpp/src/arrow/ipc/file-to-stream.cc
+1 -1
View File
@@ -35,7 +35,7 @@ Status ConvertToStream(const char* path) {
RETURN_NOT_OK(ipc::RecordBatchFileReader::Open(in_file, &reader));
io::StdoutStream sink;
std::shared_ptr<RecordBatchStreamWriter> writer;
std::shared_ptr<RecordBatchWriter> writer;
RETURN_NOT_OK(RecordBatchStreamWriter::Open(&sink, reader->schema(), &writer));
for (int i = 0; i < reader->num_record_batches(); ++i) {
std::shared_ptr<RecordBatch> chunk;
cpp/src/arrow/ipc/ipc-read-write-benchmark.cc
+1 -1
View File
@@ -109,7 +109,7 @@ static void BM_ReadRecordBatch(benchmark::State& state) { // NOLINT non-const r
std::shared_ptr<RecordBatch> result;
io::BufferReader reader(buffer);
if (!ipc::ReadRecordBatch(record_batch->schema(), 0, &reader, &result).ok()) {
if (!ipc::ReadRecordBatch(record_batch->schema(), &reader, &result).ok()) {
state.SkipWithError("Failed to read!");
}
}
cpp/src/arrow/ipc/ipc-read-write-test.cc
+28 -26
View File
@@ -29,6 +29,7 @@
#include "arrow/io/memory.h"
#include "arrow/io/test-common.h"
#include "arrow/ipc/api.h"
#include "arrow/ipc/metadata-internal.h"
#include "arrow/ipc/test-common.h"
#include "arrow/ipc/util.h"
#include "arrow/memory_pool.h"
@@ -47,21 +48,14 @@ class TestSchemaMetadata : public ::testing::Test {
public:
void SetUp() {}
void CheckRoundtrip(const Schema& schema, DictionaryMemo* memo) {
void CheckRoundtrip(const Schema& schema) {
std::shared_ptr<Buffer> buffer;
ASSERT_OK(WriteSchemaMessage(schema, memo, &buffer));
ASSERT_OK(SerializeSchema(schema, default_memory_pool(), &buffer));
std::unique_ptr<Message> message;
ASSERT_OK(Message::Open(buffer, nullptr, &message));
ASSERT_EQ(Message::SCHEMA, message->type());
DictionaryMemo empty_memo;
std::shared_ptr<Schema> schema2;
ASSERT_OK(GetSchema(message->header(), empty_memo, &schema2));
AssertSchemaEqual(schema, *schema2);
std::shared_ptr<Schema> result;
io::BufferReader reader(buffer);
ASSERT_OK(ReadSchema(&reader, &result));
AssertSchemaEqual(schema, *result);
}
};
@@ -107,9 +101,7 @@ TEST_F(TestSchemaMetadata, PrimitiveFields) {
auto f10 = field("f10", std::make_shared<BooleanType>());
Schema schema({f0, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10});
DictionaryMemo memo;
CheckRoundtrip(schema, &memo);
CheckRoundtrip(schema);
}
TEST_F(TestSchemaMetadata, NestedFields) {
@@ -121,9 +113,7 @@ TEST_F(TestSchemaMetadata, NestedFields) {
auto f1 = field("f1", type2);
Schema schema({f0, f1});
DictionaryMemo memo;
CheckRoundtrip(schema, &memo);
CheckRoundtrip(schema);
}
#define BATCH_CASES() \
@@ -137,13 +127,21 @@ static int g_file_number = 0;
class IpcTestFixture : public io::MemoryMapFixture {
public:
Status DoSchemaRoundTrip(const Schema& schema, std::shared_ptr<Schema>* result) {
std::shared_ptr<Buffer> serialized_schema;
RETURN_NOT_OK(SerializeSchema(schema, pool_, &serialized_schema));
io::BufferReader buf_reader(serialized_schema);
return ReadSchema(&buf_reader, result);
}
Status DoStandardRoundTrip(const RecordBatch& batch,
std::shared_ptr<RecordBatch>* batch_result) {
std::shared_ptr<Buffer> serialized_batch;
RETURN_NOT_OK(SerializeRecordBatch(batch, pool_, &serialized_batch));
io::BufferReader buf_reader(serialized_batch);
return ReadRecordBatch(batch.schema(), 0, &buf_reader, batch_result);
return ReadRecordBatch(batch.schema(), &buf_reader, batch_result);
}
Status DoLargeRoundTrip(const RecordBatch& batch, bool zero_data,
@@ -153,7 +151,7 @@ class IpcTestFixture : public io::MemoryMapFixture {
}
RETURN_NOT_OK(mmap_->Seek(0));
std::shared_ptr<RecordBatchFileWriter> file_writer;
std::shared_ptr<RecordBatchWriter> file_writer;
RETURN_NOT_OK(RecordBatchFileWriter::Open(mmap_.get(), batch.schema(), &file_writer));
RETURN_NOT_OK(file_writer->WriteRecordBatch(batch, true));
RETURN_NOT_OK(file_writer->Close());
@@ -182,6 +180,10 @@ class IpcTestFixture : public io::MemoryMapFixture {
ss << "test-write-row-batch-" << g_file_number++;
ASSERT_OK(io::MemoryMapFixture::InitMemoryMap(buffer_size, ss.str(), &mmap_));
std::shared_ptr<Schema> schema_result;
ASSERT_OK(DoSchemaRoundTrip(*batch.schema(), &schema_result));
ASSERT_TRUE(batch.schema()->Equals(*schema_result));
std::shared_ptr<RecordBatch> result;
ASSERT_OK(DoStandardRoundTrip(batch, &result));
CheckReadResult(*result, batch);
@@ -498,7 +500,7 @@ class TestFileFormat : public ::testing::TestWithParam<MakeRecordBatch*> {
Status RoundTripHelper(const BatchVector& in_batches, BatchVector* out_batches) {
// Write the file
std::shared_ptr<RecordBatchFileWriter> writer;
std::shared_ptr<RecordBatchWriter> writer;
RETURN_NOT_OK(
RecordBatchFileWriter::Open(sink_.get(), in_batches[0]->schema(), &writer));
@@ -565,7 +567,7 @@ class TestStreamFormat : public ::testing::TestWithParam<MakeRecordBatch*> {
Status RoundTripHelper(const RecordBatch& batch,
std::vector<std::shared_ptr<RecordBatch>>* out_batches) {
// Write the file
std::shared_ptr<RecordBatchStreamWriter> writer;
std::shared_ptr<RecordBatchWriter> writer;
RETURN_NOT_OK(RecordBatchStreamWriter::Open(sink_.get(), batch.schema(), &writer));
int num_batches = 5;
for (int i = 0; i < num_batches; ++i) {
@@ -575,10 +577,10 @@ class TestStreamFormat : public ::testing::TestWithParam<MakeRecordBatch*> {
RETURN_NOT_OK(sink_->Close());
// Open the file
auto buf_reader = std::make_shared<io::BufferReader>(buffer_);
io::BufferReader buf_reader(buffer_);
std::shared_ptr<RecordBatchStreamReader> reader;
RETURN_NOT_OK(RecordBatchStreamReader::Open(buf_reader, &reader));
std::shared_ptr<RecordBatchReader> reader;
RETURN_NOT_OK(RecordBatchStreamReader::Open(&buf_reader, &reader));
std::shared_ptr<RecordBatch> chunk;
while (true) {
cpp/src/arrow/ipc/json-integration-test.cc
+1 -1
View File
@@ -77,7 +77,7 @@ static Status ConvertJsonToArrow(const std::string& json_path,
std::cout << "Found schema: " << reader->schema()->ToString() << std::endl;
}
std::shared_ptr<ipc::RecordBatchFileWriter> writer;
std::shared_ptr<ipc::RecordBatchWriter> writer;
RETURN_NOT_OK(
ipc::RecordBatchFileWriter::Open(out_file.get(), reader->schema(), &writer));
cpp/src/arrow/ipc/json-internal.cc
+1 -1
View File
@@ -32,7 +32,7 @@
#include "arrow/array.h"
#include "arrow/builder.h"
#include "arrow/ipc/metadata.h"
#include "arrow/ipc/metadata-internal.h"
#include "arrow/memory_pool.h"
#include "arrow/status.h"
#include "arrow/table.h"
cpp/src/arrow/ipc/message.cc
+257
View File
@@ -0,0 +1,257 @@
// 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.
#include "arrow/ipc/message.h"
#include <algorithm>
#include <cstdint>
#include <memory>
#include <sstream>
#include <string>
#include <vector>
#include "arrow/array.h"
#include "arrow/buffer.h"
#include "arrow/io/interfaces.h"
#include "arrow/ipc/metadata-internal.h"
#include "arrow/ipc/util.h"
#include "arrow/status.h"
#include "arrow/tensor.h"
#include "arrow/type.h"
#include "arrow/util/logging.h"
namespace arrow {
namespace ipc {
class Message::MessageImpl {
public:
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(metadata_->data());
// Check that the metadata version is supported
if (message_->version() < kMinMetadataVersion) {
return Status::Invalid("Old metadata version not supported");
}
return Status::OK();
}
Message::Type type() const {
switch (message_->header_type()) {
case flatbuf::MessageHeader_Schema:
return Message::SCHEMA;
case flatbuf::MessageHeader_DictionaryBatch:
return Message::DICTIONARY_BATCH;
case flatbuf::MessageHeader_RecordBatch:
return Message::RECORD_BATCH;
case flatbuf::MessageHeader_Tensor:
return Message::TENSOR;
default:
return Message::NONE;
}
}
MetadataVersion version() const {
switch (message_->version()) {
case flatbuf::MetadataVersion_V1:
// Arrow 0.1
return MetadataVersion::V1;
case flatbuf::MetadataVersion_V2:
// Arrow 0.2
return MetadataVersion::V2;
case flatbuf::MetadataVersion_V3:
// Arrow >= 0.3
return MetadataVersion::V3;
// Add cases as other versions become available
default:
return MetadataVersion::V3;
}
}
const void* header() const { return message_->header(); }
std::shared_ptr<Buffer> body() const { return body_; }
std::shared_ptr<Buffer> metadata() const { return metadata_; }
private:
// 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>& 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() {}
std::shared_ptr<Buffer> Message::body() const { return impl_->body(); }
std::shared_ptr<Buffer> Message::metadata() const { return impl_->metadata(); }
Message::Type Message::type() const { return impl_->type(); }
MetadataVersion Message::metadata_version() const { return impl_->version(); }
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";
}
// ----------------------------------------------------------------------
// 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::unique_ptr<Message>* message) {
std::shared_ptr<Buffer> buffer;
RETURN_NOT_OK(file->ReadAt(offset, metadata_length, &buffer));
int32_t flatbuffer_size = *reinterpret_cast<const int32_t*>(buffer->data());
if (flatbuffer_size + static_cast<int>(sizeof(int32_t)) > metadata_length) {
std::stringstream ss;
ss << "flatbuffer size " << metadata_length << " invalid. File offset: " << offset
<< ", metadata length: " << metadata_length;
return Status::Invalid(ss.str());
}
auto metadata = SliceBuffer(buffer, 4, buffer->size() - 4);
return ReadFullMessage(metadata, file, 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));
if (buffer->size() != sizeof(int32_t)) {
*message = nullptr;
return Status::OK();
}
int32_t message_length = *reinterpret_cast<const int32_t*>(buffer->data());
if (message_length == 0) {
// Optional 0 EOS control message
*message = nullptr;
return Status::OK();
}
RETURN_NOT_OK(file->Read(message_length, &buffer));
if (buffer->size() != message_length) {
return Status::IOError("Unexpected end of stream trying to read message");
}
return ReadFullMessage(buffer, file, message);
}
// ----------------------------------------------------------------------
// Implement InputStream message reader
Status InputStreamMessageReader::ReadNextMessage(std::unique_ptr<Message>* message) {
return ReadMessage(stream_, message);
}
InputStreamMessageReader::~InputStreamMessageReader() {}
} // namespace ipc
} // namespace arrow
cpp/src/arrow/ipc/metadata.h → cpp/src/arrow/ipc/message.h
+19 -129
View File
@@ -17,8 +17,8 @@
// C++ object model and user API for interprocess schema messaging
#ifndef ARROW_IPC_METADATA_H
#define ARROW_IPC_METADATA_H
#ifndef ARROW_IPC_MESSAGE_H
#define ARROW_IPC_MESSAGE_H
#include <cstdint>
#include <memory>
@@ -26,6 +26,7 @@
#include <unordered_map>
#include <vector>
#include "arrow/status.h"
#include "arrow/util/macros.h"
#include "arrow/util/visibility.h"
@@ -33,11 +34,7 @@ namespace arrow {
class Array;
class Buffer;
class DataType;
class Field;
class Schema;
class Status;
class Tensor;
namespace io {
@@ -51,97 +48,17 @@ namespace ipc {
enum class MetadataVersion : char { V1, V2, V3 };
static constexpr const char* kArrowMagicBytes = "ARROW1";
// ARROW-109: We set this number arbitrarily to help catch user mistakes. For
// deeply nested schemas, it is expected the user will indicate explicitly the
// maximum allowed recursion depth
constexpr int kMaxNestingDepth = 64;
struct ARROW_EXPORT FieldMetadata {
int64_t length;
int64_t null_count;
int64_t offset;
};
struct ARROW_EXPORT BufferMetadata {
/// The shared memory page id where to find this. Set to -1 if unused
int32_t page;
/// The relative offset into the memory page to the starting byte of the buffer
int64_t offset;
/// Absolute length in bytes of the buffer
int64_t length;
};
struct FileBlock {
int64_t offset;
int32_t metadata_length;
int64_t body_length;
};
//----------------------------------------------------------------------
using DictionaryMap = std::unordered_map<int64_t, std::shared_ptr<Array>>;
using DictionaryTypeMap = std::unordered_map<int64_t, std::shared_ptr<Field>>;
// Memoization data structure for handling shared dictionaries
class ARROW_EXPORT DictionaryMemo {
public:
DictionaryMemo();
// Returns KeyError if dictionary not found
Status GetDictionary(int64_t id, std::shared_ptr<Array>* dictionary) const;
/// Return id for dictionary, computing new id if necessary
int64_t GetId(const std::shared_ptr<Array>& dictionary);
bool HasDictionary(const std::shared_ptr<Array>& dictionary) const;
bool HasDictionaryId(int64_t id) const;
// Add a dictionary to the memo with a particular id. Returns KeyError if
// that dictionary already exists
Status AddDictionary(int64_t id, const std::shared_ptr<Array>& dictionary);
const DictionaryMap& id_to_dictionary() const { return id_to_dictionary_; }
int size() const { return static_cast<int>(id_to_dictionary_.size()); }
private:
// Dictionary memory addresses, to track whether a dictionary has been seen
// before
std::unordered_map<intptr_t, int64_t> dictionary_to_id_;
// Map of dictionary id to dictionary array
DictionaryMap id_to_dictionary_;
DISALLOW_COPY_AND_ASSIGN(DictionaryMemo);
};
// Read interface classes. We do not fully deserialize the flatbuffers so that
// individual fields metadata can be retrieved from very large schema without
//
class Message;
// Retrieve a list of all the dictionary ids and types required by the schema for
// reconstruction. The presumption is that these will be loaded either from
// the stream or file (or they may already be somewhere else in memory)
Status GetDictionaryTypes(const void* opaque_schema, DictionaryTypeMap* id_to_field);
// Construct a complete Schema from the message. May be expensive for very
// large schemas if you are only interested in a few fields
Status ARROW_EXPORT GetSchema(const void* opaque_schema,
const DictionaryMemo& dictionary_memo,
std::shared_ptr<Schema>* out);
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:
@@ -219,15 +136,26 @@ class ARROW_EXPORT MessageReader {
/// \since 0.5.0
class ARROW_EXPORT InputStreamMessageReader : public MessageReader {
public:
explicit InputStreamMessageReader(const std::shared_ptr<io::InputStream>& stream)
: stream_(stream) {}
explicit InputStreamMessageReader(io::InputStream* stream) : stream_(stream) {}
#ifndef ARROW_NO_DEPRECATED_API
/// \deprecated Since 0.7.0
explicit InputStreamMessageReader(const std::shared_ptr<io::InputStream>& owned_stream)
: InputStreamMessageReader(owned_stream.get()) {
owned_stream_ = owned_stream;
}
#endif
~InputStreamMessageReader();
Status ReadNextMessage(std::unique_ptr<Message>* message) override;
private:
std::shared_ptr<io::InputStream> stream_;
io::InputStream* stream_;
#ifndef ARROW_NO_DEPRECATED_API
std::shared_ptr<io::InputStream> owned_stream_;
#endif
};
/// \brief Read encapulated RPC message from position in file
@@ -253,48 +181,10 @@ Status ReadMessage(const int64_t offset, const int32_t metadata_length,
/// 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::unique_ptr<Message>* message);
/// 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(const Buffer& message, io::OutputStream* file,
int32_t* message_length);
// Serialize arrow::Schema as a Flatbuffer
//
// \param[in] schema a Schema instance
// \param[inout] dictionary_memo class for tracking dictionaries and assigning
// dictionary ids
// \param[out] out the serialized arrow::Buffer
// \return Status outcome
Status ARROW_EXPORT WriteSchemaMessage(const Schema& schema,
DictionaryMemo* dictionary_memo,
std::shared_ptr<Buffer>* out);
ARROW_EXPORT
Status WriteRecordBatchMessage(const int64_t length, const int64_t body_length,
const std::vector<FieldMetadata>& nodes,
const std::vector<BufferMetadata>& buffers,
std::shared_ptr<Buffer>* out);
ARROW_EXPORT
Status WriteTensorMessage(const Tensor& tensor, const int64_t buffer_start_offset,
std::shared_ptr<Buffer>* out);
Status WriteDictionaryMessage(const int64_t id, const int64_t length,
const int64_t body_length,
const std::vector<FieldMetadata>& nodes,
const std::vector<BufferMetadata>& buffers,
std::shared_ptr<Buffer>* out);
Status WriteFileFooter(const Schema& schema, const std::vector<FileBlock>& dictionaries,
const std::vector<FileBlock>& record_batches,
DictionaryMemo* dictionary_memo, io::OutputStream* out);
Status ReadMessage(io::InputStream* stream, std::unique_ptr<Message>* message);
} // namespace ipc
} // namespace arrow
#endif // ARROW_IPC_METADATA_H
#endif // ARROW_IPC_MESSAGE_H
cpp/src/arrow/ipc/metadata.cc → cpp/src/arrow/ipc/metadata-internal.cc
+4 -289
View File
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
#include "arrow/ipc/metadata.h"
#include "arrow/ipc/metadata-internal.h"
#include <algorithm>
#include <cstdint>
@@ -29,9 +29,6 @@
#include "arrow/array.h"
#include "arrow/buffer.h"
#include "arrow/io/interfaces.h"
#include "arrow/ipc/File_generated.h"
#include "arrow/ipc/Message_generated.h"
#include "arrow/ipc/Tensor_generated.h"
#include "arrow/ipc/util.h"
#include "arrow/status.h"
#include "arrow/tensor.h"
@@ -53,12 +50,6 @@ using VectorLayoutOffset = flatbuffers::Offset<arrow::flatbuf::VectorLayout>;
using Offset = flatbuffers::Offset<void>;
using FBString = flatbuffers::Offset<flatbuffers::String>;
static constexpr flatbuf::MetadataVersion kCurrentMetadataVersion =
flatbuf::MetadataVersion_V3;
static constexpr flatbuf::MetadataVersion kMinMetadataVersion =
flatbuf::MetadataVersion_V3;
static Status IntFromFlatbuffer(const flatbuf::Int* int_data,
std::shared_ptr<DataType>* out) {
if (int_data->bitWidth() > 64) {
@@ -703,7 +694,7 @@ static Status MakeRecordBatch(FBB& fbb, int64_t length, int64_t body_length,
return Status::OK();
}
Status WriteRecordBatchMessage(const int64_t length, const int64_t body_length,
Status WriteRecordBatchMessage(int64_t length, int64_t body_length,
const std::vector<FieldMetadata>& nodes,
const std::vector<BufferMetadata>& buffers,
std::shared_ptr<Buffer>* out) {
@@ -714,7 +705,7 @@ Status WriteRecordBatchMessage(const int64_t length, const int64_t body_length,
body_length, out);
}
Status WriteTensorMessage(const Tensor& tensor, const int64_t buffer_start_offset,
Status WriteTensorMessage(const Tensor& tensor, int64_t buffer_start_offset,
std::shared_ptr<Buffer>* out) {
using TensorDimOffset = flatbuffers::Offset<flatbuf::TensorDim>;
using TensorOffset = flatbuffers::Offset<flatbuf::Tensor>;
@@ -743,8 +734,7 @@ Status WriteTensorMessage(const Tensor& tensor, const int64_t buffer_start_offse
body_length, out);
}
Status WriteDictionaryMessage(const int64_t id, const int64_t length,
const int64_t body_length,
Status WriteDictionaryMessage(int64_t id, int64_t length, int64_t body_length,
const std::vector<FieldMetadata>& nodes,
const std::vector<BufferMetadata>& buffers,
std::shared_ptr<Buffer>* out) {
@@ -802,206 +792,6 @@ Status WriteFileFooter(const Schema& schema, const std::vector<FileBlock>& dicti
return out->Write(fbb.GetBufferPointer(), size);
}
// ----------------------------------------------------------------------
// Memoization data structure for handling shared dictionaries
DictionaryMemo::DictionaryMemo() {}
// Returns KeyError if dictionary not found
Status DictionaryMemo::GetDictionary(int64_t id,
std::shared_ptr<Array>* dictionary) const {
auto it = id_to_dictionary_.find(id);
if (it == id_to_dictionary_.end()) {
std::stringstream ss;
ss << "Dictionary with id " << id << " not found";
return Status::KeyError(ss.str());
}
*dictionary = it->second;
return Status::OK();
}
int64_t DictionaryMemo::GetId(const std::shared_ptr<Array>& dictionary) {
intptr_t address = reinterpret_cast<intptr_t>(dictionary.get());
auto it = dictionary_to_id_.find(address);
if (it != dictionary_to_id_.end()) {
// Dictionary already observed, return the id
return it->second;
} else {
int64_t new_id = static_cast<int64_t>(dictionary_to_id_.size());
dictionary_to_id_[address] = new_id;
id_to_dictionary_[new_id] = dictionary;
return new_id;
}
}
bool DictionaryMemo::HasDictionary(const std::shared_ptr<Array>& dictionary) const {
intptr_t address = reinterpret_cast<intptr_t>(dictionary.get());
auto it = dictionary_to_id_.find(address);
return it != dictionary_to_id_.end();
}
bool DictionaryMemo::HasDictionaryId(int64_t id) const {
auto it = id_to_dictionary_.find(id);
return it != id_to_dictionary_.end();
}
Status DictionaryMemo::AddDictionary(int64_t id,
const std::shared_ptr<Array>& dictionary) {
if (HasDictionaryId(id)) {
std::stringstream ss;
ss << "Dictionary with id " << id << " already exists";
return Status::KeyError(ss.str());
}
intptr_t address = reinterpret_cast<intptr_t>(dictionary.get());
id_to_dictionary_[id] = dictionary;
dictionary_to_id_[address] = id;
return Status::OK();
}
//----------------------------------------------------------------------
// Message reader
class Message::MessageImpl {
public:
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(metadata_->data());
// Check that the metadata version is supported
if (message_->version() < kMinMetadataVersion) {
return Status::Invalid("Old metadata version not supported");
}
return Status::OK();
}
Message::Type type() const {
switch (message_->header_type()) {
case flatbuf::MessageHeader_Schema:
return Message::SCHEMA;
case flatbuf::MessageHeader_DictionaryBatch:
return Message::DICTIONARY_BATCH;
case flatbuf::MessageHeader_RecordBatch:
return Message::RECORD_BATCH;
case flatbuf::MessageHeader_Tensor:
return Message::TENSOR;
default:
return Message::NONE;
}
}
MetadataVersion version() const {
switch (message_->version()) {
case flatbuf::MetadataVersion_V1:
// Arrow 0.1
return MetadataVersion::V1;
case flatbuf::MetadataVersion_V2:
// Arrow 0.2
return MetadataVersion::V2;
case flatbuf::MetadataVersion_V3:
// Arrow >= 0.3
return MetadataVersion::V3;
// Add cases as other versions become available
default:
return MetadataVersion::V3;
}
}
const void* header() const { return message_->header(); }
std::shared_ptr<Buffer> body() const { return body_; }
std::shared_ptr<Buffer> metadata() const { return metadata_; }
private:
// 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>& 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() {}
std::shared_ptr<Buffer> Message::body() const { return impl_->body(); }
std::shared_ptr<Buffer> Message::metadata() const { return impl_->metadata(); }
Message::Type Message::type() const { return impl_->type(); }
MetadataVersion Message::metadata_version() const { return impl_->version(); }
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) {
@@ -1084,81 +874,6 @@ Status GetTensorMetadata(const Buffer& metadata, std::shared_ptr<DataType>* type
return TypeFromFlatbuffer(tensor->type_type(), tensor->type(), {}, type);
}
// ----------------------------------------------------------------------
// 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(const int64_t offset, const int32_t metadata_length,
io::RandomAccessFile* file, std::unique_ptr<Message>* message) {
std::shared_ptr<Buffer> buffer;
RETURN_NOT_OK(file->ReadAt(offset, metadata_length, &buffer));
int32_t flatbuffer_size = *reinterpret_cast<const int32_t*>(buffer->data());
if (flatbuffer_size + static_cast<int>(sizeof(int32_t)) > metadata_length) {
std::stringstream ss;
ss << "flatbuffer size " << metadata_length << " invalid. File offset: " << offset
<< ", metadata length: " << metadata_length;
return Status::Invalid(ss.str());
}
auto metadata = SliceBuffer(buffer, 4, buffer->size() - 4);
return ReadFullMessage(metadata, file, 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));
if (buffer->size() != sizeof(int32_t)) {
*message = nullptr;
return Status::OK();
}
int32_t message_length = *reinterpret_cast<const int32_t*>(buffer->data());
if (message_length == 0) {
// Optional 0 EOS control message
*message = nullptr;
return Status::OK();
}
RETURN_NOT_OK(file->Read(message_length, &buffer));
if (buffer->size() != message_length) {
return Status::IOError("Unexpected end of stream trying to read 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
cpp/src/arrow/ipc/metadata-internal.h
+146
View File
@@ -0,0 +1,146 @@
// 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.
// C++ object model and user API for interprocess schema messaging
#ifndef ARROW_IPC_METADATA_INTERNAL_H
#define ARROW_IPC_METADATA_INTERNAL_H
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "arrow/ipc/File_generated.h"
#include "arrow/ipc/Message_generated.h"
#include "arrow/ipc/Tensor_generated.h"
#include "arrow/ipc/dictionary.h"
#include "arrow/ipc/message.h"
#include "arrow/util/macros.h"
#include "arrow/util/visibility.h"
namespace arrow {
class Array;
class Buffer;
class DataType;
class Field;
class Schema;
class Status;
class Tensor;
namespace flatbuf = org::apache::arrow::flatbuf;
namespace io {
class InputStream;
class OutputStream;
class RandomAccessFile;
} // namespace io
namespace ipc {
static constexpr flatbuf::MetadataVersion kCurrentMetadataVersion =
flatbuf::MetadataVersion_V3;
static constexpr flatbuf::MetadataVersion kMinMetadataVersion =
flatbuf::MetadataVersion_V3;
static constexpr const char* kArrowMagicBytes = "ARROW1";
struct FieldMetadata {
int64_t length;
int64_t null_count;
int64_t offset;
};
struct BufferMetadata {
/// The shared memory page id where to find this. Set to -1 if unused
int32_t page;
/// The relative offset into the memory page to the starting byte of the buffer
int64_t offset;
/// Absolute length in bytes of the buffer
int64_t length;
};
struct FileBlock {
int64_t offset;
int32_t metadata_length;
int64_t body_length;
};
// Read interface classes. We do not fully deserialize the flatbuffers so that
// individual fields metadata can be retrieved from very large schema without
//
// Retrieve a list of all the dictionary ids and types required by the schema for
// reconstruction. The presumption is that these will be loaded either from
// the stream or file (or they may already be somewhere else in memory)
Status GetDictionaryTypes(const void* opaque_schema, DictionaryTypeMap* id_to_field);
// Construct a complete Schema from the message. May be expensive for very
// large schemas if you are only interested in a few fields
Status GetSchema(const void* opaque_schema, const DictionaryMemo& dictionary_memo,
std::shared_ptr<Schema>* out);
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);
/// Write a serialized message metadata with a length-prefix and padding to an
/// 8-byte offset
///
/// <message_size: int32><message: const void*><padding>
Status WriteMessage(const Buffer& message, io::OutputStream* file,
int32_t* message_length);
// Serialize arrow::Schema as a Flatbuffer
//
// \param[in] schema a Schema instance
// \param[inout] dictionary_memo class for tracking dictionaries and assigning
// dictionary ids
// \param[out] out the serialized arrow::Buffer
// \return Status outcome
Status WriteSchemaMessage(const Schema& schema, DictionaryMemo* dictionary_memo,
std::shared_ptr<Buffer>* out);
Status WriteRecordBatchMessage(const int64_t length, const int64_t body_length,
const std::vector<FieldMetadata>& nodes,
const std::vector<BufferMetadata>& buffers,
std::shared_ptr<Buffer>* out);
Status WriteTensorMessage(const Tensor& tensor, const int64_t buffer_start_offset,
std::shared_ptr<Buffer>* out);
Status WriteFileFooter(const Schema& schema, const std::vector<FileBlock>& dictionaries,
const std::vector<FileBlock>& record_batches,
DictionaryMemo* dictionary_memo, io::OutputStream* out);
Status WriteDictionaryMessage(const int64_t id, const int64_t length,
const int64_t body_length,
const std::vector<FieldMetadata>& nodes,
const std::vector<BufferMetadata>& buffers,
std::shared_ptr<Buffer>* out);
} // namespace ipc
} // namespace arrow
#endif // ARROW_IPC_METADATA_H
cpp/src/arrow/ipc/reader.cc
+40 -9
View File
@@ -28,7 +28,8 @@
#include "arrow/io/memory.h"
#include "arrow/ipc/File_generated.h"
#include "arrow/ipc/Message_generated.h"
#include "arrow/ipc/metadata.h"
#include "arrow/ipc/message.h"
#include "arrow/ipc/metadata-internal.h"
#include "arrow/ipc/util.h"
#include "arrow/status.h"
#include "arrow/table.h"
@@ -459,6 +460,22 @@ RecordBatchStreamReader::RecordBatchStreamReader() {
RecordBatchStreamReader::~RecordBatchStreamReader() {}
Status RecordBatchStreamReader::Open(std::unique_ptr<MessageReader> message_reader,
std::shared_ptr<RecordBatchReader>* reader) {
// Private ctor
auto result = std::shared_ptr<RecordBatchStreamReader>(new RecordBatchStreamReader());
RETURN_NOT_OK(result->impl_->Open(std::move(message_reader)));
*reader = result;
return Status::OK();
}
Status RecordBatchStreamReader::Open(io::InputStream* stream,
std::shared_ptr<RecordBatchReader>* out) {
std::unique_ptr<MessageReader> message_reader(new InputStreamMessageReader(stream));
return Open(std::move(message_reader), out);
}
#ifndef ARROW_NO_DEPRECATED_API
Status RecordBatchStreamReader::Open(std::unique_ptr<MessageReader> message_reader,
std::shared_ptr<RecordBatchStreamReader>* reader) {
// Private ctor
@@ -471,6 +488,7 @@ Status RecordBatchStreamReader::Open(const std::shared_ptr<io::InputStream>& str
std::unique_ptr<MessageReader> message_reader(new InputStreamMessageReader(stream));
return Open(std::move(message_reader), out);
}
#endif
std::shared_ptr<Schema> RecordBatchStreamReader::schema() const {
return impl_->schema();
@@ -654,33 +672,46 @@ Status RecordBatchFileReader::ReadRecordBatch(int i,
return impl_->ReadRecordBatch(i, batch);
}
static Status ReadContiguousPayload(int64_t offset, io::RandomAccessFile* file,
static Status ReadContiguousPayload(io::InputStream* file,
std::unique_ptr<Message>* message) {
std::shared_ptr<Buffer> buffer;
RETURN_NOT_OK(file->Seek(offset));
RETURN_NOT_OK(ReadMessage(file, message));
if (*message == nullptr) {
return Status::Invalid("Unable to read metadata at offset");
}
return Status::OK();
}
Status ReadRecordBatch(const std::shared_ptr<Schema>& schema, int64_t offset,
io::RandomAccessFile* file, std::shared_ptr<RecordBatch>* out) {
Status ReadSchema(io::InputStream* stream, std::shared_ptr<Schema>* out) {
std::shared_ptr<RecordBatchReader> reader;
RETURN_NOT_OK(RecordBatchStreamReader::Open(stream, &reader));
*out = reader->schema();
return Status::OK();
}
Status ReadRecordBatch(const std::shared_ptr<Schema>& schema, io::InputStream* file,
std::shared_ptr<RecordBatch>* out) {
std::unique_ptr<Message> message;
RETURN_NOT_OK(ReadContiguousPayload(offset, file, &message));
RETURN_NOT_OK(ReadContiguousPayload(file, &message));
io::BufferReader buffer_reader(message->body());
return ReadRecordBatch(*message->metadata(), schema, kMaxNestingDepth, &buffer_reader,
out);
}
// Deprecated
Status ReadRecordBatch(const std::shared_ptr<Schema>& schema, int64_t offset,
io::RandomAccessFile* file, std::shared_ptr<RecordBatch>* out) {
RETURN_NOT_OK(file->Seek(offset));
return ReadRecordBatch(schema, file, out);
}
Status ReadTensor(int64_t offset, io::RandomAccessFile* file,
std::shared_ptr<Tensor>* out) {
// Respect alignment of Tensor messages (see WriteTensor)
offset = PaddedLength(offset);
RETURN_NOT_OK(file->Seek(offset));
std::unique_ptr<Message> message;
RETURN_NOT_OK(ReadContiguousPayload(offset, file, &message));
RETURN_NOT_OK(ReadContiguousPayload(file, &message));
std::shared_ptr<DataType> type;
std::vector<int64_t> shape;
cpp/src/arrow/ipc/reader.h
+47 -18
View File
@@ -24,7 +24,7 @@
#include <memory>
#include <vector>
#include "arrow/ipc/metadata.h"
#include "arrow/ipc/message.h"
#include "arrow/util/visibility.h"
namespace arrow {
@@ -69,18 +69,28 @@ class ARROW_EXPORT RecordBatchStreamReader : public RecordBatchReader {
/// Create batch reader from generic MessageReader
///
/// \param(in) message_reader a MessageReader implementation
/// \param(out) out the created RecordBatchStreamReader object
/// \param(out) out the created RecordBatchReader object
/// \return Status
static Status Open(std::unique_ptr<MessageReader> message_reader,
std::shared_ptr<RecordBatchReader>* out);
#ifndef ARROW_NO_DEPRECATED_API
/// \deprecated Since 0.7.0
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) out the created RecordBatchStreamReader object
/// \return Status
/// \deprecated Since 0.7.0
static Status Open(const std::shared_ptr<io::InputStream>& stream,
std::shared_ptr<RecordBatchStreamReader>* out);
#endif
/// \brief Record batch stream reader from InputStream
///
/// \param(in) stream an input stream instance. Must stay alive throughout
/// lifetime of stream reader
/// \param(out) out the created RecordBatchStreamReader object
/// \return Status
static Status Open(io::InputStream* stream, std::shared_ptr<RecordBatchReader>* out);
std::shared_ptr<Schema> schema() const override;
Status ReadNextRecordBatch(std::shared_ptr<RecordBatch>* batch) override;
@@ -142,6 +152,27 @@ class ARROW_EXPORT RecordBatchFileReader {
// Generic read functions; does not copy data if the input supports zero copy reads
/// \brief Read Schema from stream serialized as a sequence of IPC messages
///
/// \param[in] stream an InputStream
/// \param[out] out the output Schema
///
/// If record batches follow the schema, it is better to use
/// RecordBatchStreamReader
ARROW_EXPORT
Status ReadSchema(io::InputStream* stream, std::shared_ptr<Schema>* out);
/// Read record batch as encapsulated IPC message with metadata size prefix and
/// header
///
/// \param(in) schema the record batch schema
/// \param(in) offset the file location of the start of the message
/// \param(in) file the file where the batch is located
/// \param(out) out the read record batch
ARROW_EXPORT
Status ReadRecordBatch(const std::shared_ptr<Schema>& schema, io::InputStream* stream,
std::shared_ptr<RecordBatch>* out);
/// \brief Read record batch from file given metadata and schema
///
/// \param(in) metadata a Message containing the record batch metadata
@@ -174,17 +205,6 @@ Status ReadRecordBatch(const Buffer& metadata, const std::shared_ptr<Schema>& sc
int max_recursion_depth, io::RandomAccessFile* file,
std::shared_ptr<RecordBatch>* out);
/// Read record batch as encapsulated IPC message with metadata size prefix and
/// header
///
/// \param(in) schema the record batch schema
/// \param(in) offset the file location of the start of the message
/// \param(in) file the file where the batch is located
/// \param(out) out the read record batch
ARROW_EXPORT
Status ReadRecordBatch(const std::shared_ptr<Schema>& schema, int64_t offset,
io::RandomAccessFile* file, std::shared_ptr<RecordBatch>* out);
/// EXPERIMENTAL: Read arrow::Tensor as encapsulated IPC message in file
///
/// \param(in) offset the file location of the start of the message
@@ -194,6 +214,15 @@ ARROW_EXPORT
Status ReadTensor(int64_t offset, io::RandomAccessFile* file,
std::shared_ptr<Tensor>* out);
#ifndef ARROW_NO_DEPRECATED_API
/// \deprecated Since 0.7.0
///
/// Deprecated in favor of more general InputStream-based API
ARROW_EXPORT
Status ReadRecordBatch(const std::shared_ptr<Schema>& schema, int64_t offset,
io::RandomAccessFile* stream, std::shared_ptr<RecordBatch>* out);
#endif
} // namespace ipc
} // namespace arrow
cpp/src/arrow/ipc/stream-to-file.cc
+4 -4
View File
@@ -30,12 +30,12 @@ namespace ipc {
// A typical usage would be:
// $ <program that produces streaming output> | stream-to-file > file.arrow
Status ConvertToFile() {
std::shared_ptr<io::InputStream> input(new io::StdinStream);
std::shared_ptr<RecordBatchStreamReader> reader;
RETURN_NOT_OK(RecordBatchStreamReader::Open(input, &reader));
io::StdinStream input;
std::shared_ptr<RecordBatchReader> reader;
RETURN_NOT_OK(RecordBatchStreamReader::Open(&input, &reader));
io::StdoutStream sink;
std::shared_ptr<RecordBatchFileWriter> writer;
std::shared_ptr<RecordBatchWriter> writer;
RETURN_NOT_OK(RecordBatchFileWriter::Open(&sink, reader->schema(), &writer));
std::shared_ptr<RecordBatch> batch;
cpp/src/arrow/ipc/writer.cc
+162 -95
View File
@@ -28,7 +28,8 @@
#include "arrow/buffer.h"
#include "arrow/io/interfaces.h"
#include "arrow/io/memory.h"
#include "arrow/ipc/metadata.h"
#include "arrow/ipc/message.h"
#include "arrow/ipc/metadata-internal.h"
#include "arrow/ipc/util.h"
#include "arrow/memory_pool.h"
#include "arrow/status.h"
@@ -542,25 +543,9 @@ Status WriteRecordBatch(const RecordBatch& batch, int64_t buffer_start_offset,
return writer.Write(batch, dst, metadata_length, body_length);
}
Status SerializeRecordBatch(const RecordBatch& batch, MemoryPool* pool,
std::shared_ptr<Buffer>* out) {
int64_t size = 0;
RETURN_NOT_OK(GetRecordBatchSize(batch, &size));
std::shared_ptr<MutableBuffer> buffer;
RETURN_NOT_OK(AllocateBuffer(pool, size, &buffer));
io::FixedSizeBufferWriter stream(buffer);
int32_t metadata_length = 0;
int64_t body_length = 0;
RETURN_NOT_OK(WriteRecordBatch(batch, 0, &stream, &metadata_length, &body_length, pool,
kMaxNestingDepth, true));
*out = buffer;
return Status::OK();
}
Status WriteRecordBatchStream(const std::vector<std::shared_ptr<RecordBatch>>& batches,
io::OutputStream* dst) {
std::shared_ptr<RecordBatchStreamWriter> writer;
std::shared_ptr<RecordBatchWriter> writer;
RETURN_NOT_OK(RecordBatchStreamWriter::Open(dst, batches[0]->schema(), &writer));
for (const auto& batch : batches) {
// allow sizes > INT32_MAX
@@ -633,32 +618,44 @@ RecordBatchWriter::~RecordBatchWriter() {}
// ----------------------------------------------------------------------
// Stream writer implementation
class RecordBatchStreamWriter::RecordBatchStreamWriterImpl {
class StreamBookKeeper {
public:
RecordBatchStreamWriterImpl()
: pool_(default_memory_pool()), position_(-1), started_(false) {}
StreamBookKeeper() : sink_(nullptr), position_(-1) {}
explicit StreamBookKeeper(io::OutputStream* sink) : sink_(sink) {}
virtual ~RecordBatchStreamWriterImpl() = default;
Status UpdatePosition() { return sink_->Tell(&position_); }
Status Open(io::OutputStream* sink, const std::shared_ptr<Schema>& schema) {
sink_ = sink;
schema_ = schema;
return UpdatePosition();
}
virtual Status Start() {
RETURN_NOT_OK(WriteSchema());
// If there are any dictionaries, write them as the next messages
RETURN_NOT_OK(WriteDictionaries());
started_ = true;
Status Align(int64_t alignment = kArrowIpcAlignment) {
// Adds padding bytes if necessary to ensure all memory blocks are written on
// 8-byte (or other alignment) boundaries.
int64_t remainder = PaddedLength(position_, alignment) - position_;
if (remainder > 0) {
return Write(kPaddingBytes, remainder);
}
return Status::OK();
}
// Write data and update position
Status Write(const uint8_t* data, int64_t nbytes) {
RETURN_NOT_OK(sink_->Write(data, nbytes));
position_ += nbytes;
return Status::OK();
}
protected:
io::OutputStream* sink_;
int64_t position_;
};
class SchemaWriter : public StreamBookKeeper {
public:
SchemaWriter(const Schema& schema, DictionaryMemo* dictionary_memo, MemoryPool* pool,
io::OutputStream* sink)
: StreamBookKeeper(sink), schema_(schema), dictionary_memo_(dictionary_memo) {}
Status WriteSchema() {
std::shared_ptr<Buffer> schema_fb;
RETURN_NOT_OK(WriteSchemaMessage(*schema_, &dictionary_memo_, &schema_fb));
RETURN_NOT_OK(WriteSchemaMessage(schema_, dictionary_memo_, &schema_fb));
int32_t metadata_length = 0;
RETURN_NOT_OK(WriteMessage(*schema_fb, sink_, &metadata_length));
@@ -667,6 +664,60 @@ class RecordBatchStreamWriter::RecordBatchStreamWriterImpl {
return Status::OK();
}
Status WriteDictionaries(std::vector<FileBlock>* dictionaries) {
const DictionaryMap& id_to_dictionary = dictionary_memo_->id_to_dictionary();
dictionaries->resize(id_to_dictionary.size());
// TODO(wesm): does sorting by id yield any benefit?
int dict_index = 0;
for (const auto& entry : id_to_dictionary) {
FileBlock* block = &(*dictionaries)[dict_index++];
block->offset = position_;
// Frame of reference in file format is 0, see ARROW-384
const int64_t buffer_start_offset = 0;
RETURN_NOT_OK(WriteDictionary(entry.first, entry.second, buffer_start_offset, sink_,
&block->metadata_length, &block->body_length, pool_));
RETURN_NOT_OK(UpdatePosition());
DCHECK(position_ % 8 == 0) << "WriteDictionary did not perform aligned writes";
}
return Status::OK();
}
Status Write(std::vector<FileBlock>* dictionaries) {
RETURN_NOT_OK(WriteSchema());
// If there are any dictionaries, write them as the next messages
return WriteDictionaries(dictionaries);
}
private:
MemoryPool* pool_;
const Schema& schema_;
DictionaryMemo* dictionary_memo_;
};
class RecordBatchStreamWriter::RecordBatchStreamWriterImpl : public StreamBookKeeper {
public:
RecordBatchStreamWriterImpl(io::OutputStream* sink,
const std::shared_ptr<Schema>& schema)
: StreamBookKeeper(sink),
schema_(schema),
pool_(default_memory_pool()),
started_(false) {}
virtual ~RecordBatchStreamWriterImpl() = default;
virtual Status Start() {
SchemaWriter schema_writer(*schema_, &dictionary_memo_, pool_, sink_);
RETURN_NOT_OK(schema_writer.Write(&dictionaries_));
started_ = true;
return Status::OK();
}
virtual Status Close() {
// Write the schema if not already written
// User is responsible for closing the OutputStream
@@ -684,33 +735,9 @@ class RecordBatchStreamWriter::RecordBatchStreamWriterImpl {
return Status::OK();
}
Status UpdatePosition() { return sink_->Tell(&position_); }
Status WriteDictionaries() {
const DictionaryMap& id_to_dictionary = dictionary_memo_.id_to_dictionary();
dictionaries_.resize(id_to_dictionary.size());
// TODO(wesm): does sorting by id yield any benefit?
int dict_index = 0;
for (const auto& entry : id_to_dictionary) {
FileBlock* block = &dictionaries_[dict_index++];
block->offset = position_;
// Frame of reference in file format is 0, see ARROW-384
const int64_t buffer_start_offset = 0;
RETURN_NOT_OK(WriteDictionary(entry.first, entry.second, buffer_start_offset, sink_,
&block->metadata_length, &block->body_length, pool_));
RETURN_NOT_OK(UpdatePosition());
DCHECK(position_ % 8 == 0) << "WriteDictionary did not perform aligned writes";
}
return Status::OK();
}
Status WriteRecordBatch(const RecordBatch& batch, bool allow_64bit, FileBlock* block) {
RETURN_NOT_OK(CheckStarted());
RETURN_NOT_OK(UpdatePosition());
block->offset = position_;
@@ -733,45 +760,22 @@ class RecordBatchStreamWriter::RecordBatchStreamWriterImpl {
&record_batches_[record_batches_.size() - 1]);
}
Status Align(int64_t alignment = kArrowIpcAlignment) {
// Adds padding bytes if necessary to ensure all memory blocks are written on
// 8-byte (or other alignment) boundaries.
int64_t remainder = PaddedLength(position_, alignment) - position_;
if (remainder > 0) {
return Write(kPaddingBytes, remainder);
}
return Status::OK();
}
// Write data and update position
Status Write(const uint8_t* data, int64_t nbytes) {
RETURN_NOT_OK(sink_->Write(data, nbytes));
position_ += nbytes;
return Status::OK();
}
void set_memory_pool(MemoryPool* pool) { pool_ = pool; }
protected:
io::OutputStream* sink_;
std::shared_ptr<Schema> schema_;
MemoryPool* pool_;
bool started_;
// When writing out the schema, we keep track of all the dictionaries we
// encounter, as they must be written out first in the stream
DictionaryMemo dictionary_memo_;
MemoryPool* pool_;
int64_t position_;
bool started_;
std::vector<FileBlock> dictionaries_;
std::vector<FileBlock> record_batches_;
};
RecordBatchStreamWriter::RecordBatchStreamWriter() {
impl_.reset(new RecordBatchStreamWriterImpl());
}
RecordBatchStreamWriter::RecordBatchStreamWriter() {}
RecordBatchStreamWriter::~RecordBatchStreamWriter() {}
@@ -784,13 +788,26 @@ void RecordBatchStreamWriter::set_memory_pool(MemoryPool* pool) {
impl_->set_memory_pool(pool);
}
Status RecordBatchStreamWriter::Open(io::OutputStream* sink,
const std::shared_ptr<Schema>& schema,
std::shared_ptr<RecordBatchWriter>* out) {
// ctor is private
auto result = std::shared_ptr<RecordBatchStreamWriter>(new RecordBatchStreamWriter());
result->impl_.reset(new RecordBatchStreamWriterImpl(sink, schema));
*out = result;
return Status::OK();
}
#ifndef ARROW_NO_DEPRECATED_API
Status RecordBatchStreamWriter::Open(io::OutputStream* sink,
const std::shared_ptr<Schema>& schema,
std::shared_ptr<RecordBatchStreamWriter>* out) {
// ctor is private
*out = std::shared_ptr<RecordBatchStreamWriter>(new RecordBatchStreamWriter());
return (*out)->impl_->Open(sink, schema);
(*out)->impl_.reset(new RecordBatchStreamWriterImpl(sink, schema));
return Status::OK();
}
#endif
Status RecordBatchStreamWriter::Close() { return impl_->Close(); }
@@ -802,6 +819,9 @@ class RecordBatchFileWriter::RecordBatchFileWriterImpl
public:
using BASE = RecordBatchStreamWriter::RecordBatchStreamWriterImpl;
RecordBatchFileWriterImpl(io::OutputStream* sink, const std::shared_ptr<Schema>& schema)
: BASE(sink, schema) {}
Status Start() override {
// It is only necessary to align to 8-byte boundary at the start of the file
RETURN_NOT_OK(Write(reinterpret_cast<const uint8_t*>(kArrowMagicBytes),
@@ -815,6 +835,8 @@ class RecordBatchFileWriter::RecordBatchFileWriterImpl
Status Close() override {
// Write metadata
RETURN_NOT_OK(UpdatePosition());
int64_t initial_position = position_;
RETURN_NOT_OK(WriteFileFooter(*schema_, dictionaries_, record_batches_,
&dictionary_memo_, sink_));
@@ -836,20 +858,31 @@ class RecordBatchFileWriter::RecordBatchFileWriterImpl
}
};
RecordBatchFileWriter::RecordBatchFileWriter() {
impl_.reset(new RecordBatchFileWriterImpl());
}
RecordBatchFileWriter::RecordBatchFileWriter() {}
RecordBatchFileWriter::~RecordBatchFileWriter() {}
Status RecordBatchFileWriter::Open(io::OutputStream* sink,
const std::shared_ptr<Schema>& schema,
std::shared_ptr<RecordBatchFileWriter>* out) {
*out = std::shared_ptr<RecordBatchFileWriter>(
new RecordBatchFileWriter()); // ctor is private
return (*out)->impl_->Open(sink, schema);
std::shared_ptr<RecordBatchWriter>* out) {
// ctor is private
auto result = std::shared_ptr<RecordBatchFileWriter>(new RecordBatchFileWriter());
result->impl_.reset(new RecordBatchFileWriterImpl(sink, schema));
*out = result;
return Status::OK();
}
#ifndef ARROW_NO_DEPRECATED_API
Status RecordBatchFileWriter::Open(io::OutputStream* sink,
const std::shared_ptr<Schema>& schema,
std::shared_ptr<RecordBatchFileWriter>* out) {
// ctor is private
*out = std::shared_ptr<RecordBatchFileWriter>(new RecordBatchFileWriter());
(*out)->impl_.reset(new RecordBatchFileWriterImpl(sink, schema));
return Status::OK();
}
#endif
Status RecordBatchFileWriter::WriteRecordBatch(const RecordBatch& batch,
bool allow_64bit) {
return impl_->WriteRecordBatch(batch, allow_64bit);
@@ -857,5 +890,39 @@ Status RecordBatchFileWriter::WriteRecordBatch(const RecordBatch& batch,
Status RecordBatchFileWriter::Close() { return impl_->Close(); }
// ----------------------------------------------------------------------
// Serialization public APIs
Status SerializeRecordBatch(const RecordBatch& batch, MemoryPool* pool,
std::shared_ptr<Buffer>* out) {
int64_t size = 0;
RETURN_NOT_OK(GetRecordBatchSize(batch, &size));
std::shared_ptr<MutableBuffer> buffer;
RETURN_NOT_OK(AllocateBuffer(pool, size, &buffer));
io::FixedSizeBufferWriter stream(buffer);
int32_t metadata_length = 0;
int64_t body_length = 0;
RETURN_NOT_OK(WriteRecordBatch(batch, 0, &stream, &metadata_length, &body_length, pool,
kMaxNestingDepth, true));
*out = buffer;
return Status::OK();
}
Status SerializeSchema(const Schema& schema, MemoryPool* pool,
std::shared_ptr<Buffer>* out) {
std::shared_ptr<io::BufferOutputStream> stream;
RETURN_NOT_OK(io::BufferOutputStream::Create(1024, pool, &stream));
DictionaryMemo memo;
SchemaWriter schema_writer(schema, &memo, pool, stream.get());
// Unused
std::vector<FileBlock> dictionary_blocks;
RETURN_NOT_OK(schema_writer.Write(&dictionary_blocks));
return stream->Finish(out);
}
} // namespace ipc
} // namespace arrow
cpp/src/arrow/ipc/writer.h
+32 -9
View File
@@ -25,7 +25,7 @@
#include <memory>
#include <vector>
#include "arrow/ipc/metadata.h"
#include "arrow/ipc/message.h"
#include "arrow/util/visibility.h"
namespace arrow {
@@ -85,8 +85,14 @@ class ARROW_EXPORT RecordBatchStreamWriter : public RecordBatchWriter {
/// \param(in) schema the schema of the record batches to be written
/// \param(out) out the created stream writer
/// \return Status indicating success or failure
static Status Open(io::OutputStream* sink, const std::shared_ptr<Schema>& schema,
std::shared_ptr<RecordBatchWriter>* out);
#ifndef ARROW_NO_DEPRECATED_API
/// \deprecated Since 0.7.0
static Status Open(io::OutputStream* sink, const std::shared_ptr<Schema>& schema,
std::shared_ptr<RecordBatchStreamWriter>* out);
#endif
Status WriteRecordBatch(const RecordBatch& batch, bool allow_64bit = false) override;
Status Close() override;
@@ -113,8 +119,14 @@ class ARROW_EXPORT RecordBatchFileWriter : public RecordBatchStreamWriter {
/// \param(in) schema the schema of the record batches to be written
/// \param(out) out the created stream writer
/// \return Status indicating success or failure
static Status Open(io::OutputStream* sink, const std::shared_ptr<Schema>& schema,
std::shared_ptr<RecordBatchWriter>* out);
#ifndef ARROW_NO_DEPRECATED_API
/// \deprecated Since 0.7.0
static Status Open(io::OutputStream* sink, const std::shared_ptr<Schema>& schema,
std::shared_ptr<RecordBatchFileWriter>* out);
#endif
Status WriteRecordBatch(const RecordBatch& batch, bool allow_64bit = false) override;
Status Close() override;
@@ -155,14 +167,6 @@ Status WriteRecordBatch(const RecordBatch& batch, int64_t buffer_start_offset,
int max_recursion_depth = kMaxNestingDepth,
bool allow_64bit = false);
/// \brief Write dictionary message to output stream
///
/// Low-level API
ARROW_EXPORT
Status WriteDictionary(int64_t dictionary_id, const std::shared_ptr<Array>& dictionary,
int64_t buffer_start_offset, io::OutputStream* dst,
int32_t* metadata_length, int64_t* body_length, MemoryPool* pool);
/// \brief Serialize record batch as encapsulated IPC message in a new buffer
///
/// \param[in] batch the record batch
@@ -173,6 +177,17 @@ ARROW_EXPORT
Status SerializeRecordBatch(const RecordBatch& batch, MemoryPool* pool,
std::shared_ptr<Buffer>* out);
/// \brief Serialize schema using stream writer as a sequence of one or more
/// IPC messages
///
/// \param[in] scheam the schema to write
/// \param[in] pool a MemoryPool to allocate memory from
/// \param[out] out the serialized schema
/// \return Status
ARROW_EXPORT
Status SerializeSchema(const Schema& schema, MemoryPool* pool,
std::shared_ptr<Buffer>* out);
/// \brief Write multiple record batches to OutputStream
/// \param[in] batches a vector of batches. Must all have same schema
/// \param[out] dst an OutputStream
@@ -205,6 +220,14 @@ ARROW_EXPORT
Status WriteLargeRecordBatch(const RecordBatch& batch, int64_t buffer_start_offset,
io::OutputStream* dst, int32_t* metadata_length,
int64_t* body_length, MemoryPool* pool);
/// \brief Write dictionary message to output stream
///
/// \deprecated Since 0.7.0 as public API
ARROW_EXPORT
Status WriteDictionary(int64_t dictionary_id, const std::shared_ptr<Array>& dictionary,
int64_t buffer_start_offset, io::OutputStream* dst,
int32_t* metadata_length, int64_t* body_length, MemoryPool* pool);
#endif
} // namespace ipc
cpp/src/arrow/python/arrow_to_python.cc
+5 -4
View File
@@ -189,22 +189,23 @@ Status DeserializeTuple(std::shared_ptr<Array> array, int64_t start_idx, int64_t
DESERIALIZE_SEQUENCE(PyTuple_New, PyTuple_SET_ITEM)
}
Status ReadSerializedObject(std::shared_ptr<io::RandomAccessFile> src,
SerializedPyObject* out) {
std::shared_ptr<ipc::RecordBatchStreamReader> reader;
Status ReadSerializedObject(io::RandomAccessFile* src, SerializedPyObject* out) {
int64_t offset;
int64_t bytes_read;
int32_t num_tensors;
// Read number of tensors
RETURN_NOT_OK(
src->Read(sizeof(int32_t), &bytes_read, reinterpret_cast<uint8_t*>(&num_tensors)));
std::shared_ptr<ipc::RecordBatchReader> reader;
RETURN_NOT_OK(ipc::RecordBatchStreamReader::Open(src, &reader));
RETURN_NOT_OK(reader->ReadNextRecordBatch(&out->batch));
RETURN_NOT_OK(src->Tell(&offset));
offset += 4; // Skip the end-of-stream message
for (int i = 0; i < num_tensors; ++i) {
std::shared_ptr<Tensor> tensor;
RETURN_NOT_OK(ipc::ReadTensor(offset, src.get(), &tensor));
RETURN_NOT_OK(ipc::ReadTensor(offset, src, &tensor));
out->tensors.push_back(tensor);
RETURN_NOT_OK(src->Tell(&offset));
}
cpp/src/arrow/python/arrow_to_python.h
+1 -2
View File
@@ -46,8 +46,7 @@ namespace py {
/// \param[out] out the reconstructed data
/// \return Status
ARROW_EXPORT
Status ReadSerializedObject(std::shared_ptr<io::RandomAccessFile> src,
SerializedPyObject* out);
Status ReadSerializedObject(io::RandomAccessFile* src, SerializedPyObject* out);
/// \brief Reconstruct Python object from Arrow-serialized representation
/// \param[in] object
cpp/src/arrow/python/python_to_arrow.cc
+10 -10
View File
@@ -53,18 +53,18 @@ namespace py {
/// scalar Python types, lists, tuples, dictionaries and tensors.
class SequenceBuilder {
public:
explicit SequenceBuilder(MemoryPool* pool = nullptr)
explicit SequenceBuilder(MemoryPool* pool ARROW_MEMORY_POOL_DEFAULT)
: pool_(pool),
types_(pool, ::arrow::int8()),
offsets_(pool, ::arrow::int32()),
types_(::arrow::int8(), pool),
offsets_(::arrow::int32(), pool),
nones_(pool),
bools_(pool, ::arrow::boolean()),
ints_(pool, ::arrow::int64()),
bytes_(pool, ::arrow::binary()),
bools_(::arrow::boolean(), pool),
ints_(::arrow::int64(), pool),
bytes_(::arrow::binary(), pool),
strings_(pool),
floats_(pool, ::arrow::float32()),
doubles_(pool, ::arrow::float64()),
tensor_indices_(pool, ::arrow::int32()),
floats_(::arrow::float32(), pool),
doubles_(::arrow::float64(), pool),
tensor_indices_(::arrow::int32(), pool),
list_offsets_({0}),
tuple_offsets_({0}),
dict_offsets_({0}) {}
@@ -184,7 +184,7 @@ class SequenceBuilder {
if (data != nullptr) {
DCHECK(data->length() == offsets.back());
std::shared_ptr<Array> offset_array;
Int32Builder builder(pool_, std::make_shared<Int32Type>());
Int32Builder builder(::arrow::int32(), pool_);
RETURN_NOT_OK(builder.Append(offsets.data(), offsets.size()));
RETURN_NOT_OK(builder.Finish(&offset_array));
std::shared_ptr<Array> list_array;
python/pyarrow/includes/libarrow.pxd
+5 -5
View File
@@ -632,8 +632,8 @@ cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil:
cdef cppclass CRecordBatchStreamReader \
" arrow::ipc::RecordBatchStreamReader"(CRecordBatchReader):
@staticmethod
CStatus Open(const shared_ptr[InputStream]& stream,
shared_ptr[CRecordBatchStreamReader]* out)
CStatus Open(const InputStream* stream,
shared_ptr[CRecordBatchReader]* out)
@staticmethod
CStatus Open2" Open"(unique_ptr[CMessageReader] message_reader,
@@ -643,13 +643,13 @@ cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil:
" arrow::ipc::RecordBatchStreamWriter"(CRecordBatchWriter):
@staticmethod
CStatus Open(OutputStream* sink, const shared_ptr[CSchema]& schema,
shared_ptr[CRecordBatchStreamWriter]* out)
shared_ptr[CRecordBatchWriter]* out)
cdef cppclass CRecordBatchFileWriter \
" arrow::ipc::RecordBatchFileWriter"(CRecordBatchWriter):
@staticmethod
CStatus Open(OutputStream* sink, const shared_ptr[CSchema]& schema,
shared_ptr[CRecordBatchFileWriter]* out)
shared_ptr[CRecordBatchWriter]* out)
cdef cppclass CRecordBatchFileReader \
" arrow::ipc::RecordBatchFileReader":
@@ -792,7 +792,7 @@ cdef extern from "arrow/python/api.h" namespace 'arrow::py' nogil:
CStatus DeserializeObject(const CSerializedPyObject& obj,
PyObject* base, PyObject** out)
CStatus ReadSerializedObject(shared_ptr[RandomAccessFile] src,
CStatus ReadSerializedObject(RandomAccessFile* src,
CSerializedPyObject* out)
void set_serialization_callbacks(object serialize_callback,
python/pyarrow/ipc.pxi
+6 -18
View File
@@ -166,18 +166,13 @@ cdef class _RecordBatchWriter:
pass
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.writer))
self.closed = False
def write_batch(self, RecordBatch batch):
@@ -215,6 +210,7 @@ cdef get_input_stream(object source, shared_ptr[InputStream]* out):
cdef class _RecordBatchReader:
cdef:
shared_ptr[CRecordBatchReader] reader
shared_ptr[InputStream] in_stream
cdef readonly:
Schema schema
@@ -223,16 +219,11 @@ cdef class _RecordBatchReader:
pass
def _open(self, source):
cdef:
shared_ptr[InputStream] in_stream
shared_ptr[CRecordBatchStreamReader] reader
get_input_stream(source, &in_stream)
get_input_stream(source, &self.in_stream)
with nogil:
check_status(CRecordBatchStreamReader.Open(in_stream, &reader))
check_status(CRecordBatchStreamReader.Open(
self.in_stream.get(), &self.reader))
self.reader = <shared_ptr[CRecordBatchReader]> reader
self.schema = Schema()
self.schema.init_schema(self.reader.get().schema())
@@ -285,16 +276,13 @@ cdef class _RecordBatchReader:
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))
&self.writer))
# Cast to base class, because has same interface
self.writer = <shared_ptr[CRecordBatchWriter]> writer
self.closed = False
python/pyarrow/serialization.pxi
+1 -1
View File
@@ -237,7 +237,7 @@ def read_serialized(source, base=None):
cdef SerializedPyObject serialized = SerializedPyObject()
serialized.base = base
with nogil:
check_status(ReadSerializedObject(stream, &serialized.data))
check_status(ReadSerializedObject(stream.get(), &serialized.data))
return serialized