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

Change-Id: I3fbacc6a0f5d02b734a04b9717d426cb1b89ace2

Author: wesm

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
   )

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")

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"
 

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

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

cpp/src/arrow/ipc/file-to-stream.cc

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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"