Test: support mock not null column (#7059)

ref #4609

dbms/src/Debug/MockStorage.cpp

@@ -150,7 +150,7 @@ std::tuple<StorageDeltaMergePtr, Names, SelectQueryInfo> MockStorage::prepareFor
     assert(!column_infos.empty());
     Names column_names;
     for (const auto & column_info : column_infos)
-        column_names.push_back(column_info.first);
+        column_names.push_back(column_info.name);
 
     auto scan_context = std::make_shared<DM::ScanContext>();
 
@@ -209,7 +209,11 @@ void MockStorage::addTableInfoForDeltaMerge(const String & name, const MockColum
     for (const auto & column : columns)
     {
         TiDB::ColumnInfo ret;
-        std::tie(ret.name, ret.tp) = column;
+        ret.name = column.name;
+        ret.tp = column.type;
+
+        if (!column.nullable)
+            ret.setNotNullFlag();
         // TODO: find a way to assign decimal field's flen.
         if (ret.tp == TiDB::TP::TypeNewDecimal)
             ret.flen = 65;
@@ -465,11 +469,14 @@ ColumnInfos mockColumnInfosToTiDBColumnInfos(const MockColumnInfoVec & mock_colu
     for (const auto & mock_column_info : mock_column_infos)
     {
         TiDB::ColumnInfo column_info;
-        std::tie(column_info.name, column_info.tp) = mock_column_info;
+        column_info.name = mock_column_info.name;
+        column_info.tp = mock_column_info.type;
         column_info.id = col_id++;
         // TODO: find a way to assign decimal field's flen.
         if (column_info.tp == TiDB::TP::TypeNewDecimal)
             column_info.flen = 65;
+        if (!mock_column_info.nullable)
+            column_info.setNotNullFlag();
         ret.push_back(std::move(column_info));
     }
     return ret;

dbms/src/Debug/MockStorage.h

@@ -31,7 +31,12 @@ using StorageDeltaMergePtr = std::shared_ptr<StorageDeltaMerge>;
 class Context;
 struct SelectQueryInfo;
 
-using MockColumnInfo = std::pair<String, TiDB::TP>;
+struct MockColumnInfo
+{
+    String name;
+    TiDB::TP type;
+    bool nullable = true;
+};
 using MockColumnInfoVec = std::vector<MockColumnInfo>;
 using TableInfo = TiDB::TableInfo;
 using CutColumnInfo = std::pair<int, int>; // <start_idx, row_num>

dbms/src/Flash/Coprocessor/MockSourceStream.h

@@ -78,7 +78,6 @@ std::pair<NamesAndTypes, std::vector<std::shared_ptr<SourceType>>> mockSourceStr
         columns_with_type_and_name = context.mockStorage()->getColumns(table_id);
 
     columns_with_type_and_name = getUsedColumns(used_columns, columns_with_type_and_name);
-
     return cutStreams<SourceType>(context, columns_with_type_and_name, max_streams, log);
 }
 

dbms/src/Flash/tests/gtest_limit_executor.cpp

@@ -38,6 +38,10 @@ class LimitExecutorTestRunner : public DB::tests::ExecutorTest
         context.addMockTable({"test", "bigtable"},
                              {{"col", TiDB::TP::TypeString}},
                              {toNullableVec<String>("col", col)});
+
+        context.addMockTable({"test", "notNull"},
+                             {{"col", TiDB::TP::TypeString, false}},
+                             {toVec<String>("col", {"a", "b", "c", "d", "e", "f", "g", "h"})});
     }
 
     std::shared_ptr<tipb::DAGRequest> buildDAGRequest(size_t limit_num)
@@ -50,6 +54,7 @@ class LimitExecutorTestRunner : public DB::tests::ExecutorTest
     const String table_name{"projection_test_table"};
     const String col_name{"limit_col"};
     const ColumnWithNullableString col0{"col0-0", {}, "col0-2", "col0-3", {}, "col0-5", "col0-6", "col0-7"};
+    const ColumnWithString col_not_null{"a", "b", "c", "d", "e", "f", "g", "h"};
 };
 
 TEST_F(LimitExecutorTestRunner, Limit)
@@ -79,6 +84,15 @@ try
 
         executeAndAssertRowsEqual(request, std::min(limit_num, col_data_num));
     }
+
+    request = context.scan("test", "notNull").limit(1).build(context);
+    expect_cols = {toVec<String>(col_name, ColumnWithString(col_not_null.begin(), col_not_null.begin() + 1))};
+
+    WRAP_FOR_TEST_BEGIN
+    ASSERT_COLUMNS_EQ_R(executeStreams(request), expect_cols);
+    WRAP_FOR_TEST_END
+
+    executeAndAssertRowsEqual(request, 1);
 }
 CATCH
 

dbms/src/TestUtils/mockExecutor.cpp

@@ -188,9 +188,9 @@ DAGRequestBuilder & DAGRequestBuilder::buildExchangeReceiver(const String & exch
     for (const auto & column : columns)
     {
         TiDB::ColumnInfo info;
-        info.tp = column.second;
-        info.name = column.first;
-        schema.push_back({exchange_name + "." + column.first, info});
+        info.name = column.name;
+        info.tp = column.type;
+        schema.push_back({exchange_name + "." + info.name, info});
     }
 
     root = mock::compileExchangeReceiver(getExecutorIndex(), schema, fine_grained_shuffle_stream_count, std::static_pointer_cast<mock::ExchangeSenderBinder>(root));
@@ -388,25 +388,26 @@ DAGRequestBuilder & DAGRequestBuilder::expand(MockVVecColumnNameVec grouping_set
     return *this;
 }
 
-void MockDAGRequestContext::addMockTable(const String & db, const String & table, const MockColumnInfoVec & mock_column_infos, size_t concurrency_hint)
+void MockDAGRequestContext::addMockTable(const String & db, const String & table, const MockColumnInfoVec & columnInfos, size_t concurrency_hint)
 {
-    auto columns = getColumnWithTypeAndName(genNamesAndTypes(mockColumnInfosToTiDBColumnInfos(mock_column_infos), "mock_table_scan"));
-    addMockTable(db, table, mock_column_infos, columns, concurrency_hint);
+    auto columns = getColumnWithTypeAndName(genNamesAndTypes(mockColumnInfosToTiDBColumnInfos(columnInfos), "mock_table_scan"));
+    addMockTable(db, table, columnInfos, columns, concurrency_hint);
 }
 
 void MockDAGRequestContext::addMockTableSchema(const String & db, const String & table, const MockColumnInfoVec & columnInfos)
 {
     mock_storage->addTableSchema(db + "." + table, columnInfos);
 }
+
 void MockDAGRequestContext::addMockTableSchema(const MockTableName & name, const MockColumnInfoVec & columnInfos)
 {
     mock_storage->addTableSchema(name.first + "." + name.second, columnInfos);
 }
 
-void MockDAGRequestContext::addMockTable(const MockTableName & name, const MockColumnInfoVec & mock_column_infos, size_t concurrency_hint)
+void MockDAGRequestContext::addMockTable(const MockTableName & name, const MockColumnInfoVec & columnInfos, size_t concurrency_hint)
 {
-    auto columns = getColumnWithTypeAndName(genNamesAndTypes(mockColumnInfosToTiDBColumnInfos(mock_column_infos), "mock_table_scan"));
-    addMockTable(name, mock_column_infos, columns, concurrency_hint);
+    auto columns = getColumnWithTypeAndName(genNamesAndTypes(mockColumnInfosToTiDBColumnInfos(columnInfos), "mock_table_scan"));
+    addMockTable(name, columnInfos, columns, concurrency_hint);
 }
 
 void MockDAGRequestContext::addMockTableConcurrencyHint(const String & db, const String & table, size_t concurrency_hint)
@@ -507,7 +508,7 @@ void MockDAGRequestContext::addExchangeReceiver(const String & name, const MockC
         {
             for (size_t col_index = 0; col_index < columns.size(); col_index++)
             {
-                if (columns[col_index].name == mock_column_info.first)
+                if (columns[col_index].name == mock_column_info.name)
                 {
                     partition_column_ids.push_back(col_index);
                     break;
@@ -582,7 +583,7 @@ void MockDAGRequestContext::assertMockInput(const MockColumnInfoVec & columnInfo
 {
     assert(columnInfos.size() == columns.size());
     for (size_t i = 0; i < columns.size(); ++i)
-        assert(columnInfos[i].first == columns[i].name);
+        assert(columnInfos[i].name == columns[i].name);
 }
 
 } // namespace DB::tests

dbms/src/TestUtils/mockExecutor.h

@@ -30,8 +30,6 @@ namespace DB
 {
 namespace tests
 {
-using MockColumnInfo = std::pair<String, TiDB::TP>;
-using MockColumnInfoVec = std::vector<MockColumnInfo>;
 using MockTableName = std::pair<String, String>;
 using MockOrderByItem = std::pair<String, bool>;
 using MockOrderByItemVec = std::vector<MockOrderByItem>;
@@ -186,6 +184,7 @@ class MockDAGRequestContext
     void addMockTable(const MockTableName & name, const MockColumnInfoVec & columnInfos, size_t concurrency_hint = 0);
     void addMockTable(const String & db, const String & table, const MockColumnInfoVec & columnInfos, ColumnsWithTypeAndName columns, size_t concurrency_hint = 0);
     void addMockTable(const MockTableName & name, const MockColumnInfoVec & columnInfos, ColumnsWithTypeAndName columns, size_t concurrency_hint = 0);
+
     void updateMockTableColumnData(const String & db, const String & table, ColumnsWithTypeAndName columns)
     {
         addMockTableColumnData(db, table, columns);