diff --git a/dbms/src/Debug/MockStorage.cpp b/dbms/src/Debug/MockStorage.cpp index a94a319f8dc..67451bb0162 100644 --- a/dbms/src/Debug/MockStorage.cpp +++ b/dbms/src/Debug/MockStorage.cpp @@ -150,7 +150,7 @@ std::tuple 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(); @@ -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; diff --git a/dbms/src/Debug/MockStorage.h b/dbms/src/Debug/MockStorage.h index 1db3ea900a6..ed80fd3ba42 100644 --- a/dbms/src/Debug/MockStorage.h +++ b/dbms/src/Debug/MockStorage.h @@ -31,7 +31,12 @@ using StorageDeltaMergePtr = std::shared_ptr; class Context; struct SelectQueryInfo; -using MockColumnInfo = std::pair; +struct MockColumnInfo +{ + String name; + TiDB::TP type; + bool nullable = true; +}; using MockColumnInfoVec = std::vector; using TableInfo = TiDB::TableInfo; using CutColumnInfo = std::pair; // diff --git a/dbms/src/Flash/Coprocessor/MockSourceStream.h b/dbms/src/Flash/Coprocessor/MockSourceStream.h index dd75b8f489f..a4ee76654f3 100644 --- a/dbms/src/Flash/Coprocessor/MockSourceStream.h +++ b/dbms/src/Flash/Coprocessor/MockSourceStream.h @@ -78,7 +78,6 @@ std::pair>> 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(context, columns_with_type_and_name, max_streams, log); } diff --git a/dbms/src/Flash/tests/gtest_limit_executor.cpp b/dbms/src/Flash/tests/gtest_limit_executor.cpp index 79dc56ee6b5..51313fb3e5c 100644 --- a/dbms/src/Flash/tests/gtest_limit_executor.cpp +++ b/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("col", col)}); + + context.addMockTable({"test", "notNull"}, + {{"col", TiDB::TP::TypeString, false}}, + {toVec("col", {"a", "b", "c", "d", "e", "f", "g", "h"})}); } std::shared_ptr 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(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 diff --git a/dbms/src/TestUtils/mockExecutor.cpp b/dbms/src/TestUtils/mockExecutor.cpp index fafd98758c3..b4cd1e0e108 100644 --- a/dbms/src/TestUtils/mockExecutor.cpp +++ b/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(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 diff --git a/dbms/src/TestUtils/mockExecutor.h b/dbms/src/TestUtils/mockExecutor.h index b2dee9a6d9e..b3fb5d522c8 100644 --- a/dbms/src/TestUtils/mockExecutor.h +++ b/dbms/src/TestUtils/mockExecutor.h @@ -30,8 +30,6 @@ namespace DB { namespace tests { -using MockColumnInfo = std::pair; -using MockColumnInfoVec = std::vector; using MockTableName = std::pair; using MockOrderByItem = std::pair; using MockOrderByItemVec = std::vector; @@ -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);