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

jepett0 · web-flow · commit 2d3a63b70c98 · 2024-08-07T11:57:30.000+03:00
diff --git a/ydb/core/kqp/provider/rewrite_io_utils.cpp b/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."));
diff --git a/ydb/core/kqp/provider/rewrite_io_utils.h b/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
 );
 
 }
diff --git a/ydb/core/kqp/provider/yql_kikimr_datasource.cpp b/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);
             }
         }
 
diff --git a/ydb/core/kqp/ut/view/input/cases/aggregates_and_window/create_view.sql b/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);
diff --git a/ydb/core/kqp/ut/view/input/cases/aggregates_and_window/drop_view.sql b/ydb/core/kqp/ut/view/input/cases/aggregates_and_window/drop_view.sql
@@ -0,0 +1 @@
+DROP VIEW `/Root/aggregates_and_window`;
diff --git a/ydb/core/kqp/ut/view/input/cases/aggregates_and_window/etalon_query.sql b/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)
+);
diff --git a/ydb/core/kqp/ut/view/input/cases/aggregates_and_window/select_from_view.sql b/ydb/core/kqp/ut/view/input/cases/aggregates_and_window/select_from_view.sql
@@ -0,0 +1,3 @@
+SELECT
+    *
+FROM `/Root/aggregates_and_window`;
diff --git a/ydb/core/kqp/ut/view/input/cases/count_episodes/create_view.sql b/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;
diff --git a/ydb/core/kqp/ut/view/input/cases/count_episodes/drop_view.sql b/ydb/core/kqp/ut/view/input/cases/count_episodes/drop_view.sql
@@ -0,0 +1 @@
+DROP VIEW `/Root/count_episodes`;
diff --git a/ydb/core/kqp/ut/view/input/cases/count_episodes/etalon_query.sql b/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
+);
diff --git a/ydb/core/kqp/ut/view/input/cases/count_episodes/select_from_view.sql b/ydb/core/kqp/ut/view/input/cases/count_episodes/select_from_view.sql
@@ -0,0 +1,3 @@
+SELECT
+    *
+FROM `/Root/count_episodes`;
diff --git a/ydb/core/kqp/ut/view/input/cases/count_episodes_with_titles/create_view.sql b/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;
diff --git a/ydb/core/kqp/ut/view/input/cases/count_episodes_with_titles/drop_view.sql b/ydb/core/kqp/ut/view/input/cases/count_episodes_with_titles/drop_view.sql
@@ -0,0 +1 @@
+DROP VIEW `/Root/count_episodes_with_titles`;
diff --git a/ydb/core/kqp/ut/view/input/cases/count_episodes_with_titles/etalon_query.sql b/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
+);
diff --git a/ydb/core/kqp/ut/view/input/cases/count_episodes_with_titles/select_from_view.sql b/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`;
diff --git a/ydb/core/kqp/ut/view/input/cases/count_rows/create_view.sql b/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`;
diff --git a/ydb/core/kqp/ut/view/input/cases/count_rows/drop_view.sql b/ydb/core/kqp/ut/view/input/cases/count_rows/drop_view.sql
@@ -0,0 +1 @@
+DROP VIEW `/Root/count_rows`;
diff --git a/ydb/core/kqp/ut/view/input/cases/count_rows/etalon_query.sql b/ydb/core/kqp/ut/view/input/cases/count_rows/etalon_query.sql
@@ -0,0 +1,7 @@
+SELECT
+    *
+FROM (
+    SELECT
+        COUNT(*)
+    FROM `/Root/episodes`
+);
diff --git a/ydb/core/kqp/ut/view/input/cases/count_rows/select_from_view.sql b/ydb/core/kqp/ut/view/input/cases/count_rows/select_from_view.sql
@@ -0,0 +1,3 @@
+SELECT
+    *
+FROM `/Root/count_rows`;

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,8 @@ TExprNode::TPtr RewriteReadFromView(`
`8`	`8`	`const TExprNode::TPtr& node,`
`9`	`9`	`TExprContext& ctx,`
`10`	`10`	`const TString& query,`
`11`		`- const TString& cluster`
	`11`	`+ const TString& cluster,`
	`12`	`+ IModuleResolver::TPtr moduleResolver`
`12`	`13`	`);`
`13`	`14`
`14`	`15`	`}`
Original file line number	Diff line number	Diff line change
`@@ -754,6 +754,7 @@ class TKikimrDataSource : public TDataProviderBase {`
`754`	`754`	`}`
`755`	`755`
`756`	`756`	`ctx.Step`
	`757`	`+ .Repeat(TExprStep::ExpandApplyForLambdas)`
`757`	`758`	`.Repeat(TExprStep::ExprEval)`
`758`	`759`	`.Repeat(TExprStep::DiscoveryIO)`
`759`	`760`	`.Repeat(TExprStep::Epochs)`
`@@ -762,7 +763,7 @@ class TKikimrDataSource : public TDataProviderBase {`
`762`	`763`	`.Repeat(TExprStep::RewriteIO);`
`763`	`764`
`764`	`765`	`const auto& query = tableDesc.Metadata->ViewPersistedData.QueryText;`
`765`		`- return RewriteReadFromView(node, ctx, query, cluster);`
	`766`	`+ return RewriteReadFromView(node, ctx, query, cluster, Types.Modules);`
`766`	`767`	`}`
`767`	`768`	`}`
`768`	`769`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+SELECT`
	`2`	`+ *`
	`3`	+FROM `/Root/aggregates_and_window`;