Enable COUNT in view queries (ydb-platform#6820)

ydb/core/kqp/host/kqp_host.cpp

@@ -1776,7 +1776,7 @@ class TKqpHost : public IKqpHost {
 
         auto queryExecutor = MakeIntrusive<TKqpQueryExecutor>(Gateway, Cluster, SessionCtx, KqpRunner);
         auto kikimrDataSource = CreateKikimrDataSource(*FuncRegistry, *TypesCtx, gatewayProxy, SessionCtx,
-            ExternalSourceFactory, IsInternalCall);
+            ExternalSourceFactory, IsInternalCall, GUCSettings);
         auto kikimrDataSink = CreateKikimrDataSink(*FuncRegistry, *TypesCtx, gatewayProxy, SessionCtx, ExternalSourceFactory, queryExecutor);
 
         FillSettings.AllResultsBytesLimit = Nothing();

ydb/core/kqp/host/kqp_translate.cpp

@@ -1,6 +1,8 @@
 #include "kqp_translate.h"
 
+#include <ydb/core/kqp/provider/yql_kikimr_results.h>
 #include <ydb/library/yql/sql/sql.h>
+#include <ydb/public/api/protos/ydb_query.pb.h>
 
 
 namespace NKikimr {

ydb/core/kqp/host/kqp_translate.h

@@ -1,7 +1,8 @@
 #pragma once
 
-#include <ydb/core/kqp/provider/yql_kikimr_results.h>
-#include <ydb/core/kqp/common/kqp.h>
+#include <ydb/core/kqp/common/simple/query_ast.h>
+#include <ydb/core/kqp/provider/yql_kikimr_provider.h>
+#include <ydb/core/protos/table_service_config.pb.h>
 #include <ydb/library/yql/core/pg_settings/guc_settings.h>
 
 namespace NKikimr {

ydb/core/kqp/provider/rewrite_io_utils.cpp

@@ -15,31 +15,25 @@ using namespace NNodes;
 
 constexpr const char* QueryGraphNodeSignature = "SavedQueryGraph";
 
-NSQLTranslation::TTranslationSettings CreateViewTranslationSettings(const TString& cluster) {
-    NSQLTranslation::TTranslationSettings settings;
-
-    settings.DefaultCluster = cluster;
-    settings.ClusterMapping[cluster] = TString(NYql::KikimrProviderName);
-    settings.Mode = NSQLTranslation::ESqlMode::LIMITED_VIEW;
-
-    return settings;
-}
-
 TExprNode::TPtr CompileViewQuery(
     const TString& query,
     TExprContext& ctx,
-    const TString& cluster
+    NKikimr::NKqp::TKqpTranslationSettingsBuilder& settingsBuilder,
+    IModuleResolver::TPtr moduleResolver
 ) {
+    auto translationSettings = settingsBuilder.Build(ctx);
+    translationSettings.Mode = NSQLTranslation::ESqlMode::LIMITED_VIEW;
+
     TAstParseResult queryAst;
-    queryAst = NSQLTranslation::SqlToYql(query, CreateViewTranslationSettings(cluster));
+    queryAst = NSQLTranslation::SqlToYql(query, translationSettings);
 
     ctx.IssueManager.AddIssues(queryAst.Issues);
     if (!queryAst.IsOk()) {
         return nullptr;
     }
 
     TExprNode::TPtr queryGraph;
-    if (!CompileExpr(*queryAst.Root, queryGraph, ctx, nullptr, nullptr)) {
+    if (!CompileExpr(*queryAst.Root, queryGraph, ctx, moduleResolver.get(), nullptr)) {
         return nullptr;
     }
 
@@ -123,7 +117,8 @@ TExprNode::TPtr RewriteReadFromView(
     const TExprNode::TPtr& node,
     TExprContext& ctx,
     const TString& query,
-    const TString& cluster
+    NKikimr::NKqp::TKqpTranslationSettingsBuilder& settingsBuilder,
+    IModuleResolver::TPtr moduleResolver
 ) {
     YQL_PROFILE_FUNC(DEBUG);
 
@@ -132,7 +127,7 @@ TExprNode::TPtr RewriteReadFromView(
 
     TExprNode::TPtr queryGraph = FindSavedQueryGraph(readNode.Ptr());
     if (!queryGraph) {
-        queryGraph = CompileViewQuery(query, ctx, cluster);
+        queryGraph = CompileViewQuery(query, ctx, settingsBuilder, 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

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <ydb/core/kqp/host/kqp_translate.h>
 #include <ydb/library/yql/ast/yql_expr.h>
 
 namespace NYql {
@@ -10,7 +11,8 @@ TExprNode::TPtr RewriteReadFromView(
     const TExprNode::TPtr& node,
     TExprContext& ctx,
     const TString& query,
-    const TString& cluster
+    NKikimr::NKqp::TKqpTranslationSettingsBuilder& settingsBuilder,
+    IModuleResolver::TPtr moduleResolver
 );
 
 }

ydb/core/kqp/provider/yql_kikimr_datasource.cpp

@@ -3,6 +3,7 @@
 #include "yql_kikimr_provider_impl.h"
 
 #include <ydb/core/kqp/common/simple/services.h>
+#include <ydb/core/kqp/host/kqp_translate.h>
 #include <ydb/library/yql/providers/common/provider/yql_data_provider_impl.h>
 #include <ydb/library/yql/providers/common/config/yql_configuration_transformer.h>
 
@@ -472,12 +473,14 @@ class TKikimrDataSource : public TDataProviderBase {
         TIntrusivePtr<IKikimrGateway> gateway,
         TIntrusivePtr<TKikimrSessionContext> sessionCtx,
         const NExternalSource::IExternalSourceFactory::TPtr& externalSourceFactory,
-        bool isInternalCall)
+        bool isInternalCall,
+        TGUCSettings::TPtr gucSettings)
         : FunctionRegistry(functionRegistry)
         , Types(types)
         , Gateway(gateway)
         , SessionCtx(sessionCtx)
         , ExternalSourceFactory(externalSourceFactory)
+        , GUCSettings(gucSettings)
         , ConfigurationTransformer(new TKikimrConfigurationTransformer(sessionCtx, types))
         , IntentDeterminationTransformer(new TKiSourceIntentDeterminationTransformer(sessionCtx))
         , LoadTableMetadataTransformer(CreateKiSourceLoadTableMetadataTransformer(gateway, sessionCtx, types, externalSourceFactory, isInternalCall))
@@ -760,6 +763,7 @@ class TKikimrDataSource : public TDataProviderBase {
                 }
 
                 ctx.Step
+                    .Repeat(TExprStep::ExpandApplyForLambdas)
                     .Repeat(TExprStep::ExprEval)
                     .Repeat(TExprStep::DiscoveryIO)
                     .Repeat(TExprStep::Epochs)
@@ -768,7 +772,15 @@ class TKikimrDataSource : public TDataProviderBase {
                     .Repeat(TExprStep::RewriteIO);
 
                 const auto& query = tableDesc.Metadata->ViewPersistedData.QueryText;
-                return RewriteReadFromView(node, ctx, query, cluster);
+                NKqp::TKqpTranslationSettingsBuilder settingsBuilder(
+                    SessionCtx->Query().Type,
+                    SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(),
+                    cluster,
+                    query,
+                    SessionCtx->Config().BindingsMode,
+                    GUCSettings
+                );
+                return RewriteReadFromView(node, ctx, query, settingsBuilder, Types.Modules);
             }
         }
 
@@ -881,6 +893,7 @@ class TKikimrDataSource : public TDataProviderBase {
     TIntrusivePtr<IKikimrGateway> Gateway;
     TIntrusivePtr<TKikimrSessionContext> SessionCtx;
     NExternalSource::IExternalSourceFactory::TPtr ExternalSourceFactory;
+    TGUCSettings::TPtr GUCSettings;
 
     TAutoPtr<IGraphTransformer> ConfigurationTransformer;
     TAutoPtr<IGraphTransformer> IntentDeterminationTransformer;
@@ -920,9 +933,10 @@ TIntrusivePtr<IDataProvider> CreateKikimrDataSource(
     TIntrusivePtr<IKikimrGateway> gateway,
     TIntrusivePtr<TKikimrSessionContext> sessionCtx,
     const NExternalSource::IExternalSourceFactory::TPtr& externalSourceFactory,
-    bool isInternalCall)
+    bool isInternalCall,
+    TGUCSettings::TPtr gucSettings)
 {
-    return new TKikimrDataSource(functionRegistry, types, gateway, sessionCtx, externalSourceFactory, isInternalCall);
+    return new TKikimrDataSource(functionRegistry, types, gateway, sessionCtx, externalSourceFactory, isInternalCall, gucSettings);
 }
 
 TAutoPtr<IGraphTransformer> CreateKiSourceLoadTableMetadataTransformer(TIntrusivePtr<IKikimrGateway> gateway,

ydb/core/kqp/provider/yql_kikimr_provider.h

@@ -549,7 +549,8 @@ TIntrusivePtr<IDataProvider> CreateKikimrDataSource(
     TIntrusivePtr<IKikimrGateway> gateway,
     TIntrusivePtr<TKikimrSessionContext> sessionCtx,
     const NKikimr::NExternalSource::IExternalSourceFactory::TPtr& sourceFactory,
-    bool isInternalCall);
+    bool isInternalCall,
+    TGUCSettings::TPtr gucSettings);
 
 TIntrusivePtr<IDataProvider> CreateKikimrDataSink(
     const NKikimr::NMiniKQL::IFunctionRegistry& functionRegistry,

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`;