Add tests for aggregates with nullable GROUP BY in Column Shards

2 files changed, 144 insertions, 3 deletions

diff --git a/ydb/core/kqp/opt/physical/kqp_opt_phy_olap_agg.cpp b/ydb/core/kqp/opt/physical/kqp_opt_phy_olap_agg.cpp
index ad312ecb30..7baeb10b8e 100644
--- a/ydb/core/kqp/opt/physical/kqp_opt_phy_olap_agg.cpp
+++ b/ydb/core/kqp/opt/physical/kqp_opt_phy_olap_agg.cpp
@@ -7,6 +7,7 @@
 #include <ydb/library/yql/core/yql_opt_utils.h>
 
 #include <vector>
+#include <unordered_set>
 
 namespace NKikimr::NKqp::NOpt {
 
@@ -15,6 +16,16 @@ using namespace NYql::NNodes;
 
 namespace {
 
+static const std::unordered_set<std::string> SupportedAggFuncs = {
+    "count",
+    "count_all",
+    "sum",
+    "min",
+    "max",
+    "avg",
+    "some"
+};
+
 struct TAggInfo {
     TAggInfo(const std::string& aggName, const std::string& colName, const std::string& opType)
         : AggName(aggName)
@@ -55,9 +66,7 @@ bool CanBePushedDown(const TExprBase& trait, TExprContext& ctx)
     }
     auto aggApply = trait.Cast<TCoAggApply>();
     auto aggName = aggApply.Name();
-    if (aggName == "count" || aggName == "count_all" || aggName == "sum"
-        || aggName == "min" || aggName == "max" || aggName == "avg" || aggName == "some")
-    {
+    if (SupportedAggFuncs.find(aggName.StringValue()) != SupportedAggFuncs.end()) {
         return true;
     }
     YQL_CLOG(DEBUG, ProviderKqp) << "Unsupported type of aggregation: " << aggName.StringValue();
diff --git a/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp b/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
index 8cf722c303..1b1e01eb56 100644
--- a/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
+++ b/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
@@ -1987,6 +1987,50 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
         TestTableWithNulls({ testCase });
     }
 
+    Y_UNIT_TEST(Aggregation_Count_GroupByNull) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    level, COUNT(id), COUNT(level), COUNT(*)
+                FROM `/Root/tableWithNulls`
+                WHERE id > 5
+                GROUP BY level
+                ORDER BY level;
+            )")
+            .SetExpectedReply("[[#;5u;0u;5u]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
+    Y_UNIT_TEST(Aggregation_Count_GroupByNullMix) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    level, COUNT(id), COUNT(level), COUNT(*)
+                FROM `/Root/tableWithNulls`
+                WHERE id >= 5
+                GROUP BY level
+                ORDER BY level;
+            )")
+            .SetExpectedReply("[[#;5u;0u;5u];[[5];1u;1u;1u]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
     Y_UNIT_TEST(Aggregation_NoPushdownOnDisabledEmitAggApply) {
         TAggregationTestCase testCase;
         testCase.SetQuery(R"(
@@ -2136,6 +2180,50 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
         TestTableWithNulls({ testCase });
     }
 
+    Y_UNIT_TEST(Aggregation_Avg_GroupByNull) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    level, AVG(id), AVG(level)
+                FROM `/Root/tableWithNulls`
+                WHERE id > 5
+                GROUP BY level
+                ORDER BY level;
+            )")
+            .SetExpectedReply("[[#;[8.];#]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
+    Y_UNIT_TEST(Aggregation_Avg_GroupByNullMix) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    level, AVG(id), AVG(level)
+                FROM `/Root/tableWithNulls`
+                WHERE id >= 5
+                GROUP BY level
+                ORDER BY level;
+            )")
+            .SetExpectedReply("[[#;[8.];#];[[5];[5.];[5.]]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
     Y_UNIT_TEST(Aggregation_Sum) {
         TAggregationTestCase testCase;
         testCase.SetQuery(R"(
@@ -2252,6 +2340,50 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
         TestTableWithNulls({ testCase });
     }
 
+    Y_UNIT_TEST(Aggregation_Sum_GroupByNull) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    level, SUM(id), SUM(level)
+                FROM `/Root/tableWithNulls`
+                WHERE id > 5
+                GROUP BY level
+                ORDER BY level;
+            )")
+            .SetExpectedReply("[[#;[40];#]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
+    Y_UNIT_TEST(Aggregation_Sum_GroupByNullMix) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    level, SUM(id), SUM(level)
+                FROM `/Root/tableWithNulls`
+                WHERE id >= 5
+                GROUP BY level
+                ORDER BY level;
+            )")
+            .SetExpectedReply("[[#;[40];#];[[5];[5];[5]]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
     Y_UNIT_TEST(Aggregation_SumL_GroupL_OrderL) {
         TAggregationTestCase testCase;
         testCase.SetQuery(R"(