Fix COUNT(*) pushdown for OLAP tables AST tree to return Stream of Uint64.

Fix COUNT(*) push-down return type from stage

2 files changed, 17 insertions, 4 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 6443ad1427..3ecedb71dd 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
@@ -270,7 +270,18 @@ TExprBase KqpPushOlapLength(TExprBase node, TExprContext& ctx, const TKqpOptimiz
         .Process(newProcessLambda)
         .Done();
 
-    return newRead;
+    auto member = Build<TCoMap>(ctx, node.Pos())
+        .Input(newRead)
+        .Lambda()
+            .Args({"row"})
+            .Body<TCoMember>()
+                .Struct("row")
+                .Name(dqPhyLength.Name())
+                .Build()
+            .Build()
+        .Done();
+
+    return member;
 }
 
 } // namespace NKikimr::NKqp::NOpt
\ No newline at end of file
diff --git a/ydb/core/kqp/ut/kqp_olap_ut.cpp b/ydb/core/kqp/ut/kqp_olap_ut.cpp
index e85c27b93f..181e8814ff 100644
--- a/ydb/core/kqp/ut/kqp_olap_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_olap_ut.cpp
@@ -19,6 +19,7 @@
 #include <ydb/core/testlib/tablet_helpers.h>
 #include <util/system/sanitizers.h>
 
+#include <fmt/format.h>
 
 namespace NKikimr {
 namespace NKqp {
@@ -1367,7 +1368,7 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
         }
     }
 
-    Y_UNIT_TEST(CountAllPushdown) {
+    Y_UNIT_TEST_TWIN(CountAllPushdown, UseLlvm) {
         // remove this return when COUNT(*) will be implemented on columnshard
         return;
 
@@ -1391,13 +1392,14 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
         }
 
         {
-            TString query = R"(
+            TString query = fmt::format(R"(
                 --!syntax_v1
                 PRAGMA Kikimr.KqpPushOlapProcess = "true";
+                PRAGMA ydb.EnableLlvm = "{}";
                 SELECT
                     COUNT(*)
                 FROM `/Root/olapStore/olapTable`
-            )";
+            )", UseLlvm ? "true" : "false");
             auto it = tableClient.StreamExecuteScanQuery(query).GetValueSync();
 
             UNIT_ASSERT_C(it.IsSuccess(), it.GetIssues().ToString());