Test: Support running raw query in Executor UT. (#6085)

ref #4609

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

dbms/src/Debug/MockStorage.h

@@ -22,6 +22,7 @@ namespace DB::tests
 {
 using MockColumnInfo = std::pair<String, TiDB::TP>;
 using MockColumnInfoVec = std::vector<MockColumnInfo>;
+using TableInfo = TiDB::TableInfo;
 using CutColumnInfo = std::pair<int, int>; // <start_idx, row_num>
 
 class MockTableIdGenerator : public ext::Singleton<MockTableIdGenerator>
@@ -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<String, Int64> name_to_id_map; /// <table_name, table_id>
     std::unordered_map<Int64, MockColumnInfoVec> table_schema; /// <table_id, columnInfo>
     std::unordered_map<Int64, ColumnsWithTypeAndName> table_columns; /// <table_id, columns>
+    std::unordered_map<String, TableInfo> table_infos;
 
     /// for mock exchange receiver
     std::unordered_map<String, String> executor_id_to_name_map; /// <executor_id, exchange name>
     std::unordered_map<String, MockColumnInfoVec> exchange_schemas; /// <exchange_name, columnInfo>
     std::unordered_map<String, ColumnsWithTypeAndName> exchange_columns; /// <exchange_name, columns>
+
+private:
+    void addTableInfo(const String & name, const MockColumnInfoVec & columns);
 };
 } // namespace DB::tests

dbms/src/Debug/dbgFuncCoprocessor.cpp

@@ -669,8 +669,6 @@ TableID findTableIdForQueryFragment(ExecutorBinderPtr root_executor, bool must_h
         {
             if (dynamic_cast<mock::ExchangeReceiverBinder *>(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)

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<String>({"banana", "banana", "banana", "banana"}), toNullableVec<String>({"apple", "apple", "apple", "banana"}), toNullableVec<String>({"banana", "banana", "banana", {}}), toNullableVec<String>({"apple", "apple", "apple", {}})};
+    ASSERT_COLUMNS_EQ_R(executeRawQuery(query, 1), cols);
+}
+CATCH
+
 } // namespace tests
 } // namespace DB

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<String>(col_name, ColumnWithData(col0.begin(), col0.begin() + 1))};
+    ASSERT_COLUMNS_EQ_R(executeRawQuery(query, 1), cols);
+}
+CATCH
+
 } // namespace tests
 } // namespace DB

dbms/src/TestUtils/ExecutorTestUtils.cpp

@@ -230,6 +230,23 @@ DB::ColumnsWithTypeAndName ExecutorTest::executeStreams(const std::shared_ptr<ti
     return readBlock(executeQuery(context.context, /*internal=*/true).in);
 }
 
+DB::ColumnsWithTypeAndName ExecutorTest::executeRawQuery(const String & query, size_t concurrency)
+{
+    DAGProperties properties;
+    // disable mpp
+    properties.is_mpp_query = false;
+    properties.start_ts = 1;
+    auto [query_tasks, func_wrap_output_stream] = compileQuery(
+        context.context,
+        query,
+        [&](const String & database_name, const String & table_name) {
+            return context.mockStorage().getTableInfo(database_name + "." + table_name);
+        },
+        properties);
+
+    return executeStreams(query_tasks[0].dag_request, concurrency);
+}
+
 void ExecutorTest::dagRequestEqual(const String & expected_string, const std::shared_ptr<tipb::DAGRequest> & actual)
 {
     ASSERT_EQ(Poco::trim(expected_string), Poco::trim(ExecutorSerializer().serialize(actual.get())));

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<BlockInputStreamPtr> streams);
+ColumnsWithTypeAndName readBlock(BlockInputStreamPtr stream);
+ColumnsWithTypeAndName readBlocks(std::vector<BlockInputStreamPtr> 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<tipb::DAGRequest> & actual);
 
     void executeInterpreter(const String & expected_string, const std::shared_ptr<tipb::DAGRequest> & request, size_t concurrency);
-
+    ColumnsWithTypeAndName executeRawQuery(const String & query, size_t concurrency);
     void executeAndAssertColumnsEqual(const std::shared_ptr<tipb::DAGRequest> & request, const ColumnsWithTypeAndName & expect_columns);
     void executeAndAssertRowsEqual(const std::shared_ptr<tipb::DAGRequest> & request, size_t expect_rows);
 

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)

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);