[NEW RBO] Added contant folding (#28398)

Co-authored-by: Pavel Velikhov <pavelvelikhov@localhost.localdomain>

Author: pavelvelikhov

ydb/core/kqp/host/kqp_runner.cpp

@@ -373,7 +373,7 @@ class TKqpRunner : public IKqpRunner {
             //.AddCommonOptimization()
 
             .Add(CreateKqpPgRewriteTransformer(OptimizeCtx, *typesCtx), "RewritePgSelect")
-            .Add(CreateKqpNewRBOTransformer(OptimizeCtx, *typesCtx, rboKqpTypeAnnTransformer, kqpTypeAnnTransformer, newRBOPhysicalPeepholeTransformer), "NewRBOTransformer")
+            .Add(CreateKqpNewRBOTransformer(OptimizeCtx, *typesCtx, rboKqpTypeAnnTransformer, kqpTypeAnnTransformer, newRBOPhysicalPeepholeTransformer, funcRegistry), "NewRBOTransformer")
             .Add(CreateKqpRBOCleanupTransformer(*typesCtx), "RBOCleanupTransformer")
 
             //.Add(CreatePhysicalDataProposalsInspector(*typesCtx), "ProvidersPhysicalOptimize")

ydb/core/kqp/opt/rbo/kqp_constant_folding_stage.cpp

@@ -0,0 +1,148 @@
+#include "kqp_rbo_rules.h"
+#include <ydb/core/kqp/common/kqp_yql.h>
+#include <yql/essentials/core/yql_expr_optimize.h>
+#include <yql/essentials/utils/log/log.h>
+#include <yql/essentials/core/services/yql_transform_pipeline.h>
+#include <ydb/library/yql/dq/opt/dq_opt_stat.h>
+#include <typeinfo>
+
+using namespace NYql;
+using namespace NYql::NNodes;
+using namespace NYql::NDq;
+
+namespace {
+
+    THashSet<TString> notAllowedDataTypeForSafeCast{"JsonDocument", "DyNumber"};
+
+    bool IsSuitableToExtractExpr(const TExprNode::TPtr &input) {
+        if (auto maybeSafeCast = TExprBase(input).Maybe<TCoSafeCast>()) {
+            auto maybeDataType = maybeSafeCast.Cast().Type().Maybe<TCoDataType>();
+            if (!maybeDataType) {
+                if (const auto maybeOptionalType = maybeSafeCast.Cast().Type().Maybe<TCoOptionalType>()) {
+                    maybeDataType = maybeOptionalType.Cast().ItemType().Maybe<TCoDataType>();
+                }
+            }
+            return (maybeDataType && !notAllowedDataTypeForSafeCast.contains(maybeDataType.Cast().Type().Value()));
+        }
+        return true;
+    }
+
+    /**
+     * Traverse a lambda and extract a list of constant expressions
+     */
+    void ExtractConstantExprs(const TExprNode::TPtr& input, TVector<std::pair<TExprNode::TPtr, TExprNode::TPtr>>& exprs, TExprContext& ctx, bool foldUdfs = true) {
+        if (!IsSuitableToExtractExpr(input)) {
+            return;
+        }
+
+        if (TCoLambda::Match(input.Get())) {
+            auto lambda = TExprBase(input).Cast<TCoLambda>();
+            return ExtractConstantExprs(lambda.Body().Ptr(), exprs, ctx);
+        }
+
+        if (IsDataOrOptionalOfData(input->GetTypeAnn()) && !NeedCalc(TExprBase(input))) {
+            return;
+        }
+
+        if (IsConstantExpr(input, foldUdfs) && !input->IsCallable("PgConst")) {
+            TNodeOnNodeOwnedMap deepClones;
+            auto inputClone = ctx.DeepCopy(*input, ctx, deepClones, false, true, true);
+            exprs.push_back(std::make_pair(input, inputClone));
+            return;
+        }
+
+        if (TCoAsStruct::Match(input.Get())) {
+            for (auto child : TExprBase(input).Cast<TCoAsStruct>()) {
+                ExtractConstantExprs(child.Item(1).Ptr(), exprs, ctx);
+            }
+            return;
+        }
+
+        if (input->IsCallable() && input->Content() != "EvaluateExpr") {
+            if (input->ChildrenSize() >= 1) {
+                for (size_t i = 0; i < input->ChildrenSize(); i++) {
+                    ExtractConstantExprs(input->Child(i), exprs, ctx);
+                }
+            }
+        }
+
+        return;
+    }
+
+}
+
+namespace NKikimr {
+namespace NKqp {
+
+void TConstantFoldingStage::RunStage(TOpRoot &root, TRBOContext &ctx) {
+    TVector<TExprNode::TPtr> lambdasWithConstExpr;
+    bool foldUdfs = ctx.KqpCtx.Config->EnableFoldUdfs;
+
+    // Iterate through all operators that contain lambdas with potential constant expression
+
+    // Internal map for remap operation
+    TNodeOnNodeOwnedMap replaces;
+
+    // Actual map used in the optimizer
+    TVector<std::pair<TExprNode::TPtr, TExprNode::TPtr>> globalExtractedExprs;
+    TVector<std::shared_ptr<IOperator>> affectedOps;
+
+    for (auto it : root) {
+        if (!it.Current->GetLambdas().empty()) {
+            auto lambdas = it.Current->GetLambdas();
+            bool affected = false;
+            for (auto l : lambdas) {
+                auto lambda = TCoLambda(l);
+                TVector<std::pair<TExprNode::TPtr, TExprNode::TPtr>> extractedExprs;
+                ExtractConstantExprs(lambda.Body().Ptr(), extractedExprs, ctx.ExprCtx, foldUdfs);
+                if (!extractedExprs.empty()) {
+                    affected = true;
+                    globalExtractedExprs.insert(globalExtractedExprs.end(), extractedExprs.begin(), extractedExprs.end());
+                }
+            }
+
+            if (affected) {
+                affectedOps.push_back(it.Current);
+            }
+        }
+    }
+
+    if (globalExtractedExprs.empty()) {
+        return;
+    }
+
+    // Build a list of eval expressions
+
+    TExprNode::TListType lambdaList;
+    TExprNode::TListType evalElements;
+    for (auto & [k, v] : globalExtractedExprs) {
+        lambdaList.push_back(k);
+        evalElements.push_back(v);
+    }
+
+    // Evaluate all the constant expressions at once
+    auto evalList = ctx.ExprCtx.NewList(root.Pos, std::move(evalElements));
+    evalList = ctx.ExprCtx.NewCallable(root.Pos, "EvaluateExpr", { evalList });
+
+    auto evaluator = TTransformationPipeline(&ctx.TypeCtx)
+        .AddServiceTransformers()
+        .AddPreTypeAnnotation()
+        .AddExpressionEvaluation(ctx.FuncRegistry).Build(false);
+
+    ctx.ExprCtx.Step.Repeat(TExprStep::ExprEval);
+    IGraphTransformer::TStatus status(IGraphTransformer::TStatus::Ok);
+    do {
+        status = evaluator->Transform(evalList, evalList, ctx.ExprCtx);
+    } while (status == IGraphTransformer::TStatus::Repeat);
+
+    // Iterate over affected operators and modify their expressions with folded expressions
+    for (size_t i=0; i<lambdaList.size(); i++) {
+        replaces[lambdaList[i].Get()] = evalList->Child(i);
+    }
+
+    for (auto op : affectedOps) {
+        op->ApplyReplaceMap(replaces, ctx);
+    }
+}
+}
+}

ydb/core/kqp/opt/rbo/kqp_operator.cpp

@@ -1,4 +1,5 @@
 #include "kqp_operator.h"
+#include <yql/essentials/core/yql_expr_optimize.h>
 
 namespace {
 using namespace NKikimr;
@@ -124,6 +125,14 @@ TExprNode::TPtr RenameMembers(TExprNode::TPtr input, const THashMap<TInfoUnit, T
     }
 }
 
+} // namespace
+
+namespace NKikimr {
+namespace NKqp {
+
+using namespace NYql;
+using namespace NNodes;
+
 TString PrintRBOExpression(TExprNode::TPtr expr, TExprContext & ctx) {
     try {
         TConvertToAstSettings settings;
@@ -142,14 +151,6 @@ TString PrintRBOExpression(TExprNode::TPtr expr, TExprContext & ctx) {
     }
 }
 
-} // namespace
-
-namespace NKikimr {
-namespace NKqp {
-
-using namespace NYql;
-using namespace NNodes;
-
 /**
  * Scan expression and retrieve all members
  */
@@ -388,6 +389,10 @@ std::pair<TExprNode::TPtr, TVector<TExprNode::TPtr>> BuildSortKeySelector(TVecto
 }
 
 
+/**
+ * Base class Operator methods
+ */
+
 void IOperator::RenameIUs(const THashMap<TInfoUnit, TInfoUnit, TInfoUnit::THashFunction> &renameMap, TExprContext &ctx) {
     Y_UNUSED(renameMap);
     Y_UNUSED(ctx);
@@ -398,11 +403,19 @@ const TTypeAnnotationNode* IOperator::GetIUType(TInfoUnit iu) {
     return structType->FindItemType(iu.GetFullName());
 }
 
+/**
+ * EmptySource operator methods
+ */
+
 TString TOpEmptySource::ToString(TExprContext& ctx) {
     Y_UNUSED(ctx); 
     return "EmptySource"; 
 }
 
+/**
+ * OpRead operator methods
+ */
+
 TOpRead::TOpRead(TExprNode::TPtr node) : IOperator(EOperator::Source, node->Pos()) {
     auto opSource = TKqpOpRead(node);
 
@@ -440,6 +453,10 @@ TString TOpRead::ToString(TExprContext& ctx) {
     return res;
 }
 
+/**
+ * OpMap operator methods
+ */
+
 TOpMap::TOpMap(std::shared_ptr<IOperator> input, TPositionHandle pos, TVector<std::pair<TInfoUnit, std::variant<TInfoUnit, TExprNode::TPtr>>> mapElements,
                bool project)
     : IUnaryOperator(EOperator::Map, pos, input), MapElements(mapElements), Project(project) {}
@@ -456,6 +473,16 @@ TVector<TInfoUnit> TOpMap::GetOutputIUs() {
     return res;
 }
 
+TVector<TExprNode::TPtr> TOpMap::GetLambdas() {
+    TVector<TExprNode::TPtr> result;
+    for (auto &[_,body] : MapElements) {
+        if (std::holds_alternative<TExprNode::TPtr>(body)) {
+            result.push_back(std::get<TExprNode::TPtr>(body));
+        }
+    }
+    return result;
+}
+
 TVector<TInfoUnit> TOpMap::GetScalarSubplanIUs(TPlanProps& props) {
     TVector<TInfoUnit> allVars;
     TVector<TInfoUnit> res;
@@ -524,6 +551,19 @@ void TOpMap::RenameIUs(const THashMap<TInfoUnit, TInfoUnit, TInfoUnit::THashFunc
     MapElements = newMapElements;
 }
 
+void TOpMap::ApplyReplaceMap(TNodeOnNodeOwnedMap map, TRBOContext & ctx) {
+    TOptimizeExprSettings settings(&ctx.TypeCtx);
+    for (size_t i=0; i<MapElements.size(); i++) {
+        auto & body = MapElements[i].second;
+        if (std::holds_alternative<TExprNode::TPtr>(body)) {
+            auto bodyLambda = std::get<TExprNode::TPtr>(body);
+            RemapExpr(bodyLambda, bodyLambda, map, ctx.ExprCtx, settings);
+            MapElements[i].second = std::variant<TInfoUnit,TExprNode::TPtr>(bodyLambda);
+        }
+    }
+}
+
+
 TString TOpMap::ToString(TExprContext& ctx) {
     auto res = TStringBuilder();
     res << "Map [";
@@ -547,6 +587,10 @@ TString TOpMap::ToString(TExprContext& ctx) {
     return res;
 }
 
+/**
+ * OpProject methods
+ */
+
 TOpProject::TOpProject(std::shared_ptr<IOperator> input, TPositionHandle pos, TVector<TInfoUnit> projectList)
     : IUnaryOperator(EOperator::Project, pos, input), ProjectList(projectList) {}
 
@@ -584,6 +628,10 @@ TString TOpProject::ToString(TExprContext& ctx) {
     return res;
 }
 
+/**
+ * OpFilter operator methods
+ */
+
 TOpFilter::TOpFilter(std::shared_ptr<IOperator> input, TPositionHandle pos, TExprNode::TPtr filterLambda)
     : IUnaryOperator(EOperator::Filter, pos, input), FilterLambda(filterLambda) {}
 
@@ -593,6 +641,15 @@ void TOpFilter::RenameIUs(const THashMap<TInfoUnit, TInfoUnit, TInfoUnit::THashF
     FilterLambda = RenameMembers(FilterLambda, renameMap, ctx);
 }
 
+TVector<TExprNode::TPtr> TOpFilter::GetLambdas() {
+    return {FilterLambda};
+}
+
+void TOpFilter::ApplyReplaceMap(TNodeOnNodeOwnedMap map, TRBOContext & ctx) {
+    TOptimizeExprSettings settings(&ctx.TypeCtx);
+    RemapExpr(FilterLambda, FilterLambda, map, ctx.ExprCtx, settings);
+}
+
 TVector<TInfoUnit> TOpFilter::GetFilterIUs(TPlanProps& props) const {
     TVector<TInfoUnit> res;
 
@@ -666,6 +723,10 @@ TString TOpFilter::ToString(TExprContext& ctx) {
     return TStringBuilder() << "Filter :" << PrintRBOExpression(FilterLambda, ctx);
 }
 
+/**
+ * OpJoin operator methods
+ */
+
 TOpJoin::TOpJoin(std::shared_ptr<IOperator> leftInput, std::shared_ptr<IOperator> rightInput, TPositionHandle pos, TString joinKind,
                  TVector<std::pair<TInfoUnit, TInfoUnit>> joinKeys)
     : IBinaryOperator(EOperator::Join, pos, leftInput, rightInput), JoinKind(joinKind), JoinKeys(joinKeys) {}
@@ -706,6 +767,10 @@ TString TOpJoin::ToString(TExprContext& ctx) {
     return res;
 }
 
+/**
+ * OpUnionAll operator methods
+ */
+
 TOpUnionAll::TOpUnionAll(std::shared_ptr<IOperator> leftInput, std::shared_ptr<IOperator> rightInput, TPositionHandle pos, bool ordered)
     : IBinaryOperator(EOperator::UnionAll, pos, leftInput, rightInput), Ordered(ordered) {}
 
@@ -718,6 +783,10 @@ TString TOpUnionAll::ToString(TExprContext& ctx) {
     return "UnionAll"; 
 }
 
+/**
+ * OpLimit operator methods
+ */
+
 TOpLimit::TOpLimit(std::shared_ptr<IOperator> input, TPositionHandle pos, TExprNode::TPtr limitCond)
     : IUnaryOperator(EOperator::Limit, pos, input), LimitCond(limitCond) {}
 
@@ -731,6 +800,10 @@ TString TOpLimit::ToString(TExprContext& ctx) {
     return TStringBuilder() << "Limit: " << PrintRBOExpression(LimitCond, ctx); 
 }
 
+/**
+ * OpAggregate operator methods
+ */
+
 TOpAggregate::TOpAggregate(std::shared_ptr<IOperator> input, TVector<TOpAggregationTraits>& aggTraitsList, TVector<TInfoUnit>& keyColumns,
                            EAggregationPhase aggPhase, bool distinctAll, TPositionHandle pos)
     : IUnaryOperator(EOperator::Aggregate, pos, input), AggregationTraitsList(aggTraitsList), KeyColumns(keyColumns),
@@ -774,6 +847,10 @@ TString TOpAggregate::ToString(TExprContext& ctx) {
     return strBuilder;
 }
 
+/**
+ * OpRoot operator methods
+ */
+
 TOpRoot::TOpRoot(std::shared_ptr<IOperator> input, TPositionHandle pos, TVector<TString> columnOrder) : 
     IUnaryOperator(EOperator::Root, pos, input), 
     ColumnOrder(columnOrder) {}