You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by we...@apache.org on 2017/08/25 03:05:12 UTC
[3/3] arrow git commit: ARROW-1408: [C++] IPC public API cleanup,
refactoring. Add SerializeSchema, ReadSchema public APIs
ARROW-1408: [C++] IPC public API cleanup, refactoring. Add SerializeSchema, ReadSchema public APIs
This is mostly moving code around. In reviewing I recommend focusing on the public headers. There were a number of places where it is more consistent to use naked pointers versus shared_ptr. Also some constructors were returning shared_ptr to subclass, where it would be simpler for clients to return a pointer to base.
This includes ARROW-1376 and ARROW-1406
Author: Wes McKinney <we...@twosigma.com>
Closes #988 from wesm/ARROW-1408 and squashes the following commits:
b156767c [Wes McKinney] Fix up glib bindings, undeprecate some APIs
4bdebfab [Wes McKinney] Add serialize methods to RecordBatch, Schema. Test round trip
ef12e0f2 [Wes McKinney] Fix a valgrind warning
73d30c98 [Wes McKinney] Better comments
8597b96f [Wes McKinney] Remove API that was never intended to be public, unlikely to be used anywhere
122a7591 [Wes McKinney] 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
b646f965 [Wes McKinney] Set device in more places
Project: http://git-wip-us.apache.org/repos/asf/arrow/repo
Commit: http://git-wip-us.apache.org/repos/asf/arrow/commit/f50f2eac
Tree: http://git-wip-us.apache.org/repos/asf/arrow/tree/f50f2eac
Diff: http://git-wip-us.apache.org/repos/asf/arrow/diff/f50f2eac
Branch: refs/heads/master
Commit: f50f2eacb499906012da377009921feb21445bb8
Parents: 750b77d
Author: Wes McKinney <we...@twosigma.com>
Authored: Thu Aug 24 23:04:58 2017 -0400
Committer: Wes McKinney <we...@twosigma.com>
Committed: Thu Aug 24 23:04:58 2017 -0400
----------------------------------------------------------------------
c_glib/arrow-glib/reader.cpp | 12 +-
c_glib/arrow-glib/writer.cpp | 30 +-
cpp/CMakeLists.txt | 8 +-
cpp/src/arrow/gpu/cuda_common.h | 11 +
cpp/src/arrow/gpu/cuda_memory.cc | 11 +-
cpp/src/arrow/gpu/cuda_memory.h | 4 +-
cpp/src/arrow/ipc/CMakeLists.txt | 3 +-
cpp/src/arrow/ipc/api.h | 3 +-
cpp/src/arrow/ipc/dictionary.cc | 85 ++
cpp/src/arrow/ipc/dictionary.h | 88 ++
cpp/src/arrow/ipc/file-to-stream.cc | 4 +-
cpp/src/arrow/ipc/ipc-read-write-benchmark.cc | 2 +-
cpp/src/arrow/ipc/ipc-read-write-test.cc | 58 +-
cpp/src/arrow/ipc/json-integration-test.cc | 6 +-
cpp/src/arrow/ipc/json-internal.cc | 2 +-
cpp/src/arrow/ipc/message.cc | 257 +++++
cpp/src/arrow/ipc/message.h | 184 +++
cpp/src/arrow/ipc/metadata-internal.cc | 916 +++++++++++++++
cpp/src/arrow/ipc/metadata-internal.h | 146 +++
cpp/src/arrow/ipc/metadata.cc | 1201 --------------------
cpp/src/arrow/ipc/metadata.h | 300 -----
cpp/src/arrow/ipc/reader.cc | 84 +-
cpp/src/arrow/ipc/reader.h | 91 +-
cpp/src/arrow/ipc/stream-to-file.cc | 8 +-
cpp/src/arrow/ipc/writer.cc | 247 ++--
cpp/src/arrow/ipc/writer.h | 33 +-
cpp/src/arrow/python/arrow_to_python.cc | 9 +-
cpp/src/arrow/python/arrow_to_python.h | 3 +-
cpp/src/arrow/python/python_to_arrow.cc | 20 +-
cpp/src/arrow/type.cc | 2 +
cpp/src/arrow/type.h | 2 +-
cpp/src/arrow/util/key_value_metadata.cc | 4 +-
python/pyarrow/__init__.py | 4 +-
python/pyarrow/includes/libarrow.pxd | 27 +-
python/pyarrow/ipc.pxi | 81 +-
python/pyarrow/ipc.py | 2 +-
python/pyarrow/serialization.pxi | 2 +-
python/pyarrow/table.pxi | 22 +
python/pyarrow/tests/test_ipc.py | 15 +
python/pyarrow/tests/test_table.py | 12 +
python/pyarrow/types.pxi | 36 +-
41 files changed, 2274 insertions(+), 1761 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/c_glib/arrow-glib/reader.cpp
----------------------------------------------------------------------
diff --git a/c_glib/arrow-glib/reader.cpp b/c_glib/arrow-glib/reader.cpp
index 25c8d67..a21b375 100644
--- a/c_glib/arrow-glib/reader.cpp
+++ b/c_glib/arrow-glib/reader.cpp
@@ -237,12 +237,15 @@ GArrowRecordBatchStreamReader *
garrow_record_batch_stream_reader_new(GArrowInputStream *stream,
GError **error)
{
+ using BaseType = arrow::ipc::RecordBatchReader;
+ using ReaderType = arrow::ipc::RecordBatchStreamReader;
+
auto arrow_input_stream = garrow_input_stream_get_raw(stream);
- std::shared_ptr<arrow::ipc::RecordBatchStreamReader> arrow_reader;
- auto status =
- arrow::ipc::RecordBatchStreamReader::Open(arrow_input_stream, &arrow_reader);
+ std::shared_ptr<BaseType> arrow_reader;
+ auto status = ReaderType::Open(arrow_input_stream, &arrow_reader);
+ auto subtype = std::dynamic_pointer_cast<ReaderType>(arrow_reader);
if (garrow_error_check(error, status, "[record-batch-stream-reader][open]")) {
- return garrow_record_batch_stream_reader_new_raw(&arrow_reader);
+ return garrow_record_batch_stream_reader_new_raw(&subtype);
} else {
return NULL;
}
@@ -354,6 +357,7 @@ garrow_record_batch_file_reader_new(GArrowSeekableInputStream *file,
GError **error)
{
auto arrow_random_access_file = garrow_seekable_input_stream_get_raw(file);
+
std::shared_ptr<arrow::ipc::RecordBatchFileReader> arrow_reader;
auto status =
arrow::ipc::RecordBatchFileReader::Open(arrow_random_access_file,
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/c_glib/arrow-glib/writer.cpp
----------------------------------------------------------------------
diff --git a/c_glib/arrow-glib/writer.cpp b/c_glib/arrow-glib/writer.cpp
index d4c36c2..c08db19 100644
--- a/c_glib/arrow-glib/writer.cpp
+++ b/c_glib/arrow-glib/writer.cpp
@@ -215,14 +215,17 @@ garrow_record_batch_stream_writer_new(GArrowOutputStream *sink,
GArrowSchema *schema,
GError **error)
{
+ using BaseType = arrow::ipc::RecordBatchWriter;
+ using WriterType = arrow::ipc::RecordBatchStreamWriter;
+
auto arrow_sink = garrow_output_stream_get_raw(sink).get();
- std::shared_ptr<arrow::ipc::RecordBatchStreamWriter> arrow_writer;
- auto status =
- arrow::ipc::RecordBatchStreamWriter::Open(arrow_sink,
- garrow_schema_get_raw(schema),
- &arrow_writer);
+ std::shared_ptr<BaseType> arrow_writer;
+ auto status = WriterType::Open(arrow_sink,
+ garrow_schema_get_raw(schema),
+ &arrow_writer);
+ auto subtype = std::dynamic_pointer_cast<WriterType>(arrow_writer);
if (garrow_error_check(error, status, "[record-batch-stream-writer][open]")) {
- return garrow_record_batch_stream_writer_new_raw(&arrow_writer);
+ return garrow_record_batch_stream_writer_new_raw(&subtype);
} else {
return NULL;
}
@@ -259,14 +262,17 @@ garrow_record_batch_file_writer_new(GArrowOutputStream *sink,
GArrowSchema *schema,
GError **error)
{
+ using BaseType = arrow::ipc::RecordBatchWriter;
+ using WriterType = arrow::ipc::RecordBatchFileWriter;
+
auto arrow_sink = garrow_output_stream_get_raw(sink);
- std::shared_ptr<arrow::ipc::RecordBatchFileWriter> arrow_writer;
- auto status =
- arrow::ipc::RecordBatchFileWriter::Open(arrow_sink.get(),
- garrow_schema_get_raw(schema),
- &arrow_writer);
+ std::shared_ptr<BaseType> arrow_writer;
+ auto status = WriterType::Open(arrow_sink.get(),
+ garrow_schema_get_raw(schema),
+ &arrow_writer);
+ auto subtype = std::dynamic_pointer_cast<WriterType>(arrow_writer);
if (garrow_error_check(error, status, "[record-batch-file-writer][open]")) {
- return garrow_record_batch_file_writer_new_raw(&arrow_writer);
+ return garrow_record_batch_file_writer_new_raw(&subtype);
} else {
return NULL;
}
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/CMakeLists.txt
----------------------------------------------------------------------
diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 6c56103..5d64333 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -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
)
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/gpu/cuda_common.h
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/gpu/cuda_common.h b/cpp/src/arrow/gpu/cuda_common.h
index fae2f6e..1d65f96 100644
--- a/cpp/src/arrow/gpu/cuda_common.h
+++ b/cpp/src/arrow/gpu/cuda_common.h
@@ -45,6 +45,17 @@ namespace gpu {
} \
} while (0)
+#define CUDADRV_RETURN_NOT_OK(STMT) \
+ do { \
+ CUresult ret = (STMT); \
+ if (ret != CUDA_SUCCESS) { \
+ std::stringstream ss; \
+ ss << "Cuda Driver API call in " << __FILE__ << " at line " << __LINE__ \
+ << " failed: " << #STMT; \
+ return Status::IOError(ss.str()); \
+ } \
+ } while (0)
+
} // namespace gpu
} // namespace arrow
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/gpu/cuda_memory.cc
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/gpu/cuda_memory.cc b/cpp/src/arrow/gpu/cuda_memory.cc
index d669619..4870813 100644
--- a/cpp/src/arrow/gpu/cuda_memory.cc
+++ b/cpp/src/arrow/gpu/cuda_memory.cc
@@ -71,12 +71,13 @@ CudaHostBuffer::~CudaHostBuffer() { CUDA_DCHECK(cudaFreeHost(mutable_data_)); }
// CudaBufferReader
CudaBufferReader::CudaBufferReader(const std::shared_ptr<CudaBuffer>& buffer)
- : io::BufferReader(buffer), cuda_buffer_(buffer) {}
+ : io::BufferReader(buffer), cuda_buffer_(buffer), gpu_number_(buffer->gpu_number()) {}
CudaBufferReader::~CudaBufferReader() {}
Status CudaBufferReader::Read(int64_t nbytes, int64_t* bytes_read, uint8_t* buffer) {
nbytes = std::min(nbytes, size_ - position_);
+ CUDA_RETURN_NOT_OK(cudaSetDevice(gpu_number_));
CUDA_RETURN_NOT_OK(
cudaMemcpy(buffer, data_ + position_, nbytes, cudaMemcpyDeviceToHost));
*bytes_read = nbytes;
@@ -95,7 +96,10 @@ Status CudaBufferReader::Read(int64_t nbytes, std::shared_ptr<Buffer>* out) {
// CudaBufferWriter
CudaBufferWriter::CudaBufferWriter(const std::shared_ptr<CudaBuffer>& buffer)
- : io::FixedSizeBufferWriter(buffer), buffer_size_(0), buffer_position_(0) {}
+ : io::FixedSizeBufferWriter(buffer),
+ gpu_number_(buffer->gpu_number()),
+ buffer_size_(0),
+ buffer_position_(0) {}
CudaBufferWriter::~CudaBufferWriter() {}
@@ -104,6 +108,7 @@ Status CudaBufferWriter::Close() { return Flush(); }
Status CudaBufferWriter::Flush() {
if (buffer_size_ > 0 && buffer_position_ > 0) {
// Only need to flush when the write has been buffered
+ CUDA_RETURN_NOT_OK(cudaSetDevice(gpu_number_));
CUDA_RETURN_NOT_OK(cudaMemcpy(mutable_data_ + position_ - buffer_position_,
host_buffer_data_, buffer_position_,
cudaMemcpyHostToDevice));
@@ -132,6 +137,7 @@ Status CudaBufferWriter::Write(const uint8_t* data, int64_t nbytes) {
if (nbytes + buffer_position_ >= buffer_size_) {
// Reach end of buffer, write everything
RETURN_NOT_OK(Flush());
+ CUDA_RETURN_NOT_OK(cudaSetDevice(gpu_number_));
CUDA_RETURN_NOT_OK(
cudaMemcpy(mutable_data_ + position_, data, nbytes, cudaMemcpyHostToDevice));
} else {
@@ -141,6 +147,7 @@ Status CudaBufferWriter::Write(const uint8_t* data, int64_t nbytes) {
}
} else {
// Unbuffered write
+ CUDA_RETURN_NOT_OK(cudaSetDevice(gpu_number_));
CUDA_RETURN_NOT_OK(
cudaMemcpy(mutable_data_ + position_, data, nbytes, cudaMemcpyHostToDevice));
}
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/gpu/cuda_memory.h
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/gpu/cuda_memory.h b/cpp/src/arrow/gpu/cuda_memory.h
index acf65e6..bd8b89a 100644
--- a/cpp/src/arrow/gpu/cuda_memory.h
+++ b/cpp/src/arrow/gpu/cuda_memory.h
@@ -97,8 +97,8 @@ class ARROW_EXPORT CudaBufferReader : public io::BufferReader {
Status Read(int64_t nbytes, std::shared_ptr<Buffer>* out) override;
private:
- // In case we need to access anything GPU-specific, like device number
std::shared_ptr<CudaBuffer> cuda_buffer_;
+ int gpu_number_;
};
/// \class CudaBufferWriter
@@ -132,6 +132,8 @@ class ARROW_EXPORT CudaBufferWriter : public io::FixedSizeBufferWriter {
int64_t num_bytes_buffered() const { return buffer_position_; }
private:
+ int gpu_number_;
+
// Pinned host buffer for buffering writes on CPU before calling cudaMalloc
int64_t buffer_size_;
int64_t buffer_position_;
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/CMakeLists.txt
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/CMakeLists.txt b/cpp/src/arrow/ipc/CMakeLists.txt
index 76e52a0..9689970 100644
--- a/cpp/src/arrow/ipc/CMakeLists.txt
+++ b/cpp/src/arrow/ipc/CMakeLists.txt
@@ -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")
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/api.h
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/api.h b/cpp/src/arrow/ipc/api.h
index 3f05e69..5cc6407 100644
--- a/cpp/src/arrow/ipc/api.h
+++ b/cpp/src/arrow/ipc/api.h
@@ -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"
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/dictionary.cc
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/dictionary.cc b/cpp/src/arrow/ipc/dictionary.cc
new file mode 100644
index 0000000..fcfb3e2
--- /dev/null
+++ b/cpp/src/arrow/ipc/dictionary.cc
@@ -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
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/dictionary.h
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/dictionary.h b/cpp/src/arrow/ipc/dictionary.h
new file mode 100644
index 0000000..744ab64
--- /dev/null
+++ b/cpp/src/arrow/ipc/dictionary.h
@@ -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
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/file-to-stream.cc
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/file-to-stream.cc b/cpp/src/arrow/ipc/file-to-stream.cc
index 4707c4f..e69444a 100644
--- a/cpp/src/arrow/ipc/file-to-stream.cc
+++ b/cpp/src/arrow/ipc/file-to-stream.cc
@@ -32,10 +32,10 @@ Status ConvertToStream(const char* path) {
std::shared_ptr<RecordBatchFileReader> reader;
RETURN_NOT_OK(io::ReadableFile::Open(path, &in_file));
- RETURN_NOT_OK(ipc::RecordBatchFileReader::Open(in_file, &reader));
+ RETURN_NOT_OK(ipc::RecordBatchFileReader::Open(in_file.get(), &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;
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/ipc-read-write-benchmark.cc
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/ipc-read-write-benchmark.cc b/cpp/src/arrow/ipc/ipc-read-write-benchmark.cc
index a6da637..9ed0abd 100644
--- a/cpp/src/arrow/ipc/ipc-read-write-benchmark.cc
+++ b/cpp/src/arrow/ipc/ipc-read-write-benchmark.cc
@@ -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!");
}
}
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/ipc-read-write-test.cc
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/ipc-read-write-test.cc b/cpp/src/arrow/ipc/ipc-read-write-test.cc
index 5aaff5f..e7b0a34 100644
--- a/cpp/src/arrow/ipc/ipc-read-write-test.cc
+++ b/cpp/src/arrow/ipc/ipc-read-write-test.cc
@@ -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());
@@ -162,7 +160,7 @@ class IpcTestFixture : public io::MemoryMapFixture {
RETURN_NOT_OK(mmap_->Tell(&offset));
std::shared_ptr<RecordBatchFileReader> file_reader;
- RETURN_NOT_OK(RecordBatchFileReader::Open(mmap_, offset, &file_reader));
+ RETURN_NOT_OK(RecordBatchFileReader::Open(mmap_.get(), offset, &file_reader));
return file_reader->ReadRecordBatch(0, result);
}
@@ -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));
@@ -517,7 +519,7 @@ class TestFileFormat : public ::testing::TestWithParam<MakeRecordBatch*> {
// Open the file
auto buf_reader = std::make_shared<io::BufferReader>(buffer_);
std::shared_ptr<RecordBatchFileReader> reader;
- RETURN_NOT_OK(RecordBatchFileReader::Open(buf_reader, footer_offset, &reader));
+ RETURN_NOT_OK(RecordBatchFileReader::Open(buf_reader.get(), footer_offset, &reader));
EXPECT_EQ(num_batches, reader->num_record_batches());
for (int i = 0; i < num_batches; ++i) {
@@ -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) {
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/json-integration-test.cc
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/json-integration-test.cc b/cpp/src/arrow/ipc/json-integration-test.cc
index 035f708..e0f01f9 100644
--- a/cpp/src/arrow/ipc/json-integration-test.cc
+++ b/cpp/src/arrow/ipc/json-integration-test.cc
@@ -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));
@@ -99,7 +99,7 @@ static Status ConvertArrowToJson(const std::string& arrow_path,
RETURN_NOT_OK(io::FileOutputStream::Open(json_path, &out_file));
std::shared_ptr<ipc::RecordBatchFileReader> reader;
- RETURN_NOT_OK(ipc::RecordBatchFileReader::Open(in_file, &reader));
+ RETURN_NOT_OK(ipc::RecordBatchFileReader::Open(in_file.get(), &reader));
if (FLAGS_verbose) {
std::cout << "Found schema: " << reader->schema()->ToString() << std::endl;
@@ -140,7 +140,7 @@ static Status ValidateArrowVsJson(const std::string& arrow_path,
RETURN_NOT_OK(io::ReadableFile::Open(arrow_path, &arrow_file));
std::shared_ptr<ipc::RecordBatchFileReader> arrow_reader;
- RETURN_NOT_OK(ipc::RecordBatchFileReader::Open(arrow_file, &arrow_reader));
+ RETURN_NOT_OK(ipc::RecordBatchFileReader::Open(arrow_file.get(), &arrow_reader));
auto json_schema = json_reader->schema();
auto arrow_schema = arrow_reader->schema();
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/json-internal.cc
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/json-internal.cc b/cpp/src/arrow/ipc/json-internal.cc
index bc2b0d1..13572ae 100644
--- a/cpp/src/arrow/ipc/json-internal.cc
+++ b/cpp/src/arrow/ipc/json-internal.cc
@@ -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"
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/message.cc
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/message.cc b/cpp/src/arrow/ipc/message.cc
new file mode 100644
index 0000000..eb06aaf
--- /dev/null
+++ b/cpp/src/arrow/ipc/message.cc
@@ -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
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/message.h
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/message.h b/cpp/src/arrow/ipc/message.h
new file mode 100644
index 0000000..dce4e27
--- /dev/null
+++ b/cpp/src/arrow/ipc/message.h
@@ -0,0 +1,184 @@
+// 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_MESSAGE_H
+#define ARROW_IPC_MESSAGE_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 {
+
+enum class MetadataVersion : char { V1, V2, V3 };
+
+// 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;
+
+// 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;
+
+/// \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();
+
+ /// \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);
+
+ /// \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;
+
+ MetadataVersion metadata_version() const;
+
+ const void* header() 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
+ Status SerializeTo(io::OutputStream* file, int64_t* output_length) const;
+
+ private:
+ // Hide serialization details from user API
+ class MessageImpl;
+ std::unique_ptr<MessageImpl> impl_;
+
+ DISALLOW_COPY_AND_ASSIGN(Message);
+};
+
+ARROW_EXPORT std::string FormatMessageType(Message::Type type);
+
+/// \brief Abstract interface for a sequence of messages
+/// \since 0.5.0
+class ARROW_EXPORT MessageReader {
+ public:
+ virtual ~MessageReader() = default;
+
+ /// \brief Read next Message from the interface
+ ///
+ /// \param[out] message an arrow::ipc::Message instance
+ /// \return Status
+ virtual Status ReadNextMessage(std::unique_ptr<Message>* message) = 0;
+};
+
+/// \brief Implementation of MessageReader that reads from InputStream
+/// \since 0.5.0
+class ARROW_EXPORT InputStreamMessageReader : public MessageReader {
+ public:
+ explicit InputStreamMessageReader(io::InputStream* stream) : stream_(stream) {}
+
+ explicit InputStreamMessageReader(const std::shared_ptr<io::InputStream>& owned_stream)
+ : InputStreamMessageReader(owned_stream.get()) {
+ owned_stream_ = owned_stream;
+ }
+
+ ~InputStreamMessageReader();
+
+ Status ReadNextMessage(std::unique_ptr<Message>* message) override;
+
+ private:
+ io::InputStream* stream_;
+ std::shared_ptr<io::InputStream> owned_stream_;
+};
+
+/// \brief Read encapulated RPC message from position in file
+///
+/// Read a length-prefixed message flatbuffer starting at the indicated file
+/// 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
+///
+/// \param[in] offset the position in the file where the message starts. The
+/// first 4 bytes after the offset are the message length
+/// \param[in] metadata_length the total number of bytes to read from file
+/// \param[in] file the seekable file interface to read from
+/// \param[out] message the message read
+/// \return Status success or failure
+ARROW_EXPORT
+Status ReadMessage(const int64_t offset, const int32_t metadata_length,
+ 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)
+ARROW_EXPORT
+Status ReadMessage(io::InputStream* stream, std::unique_ptr<Message>* message);
+
+} // namespace ipc
+} // namespace arrow
+
+#endif // ARROW_IPC_MESSAGE_H
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/metadata-internal.cc
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/metadata-internal.cc b/cpp/src/arrow/ipc/metadata-internal.cc
new file mode 100644
index 0000000..fdaab21
--- /dev/null
+++ b/cpp/src/arrow/ipc/metadata-internal.cc
@@ -0,0 +1,916 @@
+// 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/metadata-internal.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "flatbuffers/flatbuffers.h"
+
+#include "arrow/array.h"
+#include "arrow/buffer.h"
+#include "arrow/io/interfaces.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 flatbuf = org::apache::arrow::flatbuf;
+
+namespace ipc {
+
+using FBB = flatbuffers::FlatBufferBuilder;
+using DictionaryOffset = flatbuffers::Offset<flatbuf::DictionaryEncoding>;
+using FieldOffset = flatbuffers::Offset<flatbuf::Field>;
+using KeyValueOffset = flatbuffers::Offset<flatbuf::KeyValue>;
+using RecordBatchOffset = flatbuffers::Offset<flatbuf::RecordBatch>;
+using VectorLayoutOffset = flatbuffers::Offset<arrow::flatbuf::VectorLayout>;
+using Offset = flatbuffers::Offset<void>;
+using FBString = flatbuffers::Offset<flatbuffers::String>;
+
+static Status IntFromFlatbuffer(const flatbuf::Int* int_data,
+ std::shared_ptr<DataType>* out) {
+ if (int_data->bitWidth() > 64) {
+ return Status::NotImplemented("Integers with more than 64 bits not implemented");
+ }
+ if (int_data->bitWidth() < 8) {
+ return Status::NotImplemented("Integers with less than 8 bits not implemented");
+ }
+
+ switch (int_data->bitWidth()) {
+ case 8:
+ *out = int_data->is_signed() ? int8() : uint8();
+ break;
+ case 16:
+ *out = int_data->is_signed() ? int16() : uint16();
+ break;
+ case 32:
+ *out = int_data->is_signed() ? int32() : uint32();
+ break;
+ case 64:
+ *out = int_data->is_signed() ? int64() : uint64();
+ break;
+ default:
+ return Status::NotImplemented("Integers not in cstdint are not implemented");
+ }
+ return Status::OK();
+}
+
+static Status FloatFromFlatuffer(const flatbuf::FloatingPoint* float_data,
+ std::shared_ptr<DataType>* out) {
+ if (float_data->precision() == flatbuf::Precision_HALF) {
+ *out = float16();
+ } else if (float_data->precision() == flatbuf::Precision_SINGLE) {
+ *out = float32();
+ } else {
+ *out = float64();
+ }
+ return Status::OK();
+}
+
+// Forward declaration
+static Status FieldToFlatbuffer(FBB& fbb, const std::shared_ptr<Field>& field,
+ DictionaryMemo* dictionary_memo, FieldOffset* offset);
+
+static Offset IntToFlatbuffer(FBB& fbb, int bitWidth, bool is_signed) {
+ return flatbuf::CreateInt(fbb, bitWidth, is_signed).Union();
+}
+
+static Offset FloatToFlatbuffer(FBB& fbb, flatbuf::Precision precision) {
+ return flatbuf::CreateFloatingPoint(fbb, precision).Union();
+}
+
+static Status AppendChildFields(FBB& fbb, const std::shared_ptr<DataType>& type,
+ std::vector<FieldOffset>* out_children,
+ DictionaryMemo* dictionary_memo) {
+ FieldOffset field;
+ for (int i = 0; i < type->num_children(); ++i) {
+ RETURN_NOT_OK(FieldToFlatbuffer(fbb, type->child(i), dictionary_memo, &field));
+ out_children->push_back(field);
+ }
+ return Status::OK();
+}
+
+static Status ListToFlatbuffer(FBB& fbb, const std::shared_ptr<DataType>& type,
+ std::vector<FieldOffset>* out_children,
+ DictionaryMemo* dictionary_memo, Offset* offset) {
+ RETURN_NOT_OK(AppendChildFields(fbb, type, out_children, dictionary_memo));
+ *offset = flatbuf::CreateList(fbb).Union();
+ return Status::OK();
+}
+
+static Status StructToFlatbuffer(FBB& fbb, const std::shared_ptr<DataType>& type,
+ std::vector<FieldOffset>* out_children,
+ DictionaryMemo* dictionary_memo, Offset* offset) {
+ RETURN_NOT_OK(AppendChildFields(fbb, type, out_children, dictionary_memo));
+ *offset = flatbuf::CreateStruct_(fbb).Union();
+ return Status::OK();
+}
+
+// ----------------------------------------------------------------------
+// Union implementation
+
+static Status UnionFromFlatbuffer(const flatbuf::Union* union_data,
+ const std::vector<std::shared_ptr<Field>>& children,
+ std::shared_ptr<DataType>* out) {
+ UnionMode mode = union_data->mode() == flatbuf::UnionMode_Sparse ? UnionMode::SPARSE
+ : UnionMode::DENSE;
+
+ std::vector<uint8_t> type_codes;
+
+ const flatbuffers::Vector<int32_t>* fb_type_ids = union_data->typeIds();
+ if (fb_type_ids == nullptr) {
+ for (uint8_t i = 0; i < children.size(); ++i) {
+ type_codes.push_back(i);
+ }
+ } else {
+ for (int32_t id : (*fb_type_ids)) {
+ // TODO(wesm): can these values exceed 255?
+ type_codes.push_back(static_cast<uint8_t>(id));
+ }
+ }
+
+ *out = union_(children, type_codes, mode);
+ return Status::OK();
+}
+
+static Status UnionToFlatBuffer(FBB& fbb, const std::shared_ptr<DataType>& type,
+ std::vector<FieldOffset>* out_children,
+ DictionaryMemo* dictionary_memo, Offset* offset) {
+ RETURN_NOT_OK(AppendChildFields(fbb, type, out_children, dictionary_memo));
+
+ const auto& union_type = static_cast<const UnionType&>(*type);
+
+ flatbuf::UnionMode mode = union_type.mode() == UnionMode::SPARSE
+ ? flatbuf::UnionMode_Sparse
+ : flatbuf::UnionMode_Dense;
+
+ std::vector<int32_t> type_ids;
+ type_ids.reserve(union_type.type_codes().size());
+ for (uint8_t code : union_type.type_codes()) {
+ type_ids.push_back(code);
+ }
+
+ auto fb_type_ids = fbb.CreateVector(type_ids);
+
+ *offset = flatbuf::CreateUnion(fbb, mode, fb_type_ids).Union();
+ return Status::OK();
+}
+
+#define INT_TO_FB_CASE(BIT_WIDTH, IS_SIGNED) \
+ *out_type = flatbuf::Type_Int; \
+ *offset = IntToFlatbuffer(fbb, BIT_WIDTH, IS_SIGNED); \
+ break;
+
+static inline flatbuf::TimeUnit ToFlatbufferUnit(TimeUnit::type unit) {
+ switch (unit) {
+ case TimeUnit::SECOND:
+ return flatbuf::TimeUnit_SECOND;
+ case TimeUnit::MILLI:
+ return flatbuf::TimeUnit_MILLISECOND;
+ case TimeUnit::MICRO:
+ return flatbuf::TimeUnit_MICROSECOND;
+ case TimeUnit::NANO:
+ return flatbuf::TimeUnit_NANOSECOND;
+ default:
+ break;
+ }
+ return flatbuf::TimeUnit_MIN;
+}
+
+static inline TimeUnit::type FromFlatbufferUnit(flatbuf::TimeUnit unit) {
+ switch (unit) {
+ case flatbuf::TimeUnit_SECOND:
+ return TimeUnit::SECOND;
+ case flatbuf::TimeUnit_MILLISECOND:
+ return TimeUnit::MILLI;
+ case flatbuf::TimeUnit_MICROSECOND:
+ return TimeUnit::MICRO;
+ case flatbuf::TimeUnit_NANOSECOND:
+ return TimeUnit::NANO;
+ default:
+ break;
+ }
+ // cannot reach
+ return TimeUnit::SECOND;
+}
+
+static Status TypeFromFlatbuffer(flatbuf::Type type, const void* type_data,
+ const std::vector<std::shared_ptr<Field>>& children,
+ std::shared_ptr<DataType>* out) {
+ switch (type) {
+ case flatbuf::Type_NONE:
+ return Status::Invalid("Type metadata cannot be none");
+ case flatbuf::Type_Int:
+ return IntFromFlatbuffer(static_cast<const flatbuf::Int*>(type_data), out);
+ case flatbuf::Type_FloatingPoint:
+ return FloatFromFlatuffer(static_cast<const flatbuf::FloatingPoint*>(type_data),
+ out);
+ case flatbuf::Type_Binary:
+ *out = binary();
+ return Status::OK();
+ case flatbuf::Type_FixedSizeBinary: {
+ auto fw_binary = static_cast<const flatbuf::FixedSizeBinary*>(type_data);
+ *out = fixed_size_binary(fw_binary->byteWidth());
+ return Status::OK();
+ }
+ case flatbuf::Type_Utf8:
+ *out = utf8();
+ return Status::OK();
+ case flatbuf::Type_Bool:
+ *out = boolean();
+ return Status::OK();
+ case flatbuf::Type_Decimal:
+ return Status::NotImplemented("Decimal");
+ case flatbuf::Type_Date: {
+ auto date_type = static_cast<const flatbuf::Date*>(type_data);
+ if (date_type->unit() == flatbuf::DateUnit_DAY) {
+ *out = date32();
+ } else {
+ *out = date64();
+ }
+ return Status::OK();
+ }
+ case flatbuf::Type_Time: {
+ auto time_type = static_cast<const flatbuf::Time*>(type_data);
+ TimeUnit::type unit = FromFlatbufferUnit(time_type->unit());
+ int32_t bit_width = time_type->bitWidth();
+ switch (unit) {
+ case TimeUnit::SECOND:
+ case TimeUnit::MILLI:
+ if (bit_width != 32) {
+ return Status::Invalid("Time is 32 bits for second/milli unit");
+ }
+ *out = time32(unit);
+ break;
+ default:
+ if (bit_width != 64) {
+ return Status::Invalid("Time is 64 bits for micro/nano unit");
+ }
+ *out = time64(unit);
+ break;
+ }
+ return Status::OK();
+ }
+ case flatbuf::Type_Timestamp: {
+ auto ts_type = static_cast<const flatbuf::Timestamp*>(type_data);
+ TimeUnit::type unit = FromFlatbufferUnit(ts_type->unit());
+ if (ts_type->timezone() != 0 && ts_type->timezone()->Length() > 0) {
+ *out = timestamp(unit, ts_type->timezone()->str());
+ } else {
+ *out = timestamp(unit);
+ }
+ return Status::OK();
+ }
+ case flatbuf::Type_Interval:
+ return Status::NotImplemented("Interval");
+ case flatbuf::Type_List:
+ if (children.size() != 1) {
+ return Status::Invalid("List must have exactly 1 child field");
+ }
+ *out = std::make_shared<ListType>(children[0]);
+ return Status::OK();
+ case flatbuf::Type_Struct_:
+ *out = std::make_shared<StructType>(children);
+ return Status::OK();
+ case flatbuf::Type_Union:
+ return UnionFromFlatbuffer(static_cast<const flatbuf::Union*>(type_data), children,
+ out);
+ default:
+ return Status::Invalid("Unrecognized type");
+ }
+}
+
+// TODO(wesm): Convert this to visitor pattern
+static Status TypeToFlatbuffer(FBB& fbb, const std::shared_ptr<DataType>& type,
+ std::vector<FieldOffset>* children,
+ std::vector<VectorLayoutOffset>* layout,
+ flatbuf::Type* out_type, DictionaryMemo* dictionary_memo,
+ Offset* offset) {
+ if (type->id() == Type::DICTIONARY) {
+ // In this library, the dictionary "type" is a logical construct. Here we
+ // pass through to the value type, as we've already captured the index
+ // type in the DictionaryEncoding metadata in the parent field
+ const auto& dict_type = static_cast<const DictionaryType&>(*type);
+ return TypeToFlatbuffer(fbb, dict_type.dictionary()->type(), children, layout,
+ out_type, dictionary_memo, offset);
+ }
+
+ std::vector<BufferDescr> buffer_layout = type->GetBufferLayout();
+ for (const BufferDescr& descr : buffer_layout) {
+ flatbuf::VectorType vector_type;
+ switch (descr.type()) {
+ case BufferType::OFFSET:
+ vector_type = flatbuf::VectorType_OFFSET;
+ break;
+ case BufferType::DATA:
+ vector_type = flatbuf::VectorType_DATA;
+ break;
+ case BufferType::VALIDITY:
+ vector_type = flatbuf::VectorType_VALIDITY;
+ break;
+ case BufferType::TYPE:
+ vector_type = flatbuf::VectorType_TYPE;
+ break;
+ default:
+ vector_type = flatbuf::VectorType_DATA;
+ break;
+ }
+ auto offset = flatbuf::CreateVectorLayout(
+ fbb, static_cast<int16_t>(descr.bit_width()), vector_type);
+ layout->push_back(offset);
+ }
+
+ switch (type->id()) {
+ case Type::BOOL:
+ *out_type = flatbuf::Type_Bool;
+ *offset = flatbuf::CreateBool(fbb).Union();
+ break;
+ case Type::UINT8:
+ INT_TO_FB_CASE(8, false);
+ case Type::INT8:
+ INT_TO_FB_CASE(8, true);
+ case Type::UINT16:
+ INT_TO_FB_CASE(16, false);
+ case Type::INT16:
+ INT_TO_FB_CASE(16, true);
+ case Type::UINT32:
+ INT_TO_FB_CASE(32, false);
+ case Type::INT32:
+ INT_TO_FB_CASE(32, true);
+ case Type::UINT64:
+ INT_TO_FB_CASE(64, false);
+ case Type::INT64:
+ INT_TO_FB_CASE(64, true);
+ case Type::FLOAT:
+ *out_type = flatbuf::Type_FloatingPoint;
+ *offset = FloatToFlatbuffer(fbb, flatbuf::Precision_SINGLE);
+ break;
+ case Type::DOUBLE:
+ *out_type = flatbuf::Type_FloatingPoint;
+ *offset = FloatToFlatbuffer(fbb, flatbuf::Precision_DOUBLE);
+ break;
+ case Type::FIXED_SIZE_BINARY: {
+ const auto& fw_type = static_cast<const FixedSizeBinaryType&>(*type);
+ *out_type = flatbuf::Type_FixedSizeBinary;
+ *offset = flatbuf::CreateFixedSizeBinary(fbb, fw_type.byte_width()).Union();
+ } break;
+ case Type::BINARY:
+ *out_type = flatbuf::Type_Binary;
+ *offset = flatbuf::CreateBinary(fbb).Union();
+ break;
+ case Type::STRING:
+ *out_type = flatbuf::Type_Utf8;
+ *offset = flatbuf::CreateUtf8(fbb).Union();
+ break;
+ case Type::DATE32:
+ *out_type = flatbuf::Type_Date;
+ *offset = flatbuf::CreateDate(fbb, flatbuf::DateUnit_DAY).Union();
+ break;
+ case Type::DATE64:
+ *out_type = flatbuf::Type_Date;
+ *offset = flatbuf::CreateDate(fbb, flatbuf::DateUnit_MILLISECOND).Union();
+ break;
+ case Type::TIME32: {
+ const auto& time_type = static_cast<const Time32Type&>(*type);
+ *out_type = flatbuf::Type_Time;
+ *offset = flatbuf::CreateTime(fbb, ToFlatbufferUnit(time_type.unit()), 32).Union();
+ } break;
+ case Type::TIME64: {
+ const auto& time_type = static_cast<const Time64Type&>(*type);
+ *out_type = flatbuf::Type_Time;
+ *offset = flatbuf::CreateTime(fbb, ToFlatbufferUnit(time_type.unit()), 64).Union();
+ } break;
+ case Type::TIMESTAMP: {
+ const auto& ts_type = static_cast<const TimestampType&>(*type);
+ *out_type = flatbuf::Type_Timestamp;
+
+ flatbuf::TimeUnit fb_unit = ToFlatbufferUnit(ts_type.unit());
+ FBString fb_timezone = 0;
+ if (ts_type.timezone().size() > 0) {
+ fb_timezone = fbb.CreateString(ts_type.timezone());
+ }
+ *offset = flatbuf::CreateTimestamp(fbb, fb_unit, fb_timezone).Union();
+ } break;
+ case Type::LIST:
+ *out_type = flatbuf::Type_List;
+ return ListToFlatbuffer(fbb, type, children, dictionary_memo, offset);
+ case Type::STRUCT:
+ *out_type = flatbuf::Type_Struct_;
+ return StructToFlatbuffer(fbb, type, children, dictionary_memo, offset);
+ case Type::UNION:
+ *out_type = flatbuf::Type_Union;
+ return UnionToFlatBuffer(fbb, type, children, dictionary_memo, offset);
+ default:
+ *out_type = flatbuf::Type_NONE; // Make clang-tidy happy
+ std::stringstream ss;
+ ss << "Unable to convert type: " << type->ToString() << std::endl;
+ return Status::NotImplemented(ss.str());
+ }
+ return Status::OK();
+}
+
+static Status TensorTypeToFlatbuffer(FBB& fbb, const std::shared_ptr<DataType>& type,
+ flatbuf::Type* out_type, Offset* offset) {
+ switch (type->id()) {
+ case Type::UINT8:
+ INT_TO_FB_CASE(8, false);
+ case Type::INT8:
+ INT_TO_FB_CASE(8, true);
+ case Type::UINT16:
+ INT_TO_FB_CASE(16, false);
+ case Type::INT16:
+ INT_TO_FB_CASE(16, true);
+ case Type::UINT32:
+ INT_TO_FB_CASE(32, false);
+ case Type::INT32:
+ INT_TO_FB_CASE(32, true);
+ case Type::UINT64:
+ INT_TO_FB_CASE(64, false);
+ case Type::INT64:
+ INT_TO_FB_CASE(64, true);
+ case Type::HALF_FLOAT:
+ *out_type = flatbuf::Type_FloatingPoint;
+ *offset = FloatToFlatbuffer(fbb, flatbuf::Precision_HALF);
+ break;
+ case Type::FLOAT:
+ *out_type = flatbuf::Type_FloatingPoint;
+ *offset = FloatToFlatbuffer(fbb, flatbuf::Precision_SINGLE);
+ break;
+ case Type::DOUBLE:
+ *out_type = flatbuf::Type_FloatingPoint;
+ *offset = FloatToFlatbuffer(fbb, flatbuf::Precision_DOUBLE);
+ break;
+ default:
+ *out_type = flatbuf::Type_NONE; // Make clang-tidy happy
+ std::stringstream ss;
+ ss << "Unable to convert type: " << type->ToString() << std::endl;
+ return Status::NotImplemented(ss.str());
+ }
+ return Status::OK();
+}
+
+static DictionaryOffset GetDictionaryEncoding(FBB& fbb, const DictionaryType& type,
+ DictionaryMemo* memo) {
+ int64_t dictionary_id = memo->GetId(type.dictionary());
+
+ // We assume that the dictionary index type (as an integer) has already been
+ // validated elsewhere, and can safely assume we are dealing with signed
+ // integers
+ const auto& fw_index_type = static_cast<const FixedWidthType&>(*type.index_type());
+
+ auto index_type_offset = flatbuf::CreateInt(fbb, fw_index_type.bit_width(), true);
+
+ // TODO(wesm): ordered dictionaries
+ return flatbuf::CreateDictionaryEncoding(fbb, dictionary_id, index_type_offset,
+ type.ordered());
+}
+
+static Status FieldToFlatbuffer(FBB& fbb, const std::shared_ptr<Field>& field,
+ DictionaryMemo* dictionary_memo, FieldOffset* offset) {
+ auto fb_name = fbb.CreateString(field->name());
+
+ flatbuf::Type type_enum;
+ Offset type_offset;
+ Offset type_layout;
+ std::vector<FieldOffset> children;
+ std::vector<VectorLayoutOffset> layout;
+
+ RETURN_NOT_OK(TypeToFlatbuffer(fbb, field->type(), &children, &layout, &type_enum,
+ dictionary_memo, &type_offset));
+ auto fb_children = fbb.CreateVector(children);
+ auto fb_layout = fbb.CreateVector(layout);
+
+ DictionaryOffset dictionary = 0;
+ if (field->type()->id() == Type::DICTIONARY) {
+ dictionary = GetDictionaryEncoding(
+ fbb, static_cast<const DictionaryType&>(*field->type()), dictionary_memo);
+ }
+
+ // TODO: produce the list of VectorTypes
+ *offset = flatbuf::CreateField(fbb, fb_name, field->nullable(), type_enum, type_offset,
+ dictionary, fb_children, fb_layout);
+
+ return Status::OK();
+}
+
+static Status FieldFromFlatbuffer(const flatbuf::Field* field,
+ const DictionaryMemo& dictionary_memo,
+ std::shared_ptr<Field>* out) {
+ std::shared_ptr<DataType> type;
+
+ const flatbuf::DictionaryEncoding* encoding = field->dictionary();
+
+ if (encoding == nullptr) {
+ // The field is not dictionary encoded. We must potentially visit its
+ // children to fully reconstruct the data type
+ auto children = field->children();
+ std::vector<std::shared_ptr<Field>> child_fields(children->size());
+ for (int i = 0; i < static_cast<int>(children->size()); ++i) {
+ RETURN_NOT_OK(
+ FieldFromFlatbuffer(children->Get(i), dictionary_memo, &child_fields[i]));
+ }
+ RETURN_NOT_OK(
+ TypeFromFlatbuffer(field->type_type(), field->type(), child_fields, &type));
+ } else {
+ // The field is dictionary encoded. The type of the dictionary values has
+ // been determined elsewhere, and is stored in the DictionaryMemo. Here we
+ // construct the logical DictionaryType object
+
+ std::shared_ptr<Array> dictionary;
+ RETURN_NOT_OK(dictionary_memo.GetDictionary(encoding->id(), &dictionary));
+
+ std::shared_ptr<DataType> index_type;
+ RETURN_NOT_OK(IntFromFlatbuffer(encoding->indexType(), &index_type));
+ type = ::arrow::dictionary(index_type, dictionary, encoding->isOrdered());
+ }
+ *out = std::make_shared<Field>(field->name()->str(), type, field->nullable());
+ return Status::OK();
+}
+
+static Status FieldFromFlatbufferDictionary(const flatbuf::Field* field,
+ std::shared_ptr<Field>* out) {
+ // Need an empty memo to pass down for constructing children
+ DictionaryMemo dummy_memo;
+
+ // Any DictionaryEncoding set is ignored here
+
+ std::shared_ptr<DataType> type;
+ auto children = field->children();
+ std::vector<std::shared_ptr<Field>> child_fields(children->size());
+ for (int i = 0; i < static_cast<int>(children->size()); ++i) {
+ RETURN_NOT_OK(FieldFromFlatbuffer(children->Get(i), dummy_memo, &child_fields[i]));
+ }
+
+ RETURN_NOT_OK(
+ TypeFromFlatbuffer(field->type_type(), field->type(), child_fields, &type));
+
+ *out = std::make_shared<Field>(field->name()->str(), type, field->nullable());
+ return Status::OK();
+}
+
+// will return the endianness of the system we are running on
+// based the NUMPY_API function. See NOTICE.txt
+flatbuf::Endianness endianness() {
+ union {
+ uint32_t i;
+ char c[4];
+ } bint = {0x01020304};
+
+ return bint.c[0] == 1 ? flatbuf::Endianness_Big : flatbuf::Endianness_Little;
+}
+
+static Status SchemaToFlatbuffer(FBB& fbb, const Schema& schema,
+ DictionaryMemo* dictionary_memo,
+ flatbuffers::Offset<flatbuf::Schema>* out) {
+ /// Fields
+ std::vector<FieldOffset> field_offsets;
+ for (int i = 0; i < schema.num_fields(); ++i) {
+ std::shared_ptr<Field> field = schema.field(i);
+ FieldOffset offset;
+ RETURN_NOT_OK(FieldToFlatbuffer(fbb, field, dictionary_memo, &offset));
+ field_offsets.push_back(offset);
+ }
+
+ auto fb_offsets = fbb.CreateVector(field_offsets);
+
+ /// Custom metadata
+ const KeyValueMetadata* metadata = schema.metadata().get();
+
+ if (metadata != nullptr) {
+ std::vector<KeyValueOffset> key_value_offsets;
+ size_t metadata_size = metadata->size();
+ key_value_offsets.reserve(metadata_size);
+ for (size_t i = 0; i < metadata_size; ++i) {
+ const auto& key = metadata->key(i);
+ const auto& value = metadata->value(i);
+ key_value_offsets.push_back(
+ flatbuf::CreateKeyValue(fbb, fbb.CreateString(key), fbb.CreateString(value)));
+ }
+ *out = flatbuf::CreateSchema(fbb, endianness(), fb_offsets,
+ fbb.CreateVector(key_value_offsets));
+ } else {
+ *out = flatbuf::CreateSchema(fbb, endianness(), fb_offsets);
+ }
+
+ return Status::OK();
+}
+
+static Status WriteFlatbufferBuilder(FBB& fbb, std::shared_ptr<Buffer>* out) {
+ int32_t size = fbb.GetSize();
+
+ auto result = std::make_shared<PoolBuffer>();
+ RETURN_NOT_OK(result->Resize(size));
+
+ uint8_t* dst = result->mutable_data();
+ memcpy(dst, fbb.GetBufferPointer(), size);
+ *out = result;
+ return Status::OK();
+}
+
+static Status WriteFBMessage(FBB& fbb, flatbuf::MessageHeader header_type,
+ flatbuffers::Offset<void> header, int64_t body_length,
+ std::shared_ptr<Buffer>* out) {
+ auto message = flatbuf::CreateMessage(fbb, kCurrentMetadataVersion, header_type, header,
+ body_length);
+ fbb.Finish(message);
+ return WriteFlatbufferBuilder(fbb, out);
+}
+
+Status WriteSchemaMessage(const Schema& schema, DictionaryMemo* dictionary_memo,
+ std::shared_ptr<Buffer>* out) {
+ FBB fbb;
+ flatbuffers::Offset<flatbuf::Schema> fb_schema;
+ RETURN_NOT_OK(SchemaToFlatbuffer(fbb, schema, dictionary_memo, &fb_schema));
+ return WriteFBMessage(fbb, flatbuf::MessageHeader_Schema, fb_schema.Union(), 0, out);
+}
+
+using FieldNodeVector =
+ flatbuffers::Offset<flatbuffers::Vector<const flatbuf::FieldNode*>>;
+using BufferVector = flatbuffers::Offset<flatbuffers::Vector<const flatbuf::Buffer*>>;
+
+static Status WriteFieldNodes(FBB& fbb, const std::vector<FieldMetadata>& nodes,
+ FieldNodeVector* out) {
+ std::vector<flatbuf::FieldNode> fb_nodes;
+ fb_nodes.reserve(nodes.size());
+
+ for (size_t i = 0; i < nodes.size(); ++i) {
+ const FieldMetadata& node = nodes[i];
+ if (node.offset != 0) {
+ return Status::Invalid("Field metadata for IPC must have offset 0");
+ }
+ fb_nodes.emplace_back(node.length, node.null_count);
+ }
+ *out = fbb.CreateVectorOfStructs(fb_nodes);
+ return Status::OK();
+}
+
+static Status WriteBuffers(FBB& fbb, const std::vector<BufferMetadata>& buffers,
+ BufferVector* out) {
+ std::vector<flatbuf::Buffer> fb_buffers;
+ fb_buffers.reserve(buffers.size());
+
+ for (size_t i = 0; i < buffers.size(); ++i) {
+ const BufferMetadata& buffer = buffers[i];
+ fb_buffers.emplace_back(buffer.page, buffer.offset, buffer.length);
+ }
+ *out = fbb.CreateVectorOfStructs(fb_buffers);
+ return Status::OK();
+}
+
+static Status MakeRecordBatch(FBB& fbb, int64_t length, int64_t body_length,
+ const std::vector<FieldMetadata>& nodes,
+ const std::vector<BufferMetadata>& buffers,
+ RecordBatchOffset* offset) {
+ FieldNodeVector fb_nodes;
+ BufferVector fb_buffers;
+
+ RETURN_NOT_OK(WriteFieldNodes(fbb, nodes, &fb_nodes));
+ RETURN_NOT_OK(WriteBuffers(fbb, buffers, &fb_buffers));
+
+ *offset = flatbuf::CreateRecordBatch(fbb, length, fb_nodes, fb_buffers);
+ return Status::OK();
+}
+
+Status WriteRecordBatchMessage(int64_t length, int64_t body_length,
+ const std::vector<FieldMetadata>& nodes,
+ const std::vector<BufferMetadata>& buffers,
+ std::shared_ptr<Buffer>* out) {
+ FBB fbb;
+ RecordBatchOffset record_batch;
+ RETURN_NOT_OK(MakeRecordBatch(fbb, length, body_length, nodes, buffers, &record_batch));
+ return WriteFBMessage(fbb, flatbuf::MessageHeader_RecordBatch, record_batch.Union(),
+ body_length, out);
+}
+
+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>;
+
+ FBB fbb;
+
+ flatbuf::Type fb_type_type;
+ Offset fb_type;
+ RETURN_NOT_OK(TensorTypeToFlatbuffer(fbb, tensor.type(), &fb_type_type, &fb_type));
+
+ std::vector<TensorDimOffset> dims;
+ for (int i = 0; i < tensor.ndim(); ++i) {
+ FBString name = fbb.CreateString(tensor.dim_name(i));
+ dims.push_back(flatbuf::CreateTensorDim(fbb, tensor.shape()[i], name));
+ }
+
+ auto fb_shape = fbb.CreateVector(dims);
+ auto fb_strides = fbb.CreateVector(tensor.strides());
+ int64_t body_length = tensor.data()->size();
+ flatbuf::Buffer buffer(-1, buffer_start_offset, body_length);
+
+ TensorOffset fb_tensor =
+ flatbuf::CreateTensor(fbb, fb_type_type, fb_type, fb_shape, fb_strides, &buffer);
+
+ return WriteFBMessage(fbb, flatbuf::MessageHeader_Tensor, fb_tensor.Union(),
+ body_length, out);
+}
+
+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) {
+ FBB fbb;
+ RecordBatchOffset record_batch;
+ RETURN_NOT_OK(MakeRecordBatch(fbb, length, body_length, nodes, buffers, &record_batch));
+ auto dictionary_batch = flatbuf::CreateDictionaryBatch(fbb, id, record_batch).Union();
+ return WriteFBMessage(fbb, flatbuf::MessageHeader_DictionaryBatch, dictionary_batch,
+ body_length, out);
+}
+
+static flatbuffers::Offset<flatbuffers::Vector<const flatbuf::Block*>>
+FileBlocksToFlatbuffer(FBB& fbb, const std::vector<FileBlock>& blocks) {
+ std::vector<flatbuf::Block> fb_blocks;
+
+ for (const FileBlock& block : blocks) {
+ fb_blocks.emplace_back(block.offset, block.metadata_length, block.body_length);
+ }
+
+ return fbb.CreateVectorOfStructs(fb_blocks);
+}
+
+Status WriteFileFooter(const Schema& schema, const std::vector<FileBlock>& dictionaries,
+ const std::vector<FileBlock>& record_batches,
+ DictionaryMemo* dictionary_memo, io::OutputStream* out) {
+ FBB fbb;
+
+ flatbuffers::Offset<flatbuf::Schema> fb_schema;
+ RETURN_NOT_OK(SchemaToFlatbuffer(fbb, schema, dictionary_memo, &fb_schema));
+
+#ifndef NDEBUG
+ for (size_t i = 0; i < dictionaries.size(); ++i) {
+ DCHECK(BitUtil::IsMultipleOf8(dictionaries[i].offset)) << i;
+ DCHECK(BitUtil::IsMultipleOf8(dictionaries[i].metadata_length)) << i;
+ DCHECK(BitUtil::IsMultipleOf8(dictionaries[i].body_length)) << i;
+ }
+
+ for (size_t i = 0; i < record_batches.size(); ++i) {
+ DCHECK(BitUtil::IsMultipleOf8(record_batches[i].offset)) << i;
+ DCHECK(BitUtil::IsMultipleOf8(record_batches[i].metadata_length)) << i;
+ DCHECK(BitUtil::IsMultipleOf8(record_batches[i].body_length)) << i;
+ }
+#endif
+
+ auto fb_dictionaries = FileBlocksToFlatbuffer(fbb, dictionaries);
+ auto fb_record_batches = FileBlocksToFlatbuffer(fbb, record_batches);
+
+ auto footer = flatbuf::CreateFooter(fbb, kCurrentMetadataVersion, fb_schema,
+ fb_dictionaries, fb_record_batches);
+
+ fbb.Finish(footer);
+
+ int32_t size = fbb.GetSize();
+
+ return out->Write(fbb.GetBufferPointer(), size);
+}
+
+// ----------------------------------------------------------------------
+
+static Status VisitField(const flatbuf::Field* field, DictionaryTypeMap* id_to_field) {
+ const flatbuf::DictionaryEncoding* dict_metadata = field->dictionary();
+ if (dict_metadata == nullptr) {
+ // Field is not dictionary encoded. Visit children
+ auto children = field->children();
+ for (flatbuffers::uoffset_t i = 0; i < children->size(); ++i) {
+ RETURN_NOT_OK(VisitField(children->Get(i), id_to_field));
+ }
+ } else {
+ // Field is dictionary encoded. Construct the data type for the
+ // dictionary (no descendents can be dictionary encoded)
+ std::shared_ptr<Field> dictionary_field;
+ RETURN_NOT_OK(FieldFromFlatbufferDictionary(field, &dictionary_field));
+ (*id_to_field)[dict_metadata->id()] = dictionary_field;
+ }
+ return Status::OK();
+}
+
+Status GetDictionaryTypes(const void* opaque_schema, DictionaryTypeMap* id_to_field) {
+ auto schema = static_cast<const flatbuf::Schema*>(opaque_schema);
+ int num_fields = static_cast<int>(schema->fields()->size());
+ for (int i = 0; i < num_fields; ++i) {
+ RETURN_NOT_OK(VisitField(schema->fields()->Get(i), id_to_field));
+ }
+ return Status::OK();
+}
+
+Status GetSchema(const void* opaque_schema, const DictionaryMemo& dictionary_memo,
+ std::shared_ptr<Schema>* out) {
+ auto schema = static_cast<const flatbuf::Schema*>(opaque_schema);
+ int num_fields = static_cast<int>(schema->fields()->size());
+
+ std::vector<std::shared_ptr<Field>> fields(num_fields);
+ for (int i = 0; i < num_fields; ++i) {
+ const flatbuf::Field* field = schema->fields()->Get(i);
+ RETURN_NOT_OK(FieldFromFlatbuffer(field, dictionary_memo, &fields[i]));
+ }
+
+ auto metadata = std::make_shared<KeyValueMetadata>();
+ auto fb_metadata = schema->custom_metadata();
+ if (fb_metadata != nullptr) {
+ metadata->reserve(fb_metadata->size());
+ for (const auto& pair : *fb_metadata) {
+ metadata->Append(pair->key()->str(), pair->value()->str());
+ }
+ }
+
+ *out = ::arrow::schema(std::move(fields), metadata);
+ return Status::OK();
+}
+
+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 message = flatbuf::GetMessage(metadata.data());
+ auto tensor = reinterpret_cast<const flatbuf::Tensor*>(message->header());
+
+ int ndim = static_cast<int>(tensor->shape()->size());
+
+ for (int i = 0; i < ndim; ++i) {
+ auto dim = tensor->shape()->Get(i);
+
+ shape->push_back(dim->size());
+ auto fb_name = dim->name();
+ if (fb_name == 0) {
+ dim_names->push_back("");
+ } else {
+ dim_names->push_back(fb_name->str());
+ }
+ }
+
+ if (tensor->strides()->size() > 0) {
+ for (int i = 0; i < ndim; ++i) {
+ strides->push_back(tensor->strides()->Get(i));
+ }
+ }
+
+ return TypeFromFlatbuffer(tensor->type_type(), tensor->type(), {}, type);
+}
+
+// ----------------------------------------------------------------------
+// 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
+ // end on an 8-byte offset
+ int64_t start_offset;
+ RETURN_NOT_OK(file->Tell(&start_offset));
+
+ int32_t padded_message_length = static_cast<int32_t>(message.size()) + 4;
+ const int32_t remainder =
+ (padded_message_length + static_cast<int32_t>(start_offset)) % 8;
+ if (remainder != 0) {
+ padded_message_length += 8 - remainder;
+ }
+
+ // The returned message size includes the length prefix, the flatbuffer,
+ // plus padding
+ *message_length = padded_message_length;
+
+ // Write the flatbuffer size prefix including padding
+ int32_t flatbuffer_size = padded_message_length - 4;
+ RETURN_NOT_OK(
+ file->Write(reinterpret_cast<const uint8_t*>(&flatbuffer_size), sizeof(int32_t)));
+
+ // Write the flatbuffer
+ RETURN_NOT_OK(file->Write(message.data(), message.size()));
+
+ // Write any padding
+ int32_t padding = padded_message_length - static_cast<int32_t>(message.size()) - 4;
+ if (padding > 0) {
+ RETURN_NOT_OK(file->Write(kPaddingBytes, padding));
+ }
+
+ return Status::OK();
+}
+
+} // namespace ipc
+} // namespace arrow
http://git-wip-us.apache.org/repos/asf/arrow/blob/f50f2eac/cpp/src/arrow/ipc/metadata-internal.h
----------------------------------------------------------------------
diff --git a/cpp/src/arrow/ipc/metadata-internal.h b/cpp/src/arrow/ipc/metadata-internal.h
new file mode 100644
index 0000000..738ed3c
--- /dev/null
+++ b/cpp/src/arrow/ipc/metadata-internal.h
@@ -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.
+
+// Internal metadata serialization matters
+
+#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