[C++][Flight] Add basic Flight service (#124)

cpp/code/CMakeLists.txt

@@ -7,7 +7,7 @@ if (CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
 endif()
 
 # Add Arrow
-find_package(Arrow REQUIRED COMPONENTS dataset parquet)
+find_package(Arrow REQUIRED COMPONENTS dataset flight parquet)
 
 if (CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
     set(CMAKE_CXX_CLANG_TIDY "clang-tidy")
@@ -29,6 +29,7 @@ function(RECIPE TARGET)
             ${TARGET}
             arrow_shared
             arrow_dataset
+            arrow_flight
             parquet
             gtest
     )
@@ -44,3 +45,4 @@ endfunction()
 recipe(basic_arrow)
 recipe(creating_arrow_objects)
 recipe(datasets)
+recipe(flight)

cpp/code/common.cc

@@ -117,7 +117,7 @@ void PopulateMap(const arrow::Table& table,
   std::shared_ptr<arrow::StringArray> table_outputs =
       std::dynamic_pointer_cast<arrow::StringArray>(table.column(1)->chunk(0));
   for (int64_t i = 0; i < table.num_rows(); i++) {
-    values->insert({table_names->GetString(i), table_outputs->GetString(i)});
+    (*values)[table_names->GetString(i)] = table_outputs->GetString(i);
   }
 }
 

cpp/code/flight.cc

@@ -0,0 +1,294 @@
+// 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/buffer.h>
+#include <arrow/filesystem/filesystem.h>
+#include <arrow/filesystem/localfs.h>
+#include <arrow/flight/client.h>
+#include <arrow/flight/server.h>
+#include <arrow/pretty_print.h>
+#include <arrow/result.h>
+#include <arrow/status.h>
+#include <arrow/table.h>
+#include <arrow/type.h>
+#include <gtest/gtest.h>
+#include <parquet/arrow/reader.h>
+#include <parquet/arrow/writer.h>
+
+#include <algorithm>
+#include <memory>
+#include <numeric>
+#include <vector>
+
+#include "common.h"
+
+class ParquetStorageService : public arrow::flight::FlightServerBase {
+ public:
+  const arrow::flight::ActionType kActionDropDataset{"drop_dataset", "Delete a dataset."};
+
+  explicit ParquetStorageService(std::shared_ptr<arrow::fs::FileSystem> root)
+      : root_(std::move(root)) {}
+
+  arrow::Status ListFlights(const arrow::flight::ServerCallContext&,
+                            const arrow::flight::Criteria*,
+                            std::unique_ptr<arrow::flight::FlightListing>* listings) {
+    arrow::fs::FileSelector selector;
+    selector.base_dir = "/";
+    ARROW_ASSIGN_OR_RAISE(auto listing, root_->GetFileInfo(selector));
+
+    std::vector<arrow::flight::FlightInfo> flights;
+    for (const auto& file_info : listing) {
+      if (!file_info.IsFile() || file_info.extension() != "parquet") continue;
+      ARROW_ASSIGN_OR_RAISE(auto info, MakeFlightInfo(file_info));
+      flights.push_back(std::move(info));
+    }
+
+    *listings = std::unique_ptr<arrow::flight::FlightListing>(
+        new arrow::flight::SimpleFlightListing(std::move(flights)));
+    return arrow::Status::OK();
+  }
+
+  arrow::Status GetFlightInfo(const arrow::flight::ServerCallContext&,
+                              const arrow::flight::FlightDescriptor& descriptor,
+                              std::unique_ptr<arrow::flight::FlightInfo>* info) {
+    ARROW_ASSIGN_OR_RAISE(auto file_info, FileInfoFromDescriptor(descriptor));
+    ARROW_ASSIGN_OR_RAISE(auto flight_info, MakeFlightInfo(file_info));
+    *info = std::unique_ptr<arrow::flight::FlightInfo>(
+        new arrow::flight::FlightInfo(std::move(flight_info)));
+    return arrow::Status::OK();
+  }
+
+  arrow::Status DoPut(const arrow::flight::ServerCallContext&,
+                      std::unique_ptr<arrow::flight::FlightMessageReader> reader,
+                      std::unique_ptr<arrow::flight::FlightMetadataWriter>) {
+    ARROW_ASSIGN_OR_RAISE(auto file_info, FileInfoFromDescriptor(reader->descriptor()));
+    ARROW_ASSIGN_OR_RAISE(auto sink, root_->OpenOutputStream(file_info.path()));
+    std::shared_ptr<arrow::Table> table;
+    ARROW_RETURN_NOT_OK(reader->ReadAll(&table));
+
+    ARROW_RETURN_NOT_OK(parquet::arrow::WriteTable(*table, arrow::default_memory_pool(),
+                                                   sink, /*chunk_size=*/65536));
+    return arrow::Status::OK();
+  }
+
+  arrow::Status DoGet(const arrow::flight::ServerCallContext&,
+                      const arrow::flight::Ticket& request,
+                      std::unique_ptr<arrow::flight::FlightDataStream>* stream) {
+    ARROW_ASSIGN_OR_RAISE(auto input, root_->OpenInputFile(request.ticket));
+    std::unique_ptr<parquet::arrow::FileReader> reader;
+    ARROW_RETURN_NOT_OK(parquet::arrow::OpenFile(std::move(input),
+                                                 arrow::default_memory_pool(), &reader));
+
+    std::shared_ptr<arrow::Table> table;
+    ARROW_RETURN_NOT_OK(reader->ReadTable(&table));
+    // Note that we can't directly pass TableBatchReader to
+    // RecordBatchStream because TableBatchReader keeps a non-owning
+    // reference to the underlying Table, which would then get freed
+    // when we exit this function
+    std::vector<std::shared_ptr<arrow::RecordBatch>> batches;
+    arrow::TableBatchReader batch_reader(*table);
+    ARROW_RETURN_NOT_OK(batch_reader.ReadAll(&batches));
+
+    ARROW_ASSIGN_OR_RAISE(auto owning_reader, arrow::RecordBatchReader::Make(
+                                                  std::move(batches), table->schema()));
+    *stream = std::unique_ptr<arrow::flight::FlightDataStream>(
+        new arrow::flight::RecordBatchStream(owning_reader));
+
+    return arrow::Status::OK();
+  }
+
+  arrow::Status ListActions(const arrow::flight::ServerCallContext&,
+                            std::vector<arrow::flight::ActionType>* actions) override {
+    *actions = {kActionDropDataset};
+    return arrow::Status::OK();
+  }
+
+  arrow::Status DoAction(const arrow::flight::ServerCallContext&,
+                         const arrow::flight::Action& action,
+                         std::unique_ptr<arrow::flight::ResultStream>* result) {
+    if (action.type == kActionDropDataset.type) {
+      *result = std::unique_ptr<arrow::flight::ResultStream>(
+          new arrow::flight::SimpleResultStream({}));
+      return DoActionDropDataset(action.body->ToString());
+    }
+    return arrow::Status::NotImplemented("Unknown action type: ", action.type);
+  }
+
+ private:
+  arrow::Result<arrow::flight::FlightInfo> MakeFlightInfo(
+      const arrow::fs::FileInfo& file_info) {
+    ARROW_ASSIGN_OR_RAISE(auto input, root_->OpenInputFile(file_info));
+    std::unique_ptr<parquet::arrow::FileReader> reader;
+    ARROW_RETURN_NOT_OK(parquet::arrow::OpenFile(std::move(input),
+                                                 arrow::default_memory_pool(), &reader));
+
+    std::shared_ptr<arrow::Schema> schema;
+    ARROW_RETURN_NOT_OK(reader->GetSchema(&schema));
+
+    auto descriptor = arrow::flight::FlightDescriptor::Path({file_info.base_name()});
+
+    arrow::flight::FlightEndpoint endpoint;
+    endpoint.ticket.ticket = file_info.base_name();
+    arrow::flight::Location location;
+    ARROW_RETURN_NOT_OK(
+        arrow::flight::Location::ForGrpcTcp("localhost", port(), &location));
+    endpoint.locations.push_back(location);
+
+    int64_t total_records = reader->parquet_reader()->metadata()->num_rows();
+    int64_t total_bytes = file_info.size();
+
+    return arrow::flight::FlightInfo::Make(*schema, descriptor, {endpoint}, total_records,
+                                           total_bytes);
+  }
+
+  arrow::Result<arrow::fs::FileInfo> FileInfoFromDescriptor(
+      const arrow::flight::FlightDescriptor& descriptor) {
+    if (descriptor.type != arrow::flight::FlightDescriptor::PATH) {
+      return arrow::Status::Invalid("Must provide PATH-type FlightDescriptor");
+    } else if (descriptor.path.size() != 1) {
+      return arrow::Status::Invalid(
+          "Must provide PATH-type FlightDescriptor with one path component");
+    }
+    return root_->GetFileInfo(descriptor.path[0]);
+  }
+
+  arrow::Status DoActionDropDataset(const std::string& key) {
+    return root_->DeleteFile(key);
+  }
+
+  std::shared_ptr<arrow::fs::FileSystem> root_;
+};  // end ParquetStorageService
+
+arrow::Status TestPutGetDelete() {
+  StartRecipe("ParquetStorageService::StartServer");
+  auto fs = std::make_shared<arrow::fs::LocalFileSystem>();
+  ARROW_RETURN_NOT_OK(fs->CreateDir("./flight_datasets/"));
+  ARROW_RETURN_NOT_OK(fs->DeleteDirContents("./flight_datasets/"));
+  auto root = std::make_shared<arrow::fs::SubTreeFileSystem>("./flight_datasets/", fs);
+
+  arrow::flight::Location server_location;
+  ARROW_RETURN_NOT_OK(
+      arrow::flight::Location::ForGrpcTcp("0.0.0.0", 0, &server_location));
+
+  arrow::flight::FlightServerOptions options(server_location);
+  auto server = std::unique_ptr<arrow::flight::FlightServerBase>(
+      new ParquetStorageService(std::move(root)));
+  ARROW_RETURN_NOT_OK(server->Init(options));
+  rout << "Listening on port " << server->port() << std::endl;
+  EndRecipe("ParquetStorageService::StartServer");
+
+  StartRecipe("ParquetStorageService::Connect");
+  arrow::flight::Location location;
+  ARROW_RETURN_NOT_OK(
+      arrow::flight::Location::ForGrpcTcp("localhost", server->port(), &location));
+
+  std::unique_ptr<arrow::flight::FlightClient> client;
+  ARROW_RETURN_NOT_OK(arrow::flight::FlightClient::Connect(location, &client));
+  rout << "Connected to " << location.ToString() << std::endl;
+  EndRecipe("ParquetStorageService::Connect");
+
+  StartRecipe("ParquetStorageService::DoPut");
+  // Open example data file to upload
+  ARROW_ASSIGN_OR_RAISE(std::string airquality_path,
+                        FindTestDataFile("airquality.parquet"));
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::io::RandomAccessFile> input,
+                        fs->OpenInputFile(airquality_path));
+  std::unique_ptr<parquet::arrow::FileReader> reader;
+  ARROW_RETURN_NOT_OK(
+      parquet::arrow::OpenFile(std::move(input), arrow::default_memory_pool(), &reader));
+
+  auto descriptor = arrow::flight::FlightDescriptor::Path({"airquality.parquet"});
+  std::shared_ptr<arrow::Schema> schema;
+  ARROW_RETURN_NOT_OK(reader->GetSchema(&schema));
+
+  // Start the RPC call
+  std::unique_ptr<arrow::flight::FlightStreamWriter> writer;
+  std::unique_ptr<arrow::flight::FlightMetadataReader> metadata_reader;
+  ARROW_RETURN_NOT_OK(client->DoPut(descriptor, schema, &writer, &metadata_reader));
+
+  // Upload data
+  std::shared_ptr<arrow::RecordBatchReader> batch_reader;
+  std::vector<int> row_groups(reader->num_row_groups());
+  std::iota(row_groups.begin(), row_groups.end(), 0);
+  ARROW_RETURN_NOT_OK(reader->GetRecordBatchReader(row_groups, &batch_reader));
+  int64_t batches = 0;
+  while (true) {
+    ARROW_ASSIGN_OR_RAISE(auto batch, batch_reader->Next());
+    if (!batch) break;
+    ARROW_RETURN_NOT_OK(writer->WriteRecordBatch(*batch));
+    batches++;
+  }
+
+  ARROW_RETURN_NOT_OK(writer->Close());
+  rout << "Wrote " << batches << " batches" << std::endl;
+  EndRecipe("ParquetStorageService::DoPut");
+
+  StartRecipe("ParquetStorageService::GetFlightInfo");
+  std::unique_ptr<arrow::flight::FlightInfo> flight_info;
+  ARROW_RETURN_NOT_OK(client->GetFlightInfo(descriptor, &flight_info));
+  rout << flight_info->descriptor().ToString() << std::endl;
+  rout << "=== Schema ===" << std::endl;
+  std::shared_ptr<arrow::Schema> info_schema;
+  arrow::ipc::DictionaryMemo dictionary_memo;
+  ARROW_RETURN_NOT_OK(flight_info->GetSchema(&dictionary_memo, &info_schema));
+  rout << info_schema->ToString() << std::endl;
+  rout << "==============" << std::endl;
+  EndRecipe("ParquetStorageService::GetFlightInfo");
+
+  StartRecipe("ParquetStorageService::DoGet");
+  std::unique_ptr<arrow::flight::FlightStreamReader> stream;
+  ARROW_RETURN_NOT_OK(client->DoGet(flight_info->endpoints()[0].ticket, &stream));
+  std::shared_ptr<arrow::Table> table;
+  ARROW_RETURN_NOT_OK(stream->ReadAll(&table));
+  arrow::PrettyPrintOptions print_options(/*indent=*/0, /*window=*/2);
+  arrow::PrettyPrint(*table, print_options, &rout);
+  EndRecipe("ParquetStorageService::DoGet");
+
+  StartRecipe("ParquetStorageService::DoAction");
+  arrow::flight::Action action{"drop_dataset",
+                               arrow::Buffer::FromString("airquality.parquet")};
+  std::unique_ptr<arrow::flight::ResultStream> results;
+  ARROW_RETURN_NOT_OK(client->DoAction(action, &results));
+  rout << "Deleted dataset" << std::endl;
+  EndRecipe("ParquetStorageService::DoAction");
+
+  StartRecipe("ParquetStorageService::ListFlights");
+  std::unique_ptr<arrow::flight::FlightListing> listing;
+  ARROW_RETURN_NOT_OK(client->ListFlights(&listing));
+  while (true) {
+    std::unique_ptr<arrow::flight::FlightInfo> flight_info;
+    ARROW_RETURN_NOT_OK(listing->Next(&flight_info));
+    if (!flight_info) break;
+    rout << flight_info->descriptor().ToString() << std::endl;
+    rout << "=== Schema ===" << std::endl;
+    std::shared_ptr<arrow::Schema> info_schema;
+    arrow::ipc::DictionaryMemo dictionary_memo;
+    ARROW_RETURN_NOT_OK(flight_info->GetSchema(&dictionary_memo, &info_schema));
+    rout << info_schema->ToString() << std::endl;
+    rout << "==============" << std::endl;
+  }
+  rout << "End of listing" << std::endl;
+  EndRecipe("ParquetStorageService::ListFlights");
+
+  StartRecipe("ParquetStorageService::StopServer");
+  ARROW_RETURN_NOT_OK(server->Shutdown());
+  rout << "Server shut down successfully" << std::endl;
+  EndRecipe("ParquetStorageService::StopServer");
+  return arrow::Status::OK();
+}
+
+TEST(ParquetStorageServiceTest, PutGetDelete) { ASSERT_OK(TestPutGetDelete()); }