Add SOME aggregation pushdown to column shard

Add SOME aggregation pushdown to column shard

3 files changed, 140 insertions, 2 deletions

diff --git a/ydb/core/kqp/host/kqp_type_ann.cpp b/ydb/core/kqp/host/kqp_type_ann.cpp
index b902520297a..ae4907c9b88 100644
--- a/ydb/core/kqp/host/kqp_type_ann.cpp
+++ b/ydb/core/kqp/host/kqp_type_ann.cpp
@@ -924,7 +924,7 @@ TStatus AnnotateOlapAgg(const TExprNode::TPtr& node, TExprContext& ctx) {
                 return TStatus::Error;
             }
             aggTypes.push_back(ctx.MakeType<TItemExprType>(aggName->Content(), resultType));
-        } else if (opType->Content() == "min" || opType->Content() == "max") {
+        } else if (opType->Content() == "min" || opType->Content() == "max" || opType->Content() == "some") {
             auto colType = structType->FindItemType(colName->Content());
             aggTypes.push_back(ctx.MakeType<TItemExprType>(aggName->Content(), colType));
         } else {
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 a0a0fbaa481..ad312ecb30e 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
@@ -56,7 +56,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 == "min" || aggName == "max" || aggName == "avg" || aggName == "some")
     {
         return true;
     }
diff --git a/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp b/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
index 655bff01b06..8cf722c303a 100644
--- a/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
+++ b/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
@@ -2363,6 +2363,144 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
         TestAggregations({ testCase });
     }
 
+    Y_UNIT_TEST(Aggregation_Some) {
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT SOME(level) FROM `/Root/tableWithNulls` WHERE id=1
+            )")
+            .SetExpectedReply("[[[1]]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+        TestTableWithNulls({ testCase });
+    }
+
+    Y_UNIT_TEST(Aggregation_Some_Null) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT SOME(level) FROM `/Root/tableWithNulls` WHERE id > 5
+            )")
+            .SetExpectedReply("[[#]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+        TestTableWithNulls({ testCase });
+    }
+
+    Y_UNIT_TEST(Aggregation_Some_GroupBy) {
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    id, SOME(level)
+                FROM `/Root/tableWithNulls`
+                WHERE id BETWEEN 4 AND 5
+                GROUP BY id
+                ORDER BY id;
+            )")
+            .SetExpectedReply("[[[4];[4]];[[5];[5]]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
+    Y_UNIT_TEST(Aggregation_Some_NullGroupBy) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    id, SOME(level)
+                FROM `/Root/tableWithNulls`
+                WHERE id BETWEEN 6 AND 7
+                GROUP BY id
+                ORDER BY id;
+            )")
+            .SetExpectedReply("[[[6];#];[[7];#]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
+    Y_UNIT_TEST(Aggregation_Some_NullMixGroupBy) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    id, SOME(level)
+                FROM `/Root/tableWithNulls`
+                WHERE id > 4 AND id < 7
+                GROUP BY id
+                ORDER BY id;
+            )")
+            .SetExpectedReply("[[[5];[5]];[[6];#]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
+    Y_UNIT_TEST(Aggregation_Some_GroupByNullMix) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    level, SOME(id), SOME(level)
+                FROM `/Root/tableWithNulls`
+                WHERE id BETWEEN 5 AND 6
+                GROUP BY level
+                ORDER BY level;
+            )")
+            .SetExpectedReply("[[#;[6];#];[[5];[5];[5]]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
+    Y_UNIT_TEST(Aggregation_Some_GroupByNull) {
+        // Wait for KIKIMR-16831 fix
+        return;
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT
+                    level, SOME(id), SOME(level)
+                FROM `/Root/tableWithNulls`
+                WHERE id = 6
+                GROUP BY level
+                ORDER BY level;
+            )")
+            .SetExpectedReply("[[#;[6];#]]")
+#if SSA_RUNTIME_VERSION >= 2U
+            .AddExpectedPlanOptions("TKqpOlapAgg");
+#else
+            .AddExpectedPlanOptions("CombineCore");
+#endif
+
+        TestTableWithNulls({ testCase });
+    }
+
     Y_UNIT_TEST(ClickBenchSmoke) {
         TAggregationTestCase q7;
         q7.SetQuery(R"(