RFC: [FlightRPC][WIP] Substrait, transaction, cancellation for Flight…

… SQL
apache · Sep 16, 2022 · 6c48fdd · 6c48fdd
1 parent 5e6da78
commit 6c48fdd
Show file tree

Hide file tree

Showing 40 changed files with 4,445 additions and 380 deletions.
diff --git a/cpp/src/arrow/flight/integration_tests/flight_integration_test.cc b/cpp/src/arrow/flight/integration_tests/flight_integration_test.cc
@@ -55,6 +55,10 @@ TEST(FlightIntegration, Middleware) { ASSERT_OK(RunScenario("middleware")); }
 
 TEST(FlightIntegration, FlightSql) { ASSERT_OK(RunScenario("flight_sql")); }
 
+TEST(FlightIntegration, FlightSqlExtension) {
+  ASSERT_OK(RunScenario("flight_sql:extension"));
+}
+
 }  // namespace integration_tests
 }  // namespace flight
 }  // namespace arrow
diff --git a/cpp/src/arrow/flight/integration_tests/test_integration.cc b/cpp/src/arrow/flight/integration_tests/test_integration.cc
diff --git a/cpp/src/arrow/flight/server.cc b/cpp/src/arrow/flight/server.cc
@@ -353,7 +353,9 @@ RecordBatchStream::RecordBatchStream(const std::shared_ptr<RecordBatchReader>& r
   impl_.reset(new RecordBatchStreamImpl(reader, options));
 }
 
-RecordBatchStream::~RecordBatchStream() {}
+RecordBatchStream::~RecordBatchStream() {
+  ARROW_WARN_NOT_OK(impl_->Close(), "Failed to close FlightDataStream");
+}
 
 Status RecordBatchStream::Close() { return impl_->Close(); }
 

diff --git a/cpp/src/arrow/flight/sql/CMakeLists.txt b/cpp/src/arrow/flight/sql/CMakeLists.txt
@@ -89,7 +89,11 @@ if(ARROW_BUILD_TESTS OR ARROW_BUILD_EXAMPLES)
       example/sqlite_statement_batch_reader.cc
       example/sqlite_server.cc
       example/sqlite_tables_schema_batch_reader.cc)
+
   set(ARROW_FLIGHT_SQL_TEST_SRCS server_test.cc)
+  set(ARROW_FLIGHT_SQL_TEST_LIBS ${SQLite3_LIBRARIES})
+  set(ARROW_FLIGHT_SQL_ACERO_SRCS example/acero_server.cc)
+
   if(NOT MSVC AND NOT MINGW)
     # ARROW-16902: getting Protobuf generated code to have all the
     # proper dllexport/dllimport declarations is difficult, since
@@ -98,13 +102,34 @@ if(ARROW_BUILD_TESTS OR ARROW_BUILD_EXAMPLES)
     list(APPEND ARROW_FLIGHT_SQL_TEST_SRCS client_test.cc)
   endif()
 
+  if(ARROW_COMPUTE
+     AND ARROW_PARQUET
+     AND ARROW_SUBSTRAIT)
+    list(APPEND ARROW_FLIGHT_SQL_TEST_SRCS ${ARROW_FLIGHT_SQL_ACERO_SRCS} acero_test.cc)
+    if(ARROW_FLIGHT_TEST_LINKAGE STREQUAL "static")
+      list(APPEND ARROW_FLIGHT_SQL_TEST_LIBS arrow_substrait_static)
+    else()
+      list(APPEND ARROW_FLIGHT_SQL_TEST_LIBS arrow_substrait_shared)
+    endif()
+
+    if(ARROW_BUILD_EXAMPLES)
+      add_executable(acero-flight-sql-server ${ARROW_FLIGHT_SQL_ACERO_SRCS}
+                                             example/acero_main.cc)
+      target_link_libraries(acero-flight-sql-server
+                            PRIVATE ${ARROW_FLIGHT_SQL_TEST_LINK_LIBS}
+                                    ${ARROW_FLIGHT_SQL_TEST_LIBS} ${GFLAGS_LIBRARIES})
+    endif()
+  endif()
+
   add_arrow_test(flight_sql_test
                  SOURCES
                  ${ARROW_FLIGHT_SQL_TEST_SRCS}
                  ${ARROW_FLIGHT_SQL_TEST_SERVER_SRCS}
                  STATIC_LINK_LIBS
                  ${ARROW_FLIGHT_SQL_TEST_LINK_LIBS}
-                 ${SQLite3_LIBRARIES}
+                 ${ARROW_FLIGHT_SQL_TEST_LIBS}
+                 EXTRA_INCLUDES
+                 "${CMAKE_CURRENT_BINARY_DIR}/../"
                  LABELS
                  "arrow_flight_sql")
 

diff --git a/cpp/src/arrow/flight/sql/acero_test.cc b/cpp/src/arrow/flight/sql/acero_test.cc
@@ -0,0 +1,239 @@
+// 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.
+
+/// Integration test using the Acero backend
+
+#include <memory>
+#include <sstream>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "arrow/array.h"
+#include "arrow/engine/substrait/util.h"
+#include "arrow/flight/server.h"
+#include "arrow/flight/sql/client.h"
+#include "arrow/flight/sql/example/acero_server.h"
+#include "arrow/flight/sql/types.h"
+#include "arrow/flight/types.h"
+#include "arrow/stl_iterator.h"
+#include "arrow/table.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/type_fwd.h"
+#include "arrow/util/checked_cast.h"
+
+namespace arrow {
+namespace flight {
+namespace sql {
+
+using arrow::internal::checked_cast;
+
+class TestAcero : public ::testing::Test {
+ public:
+  void SetUp() override {
+    ASSERT_OK_AND_ASSIGN(auto location, Location::ForGrpcTcp("localhost", 0));
+    flight::FlightServerOptions options(location);
+
+    ASSERT_OK_AND_ASSIGN(server_, acero_example::MakeAceroServer());
+    ASSERT_OK(server_->Init(options));
+
+    ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(server_->location()));
+    client_.reset(new FlightSqlClient(std::move(client)));
+  }
+
+  void TearDown() override {
+    ASSERT_OK(client_->Close());
+    ASSERT_OK(server_->Shutdown());
+  }
+
+ protected:
+  std::unique_ptr<FlightSqlClient> client_;
+  std::unique_ptr<FlightServerBase> server_;
+};
+
+arrow::Result<std::shared_ptr<Buffer>> MakeSubstraitPlan() {
+  ARROW_ASSIGN_OR_RAISE(std::string dir_string,
+                        arrow::internal::GetEnvVar("PARQUET_TEST_DATA"));
+  ARROW_ASSIGN_OR_RAISE(auto dir,
+                        arrow::internal::PlatformFilename::FromString(dir_string));
+  ARROW_ASSIGN_OR_RAISE(auto filename, dir.Join("binary.parquet"));
+  std::string uri = std::string("file://") + filename.ToString();
+
+  // TODO(ARROW-17229): we should use a RootRel here
+  std::string json_plan = R"({
+    "relations": [
+      {
+        "rel": {
+          "read": {
+            "base_schema": {
+              "struct": {
+                "types": [
+                  {"binary": {}}
+                ]
+              },
+              "names": [
+                "foo"
+              ]
+            },
+            "local_files": {
+              "items": [
+                {
+                  "uri_file": "URI_PLACEHOLDER",
+                  "parquet": {}
+                }
+              ]
+            }
+          }
+        }
+      }
+    ]
+})";
+  std::string uri_placeholder = "URI_PLACEHOLDER";
+  json_plan.replace(json_plan.find(uri_placeholder), uri_placeholder.size(), uri);
+  return engine::SerializeJsonPlan(json_plan);
+}
+
+TEST_F(TestAcero, GetSqlInfo) {
+  FlightCallOptions call_options;
+  std::vector<int> sql_info_codes = {
+      SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_SUBSTRAIT,
+      SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_TRANSACTION,
+  };
+  ASSERT_OK_AND_ASSIGN(auto flight_info,
+                       client_->GetSqlInfo(call_options, sql_info_codes));
+  ASSERT_OK_AND_ASSIGN(auto reader,
+                       client_->DoGet(call_options, flight_info->endpoints()[0].ticket));
+  ASSERT_OK_AND_ASSIGN(auto results, reader->ToTable());
+  ASSERT_OK_AND_ASSIGN(auto batch, results->CombineChunksToBatch());
+  ASSERT_EQ(2, results->num_rows());
+  std::vector<std::pair<uint32_t, SqlInfoResult>> info;
+  const auto& ids = checked_cast<const UInt32Array&>(*batch->column(0));
+  const auto& values = checked_cast<const DenseUnionArray&>(*batch->column(1));
+  for (int64_t i = 0; i < batch->num_rows(); i++) {
+    ASSERT_OK_AND_ASSIGN(auto scalar, values.GetScalar(i));
+    if (ids.Value(i) == SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_SUBSTRAIT) {
+      ASSERT_EQ(*checked_cast<const DenseUnionScalar&>(*scalar).value,
+                BooleanScalar(true));
+    } else if (ids.Value(i) == SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_TRANSACTION) {
+      ASSERT_EQ(
+          *checked_cast<const DenseUnionScalar&>(*scalar).value,
+          Int32Scalar(
+              SqlInfoOptions::SqlSupportedTransaction::SQL_SUPPORTED_TRANSACTION_NONE));
+    } else {
+      FAIL() << "Unexpected info value: " << ids.Value(i);
+    }
+  }
+}
+
+TEST_F(TestAcero, Scan) {
+#ifdef _WIN32
+  GTEST_SKIP() << "ARROW-16392: Substrait File URI not supported for Windows";
+#endif
+
+  FlightCallOptions call_options;
+  ASSERT_OK_AND_ASSIGN(auto serialized_plan, MakeSubstraitPlan());
+
+  SubstraitPlan plan{serialized_plan->ToString(), /*version=*/"0.6.0"};
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<FlightInfo> info,
+                       client_->ExecuteSubstrait(call_options, plan));
+  ipc::DictionaryMemo memo;
+  ASSERT_OK_AND_ASSIGN(auto schema, info->GetSchema(&memo));
+  // TODO(ARROW-17229): the scanner "special" fields are still included, strip them
+  // manually
+  auto fixed_schema = arrow::schema({schema->fields()[0]});
+  ASSERT_NO_FATAL_FAILURE(
+      AssertSchemaEqual(fixed_schema, arrow::schema({field("foo", binary())})));
+
+  ASSERT_EQ(1, info->endpoints().size());
+  ASSERT_EQ(0, info->endpoints()[0].locations.size());
+  ASSERT_OK_AND_ASSIGN(auto reader,
+                       client_->DoGet(call_options, info->endpoints()[0].ticket));
+  ASSERT_OK_AND_ASSIGN(auto reader_schema, reader->GetSchema());
+  ASSERT_NO_FATAL_FAILURE(AssertSchemaEqual(schema, reader_schema));
+  ASSERT_OK_AND_ASSIGN(auto table, reader->ToTable());
+  ASSERT_GT(table->num_rows(), 0);
+}
+
+TEST_F(TestAcero, Update) {
+#ifdef _WIN32
+  GTEST_SKIP() << "ARROW-16392: Substrait File URI not supported for Windows";
+#endif
+
+  FlightCallOptions call_options;
+  ASSERT_OK_AND_ASSIGN(auto serialized_plan, MakeSubstraitPlan());
+  SubstraitPlan plan{serialized_plan->ToString(), /*version=*/"0.6.0"};
+  EXPECT_RAISES_WITH_MESSAGE_THAT(NotImplemented,
+                                  ::testing::HasSubstr("Updates are unsupported"),
+                                  client_->ExecuteSubstraitUpdate(call_options, plan));
+}
+
+TEST_F(TestAcero, Prepare) {
+#ifdef _WIN32
+  GTEST_SKIP() << "ARROW-16392: Substrait File URI not supported for Windows";
+#endif
+
+  FlightCallOptions call_options;
+  ASSERT_OK_AND_ASSIGN(auto serialized_plan, MakeSubstraitPlan());
+  SubstraitPlan plan{serialized_plan->ToString(), /*version=*/"0.6.0"};
+  ASSERT_OK_AND_ASSIGN(auto prepared_statement,
+                       client_->PrepareSubstrait(call_options, plan));
+  ASSERT_NE(prepared_statement->dataset_schema(), nullptr);
+  ASSERT_EQ(prepared_statement->parameter_schema(), nullptr);
+
+  auto fixed_schema = arrow::schema({prepared_statement->dataset_schema()->fields()[0]});
+  ASSERT_NO_FATAL_FAILURE(
+      AssertSchemaEqual(fixed_schema, arrow::schema({field("foo", binary())})));
+
+  EXPECT_RAISES_WITH_MESSAGE_THAT(NotImplemented,
+                                  ::testing::HasSubstr("Updates are unsupported"),
+                                  prepared_statement->ExecuteUpdate());
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<FlightInfo> info, prepared_statement->Execute());
+  ASSERT_EQ(1, info->endpoints().size());
+  ASSERT_EQ(0, info->endpoints()[0].locations.size());
+  ASSERT_OK_AND_ASSIGN(auto reader,
+                       client_->DoGet(call_options, info->endpoints()[0].ticket));
+  ASSERT_OK_AND_ASSIGN(auto reader_schema, reader->GetSchema());
+  ASSERT_NO_FATAL_FAILURE(
+      AssertSchemaEqual(prepared_statement->dataset_schema(), reader_schema));
+  ASSERT_OK_AND_ASSIGN(auto table, reader->ToTable());
+  ASSERT_GT(table->num_rows(), 0);
+
+  ASSERT_OK(prepared_statement->Close());
+}
+
+TEST_F(TestAcero, Transactions) {
+#ifdef _WIN32
+  GTEST_SKIP() << "ARROW-16392: Substrait File URI not supported for Windows";
+#endif
+
+  FlightCallOptions call_options;
+  ASSERT_OK_AND_ASSIGN(auto serialized_plan, MakeSubstraitPlan());
+  Transaction handle("fake-id");
+  SubstraitPlan plan{serialized_plan->ToString(), /*version=*/"0.6.0"};
+
+  EXPECT_RAISES_WITH_MESSAGE_THAT(NotImplemented,
+                                  ::testing::HasSubstr("Transactions are unsupported"),
+                                  client_->ExecuteSubstrait(call_options, plan, handle));
+  EXPECT_RAISES_WITH_MESSAGE_THAT(NotImplemented,
+                                  ::testing::HasSubstr("Transactions are unsupported"),
+                                  client_->PrepareSubstrait(call_options, plan, handle));
+}
+
+}  // namespace sql
+}  // namespace flight
+}  // namespace arrow