From a57e4b1fb3d15046f81289ab4d5e80c4d585f2c3 Mon Sep 17 00:00:00 2001 From: yanweiqi <592838129@qq.com> Date: Mon, 17 Oct 2022 20:47:53 +0800 Subject: [PATCH] Test: Support running raw query in Executor UT. (#6085) ref pingcap/tiflash#4609 --- dbms/src/Debug/MockStorage.cpp | 25 +++++++++++++++++++ dbms/src/Debug/MockStorage.h | 7 ++++++ dbms/src/Debug/dbgFuncCoprocessor.cpp | 2 -- dbms/src/Flash/tests/gtest_join_executor.cpp | 11 ++++++++ dbms/src/Flash/tests/gtest_limit_executor.cpp | 9 +++++++ dbms/src/TestUtils/ExecutorTestUtils.cpp | 17 +++++++++++++ dbms/src/TestUtils/ExecutorTestUtils.h | 6 ++--- dbms/src/TestUtils/mockExecutor.cpp | 25 ++++--------------- dbms/src/TestUtils/mockExecutor.h | 4 +-- 9 files changed, 79 insertions(+), 27 deletions(-) diff --git a/dbms/src/Debug/MockStorage.cpp b/dbms/src/Debug/MockStorage.cpp index 834a494311c..dbcf38c831b 100644 --- a/dbms/src/Debug/MockStorage.cpp +++ b/dbms/src/Debug/MockStorage.cpp @@ -19,6 +19,7 @@ void MockStorage::addTableSchema(const String & name, const MockColumnInfoVec & { name_to_id_map[name] = MockTableIdGenerator::instance().nextTableId(); table_schema[getTableId(name)] = columnInfos; + addTableInfo(name, columnInfos); } void MockStorage::addTableData(const String & name, const ColumnsWithTypeAndName & columns) @@ -150,4 +151,28 @@ ColumnsWithTypeAndName MockStorage::getColumnsForMPPTableScan(Int64 table_id, In } throw Exception(fmt::format("Failed to get table columns by table_id '{}'", table_id)); } + +void MockStorage::addTableInfo(const String & name, const MockColumnInfoVec & columns) +{ + TableInfo table_info; + table_info.name = name; + table_info.id = getTableId(name); + int i = 0; + for (const auto & column : columns) + { + TiDB::ColumnInfo ret; + std::tie(ret.name, ret.tp) = column; + // TODO: find a way to assign decimal field's flen. + if (ret.tp == TiDB::TP::TypeNewDecimal) + ret.flen = 65; + ret.id = i++; + table_info.columns.push_back(std::move(ret)); + } + table_infos[name] = table_info; +} + +TableInfo MockStorage::getTableInfo(const String & name) +{ + return table_infos[name]; +} } // namespace DB::tests diff --git a/dbms/src/Debug/MockStorage.h b/dbms/src/Debug/MockStorage.h index 5d92c9bf3ec..46e8331602f 100644 --- a/dbms/src/Debug/MockStorage.h +++ b/dbms/src/Debug/MockStorage.h @@ -22,6 +22,7 @@ namespace DB::tests { using MockColumnInfo = std::pair; using MockColumnInfoVec = std::vector; +using TableInfo = TiDB::TableInfo; using CutColumnInfo = std::pair; // class MockTableIdGenerator : public ext::Singleton @@ -73,15 +74,21 @@ class MockStorage /// for MPP Tasks, it will split data by partition num, then each MPP service will have a subset of mock data. ColumnsWithTypeAndName getColumnsForMPPTableScan(Int64 table_id, Int64 partition_id, Int64 partition_num); + TableInfo getTableInfo(const String & name); + private: /// for mock table scan std::unordered_map name_to_id_map; /// std::unordered_map table_schema; /// std::unordered_map table_columns; /// + std::unordered_map table_infos; /// for mock exchange receiver std::unordered_map executor_id_to_name_map; /// std::unordered_map exchange_schemas; /// std::unordered_map exchange_columns; /// + +private: + void addTableInfo(const String & name, const MockColumnInfoVec & columns); }; } // namespace DB::tests diff --git a/dbms/src/Debug/dbgFuncCoprocessor.cpp b/dbms/src/Debug/dbgFuncCoprocessor.cpp index 0685930d9d7..4745e7130b7 100644 --- a/dbms/src/Debug/dbgFuncCoprocessor.cpp +++ b/dbms/src/Debug/dbgFuncCoprocessor.cpp @@ -669,8 +669,6 @@ TableID findTableIdForQueryFragment(ExecutorBinderPtr root_executor, bool must_h { if (dynamic_cast(c.get())) continue; - if (non_exchange_child != nullptr) - throw Exception("More than one non-exchange child, should not happen"); non_exchange_child = c; } if (non_exchange_child == nullptr) diff --git a/dbms/src/Flash/tests/gtest_join_executor.cpp b/dbms/src/Flash/tests/gtest_join_executor.cpp index 3f504ed29b6..de98df0c997 100644 --- a/dbms/src/Flash/tests/gtest_join_executor.cpp +++ b/dbms/src/Flash/tests/gtest_join_executor.cpp @@ -537,5 +537,16 @@ try } CATCH + +// Currently only support join with `using` +TEST_F(JoinExecutorTestRunner, RawQuery) +try +{ + String query = "select * from test_db.l_table left outer join test_db.r_table_2 using join_c"; + auto cols = {toNullableVec({"banana", "banana", "banana", "banana"}), toNullableVec({"apple", "apple", "apple", "banana"}), toNullableVec({"banana", "banana", "banana", {}}), toNullableVec({"apple", "apple", "apple", {}})}; + ASSERT_COLUMNS_EQ_R(executeRawQuery(query, 1), cols); +} +CATCH + } // namespace tests } // namespace DB diff --git a/dbms/src/Flash/tests/gtest_limit_executor.cpp b/dbms/src/Flash/tests/gtest_limit_executor.cpp index 239df250306..140d58f1a2c 100644 --- a/dbms/src/Flash/tests/gtest_limit_executor.cpp +++ b/dbms/src/Flash/tests/gtest_limit_executor.cpp @@ -77,5 +77,14 @@ try } CATCH +TEST_F(ExecutorLimitTestRunner, RawQuery) +try +{ + String query = "select * from test_db.projection_test_table limit 1"; + auto cols = {toNullableVec(col_name, ColumnWithData(col0.begin(), col0.begin() + 1))}; + ASSERT_COLUMNS_EQ_R(executeRawQuery(query, 1), cols); +} +CATCH + } // namespace tests } // namespace DB diff --git a/dbms/src/TestUtils/ExecutorTestUtils.cpp b/dbms/src/TestUtils/ExecutorTestUtils.cpp index 1e1a62b5d5d..437469940e4 100644 --- a/dbms/src/TestUtils/ExecutorTestUtils.cpp +++ b/dbms/src/TestUtils/ExecutorTestUtils.cpp @@ -230,6 +230,23 @@ DB::ColumnsWithTypeAndName ExecutorTest::executeStreams(const std::shared_ptr & actual) { ASSERT_EQ(Poco::trim(expected_string), Poco::trim(ExecutorSerializer().serialize(actual.get()))); diff --git a/dbms/src/TestUtils/ExecutorTestUtils.h b/dbms/src/TestUtils/ExecutorTestUtils.h index 67ec19007b7..2dbb51f7e61 100644 --- a/dbms/src/TestUtils/ExecutorTestUtils.h +++ b/dbms/src/TestUtils/ExecutorTestUtils.h @@ -29,8 +29,8 @@ namespace DB::tests { TiDB::TP dataTypeToTP(const DataTypePtr & type); -DB::ColumnsWithTypeAndName readBlock(BlockInputStreamPtr stream); -DB::ColumnsWithTypeAndName readBlocks(std::vector streams); +ColumnsWithTypeAndName readBlock(BlockInputStreamPtr stream); +ColumnsWithTypeAndName readBlocks(std::vector streams); Block mergeBlocks(Blocks blocks); @@ -69,7 +69,7 @@ class ExecutorTest : public ::testing::Test static void dagRequestEqual(const String & expected_string, const std::shared_ptr & actual); void executeInterpreter(const String & expected_string, const std::shared_ptr & request, size_t concurrency); - + ColumnsWithTypeAndName executeRawQuery(const String & query, size_t concurrency); void executeAndAssertColumnsEqual(const std::shared_ptr & request, const ColumnsWithTypeAndName & expect_columns); void executeAndAssertRowsEqual(const std::shared_ptr & request, size_t expect_rows); diff --git a/dbms/src/TestUtils/mockExecutor.cpp b/dbms/src/TestUtils/mockExecutor.cpp index 19e12060447..c98547b164f 100644 --- a/dbms/src/TestUtils/mockExecutor.cpp +++ b/dbms/src/TestUtils/mockExecutor.cpp @@ -151,31 +151,16 @@ QueryTasks DAGRequestBuilder::buildMPPTasks(MockDAGRequestContext & mock_context return query_tasks; } -DAGRequestBuilder & DAGRequestBuilder::mockTable(const String & db, const String & table, Int64 table_id, const MockColumnInfoVec & columns) +DAGRequestBuilder & DAGRequestBuilder::mockTable(const String & db, const String & table, TableInfo & table_info, const MockColumnInfoVec & columns) { assert(!columns.empty()); - TableInfo table_info; - table_info.name = db + "." + table; - table_info.id = table_id; - int i = 0; - for (const auto & column : columns) - { - TiDB::ColumnInfo ret; - ret.tp = column.second; - ret.name = column.first; - // TODO: find a way to assign decimal field's flen. - if (ret.tp == TiDB::TP::TypeNewDecimal) - ret.flen = 65; - ret.id = i++; - table_info.columns.push_back(std::move(ret)); - } root = mock::compileTableScan(getExecutorIndex(), table_info, db, table, false); return *this; } -DAGRequestBuilder & DAGRequestBuilder::mockTable(const MockTableName & name, Int64 table_id, const MockColumnInfoVec & columns) +DAGRequestBuilder & DAGRequestBuilder::mockTable(const MockTableName & name, TableInfo & table_info, const MockColumnInfoVec & columns) { - return mockTable(name.first, name.second, table_id, columns); + return mockTable(name.first, name.second, table_info, columns); } DAGRequestBuilder & DAGRequestBuilder::exchangeReceiver(const MockColumnInfoVec & columns, uint64_t fine_grained_shuffle_stream_count) @@ -412,8 +397,8 @@ void MockDAGRequestContext::addExchangeReceiver(const String & name, MockColumnI DAGRequestBuilder MockDAGRequestContext::scan(const String & db_name, const String & table_name) { - auto table_id = mock_storage.getTableId(db_name + "." + table_name); - return DAGRequestBuilder(index, collation).mockTable({db_name, table_name}, table_id, mock_storage.getTableSchema(db_name + "." + table_name)); + auto table_info = mock_storage.getTableInfo(db_name + "." + table_name); + return DAGRequestBuilder(index, collation).mockTable({db_name, table_name}, table_info, mock_storage.getTableSchema(db_name + "." + table_name)); } DAGRequestBuilder MockDAGRequestContext::receive(const String & exchange_name, uint64_t fine_grained_shuffle_stream_count) diff --git a/dbms/src/TestUtils/mockExecutor.h b/dbms/src/TestUtils/mockExecutor.h index 1160d35f411..adf99ee0d23 100644 --- a/dbms/src/TestUtils/mockExecutor.h +++ b/dbms/src/TestUtils/mockExecutor.h @@ -80,8 +80,8 @@ class DAGRequestBuilder QueryTasks buildMPPTasks(MockDAGRequestContext & mock_context); QueryTasks buildMPPTasks(MockDAGRequestContext & mock_context, const DAGProperties & properties); - DAGRequestBuilder & mockTable(const String & db, const String & table, Int64 table_id, const MockColumnInfoVec & columns); - DAGRequestBuilder & mockTable(const MockTableName & name, Int64 table_id, const MockColumnInfoVec & columns); + DAGRequestBuilder & mockTable(const String & db, const String & table, TableInfo & table_info, const MockColumnInfoVec & columns); + DAGRequestBuilder & mockTable(const MockTableName & name, TableInfo & table_info, const MockColumnInfoVec & columns); DAGRequestBuilder & exchangeReceiver(const MockColumnInfoVec & columns, uint64_t fine_grained_shuffle_stream_count = 0);