Add gtest for Limit, TopN, Projection (#5187)

dbms/src/Debug/astToExecutor.cpp

@@ -1629,7 +1629,6 @@ ExecutorPtr compileProject(ExecutorPtr input, size_t & executor_index, ASTPtr se
             }
         }
     }
-
     auto project = std::make_shared<mock::Project>(executor_index, output_schema, std::move(exprs));
     project->children.push_back(input);
     return project;

dbms/src/Flash/tests/gtest_executor.cpp

@@ -227,4 +227,4 @@ try
 CATCH
 
 } // namespace tests
-} // namespace DB
+} // namespace DB

dbms/src/Flash/tests/gtest_limit_executor.cpp

@@ -0,0 +1,77 @@
+// Copyright 2022 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// 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 <TestUtils/ExecutorTestUtils.h>
+#include <TestUtils/mockExecutor.h>
+
+namespace DB
+{
+namespace tests
+{
+
+class ExecutorLimitTestRunner : public DB::tests::ExecutorTest
+{
+public:
+    using ColDataType = std::optional<typename TypeTraits<String>::FieldType>;
+    using ColumnWithData = std::vector<ColDataType>;
+
+    void initializeContext() override
+    {
+        ExecutorTest::initializeContext();
+
+        context.addMockTable({db_name, table_name},
+                             {{col_name, TiDB::TP::TypeString}},
+                             {toNullableVec<String>(col_name, col0)});
+    }
+
+    std::shared_ptr<tipb::DAGRequest> getRequest(size_t limit_num)
+    {
+        return context.scan(db_name, table_name).limit(limit_num).build(context);
+    }
+
+    /// Prepare some names
+    const String db_name{"test_db"};
+    const String table_name{"projection_test_table"};
+    const String col_name{"limit_col"};
+    const ColumnWithData col0{"col0-0", {}, "col0-2", "col0-3", {}, "col0-5", "col0-6", "col0-7"};
+};
+
+TEST_F(ExecutorLimitTestRunner, Limit)
+try
+{
+    std::shared_ptr<tipb::DAGRequest> request;
+    ColumnsWithTypeAndName expect_cols;
+
+    /// Check limit result with various parameters
+    const size_t col_data_num = col0.size();
+    for (size_t limit_num = 0; limit_num <= col_data_num + 3; ++limit_num)
+    {
+        if (limit_num == col_data_num + 3)
+            limit_num = INT_MAX;
+        request = getRequest(limit_num);
+
+        if (limit_num == 0)
+            expect_cols = {};
+        else if (limit_num > col_data_num)
+            expect_cols = {toNullableVec<String>(col_name, ColumnWithData(col0.begin(), col0.end()))};
+        else
+            expect_cols = {toNullableVec<String>(col_name, ColumnWithData(col0.begin(), col0.begin() + limit_num))};
+
+        executeStreams(request, expect_cols);
+    }
+}
+CATCH
+
+} // namespace tests
+} // namespace DB

dbms/src/Flash/tests/gtest_projection_executor.cpp

@@ -0,0 +1,225 @@
+// Copyright 2022 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// 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 <TestUtils/ExecutorTestUtils.h>
+#include <TestUtils/mockExecutor.h>
+
+namespace DB
+{
+namespace tests
+{
+
+class ExecutorProjectionTestRunner : public DB::tests::ExecutorTest
+{
+public:
+    using ColDataString = std::vector<std::optional<typename TypeTraits<String>::FieldType>>;
+    using ColDataInt32 = std::vector<std::optional<typename TypeTraits<Int32>::FieldType>>;
+
+    void initializeContext() override
+    {
+        ExecutorTest::initializeContext();
+
+        context.addMockTable({db_name, table_name},
+                             {{col_names[0], TiDB::TP::TypeString},
+                              {col_names[1], TiDB::TP::TypeString},
+                              {col_names[2], TiDB::TP::TypeString},
+                              {col_names[3], TiDB::TP::TypeLong},
+                              {col_names[4], TiDB::TP::TypeLong}},
+                             {toNullableVec<String>(col_names[0], col0),
+                              toNullableVec<String>(col_names[1], col1),
+                              toNullableVec<String>(col_names[2], col2),
+                              toNullableVec<Int32>(col_names[3], col3),
+                              toNullableVec<Int32>(col_names[4], col4)});
+    }
+
+    template <typename T>
+    std::shared_ptr<tipb::DAGRequest> getRequest(T param, const String & sort_col)
+    {
+        /// topN is introduced, so that we can get stable results in concurrency environment.
+        return context.scan(db_name, table_name).project(param).topN(sort_col, false, 100).build(context);
+    };
+
+    void executeWithConcurrency(const std::shared_ptr<tipb::DAGRequest> & request, const ColumnsWithTypeAndName & expect_columns)
+    {
+        for (size_t i = 1; i < 10; i += 2)
+        {
+            executeStreams(request, expect_columns, i);
+        }
+    }
+
+    /// Prepare column data
+    const ColDataString col0{"col0-0", "col0-1", "", "col0-2", {}, "col0-3", ""};
+    const ColDataString col1{"col1-0", {}, "", "col1-1", "", "col1-2", "col1-3"};
+    const ColDataString col2{"", "col2-0", "col2-1", {}, "col2-3", {}, "col2-4"};
+    const ColDataInt32 col3{1, {}, 0, -111111, {}, 0, 9999};
+
+    /** Each value in col4 should be different from each other so that topn 
+     *  could sort the columns into an unique result, or multi-results could
+     *  be right.
+     */
+    const ColDataInt32 col4{0, 5, -123, -234, {}, 24353, 9999};
+
+    /// Results after sorted by col4
+    const ColDataString col0_sorted_asc{{}, "col0-2", "", "col0-0", "col0-1", "", "col0-3"};
+    const ColDataString col1_sorted_asc{"", "col1-1", "", "col1-0", {}, "col1-3", "col1-2"};
+    const ColDataString col2_sorted_asc{"col2-3", {}, "col2-1", "", "col2-0", "col2-4", {}};
+    const ColDataInt32 col3_sorted_asc{{}, -111111, 0, 1, {}, 9999, 0};
+    const ColDataInt32 col4_sorted_asc{{}, -234, -123, 0, 5, 9999, 24353};
+
+    /// Prepare some names
+    std::vector<String> col_names{"col0", "col1", "col2", "col3", "col4"};
+    const String db_name{"test_db"};
+    const String table_name{"projection_test_table"};
+};
+
+TEST_F(ExecutorProjectionTestRunner, Projection)
+try
+{
+    /// Check single column
+    auto request = getRequest<MockColumnNames>({col_names[4]}, col_names[4]);
+    executeWithConcurrency(request, {toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    /// Check multi columns
+    request = getRequest<MockColumnNames>({col_names[0], col_names[4]}, col_names[4]);
+    executeWithConcurrency(request,
+                           {
+                               toNullableVec<String>(col_names[0], col0_sorted_asc),
+                               toNullableVec<Int32>(col_names[4], col4_sorted_asc),
+                           });
+
+    /// Check multi columns
+    request = getRequest<MockColumnNames>({col_names[0], col_names[1], col_names[4]}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<String>(col_names[0], col0_sorted_asc),
+                            toNullableVec<String>(col_names[1], col1_sorted_asc),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    /// Check duplicate columns
+    request = getRequest<MockColumnNames>({col_names[4], col_names[4], col_names[4]}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<Int32>(col_names[4], col4_sorted_asc),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    {
+        /// Check large number of columns
+        const size_t col_num = 100;
+        MockColumnNamesVec projection_input;
+        ColumnsWithTypeAndName columns;
+        auto expect_column = toNullableVec<Int32>(col_names[4], col4_sorted_asc);
+
+        for (size_t i = 0; i < col_num; ++i)
+        {
+            projection_input.push_back(col_names[4]);
+            columns.push_back(expect_column);
+        }
+
+        request = getRequest<MockColumnNamesVec>(projection_input, col_names[4]);
+        executeWithConcurrency(request, columns);
+    }
+}
+CATCH
+
+TEST_F(ExecutorProjectionTestRunner, ProjectionFunction)
+try
+{
+    std::shared_ptr<tipb::DAGRequest> request;
+
+    /// Test "equal" function
+
+    /// Data type: TypeString
+    request = getRequest<MockAsts>({eq(col(col_names[0]), col(col_names[0])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 1, 1, 1, 1, 1, 1}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    request = getRequest<MockAsts>({eq(col(col_names[0]), col(col_names[1])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 0, 1, 0, {}, 0, 0}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    /// Data type: TypeLong
+    request = getRequest<MockAsts>({eq(col(col_names[3]), col(col_names[4])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 0, 0, 0, {}, 1, 0}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+
+    /// Test "greater" function
+
+    /// Data type: TypeString
+    request = getRequest<MockAsts>({gt(col(col_names[0]), col(col_names[1])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 0, 0, 0, {}, 0, 0}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    request = getRequest<MockAsts>({gt(col(col_names[1]), col(col_names[0])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 1, 0, 1, {}, 1, 1}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    /// Data type: TypeLong
+    request = getRequest<MockAsts>({gt(col(col_names[3]), col(col_names[4])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 0, 1, 1, {}, 0, 0}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    request = getRequest<MockAsts>({gt(col(col_names[4]), col(col_names[3])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 1, 0, 0, {}, 0, 1}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+
+    /// Test "and" function
+
+    /// Data type: TypeString
+    request = getRequest<MockAsts>({And(col(col_names[0]), col(col_names[0])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 0, 0, 0, 0, 0, 0}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    request = getRequest<MockAsts>({And(col(col_names[0]), col(col_names[1])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({0, 0, 0, 0, 0, 0, 0}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    /// Data type: TypeLong
+    request = getRequest<MockAsts>({And(col(col_names[3]), col(col_names[4])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 1, 0, 0, {}, 1, 0}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    /// Test "not" function
+
+    /// Data type: TypeString
+    request = getRequest<MockAsts>({NOT(col(col_names[0])), NOT(col(col_names[1])), NOT(col(col_names[2])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 1, 1, 1, 1, 1, 1}),
+                            toNullableVec<UInt64>({1, 1, 1, 1, {}, 1, 1}),
+                            toNullableVec<UInt64>({1, {}, 1, 1, 1, 1, {}}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    /// Data type: TypeLong
+    request = getRequest<MockAsts>({NOT(col(col_names[3])), NOT(col(col_names[4])), col(col_names[4])}, col_names[4]);
+    executeWithConcurrency(request,
+                           {toNullableVec<UInt64>({{}, 0, 1, 0, {}, 0, 1}),
+                            toNullableVec<UInt64>({{}, 0, 0, 1, 0, 0, 0}),
+                            toNullableVec<Int32>(col_names[4], col4_sorted_asc)});
+
+    /// TODO more functions...
+}
+CATCH
+
+} // namespace tests
+} // namespace DB