pingcap · ti-chi-bot · Dec 22, 2022 · Dec 19, 2022 · Dec 20, 2022 · Dec 20, 2022
diff --git a/dbms/src/Debug/MockComputeServerManager.cpp b/dbms/src/Debug/MockComputeServerManager.cpp
@@ -13,6 +13,8 @@
 // limitations under the License.
 #include <Common/FmtUtils.h>
 #include <Debug/MockComputeServerManager.h>
+#include <Flash/Mpp/MPPTaskManager.h>
+#include <Storages/Transaction/TMTContext.h>
 #include <TestUtils/TiFlashTestEnv.h>
 
 #include <chrono>

diff --git a/dbms/src/Debug/MockComputeServerManager.h b/dbms/src/Debug/MockComputeServerManager.h
@@ -15,10 +15,8 @@
 #pragma once
 
 #include <Debug/MockStorage.h>
-#include <Flash/Mpp/MPPTaskManager.h>
+#include <Flash/Mpp/MPPTaskId.h>
 #include <Server/FlashGrpcServerHolder.h>
-#include <Storages/Transaction/TMTContext.h>
-#include <TestUtils/TiFlashTestEnv.h>
 
 namespace DB::tests
 {

diff --git a/dbms/src/Debug/MockExecutor/AggregationBinder.cpp b/dbms/src/Debug/MockExecutor/AggregationBinder.cpp
@@ -12,11 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <AggregateFunctions/AggregateFunctionFactory.h>
 #include <AggregateFunctions/AggregateFunctionUniq.h>
 #include <Debug/MockExecutor/AggregationBinder.h>
+#include <Debug/MockExecutor/AstToPB.h>
 #include <Debug/MockExecutor/ExchangeReceiverBinder.h>
 #include <Debug/MockExecutor/ExchangeSenderBinder.h>
-#include <Debug/MockExecutor/ExecutorBinder.h>
+#include <Debug/MockExecutor/FuncSigMap.h>
+#include <Parsers/ASTIdentifier.h>
 #include <fmt/core.h>
 
 namespace DB::mock

diff --git a/dbms/src/Debug/MockExecutor/AggregationBinder.h b/dbms/src/Debug/MockExecutor/AggregationBinder.h
@@ -14,12 +14,14 @@
 
 #pragma once
 
-#include <Debug/MockExecutor/ExchangeReceiverBinder.h>
-#include <Debug/MockExecutor/ExchangeSenderBinder.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
+#include <Parsers/ASTFunction.h>
 
 namespace DB::mock
 {
+class ExchangeSenderBinder;
+class ExchangeReceiverBinder;
+
 class AggregationBinder : public ExecutorBinder
 {
 public:

diff --git a/dbms/src/Debug/MockExecutor/AstToPB.cpp b/dbms/src/Debug/MockExecutor/AstToPB.cpp
@@ -12,7 +12,25 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <AggregateFunctions/AggregateFunctionFactory.h>
+#include <DataTypes/FieldToDataType.h>
 #include <Debug/MockExecutor/AstToPB.h>
+#include <Debug/MockExecutor/AstToPBUtils.h>
+#include <Debug/MockExecutor/FuncSigMap.h>
+#include <Flash/Coprocessor/DAGCodec.h>
+#include <Flash/Coprocessor/DAGUtils.h>
+#include <Functions/FunctionFactory.h>
+#include <Interpreters/Context.h>
+#include <Interpreters/convertFieldToType.h>
+#include <Parsers/ASTFunction.h>
+#include <Parsers/ASTIdentifier.h>
+#include <Parsers/ASTLiteral.h>
+#include <Storages/Transaction/TypeMapping.h>
+#include <Storages/Transaction/Types.h>
+#include <common/logger_useful.h>
+#include <common/types.h>
+#include <tipb/executor.pb.h>
+#include <tipb/select.pb.h>
 
 namespace DB
 {

diff --git a/dbms/src/Debug/MockExecutor/AstToPB.h b/dbms/src/Debug/MockExecutor/AstToPB.h
@@ -14,24 +14,8 @@
 
 #pragma once
 
-#include <AggregateFunctions/AggregateFunctionFactory.h>
-#include <DataTypes/FieldToDataType.h>
-#include <Debug/MockExecutor/AstToPBUtils.h>
-#include <Debug/MockExecutor/FuncSigMap.h>
-#include <Flash/Coprocessor/DAGCodec.h>
-#include <Flash/Coprocessor/DAGUtils.h>
-#include <Functions/FunctionFactory.h>
-#include <Interpreters/Context.h>
-#include <Interpreters/convertFieldToType.h>
-#include <Parsers/ASTFunction.h>
-#include <Parsers/ASTIdentifier.h>
-#include <Parsers/ASTLiteral.h>
-#include <Storages/Transaction/TypeMapping.h>
-#include <Storages/Transaction/Types.h>
-#include <common/logger_useful.h>
-#include <common/types.h>
-#include <tipb/executor.pb.h>
-#include <tipb/select.pb.h>
+#include <Flash/Coprocessor/ChunkCodec.h>
+#include <Storages/Transaction/TiDB.h>
 
 namespace DB
 {
@@ -41,6 +25,11 @@ extern const int BAD_ARGUMENTS;
 extern const int LOGICAL_ERROR;
 extern const int NO_SUCH_COLUMN_IN_TABLE;
 } // namespace ErrorCodes
+
+class ASTFunction;
+class ASTIdentifier;
+class Context;
+
 struct MPPCtx
 {
     Timestamp start_ts;

diff --git a/dbms/src/Debug/MockExecutor/ExchangeReceiverBinder.cpp b/dbms/src/Debug/MockExecutor/ExchangeReceiverBinder.cpp
@@ -12,8 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Debug/MockComputeServerManager.h>
+#include <Debug/MockExecutor/AstToPB.h>
+#include <Debug/MockExecutor/AstToPBUtils.h>
 #include <Debug/MockExecutor/ExchangeReceiverBinder.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
+#include <Storages/Transaction/TiDB.h>
+#include <kvproto/mpp.pb.h>
 
 namespace DB::mock
 {

diff --git a/dbms/src/Debug/MockExecutor/ExchangeReceiverBinder.h b/dbms/src/Debug/MockExecutor/ExchangeReceiverBinder.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <Debug/MockExecutor/AstToPB.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
 
 namespace DB::mock

diff --git a/dbms/src/Debug/MockExecutor/ExchangeSenderBinder.cpp b/dbms/src/Debug/MockExecutor/ExchangeSenderBinder.cpp
@@ -12,8 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Debug/MockComputeServerManager.h>
+#include <Debug/MockExecutor/AstToPB.h>
+#include <Debug/MockExecutor/AstToPBUtils.h>
 #include <Debug/MockExecutor/ExchangeSenderBinder.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
+#include <Flash/Coprocessor/DAGCodec.h>
+#include <Interpreters/Context.h>
 
 namespace DB::mock
 {

diff --git a/dbms/src/Debug/MockExecutor/ExchangeSenderBinder.h b/dbms/src/Debug/MockExecutor/ExchangeSenderBinder.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <Debug/MockExecutor/AstToPB.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
 
 namespace DB::mock

diff --git a/dbms/src/Debug/MockExecutor/ExecutorBinder.h b/dbms/src/Debug/MockExecutor/ExecutorBinder.h
@@ -14,18 +14,20 @@
 
 #pragma once
 
+#include <Common/typeid_cast.h>
 #include <Debug/DAGProperties.h>
-#include <Debug/MockComputeServerManager.h>
 #include <Debug/MockExecutor/AstToPB.h>
-#include <Debug/MockExecutor/AstToPBUtils.h>
+#include <Flash/Coprocessor/ChunkCodec.h>
+#include <Parsers/IAST.h>
+#include <kvproto/mpp.pb.h>
+#include <tipb/executor.pb.h>
 
 
 namespace DB::mock
 {
 class ExchangeSenderBinder;
 class ExchangeReceiverBinder;
 
-
 // Convert CH AST to tipb::Executor
 // Used in integration test framework and Unit test framework.
 class ExecutorBinder
@@ -45,7 +47,7 @@ class ExecutorBinder
         index_++;
     }
 
-    std::vector<std::shared_ptr<ExecutorBinder>> getChildren()
+    std::vector<std::shared_ptr<ExecutorBinder>> getChildren() const
     {
         return children;
     }

diff --git a/dbms/src/Debug/MockExecutor/JoinBinder.cpp b/dbms/src/Debug/MockExecutor/JoinBinder.cpp
@@ -12,10 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Debug/MockExecutor/AstToPB.h>
+#include <Debug/MockExecutor/AstToPBUtils.h>
 #include <Debug/MockExecutor/ExchangeReceiverBinder.h>
 #include <Debug/MockExecutor/ExchangeSenderBinder.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
 #include <Debug/MockExecutor/JoinBinder.h>
+#include <Flash/Coprocessor/DAGCodec.h>
+#include <Parsers/ASTIdentifier.h>
 #include <Parsers/ASTTablesInSelectQuery.h>
 
 namespace DB::mock

diff --git a/dbms/src/Debug/MockExecutor/JoinBinder.h b/dbms/src/Debug/MockExecutor/JoinBinder.h
@@ -14,12 +14,12 @@
 
 #pragma once
 
-#include <Debug/MockExecutor/ExchangeReceiverBinder.h>
-#include <Debug/MockExecutor/ExchangeSenderBinder.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
 
 namespace DB::mock
 {
+class ExchangeSenderBinder;
+class ExchangeReceiverBinder;
 class JoinBinder : public ExecutorBinder
 {
 public:

diff --git a/dbms/src/Debug/MockExecutor/LimitBinder.cpp b/dbms/src/Debug/MockExecutor/LimitBinder.cpp
@@ -14,6 +14,7 @@
 
 #include <Debug/MockExecutor/ExecutorBinder.h>
 #include <Debug/MockExecutor/LimitBinder.h>
+#include <Parsers/ASTLiteral.h>
 
 namespace DB::mock
 {

diff --git a/dbms/src/Debug/MockExecutor/ProjectBinder.cpp b/dbms/src/Debug/MockExecutor/ProjectBinder.cpp
@@ -12,9 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <AggregateFunctions/AggregateFunctionFactory.h>
+#include <Debug/MockExecutor/AstToPB.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
 #include <Debug/MockExecutor/ProjectBinder.h>
+#include <Flash/Coprocessor/DAGCodec.h>
 #include <Parsers/ASTAsterisk.h>
+#include <Parsers/ASTFunction.h>
 
 namespace DB::mock
 {

diff --git a/dbms/src/Debug/MockExecutor/SelectionBinder.cpp b/dbms/src/Debug/MockExecutor/SelectionBinder.cpp
@@ -11,7 +11,7 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-#include <Debug/MockExecutor/ExecutorBinder.h>
+#include <Debug/MockExecutor/AstToPB.h>
 #include <Debug/MockExecutor/SelectionBinder.h>
 
 namespace DB::mock

diff --git a/dbms/src/Debug/MockExecutor/SelectionBinder.h b/dbms/src/Debug/MockExecutor/SelectionBinder.h
@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include <Debug/MockExecutor/AstToPB.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
 
 namespace DB::mock

diff --git a/dbms/src/Debug/MockExecutor/SortBinder.cpp b/dbms/src/Debug/MockExecutor/SortBinder.cpp
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Debug/MockExecutor/AstToPB.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
 #include <Debug/MockExecutor/SortBinder.h>
 #include <Parsers/ASTOrderByElement.h>

diff --git a/dbms/src/Debug/MockExecutor/TableScanBinder.cpp b/dbms/src/Debug/MockExecutor/TableScanBinder.cpp
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Debug/MockExecutor/AstToPBUtils.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
 #include <Debug/MockExecutor/TableScanBinder.h>
 #include <Storages/MutableSupport.h>

diff --git a/dbms/src/Debug/MockExecutor/TopNBinder.cpp b/dbms/src/Debug/MockExecutor/TopNBinder.cpp
@@ -12,8 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Debug/MockExecutor/AstToPB.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
 #include <Debug/MockExecutor/TopNBinder.h>
+#include <Parsers/ASTLiteral.h>
 #include <Parsers/ASTOrderByElement.h>
 
 namespace DB::mock

diff --git a/dbms/src/Debug/MockExecutor/WindowBinder.cpp b/dbms/src/Debug/MockExecutor/WindowBinder.cpp
@@ -12,8 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Debug/MockExecutor/AstToPB.h>
 #include <Debug/MockExecutor/ExecutorBinder.h>
+#include <Debug/MockExecutor/FuncSigMap.h>
 #include <Debug/MockExecutor/WindowBinder.h>
+#include <Parsers/ASTFunction.h>
 
 namespace DB::mock
 {

diff --git a/dbms/src/Debug/MockExecutor/WindowBinder.h b/dbms/src/Debug/MockExecutor/WindowBinder.h
@@ -27,7 +27,6 @@ struct MockWindowFrame
     std::optional<MockWindowFrameBound> end;
     // TODO: support calcFuncs
 };
-
 using ASTPartitionByElement = ASTOrderByElement;
 
 class WindowBinder : public ExecutorBinder

diff --git a/dbms/src/Debug/MockStorage.cpp b/dbms/src/Debug/MockStorage.cpp
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 #include <Debug/MockStorage.h>
+#include <Flash/Coprocessor/TiDBTableScan.h>
 
 namespace DB::tests
 {
@@ -22,8 +23,11 @@ void MockStorage::addTableSchema(const String & name, const MockColumnInfoVec &
     addTableInfo(name, columnInfos);
 }
 
-void MockStorage::addTableData(const String & name, const ColumnsWithTypeAndName & columns)
+void MockStorage::addTableData(const String & name, ColumnsWithTypeAndName & columns)
 {
+    for (size_t i = 0; i < columns.size(); ++i)
+        columns[i].column_id = i;
+
     table_columns[getTableId(name)] = columns;
 }
 
@@ -123,11 +127,12 @@ CutColumnInfo getCutColumnInfo(size_t rows, Int64 partition_id, Int64 partition_
     return {start, cur_rows};
 }
 
-ColumnsWithTypeAndName MockStorage::getColumnsForMPPTableScan(Int64 table_id, Int64 partition_id, Int64 partition_num)
+ColumnsWithTypeAndName MockStorage::getColumnsForMPPTableScan(const TiDBTableScan & table_scan, Int64 partition_id, Int64 partition_num)
 {
+    auto table_id = table_scan.getLogicalTableID();
     if (tableExists(table_id))
     {
-        auto columns_with_type_and_name = table_columns[table_id];
+        auto columns_with_type_and_name = table_columns[table_scan.getLogicalTableID()];
         size_t rows = 0;
         for (const auto & col : columns_with_type_and_name)
         {
@@ -141,11 +146,23 @@ ColumnsWithTypeAndName MockStorage::getColumnsForMPPTableScan(Int64 table_id, In
         ColumnsWithTypeAndName res;
         for (const auto & column_with_type_and_name : columns_with_type_and_name)
         {
-            res.push_back(
-                ColumnWithTypeAndName(
-                    column_with_type_and_name.column->cut(cut_info.first, cut_info.second),
-                    column_with_type_and_name.type,
-                    column_with_type_and_name.name));
+            bool contains = false;
+            for (const auto & column : table_scan.getColumns())
+            {
+                if (column.id == column_with_type_and_name.column_id)
+                {
+                    contains = true;
+                    break;
+                }
+            }
+            if (contains)
+            {
+                res.push_back(
+                    ColumnWithTypeAndName(
+                        column_with_type_and_name.column->cut(cut_info.first, cut_info.second),
+                        column_with_type_and_name.type,
+                        column_with_type_and_name.name));
+            }
         }
         return res;
     }