diff --git a/Tests/ArrowIpcIntegrationTest.cpp b/Tests/ArrowIpcIntegrationTest.cpp index 562a087173..9e9d9210f0 100644 --- a/Tests/ArrowIpcIntegrationTest.cpp +++ b/Tests/ArrowIpcIntegrationTest.cpp @@ -48,6 +48,7 @@ using namespace std::literals; #include "QueryEngine/CompilationOptions.h" #include "Shared/ArrowUtil.h" #include "Shared/ThriftClient.h" +#include "Shared/scope.h" #include "gen-cpp/OmniSci.h" @@ -379,10 +380,9 @@ TEST_F(ArrowIpcBasic, IpcWire) { } } -TEST_F(ArrowIpcBasic, IpcCpu) { - auto data_frame = - execute_arrow_ipc("SELECT * FROM arrow_ipc_test;", ExecutorDeviceType::CPU); +namespace { +void check_cpu_dataframe(TDataFrame& data_frame) { ASSERT_TRUE(data_frame.df_size > 0); auto df = @@ -439,6 +439,29 @@ TEST_F(ArrowIpcBasic, IpcCpu) { ASSERT_EQ(str, truth_strings.GetString(i)); } } +} + +} // namespace + +TEST_F(ArrowIpcBasic, IpcCpu) { + auto data_frame = + execute_arrow_ipc("SELECT * FROM arrow_ipc_test;", ExecutorDeviceType::CPU); + + check_cpu_dataframe(data_frame); + + deallocate_df(data_frame, ExecutorDeviceType::CPU); +} + +TEST_F(ArrowIpcBasic, IpcCpuWithCpuExecution) { + g_client->set_execution_mode(g_session_id, TExecuteMode::CPU); + ScopeGuard reset_execution_mode = [&] { + g_client->set_execution_mode(g_session_id, TExecuteMode::GPU); + }; + + auto data_frame = + execute_arrow_ipc("SELECT * FROM arrow_ipc_test;", ExecutorDeviceType::CPU); + + check_cpu_dataframe(data_frame); deallocate_df(data_frame, ExecutorDeviceType::CPU); } @@ -550,6 +573,22 @@ TEST_F(ArrowIpcBasic, IpcGpu) { deallocate_df(data_frame, ExecutorDeviceType::GPU); } +TEST_F(ArrowIpcBasic, IpcGpuWithCpuQuery) { + const size_t device_id = 0; + if (g_cpu_only) { + LOG(ERROR) << "Test not valid in CPU mode."; + return; + } + + g_client->set_execution_mode(g_session_id, TExecuteMode::CPU); + ScopeGuard reset_execution_mode = [&] { + g_client->set_execution_mode(g_session_id, TExecuteMode::GPU); + }; + + EXPECT_ANY_THROW(execute_arrow_ipc( + "SELECT * FROM arrow_ipc_test;", ExecutorDeviceType::GPU, device_id)); +} + TEST_F(ArrowIpcBasic, EmptyResultSet) { char const* drop_flights = "DROP TABLE IF EXISTS flights;"; run_ddl_statement(drop_flights); diff --git a/ThriftHandler/DBHandler.cpp b/ThriftHandler/DBHandler.cpp index 39b5e32f75..408ebbe982 100644 --- a/ThriftHandler/DBHandler.cpp +++ b/ThriftHandler/DBHandler.cpp @@ -1398,7 +1398,7 @@ int64_t DBHandler::process_geo_copy_from(const TSessionId& session_id) { void DBHandler::sql_execute_df(TDataFrame& _return, const TSessionId& session, const std::string& query_str, - const TDeviceType::type device_type, + const TDeviceType::type results_device_type, const int32_t device_id, const int32_t first_n, const TArrowTransport::type transport_method) { @@ -1406,8 +1406,9 @@ void DBHandler::sql_execute_df(TDataFrame& _return, auto query_state = create_query_state(session_ptr, query_str); auto stdlog = STDLOG(session_ptr, query_state); - if (device_type == TDeviceType::GPU) { - const auto executor_device_type = session_ptr->get_executor_device_type(); + const auto executor_device_type = session_ptr->get_executor_device_type(); + + if (results_device_type == TDeviceType::GPU) { if (executor_device_type != ExecutorDeviceType::GPU) { THROW_MAPD_EXCEPTION(std::string("GPU mode is not allowed in this session")); } @@ -1453,8 +1454,10 @@ void DBHandler::sql_execute_df(TDataFrame& _return, query_ra, query_state_proxy, *session_ptr, - device_type == TDeviceType::CPU ? ExecutorDeviceType::CPU - : ExecutorDeviceType::GPU, + executor_device_type, + results_device_type == TDeviceType::CPU + ? ExecutorDeviceType::CPU + : ExecutorDeviceType::GPU, static_cast(device_id), first_n, transport_method); @@ -5622,13 +5625,12 @@ void DBHandler::execute_rel_alg_df(TDataFrame& _return, const std::string& query_ra, QueryStateProxy query_state_proxy, const Catalog_Namespace::SessionInfo& session_info, - const ExecutorDeviceType device_type, + const ExecutorDeviceType executor_device_type, + const ExecutorDeviceType results_device_type, const size_t device_id, const int32_t first_n, const TArrowTransport::type transport_method) const { const auto& cat = session_info.getCatalog(); - CHECK(device_type == ExecutorDeviceType::CPU || - session_info.get_executor_device_type() == ExecutorDeviceType::GPU); auto executor = Executor::getExecutor(Executor::UNITARY_EXECUTOR_ID, jit_debug_ ? "/tmp" : "", jit_debug_ ? "mapdquery" : "", @@ -5639,14 +5641,15 @@ void DBHandler::execute_rel_alg_df(TDataFrame& _return, query_state_proxy.getQueryState().shared_from_this()); const auto& query_hints = ra_executor.getParsedQueryHints(); const bool cpu_mode_enabled = query_hints.isHintRegistered(QueryHint::kCpuMode); - CompilationOptions co = {cpu_mode_enabled ? ExecutorDeviceType::CPU : device_type, - /*hoist_literals=*/true, - ExecutorOptLevel::Default, - g_enable_dynamic_watchdog, - /*allow_lazy_fetch=*/true, - /*filter_on_deleted_column=*/true, - ExecutorExplainType::Default, - intel_jit_profile_}; + CompilationOptions co = { + cpu_mode_enabled ? ExecutorDeviceType::CPU : executor_device_type, + /*hoist_literals=*/true, + ExecutorOptLevel::Default, + g_enable_dynamic_watchdog, + /*allow_lazy_fetch=*/true, + /*filter_on_deleted_column=*/true, + ExecutorExplainType::Default, + intel_jit_profile_}; ExecutionOptions eo = { g_enable_columnar_output, allow_multifrag_, @@ -5681,7 +5684,7 @@ void DBHandler::execute_rel_alg_df(TDataFrame& _return, const auto converter = std::make_unique(rs, data_mgr_, - device_type, + results_device_type, device_id, getTargetNames(result.getTargetsMeta()), first_n, @@ -5696,7 +5699,7 @@ void DBHandler::execute_rel_alg_df(TDataFrame& _return, std::string(arrow_result.df_handle.begin(), arrow_result.df_handle.end()); _return.df_buffer = std::string(arrow_result.df_buffer.begin(), arrow_result.df_buffer.end()); - if (device_type == ExecutorDeviceType::GPU) { + if (results_device_type == ExecutorDeviceType::GPU) { std::lock_guard map_lock(handle_to_dev_ptr_mutex_); CHECK(!ipc_handle_to_dev_ptr_.count(_return.df_handle)); ipc_handle_to_dev_ptr_.insert( diff --git a/ThriftHandler/DBHandler.h b/ThriftHandler/DBHandler.h index ff853775a1..45496cdcac 100644 --- a/ThriftHandler/DBHandler.h +++ b/ThriftHandler/DBHandler.h @@ -763,7 +763,8 @@ class DBHandler : public OmniSciIf { const std::string& query_ra, QueryStateProxy query_state_proxy, const Catalog_Namespace::SessionInfo& session_info, - const ExecutorDeviceType device_type, + const ExecutorDeviceType executor_device_type, + const ExecutorDeviceType results_device_type, const size_t device_id, const int32_t first_n, const TArrowTransport::type transport_method) const;