Skip to content

Commit 2ca83de

Browse files
rvu1024rvu1024
authored
AssumeConstraints callable (#8349)
1 parent f8f5387 commit 2ca83de

File tree

9 files changed

+110
-19
lines changed

9 files changed

+110
-19
lines changed

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

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
#include <ydb/library/yql/utils/log/log.h>
1414
#include <ydb/library/yql/parser/pg_catalog/catalog.h>
1515

16+
#include <library/cpp/yson/node/node_io.h>
17+
1618
#include <util/generic/map.h>
1719
#include <util/string/cast.h>
1820
#include <util/generic/xrange.h>
@@ -2266,7 +2268,7 @@ TExprNode::TPtr HasNullOverVariant(const TExprNode::TPtr& node, TExprContext& ct
22662268
}
22672269

22682270
constexpr std::initializer_list<std::string_view> FlowPriority = {
2269-
"AssumeSorted", "AssumeUnique", "AssumeDistinct",
2271+
"AssumeSorted", "AssumeUnique", "AssumeDistinct", "AssumeConstraints",
22702272
"Map", "OrderedMap", "MapNext",
22712273
"Filter", "OrderedFilter",
22722274
"FlatMap", "OrderedFlatMap",
@@ -5944,6 +5946,22 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) {
59445946
return ctx.ChangeChild(*node, 0, node->Head().HeadPtr());
59455947
}
59465948

5949+
if (node->Head().IsCallable("AssumeConstraints") && node->Head().GetConstraint<TSortedConstraintNode>()) {
5950+
TConstraintSet constrSet = node->Head().GetConstraintSet();
5951+
constrSet.RemoveConstraint<TSortedConstraintNode>();
5952+
YQL_CLOG(DEBUG, Core) << node->Content() << " over " << node->Head().Content();
5953+
auto res = ctx.ChangeChild(*node, 0, node->Head().HeadPtr());
5954+
if (constrSet) {
5955+
res = ctx.Builder(node->Head().Pos())
5956+
.Callable("AssumeConstraints")
5957+
.Add(0, std::move(res))
5958+
.Atom(1, NYT::NodeToYsonString(constrSet.ToYson(), NYson::EYsonFormat::Text), TNodeFlags::MultilineContent)
5959+
.Seal()
5960+
.Build();
5961+
}
5962+
return res;
5963+
}
5964+
59475965
return node;
59485966
};
59495967

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1625,7 +1625,7 @@
16251625
{"Index": 5, "Name": "LeftRenames", "Type": "TCoAtomList"},
16261626
{"Index": 6, "Name": "RightRenames", "Type": "TCoAtomList"},
16271627
{"Index": 7, "Name": "LeftKeysColumnNames", "Type": "TCoAtomList"},
1628-
{"Index": 8, "Name": "RightKeysColumnNames", "Type": "TCoAtomList"},
1628+
{"Index": 8, "Name": "RightKeysColumnNames", "Type": "TCoAtomList"},
16291629
{"Index": 9, "Name": "Flags", "Type": "TCoAtomList"}
16301630
]
16311631
},
@@ -1641,7 +1641,7 @@
16411641
{"Index": 4, "Name": "LeftRenames", "Type": "TCoAtomList"},
16421642
{"Index": 5, "Name": "RightRenames", "Type": "TCoAtomList"},
16431643
{"Index": 6, "Name": "LeftKeysColumnNames", "Type": "TCoAtomList"},
1644-
{"Index": 7, "Name": "RightKeysColumnNames", "Type": "TCoAtomList"},
1644+
{"Index": 7, "Name": "RightKeysColumnNames", "Type": "TCoAtomList"},
16451645
{"Index": 8, "Name": "Flags", "Type": "TCoAtomList"}
16461646
]
16471647
},
@@ -1863,6 +1863,14 @@
18631863
"Base": "TCoInputBase",
18641864
"Match": {"Type": "Callable", "Name": "AssumeAllMembersNullableAtOnce"}
18651865
},
1866+
{
1867+
"Name": "TCoAssumeConstraints",
1868+
"Base": "TCoInputBase",
1869+
"Match": {"Type": "Callable", "Name": "AssumeConstraints"},
1870+
"Children": [
1871+
{"Index": 1, "Name": "Value", "Type": "TCoAtom"}
1872+
]
1873+
},
18661874
{
18671875
"Name": "TCoListItemType",
18681876
"Base": "TCallable",

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6980,7 +6980,7 @@ template <bool Asc, bool Equals>
69806980
TExprNode::TPtr AggrComparePg(const TExprNode& node, TExprContext& ctx) {
69816981
YQL_CLOG(DEBUG, CorePeepHole) << "Expand '" << node.Content() << "' over Pg.";
69826982
auto op = Asc ? ( Equals ? "<=" : "<") : ( Equals ? ">=" : ">");
6983-
auto finalPart = (Equals == Asc) ? MakeBool<true>(node.Pos(), ctx) :
6983+
auto finalPart = (Equals == Asc) ? MakeBool<true>(node.Pos(), ctx) :
69846984
ctx.NewCallable(node.Pos(), "Exists", { node.TailPtr() });
69856985
if (!Asc) {
69866986
finalPart = ctx.NewCallable(node.Pos(), "Not", { finalPart });
@@ -8445,6 +8445,7 @@ struct TPeepHoleRules {
84458445
{"AssumeUnique", &DropAssume},
84468446
{"AssumeDistinct", &DropAssume},
84478447
{"AssumeChopped", &DropAssume},
8448+
{"AssumeConstraints", &DropAssume},
84488449
{"EmptyFrom", &DropEmptyFrom},
84498450
{"Top", &OptimizeTopOrSort<false, true>},
84508451
{"TopSort", &OptimizeTopOrSort<true, true>},
@@ -8493,7 +8494,7 @@ struct TPeepHoleRules {
84938494
const TExtPeepHoleOptimizerMap BlockStageExtFinalRules = {
84948495
{"BlockExtend", &ExpandBlockExtend},
84958496
{"BlockOrderedExtend", &ExpandBlockExtend},
8496-
{"ReplicateScalars", &ExpandReplicateScalars}
8497+
{"ReplicateScalars", &ExpandReplicateScalars}
84978498
};
84988499

84998500
static const TPeepHoleRules& Instance() {

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

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1232,7 +1232,7 @@ namespace NTypeAnnImpl {
12321232

12331233
YQL_ENSURE(IsSameAnnotation(*memberType, *resultType) ||
12341234
IsSameAnnotation(*ctx.Expr.MakeType<TOptionalExprType>(memberType), *resultType));
1235-
1235+
12361236
output = ctx.Expr.Builder(input->Pos())
12371237
.Callable("IfStrict")
12381238
.Callable(0, "Exists")
@@ -2749,8 +2749,8 @@ namespace NTypeAnnImpl {
27492749
if (!(*dataTypeOne == *dataTypeTwo)) {
27502750
ctx.Expr.AddError(TIssue(
27512751
ctx.Expr.GetPosition(input->Pos()),
2752-
TStringBuilder() << "Cannot calculate with different decimals: "
2753-
<< static_cast<const TTypeAnnotationNode&>(*dataType[0]) << " != "
2752+
TStringBuilder() << "Cannot calculate with different decimals: "
2753+
<< static_cast<const TTypeAnnotationNode&>(*dataType[0]) << " != "
27542754
<< static_cast<const TTypeAnnotationNode&>(*dataType[1])
27552755
));
27562756

@@ -6193,18 +6193,18 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
61936193
};
61946194

61956195
auto mergeMembers = [&ctx, &buildJustMember, &input, &left, &right, &mergeLambda](const TStringBuf& name, bool hasLeft, bool hasRight) -> TExprNode::TPtr {
6196-
auto leftMaybe = hasLeft ?
6196+
auto leftMaybe = hasLeft ?
61976197
buildJustMember(left, name) :
61986198
ctx.Expr.NewCallable(input->Pos(), "Nothing", {
61996199
ExpandType(input->Pos(), *ctx.Expr.MakeType<TOptionalExprType>(right->GetTypeAnn()->Cast<TStructExprType>()->FindItemType(name)), ctx.Expr)
62006200
});
6201-
6202-
auto rightMaybe = hasRight ?
6201+
6202+
auto rightMaybe = hasRight ?
62036203
buildJustMember(right, name) :
62046204
ctx.Expr.NewCallable(input->Pos(), "Nothing", {
62056205
ExpandType(input->Pos(), *ctx.Expr.MakeType<TOptionalExprType>(left->GetTypeAnn()->Cast<TStructExprType>()->FindItemType(name)), ctx.Expr)
62066206
});
6207-
6207+
62086208
return ctx.Expr.Builder(input->Pos())
62096209
.List()
62106210
.Atom(0, name)
@@ -6492,7 +6492,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
64926492
.With(1, result)
64936493
.Seal()
64946494
.Build();
6495-
}
6495+
}
64966496
}
64976497
} else if (collection->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Tuple) {
64986498
for (size_t idx = 0; idx < collection->GetTypeAnn()->Cast<TTupleExprType>()->GetSize(); idx++) {
@@ -6523,7 +6523,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
65236523
return IGraphTransformer::TStatus::Error;
65246524
}
65256525

6526-
if (result) {
6526+
if (result) {
65276527
output = result;
65286528
} else {
65296529
output = ctx.Expr.Builder(input->Pos()).Callable("Null").Seal().Build();
@@ -12170,6 +12170,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1217012170
Functions["AssumeUniqueHint"] = &AssumeConstraintWrapper<false>;
1217112171
Functions["AssumeDistinctHint"] = &AssumeConstraintWrapper<false>;
1217212172
Functions["AssumeChopped"] = &AssumeConstraintWrapper<true>;
12173+
Functions["AssumeConstraints"] = &AssumeConstraintsWrapper;
1217312174
Functions["AssumeAllMembersNullableAtOnce"] = &AssumeAllMembersNullableAtOnceWrapper;
1217412175
Functions["AssumeStrict"] = &AssumeStrictWrapper;
1217512176
Functions["AssumeNonStrict"] = &AssumeStrictWrapper;

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

Lines changed: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111

1212
#include <ydb/library/yql/parser/pg_catalog/catalog.h>
1313

14+
#include <library/cpp/yson/node/node_io.h>
15+
1416
#include <util/generic/algorithm.h>
1517
#include <util/string/join.h>
1618

@@ -1432,10 +1434,10 @@ namespace {
14321434
}
14331435

14341436
IGraphTransformer::TStatus ListTopSortWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
1435-
if (!EnsureMinMaxArgsCount(*input, 2, 3, ctx.Expr)) {
1437+
if (!EnsureMinMaxArgsCount(*input, 2, 3, ctx.Expr)) {
14361438
return IGraphTransformer::TStatus::Error;
14371439
}
1438-
1440+
14391441
TStringBuf newName = input->Content();
14401442
newName.Skip(4);
14411443
bool desc = false;
@@ -1457,7 +1459,7 @@ namespace {
14571459
.Seal()
14581460
.Build();
14591461
}
1460-
1462+
14611463
return OptListWrapperImpl<4U, 4U>(ctx.Expr.Builder(input->Pos())
14621464
.Callable(newName)
14631465
.Add(0, input->ChildPtr(0))
@@ -4412,6 +4414,37 @@ namespace {
44124414
template IGraphTransformer::TStatus AssumeConstraintWrapper<true>(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
44134415
template IGraphTransformer::TStatus AssumeConstraintWrapper<false>(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
44144416

4417+
IGraphTransformer::TStatus AssumeConstraintsWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TContext& ctx) {
4418+
if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
4419+
return IGraphTransformer::TStatus::Error;
4420+
}
4421+
4422+
if (!EnsureAnySeqType(input->Head(), ctx.Expr)) {
4423+
return IGraphTransformer::TStatus::Error;
4424+
}
4425+
4426+
if (!EnsureAtom(input->Tail(), ctx.Expr)) {
4427+
return IGraphTransformer::TStatus::Error;
4428+
}
4429+
4430+
NYT::TNode node;
4431+
try {
4432+
node = NYT::NodeFromYsonString(input->Tail().Content());
4433+
} catch (...) {
4434+
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), TStringBuilder() <<
4435+
"Bad constraints yson-value: " << CurrentExceptionMessage()));
4436+
return IGraphTransformer::TStatus::Error;
4437+
}
4438+
if (!node.IsMap()) {
4439+
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), TStringBuilder() <<
4440+
"Expected yson-map as serialized constraints value, actual " << node.GetType()));
4441+
return IGraphTransformer::TStatus::Error;
4442+
}
4443+
4444+
input->SetTypeAnn(input->Head().GetTypeAnn());
4445+
return IGraphTransformer::TStatus::Ok;
4446+
}
4447+
44154448
IGraphTransformer::TStatus AssumeColumnOrderWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx) {
44164449
if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
44174450
return IGraphTransformer::TStatus::Error;
@@ -6503,7 +6536,7 @@ namespace {
65036536
return IGraphTransformer::TStatus::Repeat;
65046537
}
65056538

6506-
const TTypeAnnotationNode* outputType = ctx.Expr.MakeType<TDataExprType>(input->IsCallable("PercentRank") ?
6539+
const TTypeAnnotationNode* outputType = ctx.Expr.MakeType<TDataExprType>(input->IsCallable("PercentRank") ?
65076540
EDataSlot::Double : EDataSlot::Uint64);
65086541
if (!isAnsi && keyType->GetKind() == ETypeAnnotationKind::Optional) {
65096542
outputType = ctx.Expr.MakeType<TOptionalExprType>(outputType);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ namespace NTypeAnnImpl {
9494
IGraphTransformer::TStatus ExtractMembersWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
9595
template<bool Chopped>
9696
IGraphTransformer::TStatus AssumeConstraintWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
97+
IGraphTransformer::TStatus AssumeConstraintsWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
9798
IGraphTransformer::TStatus AssumeColumnOrderWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
9899
IGraphTransformer::TStatus AggregationTraitsWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
99100
IGraphTransformer::TStatus MultiAggregateWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ PEERDIR(
3535
ydb/library/yql/providers/common/schema/expr
3636
ydb/library/yql/parser/pg_catalog
3737
ydb/library/yql/parser/pg_wrapper/interface
38+
library/cpp/yson/node
3839
)
3940

4041
YQL_LAST_ABI_VERSION()

‎ydb/library/yql/core/yql_expr_constraint.cpp

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
#include <ydb/library/yql/core/yql_join.h>
88
#include <ydb/library/yql/utils/log/profile.h>
99

10+
#include <library/cpp/yson/node/node_io.h>
11+
1012
#include <util/generic/scope.h>
1113
#include <util/generic/utility.h>
1214
#include <util/generic/algorithm.h>
@@ -85,6 +87,7 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
8587
Functions["AssumeDistinct"] = &TCallableConstraintTransformer::AssumeUniqueWrap<true, true>;
8688
Functions["AssumeUniqueHint"] = &TCallableConstraintTransformer::AssumeUniqueWrap<false, false>;
8789
Functions["AssumeDistinctHint"] = &TCallableConstraintTransformer::AssumeUniqueWrap<true, false>;
90+
Functions["AssumeConstraints"] = &TCallableConstraintTransformer::AssumeConstraintsWrap;
8891
Functions["AssumeChopped"] = &TCallableConstraintTransformer::AssumeChoppedWrap;
8992
Functions["AssumeColumnOrder"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
9093
Functions["AssumeAllMembersNullableAtOnce"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
@@ -387,6 +390,31 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
387390
return FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
388391
}
389392

393+
TStatus AssumeConstraintsWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
394+
TConstraintSet set;
395+
try {
396+
set = ctx.MakeConstraintSet(NYT::NodeFromYsonString(input->Tail().Content()));
397+
} catch (...) {
398+
ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() <<
399+
"Bad constraints yson-value: " << CurrentExceptionMessage()));
400+
return IGraphTransformer::TStatus::Error;
401+
}
402+
for (auto constraint: set.GetAllConstraints()) {
403+
if (!constraint->IsApplicableToType(*input->GetTypeAnn())) {
404+
ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << *constraint
405+
<< " is not applicable to " << *input->GetTypeAnn()));
406+
return IGraphTransformer::TStatus::Error;
407+
}
408+
}
409+
for (auto constr: input->Head().GetAllConstraints()) {
410+
if (!constr->GetName().starts_with("PartOf") && !set.GetConstraint(constr->GetName())) {
411+
set.AddConstraint(constr);
412+
}
413+
}
414+
input->SetConstraints(set);
415+
return IGraphTransformer::TStatus::Ok;
416+
}
417+
390418
template<bool Distinct, bool Strict>
391419
TStatus AssumeUniqueWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
392420
typename TUniqueConstraintNodeBase<Distinct>::TContentType content;

‎ydb/library/yql/core/yql_opt_utils.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ TExprNode::TPtr MakePgBool(TPositionHandle position, bool value, TExprContext& c
103103
TExprNode::TPtr MakeIdentityLambda(TPositionHandle position, TExprContext& ctx);
104104

105105
constexpr std::initializer_list<std::string_view> SkippableCallables = {"Unordered", "AssumeSorted", "AssumeUnique", "AssumeDistinct",
106-
"AssumeChopped", "AssumeColumnOrder", "AssumeAllMembersNullableAtOnce"};
106+
"AssumeChopped", "AssumeColumnOrder", "AssumeAllMembersNullableAtOnce", "AssumeConstraints"};
107107

108108
const TExprNode& SkipCallables(const TExprNode& node, const std::initializer_list<std::string_view>& skipCallables);
109109

0 commit comments

Comments
 (0)