AssumeConstraints callable (#8349)

ydb/library/yql/core/common_opt/yql_co_simple1.cpp

@@ -13,6 +13,8 @@
 #include <ydb/library/yql/utils/log/log.h>
 #include <ydb/library/yql/parser/pg_catalog/catalog.h>
 
+#include <library/cpp/yson/node/node_io.h>
+
 #include <util/generic/map.h>
 #include <util/string/cast.h>
 #include <util/generic/xrange.h>
@@ -2266,7 +2268,7 @@ TExprNode::TPtr HasNullOverVariant(const TExprNode::TPtr& node, TExprContext& ct
 }
 
 constexpr std::initializer_list<std::string_view> FlowPriority = {
-    "AssumeSorted", "AssumeUnique", "AssumeDistinct",
+    "AssumeSorted", "AssumeUnique", "AssumeDistinct", "AssumeConstraints",
     "Map", "OrderedMap", "MapNext",
     "Filter", "OrderedFilter",
     "FlatMap", "OrderedFlatMap",
@@ -5944,6 +5946,22 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) {
             return ctx.ChangeChild(*node, 0, node->Head().HeadPtr());
         }
 
+        if (node->Head().IsCallable("AssumeConstraints") && node->Head().GetConstraint<TSortedConstraintNode>()) {
+            TConstraintSet constrSet = node->Head().GetConstraintSet();
+            constrSet.RemoveConstraint<TSortedConstraintNode>();
+            YQL_CLOG(DEBUG, Core) << node->Content() << " over " << node->Head().Content();
+            auto res = ctx.ChangeChild(*node, 0, node->Head().HeadPtr());
+            if (constrSet) {
+                res = ctx.Builder(node->Head().Pos())
+                    .Callable("AssumeConstraints")
+                        .Add(0, std::move(res))
+                        .Atom(1, NYT::NodeToYsonString(constrSet.ToYson(), NYson::EYsonFormat::Text), TNodeFlags::MultilineContent)
+                    .Seal()
+                    .Build();
+            }
+            return res;
+        }
+
         return node;
     };
 

ydb/library/yql/core/expr_nodes/yql_expr_nodes.json

@@ -1625,7 +1625,7 @@
                 {"Index": 5, "Name": "LeftRenames", "Type": "TCoAtomList"},
                 {"Index": 6, "Name": "RightRenames", "Type": "TCoAtomList"},
                 {"Index": 7, "Name": "LeftKeysColumnNames", "Type": "TCoAtomList"},
-                {"Index": 8, "Name": "RightKeysColumnNames", "Type": "TCoAtomList"},  
+                {"Index": 8, "Name": "RightKeysColumnNames", "Type": "TCoAtomList"},
                 {"Index": 9, "Name": "Flags", "Type": "TCoAtomList"}
             ]
         },
@@ -1641,7 +1641,7 @@
                 {"Index": 4, "Name": "LeftRenames", "Type": "TCoAtomList"},
                 {"Index": 5, "Name": "RightRenames", "Type": "TCoAtomList"},
                 {"Index": 6, "Name": "LeftKeysColumnNames", "Type": "TCoAtomList"},
-                {"Index": 7, "Name": "RightKeysColumnNames", "Type": "TCoAtomList"}, 
+                {"Index": 7, "Name": "RightKeysColumnNames", "Type": "TCoAtomList"},
                 {"Index": 8, "Name": "Flags", "Type": "TCoAtomList"}
             ]
         },
@@ -1863,6 +1863,14 @@
             "Base": "TCoInputBase",
             "Match": {"Type": "Callable", "Name": "AssumeAllMembersNullableAtOnce"}
         },
+        {
+            "Name": "TCoAssumeConstraints",
+            "Base": "TCoInputBase",
+            "Match": {"Type": "Callable", "Name": "AssumeConstraints"},
+            "Children": [
+                {"Index": 1, "Name": "Value", "Type": "TCoAtom"}
+            ]
+        },
         {
             "Name": "TCoListItemType",
             "Base": "TCallable",

ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp

@@ -6980,7 +6980,7 @@ template <bool Asc, bool Equals>
 TExprNode::TPtr AggrComparePg(const TExprNode& node, TExprContext& ctx) {
     YQL_CLOG(DEBUG, CorePeepHole) << "Expand '" << node.Content() << "' over Pg.";
     auto op = Asc ? ( Equals ? "<=" : "<") : ( Equals ? ">=" : ">");
-    auto finalPart = (Equals == Asc) ? MakeBool<true>(node.Pos(), ctx) : 
+    auto finalPart = (Equals == Asc) ? MakeBool<true>(node.Pos(), ctx) :
         ctx.NewCallable(node.Pos(), "Exists", { node.TailPtr() });
     if (!Asc) {
         finalPart = ctx.NewCallable(node.Pos(), "Not", { finalPart });
@@ -8445,6 +8445,7 @@ struct TPeepHoleRules {
         {"AssumeUnique", &DropAssume},
         {"AssumeDistinct", &DropAssume},
         {"AssumeChopped", &DropAssume},
+        {"AssumeConstraints", &DropAssume},
         {"EmptyFrom", &DropEmptyFrom},
         {"Top", &OptimizeTopOrSort<false, true>},
         {"TopSort", &OptimizeTopOrSort<true, true>},
@@ -8493,7 +8494,7 @@ struct TPeepHoleRules {
     const TExtPeepHoleOptimizerMap BlockStageExtFinalRules = {
         {"BlockExtend", &ExpandBlockExtend},
         {"BlockOrderedExtend", &ExpandBlockExtend},
-        {"ReplicateScalars", &ExpandReplicateScalars} 
+        {"ReplicateScalars", &ExpandReplicateScalars}
     };
 
     static const TPeepHoleRules& Instance() {

ydb/library/yql/core/type_ann/type_ann_core.cpp

@@ -1232,7 +1232,7 @@ namespace NTypeAnnImpl {
 
         YQL_ENSURE(IsSameAnnotation(*memberType, *resultType) ||
                    IsSameAnnotation(*ctx.Expr.MakeType<TOptionalExprType>(memberType), *resultType));
-        
+
         output = ctx.Expr.Builder(input->Pos())
             .Callable("IfStrict")
                 .Callable(0, "Exists")
@@ -2749,8 +2749,8 @@ namespace NTypeAnnImpl {
                 if (!(*dataTypeOne == *dataTypeTwo)) {
                     ctx.Expr.AddError(TIssue(
                         ctx.Expr.GetPosition(input->Pos()),
-                        TStringBuilder() << "Cannot calculate with different decimals: " 
-                            << static_cast<const TTypeAnnotationNode&>(*dataType[0]) << " != " 
+                        TStringBuilder() << "Cannot calculate with different decimals: "
+                            << static_cast<const TTypeAnnotationNode&>(*dataType[0]) << " != "
                             << static_cast<const TTypeAnnotationNode&>(*dataType[1])
                     ));
 
@@ -6193,18 +6193,18 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         };
 
         auto mergeMembers = [&ctx, &buildJustMember, &input, &left, &right, &mergeLambda](const TStringBuf& name, bool hasLeft, bool hasRight) -> TExprNode::TPtr {
-            auto leftMaybe = hasLeft ? 
+            auto leftMaybe = hasLeft ?
                 buildJustMember(left, name) :
                 ctx.Expr.NewCallable(input->Pos(), "Nothing", {
                     ExpandType(input->Pos(), *ctx.Expr.MakeType<TOptionalExprType>(right->GetTypeAnn()->Cast<TStructExprType>()->FindItemType(name)), ctx.Expr)
                 });
-            
-            auto rightMaybe = hasRight ? 
+
+            auto rightMaybe = hasRight ?
                 buildJustMember(right, name) :
                 ctx.Expr.NewCallable(input->Pos(), "Nothing", {
                     ExpandType(input->Pos(), *ctx.Expr.MakeType<TOptionalExprType>(left->GetTypeAnn()->Cast<TStructExprType>()->FindItemType(name)), ctx.Expr)
                 });
-            
+
             return ctx.Expr.Builder(input->Pos())
                 .List()
                     .Atom(0, name)
@@ -6492,7 +6492,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
                             .With(1, result)
                         .Seal()
                     .Build();
-                }                        
+                }
             }
         } else if (collection->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Tuple) {
             for (size_t idx = 0; idx < collection->GetTypeAnn()->Cast<TTupleExprType>()->GetSize(); idx++) {
@@ -6523,7 +6523,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
             return IGraphTransformer::TStatus::Error;
         }
 
-        if (result) { 
+        if (result) {
             output = result;
         } else {
             output = ctx.Expr.Builder(input->Pos()).Callable("Null").Seal().Build();
@@ -12170,6 +12170,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["AssumeUniqueHint"] = &AssumeConstraintWrapper<false>;
         Functions["AssumeDistinctHint"] = &AssumeConstraintWrapper<false>;
         Functions["AssumeChopped"] = &AssumeConstraintWrapper<true>;
+        Functions["AssumeConstraints"] = &AssumeConstraintsWrapper;
         Functions["AssumeAllMembersNullableAtOnce"] = &AssumeAllMembersNullableAtOnceWrapper;
         Functions["AssumeStrict"] = &AssumeStrictWrapper;
         Functions["AssumeNonStrict"] = &AssumeStrictWrapper;

ydb/library/yql/core/type_ann/type_ann_list.cpp

@@ -11,6 +11,8 @@
 
 #include <ydb/library/yql/parser/pg_catalog/catalog.h>
 
+#include <library/cpp/yson/node/node_io.h>
+
 #include <util/generic/algorithm.h>
 #include <util/string/join.h>
 
@@ -1432,10 +1434,10 @@ namespace {
     }
 
     IGraphTransformer::TStatus ListTopSortWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
-        if (!EnsureMinMaxArgsCount(*input, 2, 3, ctx.Expr)) { 
+        if (!EnsureMinMaxArgsCount(*input, 2, 3, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
-        
+
         TStringBuf newName = input->Content();
         newName.Skip(4);
         bool desc = false;
@@ -1457,7 +1459,7 @@ namespace {
                 .Seal()
             .Build();
         }
-        
+
         return OptListWrapperImpl<4U, 4U>(ctx.Expr.Builder(input->Pos())
             .Callable(newName)
                 .Add(0, input->ChildPtr(0))
@@ -4412,6 +4414,37 @@ namespace {
     template IGraphTransformer::TStatus AssumeConstraintWrapper<true>(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     template IGraphTransformer::TStatus AssumeConstraintWrapper<false>(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 
+    IGraphTransformer::TStatus AssumeConstraintsWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TContext& ctx) {
+        if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (!EnsureAnySeqType(input->Head(), ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (!EnsureAtom(input->Tail(), ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        NYT::TNode node;
+        try {
+            node = NYT::NodeFromYsonString(input->Tail().Content());
+        } catch (...) {
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), TStringBuilder() <<
+                "Bad constraints yson-value: " << CurrentExceptionMessage()));
+            return IGraphTransformer::TStatus::Error;
+        }
+        if (!node.IsMap()) {
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), TStringBuilder() <<
+                "Expected yson-map as serialized constraints value, actual " << node.GetType()));
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        input->SetTypeAnn(input->Head().GetTypeAnn());
+        return IGraphTransformer::TStatus::Ok;
+    }
+
     IGraphTransformer::TStatus AssumeColumnOrderWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx) {
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
@@ -6503,7 +6536,7 @@ namespace {
             return IGraphTransformer::TStatus::Repeat;
         }
 
-        const TTypeAnnotationNode* outputType = ctx.Expr.MakeType<TDataExprType>(input->IsCallable("PercentRank") ? 
+        const TTypeAnnotationNode* outputType = ctx.Expr.MakeType<TDataExprType>(input->IsCallable("PercentRank") ?
             EDataSlot::Double : EDataSlot::Uint64);
         if (!isAnsi && keyType->GetKind() == ETypeAnnotationKind::Optional) {
             outputType = ctx.Expr.MakeType<TOptionalExprType>(outputType);

ydb/library/yql/core/type_ann/type_ann_list.h

@@ -94,6 +94,7 @@ namespace NTypeAnnImpl {
     IGraphTransformer::TStatus ExtractMembersWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     template<bool Chopped>
     IGraphTransformer::TStatus AssumeConstraintWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
+    IGraphTransformer::TStatus AssumeConstraintsWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus AssumeColumnOrderWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
     IGraphTransformer::TStatus AggregationTraitsWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus MultiAggregateWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);

ydb/library/yql/core/type_ann/ya.make

@@ -35,6 +35,7 @@ PEERDIR(
     ydb/library/yql/providers/common/schema/expr
     ydb/library/yql/parser/pg_catalog
     ydb/library/yql/parser/pg_wrapper/interface
+    library/cpp/yson/node
 )
 
 YQL_LAST_ABI_VERSION()

ydb/library/yql/core/yql_expr_constraint.cpp

@@ -7,6 +7,8 @@
 #include <ydb/library/yql/core/yql_join.h>
 #include <ydb/library/yql/utils/log/profile.h>
 
+#include <library/cpp/yson/node/node_io.h>
+
 #include <util/generic/scope.h>
 #include <util/generic/utility.h>
 #include <util/generic/algorithm.h>
@@ -85,6 +87,7 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
         Functions["AssumeDistinct"] = &TCallableConstraintTransformer::AssumeUniqueWrap<true, true>;
         Functions["AssumeUniqueHint"] = &TCallableConstraintTransformer::AssumeUniqueWrap<false, false>;
         Functions["AssumeDistinctHint"] = &TCallableConstraintTransformer::AssumeUniqueWrap<true, false>;
+        Functions["AssumeConstraints"] = &TCallableConstraintTransformer::AssumeConstraintsWrap;
         Functions["AssumeChopped"] = &TCallableConstraintTransformer::AssumeChoppedWrap;
         Functions["AssumeColumnOrder"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
         Functions["AssumeAllMembersNullableAtOnce"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
@@ -387,6 +390,31 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
         return FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
     }
 
+    TStatus AssumeConstraintsWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
+        TConstraintSet set;
+        try {
+            set = ctx.MakeConstraintSet(NYT::NodeFromYsonString(input->Tail().Content()));
+        } catch (...) {
+            ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() <<
+                "Bad constraints yson-value: " << CurrentExceptionMessage()));
+            return IGraphTransformer::TStatus::Error;
+        }
+        for (auto constraint: set.GetAllConstraints()) {
+            if (!constraint->IsApplicableToType(*input->GetTypeAnn())) {
+                ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << *constraint
+                    << " is not applicable to " << *input->GetTypeAnn()));
+                return IGraphTransformer::TStatus::Error;
+            }
+        }
+        for (auto constr: input->Head().GetAllConstraints()) {
+            if (!constr->GetName().starts_with("PartOf") && !set.GetConstraint(constr->GetName())) {
+                set.AddConstraint(constr);
+            }
+        }
+        input->SetConstraints(set);
+        return IGraphTransformer::TStatus::Ok;
+    }
+
     template<bool Distinct, bool Strict>
     TStatus AssumeUniqueWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
         typename TUniqueConstraintNodeBase<Distinct>::TContentType content;

ydb/library/yql/core/yql_opt_utils.h

@@ -103,7 +103,7 @@ TExprNode::TPtr MakePgBool(TPositionHandle position, bool value, TExprContext& c
 TExprNode::TPtr MakeIdentityLambda(TPositionHandle position, TExprContext& ctx);
 
 constexpr std::initializer_list<std::string_view> SkippableCallables = {"Unordered", "AssumeSorted", "AssumeUnique", "AssumeDistinct",
-    "AssumeChopped", "AssumeColumnOrder", "AssumeAllMembersNullableAtOnce"};
+    "AssumeChopped", "AssumeColumnOrder", "AssumeAllMembersNullableAtOnce", "AssumeConstraints"};
 
 const TExprNode& SkipCallables(const TExprNode& node, const std::initializer_list<std::string_view>& skipCallables);