stable-24-2: Enable COUNT in view queries (#6820) (#7501)

ydb/core/kqp/provider/rewrite_io_utils.cpp

@@ -28,7 +28,8 @@ NSQLTranslation::TTranslationSettings CreateViewTranslationSettings(const TStrin
 TExprNode::TPtr CompileViewQuery(
     const TString& query,
     TExprContext& ctx,
-    const TString& cluster
+    const TString& cluster,
+    IModuleResolver::TPtr moduleResolver
 ) {
     TAstParseResult queryAst;
     queryAst = NSQLTranslation::SqlToYql(query, CreateViewTranslationSettings(cluster));
@@ -39,7 +40,7 @@ TExprNode::TPtr CompileViewQuery(
     }
 
     TExprNode::TPtr queryGraph;
-    if (!CompileExpr(*queryAst.Root, queryGraph, ctx, nullptr, nullptr)) {
+    if (!CompileExpr(*queryAst.Root, queryGraph, ctx, moduleResolver.get(), nullptr)) {
         return nullptr;
     }
 
@@ -123,14 +124,15 @@ TExprNode::TPtr RewriteReadFromView(
     const TExprNode::TPtr& node,
     TExprContext& ctx,
     const TString& query,
-    const TString& cluster
+    const TString& cluster,
+    IModuleResolver::TPtr moduleResolver
 ) {
     const TCoRead readNode(node->ChildPtr(0));
     const auto worldBeforeThisRead = readNode.World().Ptr();
 
     TExprNode::TPtr queryGraph = FindSavedQueryGraph(readNode.Ptr());
     if (!queryGraph) {
-        queryGraph = CompileViewQuery(query, ctx, cluster);
+        queryGraph = CompileViewQuery(query, ctx, cluster, moduleResolver);
         if (!queryGraph) {
             ctx.AddError(TIssue(ctx.GetPosition(readNode.Pos()),
                          "The query stored in the view cannot be compiled."));

ydb/core/kqp/provider/rewrite_io_utils.h

@@ -8,7 +8,8 @@ TExprNode::TPtr RewriteReadFromView(
     const TExprNode::TPtr& node,
     TExprContext& ctx,
     const TString& query,
-    const TString& cluster
+    const TString& cluster,
+    IModuleResolver::TPtr moduleResolver
 );
 
 }

ydb/core/kqp/provider/yql_kikimr_datasource.cpp

@@ -754,6 +754,7 @@ class TKikimrDataSource : public TDataProviderBase {
                 }
 
                 ctx.Step
+                    .Repeat(TExprStep::ExpandApplyForLambdas)
                     .Repeat(TExprStep::ExprEval)
                     .Repeat(TExprStep::DiscoveryIO)
                     .Repeat(TExprStep::Epochs)
@@ -762,7 +763,7 @@ class TKikimrDataSource : public TDataProviderBase {
                     .Repeat(TExprStep::RewriteIO);
 
                 const auto& query = tableDesc.Metadata->ViewPersistedData.QueryText;
-                return RewriteReadFromView(node, ctx, query, cluster);
+                return RewriteReadFromView(node, ctx, query, cluster, Types.Modules);
             }
         }
 

ydb/core/kqp/ut/view/input/cases/aggregates_and_window/create_view.sql

@@ -0,0 +1,43 @@
+CREATE VIEW `/Root/aggregates_and_window` WITH (security_invoker = TRUE) AS
+    SELECT
+        series.title AS series,
+        series_stats.seasons_with_episode_count_greater_than_average AS seasons_with_episode_count_greater_than_average
+    FROM (
+        SELECT
+            series_id,
+            SUM(
+                CASE
+                    WHEN episode_count > average_episodes_in_season
+                        THEN 1
+                    ELSE 0
+                END
+            ) AS seasons_with_episode_count_greater_than_average
+        FROM (
+            SELECT
+                series_id,
+                season_id,
+                episode_count,
+                AVG(episode_count) OVER average_episodes_in_season_window AS average_episodes_in_season
+            FROM (
+                SELECT
+                    series_id,
+                    season_id,
+                    COUNT(*) AS episode_count
+                FROM `/Root/episodes`
+                GROUP BY
+                    series_id,
+                    season_id
+            )
+            WINDOW
+                average_episodes_in_season_window AS (
+                    PARTITION BY
+                        series_id
+                )
+        )
+        GROUP BY
+            series_id
+    )
+        AS series_stats
+    JOIN `/Root/series`
+        AS series
+    USING (series_id);

ydb/core/kqp/ut/view/input/cases/aggregates_and_window/drop_view.sql

@@ -0,0 +1 @@
+DROP VIEW `/Root/aggregates_and_window`;

ydb/core/kqp/ut/view/input/cases/aggregates_and_window/etalon_query.sql

@@ -0,0 +1,46 @@
+SELECT
+    *
+FROM (
+    SELECT
+        series.title AS series,
+        series_stats.seasons_with_episode_count_greater_than_average AS seasons_with_episode_count_greater_than_average
+    FROM (
+        SELECT
+            series_id,
+            SUM(
+                CASE
+                    WHEN episode_count > average_episodes_in_season
+                        THEN 1
+                    ELSE 0
+                END
+            ) AS seasons_with_episode_count_greater_than_average
+        FROM (
+            SELECT
+                series_id,
+                season_id,
+                episode_count,
+                AVG(episode_count) OVER average_episodes_in_season_window AS average_episodes_in_season
+            FROM (
+                SELECT
+                    series_id,
+                    season_id,
+                    COUNT(*) AS episode_count
+                FROM `/Root/episodes`
+                GROUP BY
+                    series_id,
+                    season_id
+            )
+            WINDOW
+                average_episodes_in_season_window AS (
+                    PARTITION BY
+                        series_id
+                )
+        )
+        GROUP BY
+            series_id
+    )
+        AS series_stats
+    JOIN `/Root/series`
+        AS series
+    USING (series_id)
+);

ydb/core/kqp/ut/view/input/cases/aggregates_and_window/select_from_view.sql

@@ -0,0 +1,3 @@
+SELECT
+    *
+FROM `/Root/aggregates_and_window`;

ydb/core/kqp/ut/view/input/cases/count_episodes/create_view.sql

@@ -0,0 +1,9 @@
+CREATE VIEW `/Root/count_episodes` WITH (security_invoker = TRUE) AS
+    SELECT
+        series_id,
+        season_id,
+        COUNT(*)
+    FROM `/Root/episodes`
+    GROUP BY
+        series_id,
+        season_id;

ydb/core/kqp/ut/view/input/cases/count_episodes/drop_view.sql

@@ -0,0 +1 @@
+DROP VIEW `/Root/count_episodes`;

ydb/core/kqp/ut/view/input/cases/count_episodes/etalon_query.sql

@@ -0,0 +1,12 @@
+SELECT
+    *
+FROM (
+    SELECT
+        series_id,
+        season_id,
+        COUNT(*)
+    FROM `/Root/episodes`
+    GROUP BY
+        series_id,
+        season_id
+);

ydb/core/kqp/ut/view/input/cases/count_episodes/select_from_view.sql

@@ -0,0 +1,3 @@
+SELECT
+    *
+FROM `/Root/count_episodes`;

ydb/core/kqp/ut/view/input/cases/count_episodes_with_titles/create_view.sql

@@ -0,0 +1,22 @@
+CREATE VIEW `/Root/count_episodes_with_titles` WITH (security_invoker = TRUE) AS
+    SELECT
+        series.title AS series,
+        seasons.title AS season,
+        episodes.episode_count AS episode_count
+    FROM (
+        SELECT
+            series_id,
+            season_id,
+            COUNT(*) AS episode_count
+        FROM `/Root/episodes`
+        GROUP BY
+            series_id,
+            season_id
+    )
+        AS episodes
+    JOIN `/Root/series`
+        AS series
+    ON episodes.series_id == series.series_id
+    JOIN `/Root/seasons`
+        AS seasons
+    ON episodes.series_id == seasons.series_id AND episodes.season_id == seasons.season_id;

ydb/core/kqp/ut/view/input/cases/count_episodes_with_titles/drop_view.sql

@@ -0,0 +1 @@
+DROP VIEW `/Root/count_episodes_with_titles`;

ydb/core/kqp/ut/view/input/cases/count_episodes_with_titles/etalon_query.sql

@@ -0,0 +1,25 @@
+SELECT
+    *
+FROM (
+    SELECT
+        series.title AS series,
+        seasons.title AS season,
+        episodes.episode_count AS episode_count
+    FROM (
+        SELECT
+            series_id,
+            season_id,
+            COUNT(*) AS episode_count
+        FROM `/Root/episodes`
+        GROUP BY
+            series_id,
+            season_id
+    )
+        AS episodes
+    JOIN `/Root/series`
+        AS series
+    ON episodes.series_id == series.series_id
+    JOIN `/Root/seasons`
+        AS seasons
+    ON episodes.series_id == seasons.series_id AND episodes.season_id == seasons.season_id
+);

ydb/core/kqp/ut/view/input/cases/count_episodes_with_titles/select_from_view.sql

@@ -0,0 +1,3 @@
+SELECT
+    *
+FROM `/Root/count_episodes_with_titles`;

ydb/core/kqp/ut/view/input/cases/count_rows/create_view.sql

@@ -0,0 +1,4 @@
+CREATE VIEW `/Root/count_rows` WITH (security_invoker = TRUE) AS
+    SELECT
+        COUNT(*)
+    FROM `/Root/episodes`;

ydb/core/kqp/ut/view/input/cases/count_rows/drop_view.sql

@@ -0,0 +1 @@
+DROP VIEW `/Root/count_rows`;

ydb/core/kqp/ut/view/input/cases/count_rows/etalon_query.sql

@@ -0,0 +1,7 @@
+SELECT
+    *
+FROM (
+    SELECT
+        COUNT(*)
+    FROM `/Root/episodes`
+);

ydb/core/kqp/ut/view/input/cases/count_rows/select_from_view.sql

@@ -0,0 +1,3 @@
+SELECT
+    *
+FROM `/Root/count_rows`;