ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL#13492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lidavidm merged 1 commit into from
Sep 16, 2022
Merged

ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL #13492

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions cpp/src/arrow/flight/integration_tests/flight_integration_test.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
656 changes: 597 additions & 59 deletions cpp/src/arrow/flight/integration_tests/test_integration.cc
View file
Open in desktop

Large diffs are not rendered by default.

4 changes: 3 additions & 1 deletion cpp/src/arrow/flight/server.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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(); }

Expand Down
27 changes: 26 additions & 1 deletion cpp/src/arrow/flight/sql/CMakeLists.txt
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand All@@ -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")

Expand Down
239 changes: 239 additions & 0 deletions cpp/src/arrow/flight/sql/acero_test.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -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
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL#13492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

lidavidm
merged 1 commit into

from

Sep 16, 2022
 Merged


ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL
#13492
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


4 changes: 4 additions & 0 deletions
cpp/src/arrow/flight/integration_tests/flight_integration_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

656 changes: 597 additions & 59 deletions
cpp/src/arrow/flight/integration_tests/test_integration.cc

View file


Open in desktop


Large diffs are not rendered by default.


4 changes: 3 additions & 1 deletion
cpp/src/arrow/flight/server.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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(); }

Expand Down

27 changes: 26 additions & 1 deletion
cpp/src/arrow/flight/sql/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand All@@ -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")

Expand Down

239 changes: 239 additions & 0 deletions
cpp/src/arrow/flight/sql/acero_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL#13492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

lidavidm
merged 1 commit into

from

Sep 16, 2022
 Merged


ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL
#13492
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


4 changes: 4 additions & 0 deletions
cpp/src/arrow/flight/integration_tests/flight_integration_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

656 changes: 597 additions & 59 deletions
cpp/src/arrow/flight/integration_tests/test_integration.cc

View file


Open in desktop


Large diffs are not rendered by default.


4 changes: 3 additions & 1 deletion
cpp/src/arrow/flight/server.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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(); }

Expand Down

27 changes: 26 additions & 1 deletion
cpp/src/arrow/flight/sql/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand All@@ -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")

Expand Down

239 changes: 239 additions & 0 deletions
cpp/src/arrow/flight/sql/acero_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL#13492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

lidavidm
merged 1 commit into

from

Sep 16, 2022
 Merged


ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL
#13492
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


4 changes: 4 additions & 0 deletions
cpp/src/arrow/flight/integration_tests/flight_integration_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

656 changes: 597 additions & 59 deletions
cpp/src/arrow/flight/integration_tests/test_integration.cc

View file


Open in desktop


Large diffs are not rendered by default.


4 changes: 3 additions & 1 deletion
cpp/src/arrow/flight/server.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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(); }

Expand Down

27 changes: 26 additions & 1 deletion
cpp/src/arrow/flight/sql/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand All@@ -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")

Expand Down

239 changes: 239 additions & 0 deletions
cpp/src/arrow/flight/sql/acero_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL#13492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

lidavidm
merged 1 commit into

from

Sep 16, 2022
 Merged


ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL
#13492
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


4 changes: 4 additions & 0 deletions
cpp/src/arrow/flight/integration_tests/flight_integration_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

656 changes: 597 additions & 59 deletions
cpp/src/arrow/flight/integration_tests/test_integration.cc

View file


Open in desktop


Large diffs are not rendered by default.


4 changes: 3 additions & 1 deletion
cpp/src/arrow/flight/server.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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(); }

Expand Down

27 changes: 26 additions & 1 deletion
cpp/src/arrow/flight/sql/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand All@@ -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")

Expand Down

239 changes: 239 additions & 0 deletions
cpp/src/arrow/flight/sql/acero_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL#13492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

lidavidm
merged 1 commit into

from

Sep 16, 2022
 Merged


ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL
#13492
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


4 changes: 4 additions & 0 deletions
cpp/src/arrow/flight/integration_tests/flight_integration_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

656 changes: 597 additions & 59 deletions
cpp/src/arrow/flight/integration_tests/test_integration.cc

View file


Open in desktop


Large diffs are not rendered by default.


4 changes: 3 additions & 1 deletion
cpp/src/arrow/flight/server.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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(); }

Expand Down

27 changes: 26 additions & 1 deletion
cpp/src/arrow/flight/sql/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand All@@ -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")

Expand Down

239 changes: 239 additions & 0 deletions
cpp/src/arrow/flight/sql/acero_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL#13492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

lidavidm
merged 1 commit into

from

Sep 16, 2022
 Merged


ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL
#13492
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


4 changes: 4 additions & 0 deletions
cpp/src/arrow/flight/integration_tests/flight_integration_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

656 changes: 597 additions & 59 deletions
cpp/src/arrow/flight/integration_tests/test_integration.cc

View file


Open in desktop


Large diffs are not rendered by default.


4 changes: 3 additions & 1 deletion
cpp/src/arrow/flight/server.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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(); }

Expand Down

27 changes: 26 additions & 1 deletion
cpp/src/arrow/flight/sql/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand All@@ -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")

Expand Down

239 changes: 239 additions & 0 deletions
cpp/src/arrow/flight/sql/acero_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL#13492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

lidavidm
merged 1 commit into

from

Sep 16, 2022
 Merged


ARROW-17688: [C++][Java][FlightRPC] Substrait, transaction, cancellation for Flight SQL
#13492
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


4 changes: 4 additions & 0 deletions
cpp/src/arrow/flight/integration_tests/flight_integration_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

656 changes: 597 additions & 59 deletions
cpp/src/arrow/flight/integration_tests/test_integration.cc

View file


Open in desktop


Large diffs are not rendered by default.


4 changes: 3 additions & 1 deletion
cpp/src/arrow/flight/server.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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(); }

Expand Down

27 changes: 26 additions & 1 deletion
cpp/src/arrow/flight/sql/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand All@@ -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")

Expand Down

239 changes: 239 additions & 0 deletions
cpp/src/arrow/flight/sql/acero_test.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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
Loading