Skip to content

Commit

Permalink
Convert count(const not-null) to count() (#5898)
Browse files Browse the repository at this point in the history
close #5613, close #5793
  • Loading branch information
Lloyd-Pottiger authored Sep 16, 2022
1 parent 892f852 commit a462d5b
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 9 deletions.
11 changes: 10 additions & 1 deletion dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@
#include <Parsers/ASTIdentifier.h>
#include <Storages/Transaction/TypeMapping.h>
#include <WindowFunctions/WindowFunctionFactory.h>


namespace DB
{
namespace ErrorCodes
Expand Down Expand Up @@ -332,11 +334,18 @@ void DAGExpressionAnalyzer::buildCommonAggFunc(
Names arg_names;
DataTypes arg_types;
TiDB::TiDBCollators arg_collators;

for (Int32 i = 0; i < child_size; ++i)
{
fillArgumentDetail(actions, expr.children(i), arg_names, arg_types, arg_collators);
}

// For count(not null column), we can transform it to count() to avoid the cost of convertToFullColumn.
if (expr.tp() == tipb::ExprType::Count && !expr.has_distinct() && child_size == 1 && !arg_types[0]->isNullable())
{
arg_names.clear();
arg_types.clear();
arg_collators.clear();
}
appendAggDescription(arg_names, arg_types, arg_collators, agg_func_name, aggregate_descriptions, aggregated_columns, empty_input_as_null);
}

Expand Down
34 changes: 26 additions & 8 deletions dbms/src/Flash/tests/gtest_aggregation_executor.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ namespace tests
types_col_name[a], types_col_name[b] \
}

class ExecutorAggTestRunner : public DB::tests::ExecutorTest
class ExecutorAggTestRunner : public ExecutorTest
{
public:
using ColStringNullableType = std::optional<typename TypeTraits<String>::FieldType>;
Expand All @@ -57,6 +57,8 @@ class ExecutorAggTestRunner : public DB::tests::ExecutorTest
using ColumnWithNullableDecimal = std::vector<ColDecimalNullableType>;
using ColumnWithUInt64 = std::vector<ColUInt64Type>;

virtual ~ExecutorAggTestRunner() = default;

void initializeContext() override
{
ExecutorTest::initializeContext();
Expand Down Expand Up @@ -94,12 +96,14 @@ class ExecutorAggTestRunner : public DB::tests::ExecutorTest
{{col_name[0], TiDB::TP::TypeLong},
{col_name[1], TiDB::TP::TypeString},
{col_name[2], TiDB::TP::TypeString},
{col_name[3], TiDB::TP::TypeDouble}},
{col_name[3], TiDB::TP::TypeDouble},
{col_name[4], TiDB::TP::TypeLong}},
/* columns= */
{toNullableVec<Int32>(col_name[0], col_age),
toNullableVec<String>(col_name[1], col_gender),
toNullableVec<String>(col_name[2], col_country),
toNullableVec<Float64>(col_name[3], col_salary)});
toNullableVec<Float64>(col_name[3], col_salary),
toVec<UInt64>(col_name[4], col_pr)});

context.addMockTable({"aggnull_test", "t1"},
{{"s1", TiDB::TP::TypeString}, {"s2", TiDB::TP::TypeString}},
Expand Down Expand Up @@ -138,7 +142,7 @@ class ExecutorAggTestRunner : public DB::tests::ExecutorTest

/// Prepare some data and names for aggregation functions
const String table_name{"clerk"};
const std::vector<String> col_name{"age", "gender", "country", "salary"};
const std::vector<String> col_name{"age", "gender", "country", "salary", "pr"};
ColumnWithNullableInt32 col_age{30, {}, 27, 32, 25, 36, {}, 22, 34};
ColumnWithNullableString col_gender{
"male",
Expand All @@ -153,6 +157,7 @@ class ExecutorAggTestRunner : public DB::tests::ExecutorTest
};
ColumnWithNullableString col_country{"russia", "korea", "usa", "usa", "usa", "china", "china", "china", "china"};
ColumnWithNullableFloat64 col_salary{1000.1, 1300.2, 0.3, {}, -200.4, 900.5, -999.6, 2000.7, -300.8};
ColumnWithUInt64 col_pr{1, 2, 0, 3290124, 968933, 3125, 31236, 4327, 80000};
};

/// Guarantee the correctness of group by
Expand Down Expand Up @@ -282,17 +287,30 @@ try
std::shared_ptr<tipb::DAGRequest> request;
auto agg_func0 = Count(col(col_name[0])); /// select count(age) from clerk group by country;
auto agg_func1 = Count(col(col_name[1])); /// select count(gender) from clerk group by country, gender;
std::vector<MockAstVec> agg_funcs = {{agg_func0}, {agg_func1}};
auto agg_func2 = Count(lit(Field(static_cast<UInt64>(1)))); /// select count(1) from clerk;
auto agg_func3 = Count(lit(Field())); /// select count(NULL) from clerk;
auto agg_func4 = Count(lit(Field(static_cast<UInt64>(1)))); /// select count(1) from clerk group by country;
auto agg_func5 = Count(lit(Field())); /// select count(NULL) from clerk group by country;
auto agg_func6 = Count(col(col_name[4])); /// select count(pr) from clerk group by country;
std::vector<MockAstVec> agg_funcs = {{agg_func0}, {agg_func1}, {agg_func2}, {agg_func3}, {agg_func4}, {agg_func5}, {agg_func6}};

auto group_by_expr0 = col(col_name[2]);
auto group_by_expr10 = col(col_name[2]);
auto group_by_expr11 = col(col_name[1]);
auto group_by_expr4 = col(col_name[2]);
auto group_by_expr5 = col(col_name[2]);
auto group_by_expr6 = col(col_name[2]);

std::vector<ColumnsWithTypeAndName> expect_cols{
{toVec<UInt64>("count(age)", ColumnWithUInt64{3, 3, 1, 0})},
{toVec<UInt64>("count(gender)", ColumnWithUInt64{2, 2, 2, 1, 1, 1})}};
std::vector<MockAstVec> group_by_exprs{{group_by_expr0}, {group_by_expr10, group_by_expr11}};
std::vector<MockColumnNameVec> projections{{"count(age)"}, {"count(gender)"}};
{toVec<UInt64>("count(gender)", ColumnWithUInt64{2, 2, 2, 1, 1, 1})},
{toVec<UInt64>("count(1)", ColumnWithUInt64{9})},
{toVec<UInt64>("count(NULL)", ColumnWithUInt64{0})},
{toVec<UInt64>("count(1)", ColumnWithUInt64{4, 3, 1, 1})},
{toVec<UInt64>("count(NULL)", ColumnWithUInt64{0, 0, 0, 0})},
{toVec<UInt64>("count(pr)", ColumnWithUInt64{4, 3, 1, 1})}};
std::vector<MockAstVec> group_by_exprs{{group_by_expr0}, {group_by_expr10, group_by_expr11}, {}, {}, {group_by_expr4}, {group_by_expr5}, {group_by_expr6}};
std::vector<MockColumnNameVec> projections{{"count(age)"}, {"count(gender)"}, {"count(1)"}, {"count(NULL)"}, {"count(1)"}, {"count(NULL)"}, {"count(pr)"}};
size_t test_num = expect_cols.size();

/// Start to test
Expand Down

0 comments on commit a462d5b

Please sign in to comment.