[CBO] Make optimizer get only columns stats that are used in request

ydb/core/kqp/host/kqp_host.cpp

@@ -1762,7 +1762,7 @@ class TKqpHost : public IKqpHost {
 
     void Init(EKikimrQueryType queryType) {
         TransformCtx = MakeIntrusive<TKqlTransformContext>(Config, SessionCtx->QueryPtr(), SessionCtx->TablesPtr());
-        KqpRunner = CreateKqpRunner(Gateway, Cluster, TypesCtx, SessionCtx, TransformCtx, *FuncRegistry);
+        KqpRunner = CreateKqpRunner(Gateway, Cluster, TypesCtx, SessionCtx, TransformCtx, *FuncRegistry, ActorSystem);
 
         ExprCtx->NodesAllocationLimit = SessionCtx->Config()._KqpExprNodesAllocationLimit.Get().GetRef();
         ExprCtx->StringsAllocationLimit = SessionCtx->Config()._KqpExprStringsAllocationLimit.Get().GetRef();

ydb/core/kqp/host/kqp_host_impl.h

@@ -265,7 +265,7 @@ class IKqpRunner : public TThrRefBase {
 
 TIntrusivePtr<IKqpRunner> CreateKqpRunner(TIntrusivePtr<IKqpGateway> gateway, const TString& cluster,
     const TIntrusivePtr<NYql::TTypeAnnotationContext>& typesCtx, const TIntrusivePtr<NYql::TKikimrSessionContext>& sessionCtx,
-    const TIntrusivePtr<TKqlTransformContext>& transformCtx, const NMiniKQL::IFunctionRegistry& funcRegistry);
+    const TIntrusivePtr<TKqlTransformContext>& transformCtx, const NMiniKQL::IFunctionRegistry& funcRegistry, TActorSystem* actorSystem);
 
 TAutoPtr<NYql::IGraphTransformer> CreateKqpExplainPreparedTransformer(TIntrusivePtr<IKqpGateway> gateway,
     const TString& cluster, TIntrusivePtr<TKqlTransformContext> transformCtx, const NMiniKQL::IFunctionRegistry* funcRegistry,

ydb/core/kqp/host/kqp_runner.cpp

@@ -5,6 +5,7 @@
 #include <ydb/core/kqp/opt/kqp_opt.h>
 #include <ydb/core/kqp/opt/logical/kqp_opt_log.h>
 #include <ydb/core/kqp/opt/kqp_statistics_transformer.h>
+#include <ydb/core/kqp/opt/kqp_column_statistics_requester.h>
 #include <ydb/core/kqp/opt/kqp_constant_folding_transformer.h>
 #include <ydb/core/kqp/opt/logical/kqp_opt_cbo.h>
 
@@ -137,7 +138,8 @@ class TKqpRunner : public IKqpRunner {
 public:
     TKqpRunner(TIntrusivePtr<IKqpGateway> gateway, const TString& cluster,
         const TIntrusivePtr<TTypeAnnotationContext>& typesCtx, const TIntrusivePtr<TKikimrSessionContext>& sessionCtx,
-        const TIntrusivePtr<TKqlTransformContext>& transformCtx, const NMiniKQL::IFunctionRegistry& funcRegistry)
+        const TIntrusivePtr<TKqlTransformContext>& transformCtx, const NMiniKQL::IFunctionRegistry& funcRegistry,
+        TActorSystem* actorSystem)
         : Gateway(gateway)
         , Cluster(cluster)
         , TypesCtx(*typesCtx)
@@ -149,6 +151,7 @@ class TKqpRunner : public IKqpRunner {
             sessionCtx->TablesPtr(), sessionCtx->GetUserRequestContext()))
         , BuildQueryCtx(MakeIntrusive<TKqpBuildQueryContext>())
         , Pctx(TKqpProviderContext(*OptimizeCtx, Config->CostBasedOptimizationLevel.Get().GetOrElse(Config->DefaultCostBasedOptimizationLevel)))
+        , ActorSystem(actorSystem)
     {
         CreateGraphTransformer(typesCtx, sessionCtx, funcRegistry);
     }
@@ -297,6 +300,7 @@ class TKqpRunner : public IKqpRunner {
             .AddPostTypeAnnotation(/* forSubgraph */ true)
             .AddCommonOptimization()
             .Add(CreateKqpConstantFoldingTransformer(OptimizeCtx, *typesCtx, Config), "ConstantFolding")
+            .Add(CreateKqpColumnStatisticsRequester(Config, *typesCtx, SessionCtx->Tables(), Cluster, ActorSystem), "ColumnStatisticsRequester")
             .Add(CreateKqpStatisticsTransformer(OptimizeCtx, *typesCtx, Config, Pctx), "Statistics")
             .Add(CreateKqpLogOptTransformer(OptimizeCtx, *typesCtx, Config), "LogicalOptimize")
             .Add(CreateLogicalDataProposalsInspector(*typesCtx), "ProvidersLogicalOptimize")
@@ -324,7 +328,7 @@ class TKqpRunner : public IKqpRunner {
                     Config),
                 "BuildPhysicalTxs")
             .Build(false));
-
+        
         auto physicalBuildQueryTransformer = TTransformationPipeline(typesCtx)
             .AddServiceTransformers()
             .Add(Log("PhysicalBuildQuery"), "LogPhysicalBuildQuery")
@@ -399,15 +403,17 @@ class TKqpRunner : public IKqpRunner {
     TKqpProviderContext Pctx;
 
     TAutoPtr<IGraphTransformer> Transformer;
+
+    TActorSystem* ActorSystem; 
 };
 
 } // namespace
 
 TIntrusivePtr<IKqpRunner> CreateKqpRunner(TIntrusivePtr<IKqpGateway> gateway, const TString& cluster,
     const TIntrusivePtr<TTypeAnnotationContext>& typesCtx, const TIntrusivePtr<TKikimrSessionContext>& sessionCtx,
-    const TIntrusivePtr<TKqlTransformContext>& transformCtx, const NMiniKQL::IFunctionRegistry& funcRegistry)
+    const TIntrusivePtr<TKqlTransformContext>& transformCtx, const NMiniKQL::IFunctionRegistry& funcRegistry, TActorSystem* actorSystem)
 {
-    return new TKqpRunner(gateway, cluster, typesCtx, sessionCtx, transformCtx, funcRegistry);
+    return new TKqpRunner(gateway, cluster, typesCtx, sessionCtx, transformCtx, funcRegistry, actorSystem);
 }
 
 } // namespace NKqp

ydb/core/kqp/opt/kqp_column_statistics_requester.cpp

@@ -0,0 +1,253 @@
+#include "kqp_column_statistics_requester.h"
+
+#include <ydb/library/yql/core/yql_expr_optimize.h>
+#include <ydb/core/statistics/service/service.h>
+#include <ydb/core/statistics/events.h>
+#include <ydb/core/kqp/gateway/actors/kqp_ic_gateway_actors.h>
+#include <ydb/library/yql/core/yql_statistics.h>
+#include <ydb/library/yql/providers/dq/common/yql_dq_settings.h>
+#include <ydb/library/yql/dq/opt/dq_opt_stat.h>
+#include <ydb/library/yql/utils/log/log.h>
+
+namespace NKikimr::NKqp {
+
+using namespace NThreading;
+using namespace NYql;
+
+void TKqpColumnStatisticsRequester::PropagateTableToLambdaArgument(const TExprNode::TPtr& input) {
+    if (input->ChildrenSize() < 2) {
+        return;
+    }
+
+    auto callableInput = input->ChildRef(0);
+
+
+    for (size_t i = 1; i < input->ChildrenSize(); ++i) {
+        auto maybeLambda = TExprBase(input->ChildRef(i));
+        if (!maybeLambda.Maybe<TCoLambda>()) {
+            continue;
+        }
+
+        auto lambda = maybeLambda.Cast<TCoLambda>();
+        if (!lambda.Args().Size()){
+            continue;
+        }
+
+        if (callableInput->IsList()){
+            for (size_t j = 0; j < callableInput->ChildrenSize(); ++j){
+                KqpTableByExprNode[lambda.Args().Arg(j).Ptr()] = KqpTableByExprNode[callableInput->Child(j)];  
+            }
+        } else {
+            KqpTableByExprNode[lambda.Args().Arg(0).Ptr()] = KqpTableByExprNode[callableInput.Get()];
+        }
+    }
+}
+
+IGraphTransformer::TStatus TKqpColumnStatisticsRequester::DoTransform(TExprNode::TPtr input, TExprNode::TPtr& output, TExprContext& ctx) {
+    Y_UNUSED(ctx);
+
+    output = input;
+    auto optLvl = Config->CostBasedOptimizationLevel.Get().GetOrElse(TDqSettings::TDefault::CostBasedOptimizationLevel);
+    auto enableColumnStats = Config->FeatureFlags.GetEnableColumnStatistics();
+    if (!(optLvl > 0 && enableColumnStats)) {
+        return IGraphTransformer::TStatus::Ok;
+    }
+
+    VisitExprLambdasLast(
+        input, 
+        [&](const TExprNode::TPtr& input) {
+            BeforeLambdas(input) || BeforeLambdasUnmatched(input);
+
+            if (input->IsCallable()) {
+                PropagateTableToLambdaArgument(input);
+            }
+
+            return true;
+        },
+        [&](const TExprNode::TPtr& input) {
+            return AfterLambdas(input) || AfterLambdasUnmatched(input);
+        }
+    );
+
+    if (ColumnsByTableName.empty()) {
+        return IGraphTransformer::TStatus::Ok;
+    }
+
+    struct TTableMeta {
+        TString TableName;
+        THashMap<ui32, TString> ColumnNameByTag;
+    };
+    THashMap<TPathId, TTableMeta> tableMetaByPathId;
+
+    // TODO: Add other statistics, not only COUNT_MIN_SKETCH.
+    auto getStatisticsRequest = MakeHolder<NStat::TEvStatistics::TEvGetStatistics>();
+    getStatisticsRequest->StatType = NKikimr::NStat::EStatType::COUNT_MIN_SKETCH;  
+
+    for (const auto& [table, columns]: ColumnsByTableName) {
+        auto tableMeta = Tables.GetTable(Cluster, table).Metadata;
+        auto& columnsMeta = tableMeta->Columns;
+
+        auto pathId = TPathId(tableMeta->PathId.OwnerId(), tableMeta->PathId.TableId());
+        for (const auto& column: columns) {
+            if (TypesCtx.ColumnStatisticsByTableName.contains(table) && TypesCtx.ColumnStatisticsByTableName[table]->Data.contains(column)) {
+                continue;
+            }
+
+            if (!columns.contains(column)) {
+                YQL_CLOG(DEBUG, ProviderKikimr) << "Table: " + table + " doesn't contain " + column + " to request for column statistics";
+            }
+
+            NKikimr::NStat::TRequest req;
+            req.ColumnTag = columnsMeta[column].Id;
+            req.PathId = pathId;
+            getStatisticsRequest->StatRequests.push_back(req);
+
+            tableMetaByPathId[pathId].TableName = table;
+            tableMetaByPathId[pathId].ColumnNameByTag[req.ColumnTag.value()] = column;
+        }
+    }
+
+    if (getStatisticsRequest->StatRequests.empty()) {
+        return IGraphTransformer::TStatus::Ok;
+    }
+
+    using TRequest = NStat::TEvStatistics::TEvGetStatistics;
+    using TResponse = NStat::TEvStatistics::TEvGetStatisticsResult;
+    struct TResult : public NYql::IKikimrGateway::TGenericResult {
+        THashMap<TString, TOptimizerStatistics::TColumnStatMap> columnStatisticsByTableName;
+    };
+
+    auto promise = NewPromise<TResult>();
+    auto callback = [tableMetaByPathId = std::move(tableMetaByPathId)]
+    (TPromise<TResult> promise, NStat::TEvStatistics::TEvGetStatisticsResult&& response) mutable {
+        if (!response.Success) {
+            promise.SetValue(NYql::NCommon::ResultFromError<TResult>("can't get column statistics!"));
+        }
+
+        THashMap<TString, TOptimizerStatistics::TColumnStatMap> columnStatisticsByTableName;
+
+        for (auto&& stat: response.StatResponses) {
+            auto meta = tableMetaByPathId[stat.Req.PathId];
+            auto columnName = meta.ColumnNameByTag[stat.Req.ColumnTag.value()];
+            auto& columnStatistics = columnStatisticsByTableName[meta.TableName].Data[columnName];
+            columnStatistics.CountMinSketch = std::move(stat.CountMinSketch.CountMin);
+        }
+
+        promise.SetValue(TResult{.columnStatisticsByTableName = std::move(columnStatisticsByTableName)});
+    };
+    auto statServiceId = NStat::MakeStatServiceID(ActorSystem->NodeId);
+    IActor* requestHandler = 
+        new TActorRequestHandler<TRequest, TResponse, TResult>(statServiceId, getStatisticsRequest.Release(), promise, callback);
+    auto actorId = ActorSystem
+        ->Register(requestHandler, TMailboxType::HTSwap, ActorSystem->AppData<TAppData>()->UserPoolId);
+    Y_UNUSED(actorId);
+
+    auto res = promise.GetFuture().GetValueSync();
+    if (!res.Issues().Empty()) {
+        TStringStream ss;
+        res.Issues().PrintTo(ss);
+        YQL_CLOG(DEBUG, ProviderKikimr) << "Can't load columns statistics for request: " << ss.Str();
+        return IGraphTransformer::TStatus::Ok;
+    }
+
+    for (auto&& [tableName, columnStatistics]:  res.columnStatisticsByTableName) {
+        TypesCtx.ColumnStatisticsByTableName.insert(
+            {std::move(tableName), new TOptimizerStatistics::TColumnStatMap(std::move(columnStatistics))}
+        );
+    }
+
+    return IGraphTransformer::TStatus::Ok;
+}
+
+bool TKqpColumnStatisticsRequester::BeforeLambdas(const TExprNode::TPtr& input) {
+    bool matched = true;
+
+    if (TKqpTable::Match(input.Get())) {
+        KqpTableByExprNode[input.Get()] = input.Get();
+    } else if (auto maybeStreamLookup = TExprBase(input).Maybe<TKqpCnStreamLookup>()) {
+        KqpTableByExprNode[input.Get()] = maybeStreamLookup.Cast().Table().Ptr();
+    } else {
+        matched = false;
+    }
+
+    return matched;
+}
+
+bool TKqpColumnStatisticsRequester::BeforeLambdasUnmatched(const TExprNode::TPtr& input) {
+    for (const auto& node: input->Children()) {
+        if (KqpTableByExprNode.contains(node)) {
+            KqpTableByExprNode[input.Get()] = KqpTableByExprNode[node];
+            return true;
+        }
+    }
+
+    return true;
+}
+
+bool TKqpColumnStatisticsRequester::AfterLambdas(const TExprNode::TPtr& input) {
+    bool matched = true;
+
+    if (
+        TCoFilterBase::Match(input.Get()) ||
+        TCoFlatMapBase::Match(input.Get()) && IsPredicateFlatMap(TExprBase(input).Cast<TCoFlatMapBase>().Lambda().Body().Ref())
+    ) {
+        std::shared_ptr<TOptimizerStatistics> dummyStats = nullptr;
+        auto computer = NDq::TPredicateSelectivityComputer(dummyStats, true);
+
+        if (TCoFilterBase::Match(input.Get())) {
+            computer.Compute(TExprBase(input).Cast<TCoFilterBase>().Lambda().Body());
+        } else if (TCoFlatMapBase::Match(input.Get())) {
+            computer.Compute(TExprBase(input).Cast<TCoFlatMapBase>().Lambda().Body());
+        } else {
+            Y_ENSURE(false);
+        }
+
+        auto columnStatsUsedMembers = computer.GetColumnStatsUsedMembers();
+        for (const auto& item: columnStatsUsedMembers.Data) {
+            auto exprNode = TExprBase(item.Member).Ptr();
+            if (!KqpTableByExprNode.contains(exprNode) || KqpTableByExprNode[exprNode] == nullptr) {
+                continue;
+            }
+
+            auto table = TExprBase(KqpTableByExprNode[exprNode]).Cast<TKqpTable>().Path().StringValue();
+            auto column = item.Member.Name().StringValue();
+            size_t pointPos = column.find('.'); // table.column
+            if (pointPos != TString::npos) {
+                column = column.substr(pointPos + 1);
+            }
+
+            ColumnsByTableName[table].insert(std::move(column));
+        }
+    } else {
+        matched = false;
+    }
+
+    return matched;
+}
+
+bool TKqpColumnStatisticsRequester::AfterLambdasUnmatched(const TExprNode::TPtr& input) {
+    if (KqpTableByExprNode.contains(input.Get())) {
+        return true;
+    }
+
+    for (const auto& node: input->Children()) {
+        if (KqpTableByExprNode.contains(node)) {
+            KqpTableByExprNode[input.Get()] = KqpTableByExprNode[node];
+            return true;
+        }
+    }
+
+    return true;
+}
+
+TAutoPtr<IGraphTransformer> CreateKqpColumnStatisticsRequester(
+    const TKikimrConfiguration::TPtr& config,
+    TTypeAnnotationContext& typesCtx,
+    TKikimrTablesData& tables,
+    TString cluster,
+    TActorSystem* actorSystem
+) {
+    return THolder<IGraphTransformer>(new TKqpColumnStatisticsRequester(config, typesCtx, tables, cluster, actorSystem));
+}
+
+} // end of NKikimr::NKqp