Skip to content

Commit 6f4d67d

Browse files
nepalnepal
authored
[YQL-17103] Basic support of pg types in predicate extractor library (#709)
* initial * fix typo * fix minus inf for pg types * more tests + canonize * temporarily disable KqpNotNullColumns::SecondaryIndexWithNotNullDataColumnPg
1 parent fdedcee commit 6f4d67d

File tree

15 files changed

+447
-55
lines changed

15 files changed

+447
-55
lines changed

‎ydb/core/kqp/ut/opt/kqp_not_null_ut.cpp

+3
Original file line numberDiff line numberDiff line change
@@ -1653,6 +1653,8 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
16531653
}
16541654
}
16551655

1656+
#if 0
1657+
// TODO: fix TxPlanSerializer with PG keys
16561658
Y_UNIT_TEST(SecondaryIndexWithNotNullDataColumnPg) {
16571659
auto settings = TKikimrSettings()
16581660
.SetWithSampleTables(false)
@@ -1757,6 +1759,7 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
17571759
result.GetIssues().ToString());
17581760
}
17591761
}
1762+
#endif
17601763

17611764
Y_UNIT_TEST_TWIN(JoinBothTablesWithNotNullPk, StreamLookup) {
17621765
NKikimrConfig::TAppConfig appConfig;

‎ydb/library/yql/core/extract_predicate/extract_predicate_impl.cpp

+53-13
Original file line numberDiff line numberDiff line change
@@ -32,13 +32,16 @@ TExprNode::TPtr BuildMultiplyLimit(TMaybe<size_t> limit, TExprContext& ctx, TPos
3232
}
3333
}
3434

35-
const TTypeAnnotationNode* GetBaseDataType(const TTypeAnnotationNode *type) {
35+
const TTypeAnnotationNode* GetBasePgOrDataType(const TTypeAnnotationNode* type) {
36+
if (type && type->GetKind() == ETypeAnnotationKind::Pg) {
37+
return type;
38+
}
3639
type = RemoveAllOptionals(type);
3740
return (type && type->GetKind() == ETypeAnnotationKind::Data) ? type : nullptr;
3841
}
3942

40-
bool IsDataOrMultiOptionalOfData(const TTypeAnnotationNode* type) {
41-
return GetBaseDataType(type) != nullptr;
43+
bool IsPgOrDataOrMultiOptionalOfData(const TTypeAnnotationNode* type) {
44+
return GetBasePgOrDataType(type) != nullptr;
4245
}
4346

4447
TMaybe<size_t> GetSqlInCollectionSize(const TExprNode::TPtr& collection) {
@@ -171,13 +174,13 @@ bool IsRoundingSupported(const TExprNode& op, bool negated) {
171174

172175
for (size_t i = 0; i < keyTypes.size(); ++i) {
173176
result.emplace_back();
174-
result.back().KeyBaseType = GetBaseDataType(keyTypes[i]);
175-
result.back().ValueBaseType = GetBaseDataType(valueTypes[i]);
177+
result.back().KeyBaseType = GetBasePgOrDataType(keyTypes[i]);
178+
result.back().ValueBaseType = GetBasePgOrDataType(valueTypes[i]);
176179
}
177180
} else {
178181
result.emplace_back();
179-
result.back().KeyBaseType = GetBaseDataType(keyType);
180-
result.back().ValueBaseType = GetBaseDataType(valueType);
182+
result.back().KeyBaseType = GetBasePgOrDataType(keyType);
183+
result.back().ValueBaseType = GetBasePgOrDataType(valueType);
181184
}
182185

183186
return result;
@@ -204,6 +207,16 @@ bool IsRoundingSupported(const TExprNode& op, bool negated) {
204207
for (auto& item : types) {
205208
YQL_ENSURE(item.KeyBaseType);
206209
YQL_ENSURE(item.ValueBaseType);
210+
const auto keyKind = item.KeyBaseType->GetKind();
211+
const auto valueKind = item.ValueBaseType->GetKind();
212+
if (keyKind != ETypeAnnotationKind::Data || valueKind != ETypeAnnotationKind::Data) {
213+
YQL_ENSURE(keyKind == valueKind);
214+
YQL_ENSURE(keyKind == ETypeAnnotationKind::Pg);
215+
if (!IsSameAnnotation(*item.KeyBaseType, *item.ValueBaseType)) {
216+
return false;
217+
}
218+
continue;
219+
}
207220

208221
EDataSlot valueSlot = item.ValueBaseType->Cast<TDataExprType>()->GetSlot();
209222
EDataSlot keySlot = item.KeyBaseType->Cast<TDataExprType>()->GetSlot();
@@ -225,7 +238,7 @@ bool IsRoundingSupported(const TExprNode& op, bool negated) {
225238
}
226239

227240
bool IsSupportedMemberNode(const TExprNode& node, const TExprNode& row) {
228-
return node.IsCallable("Member") && node.Child(0) == &row && IsDataOrMultiOptionalOfData(node.GetTypeAnn());
241+
return node.IsCallable("Member") && node.Child(0) == &row && IsPgOrDataOrMultiOptionalOfData(node.GetTypeAnn());
229242
}
230243

231244
bool IsValidForRange(const TExprNode& node, const TExprNode* otherNode, const TExprNode& row) {
@@ -572,14 +585,14 @@ bool IsListOfMembers(const TExprNode& node) {
572585
}
573586

574587
return AllOf(node.ChildrenList(), [](const auto& child) {
575-
return child->IsCallable("Member") && IsDataOrMultiOptionalOfData(child->GetTypeAnn());
588+
return child->IsCallable("Member") && IsPgOrDataOrMultiOptionalOfData(child->GetTypeAnn());
576589
});
577590
}
578591

579592
bool IsMemberBinOpNode(const TExprNode& node) {
580593
YQL_ENSURE(node.ChildrenSize() == 2);
581-
return node.Head().IsCallable("Member") && IsDataOrMultiOptionalOfData(node.Head().GetTypeAnn()) ||
582-
node.Tail().IsCallable("Member") && IsDataOrMultiOptionalOfData(node.Tail().GetTypeAnn());
594+
return node.Head().IsCallable("Member") && IsPgOrDataOrMultiOptionalOfData(node.Head().GetTypeAnn()) ||
595+
node.Tail().IsCallable("Member") && IsPgOrDataOrMultiOptionalOfData(node.Tail().GetTypeAnn());
583596
}
584597

585598
bool IsMemberListBinOpNode(const TExprNode& node) {
@@ -741,6 +754,33 @@ TExprNode::TPtr OptimizeNodeForRangeExtraction(const TExprNode::TPtr& node, cons
741754
}
742755
}
743756

757+
if (node->IsCallable("FromPg") && node->Head().IsCallable("PgResolvedOp")) {
758+
auto opNode = node->HeadPtr();
759+
TStringBuf op = opNode->Head().Content();
760+
if (SupportedBinOps.contains(op) || op == "<>" || op == "=") {
761+
auto left = opNode->ChildPtr(2);
762+
auto right = opNode->ChildPtr(3);
763+
if (IsSameAnnotation(*left->GetTypeAnn(), *right->GetTypeAnn())) {
764+
TStringBuf newOp;
765+
if (op == "=") {
766+
newOp = "==";
767+
} else if (op == "<>") {
768+
newOp = "!=";
769+
} else {
770+
newOp = op;
771+
}
772+
YQL_ENSURE(opNode->ChildrenSize() == 4);
773+
YQL_CLOG(DEBUG, Core) << "Replace PgResolvedOp(" << op << ") with corresponding plain binary operation";
774+
return ctx.Builder(node->Pos())
775+
.Callable(newOp)
776+
.Add(0, opNode->ChildPtr(2))
777+
.Add(1, opNode->ChildPtr(3))
778+
.Seal()
779+
.Build();
780+
}
781+
}
782+
}
783+
744784
return node;
745785
}
746786

@@ -2177,8 +2217,8 @@ TPredicateRangeExtractor::TBuildResult TPredicateRangeExtractor::BuildComputeNod
21772217
for (size_t i = 0; i < effectiveIndexKeys.size(); ++i) {
21782218
TMaybe<ui32> idx = RowType->FindItem(effectiveIndexKeys[i]);
21792219
if (idx) {
2180-
auto keyBaseType = RemoveAllOptionals(RowType->GetItems()[*idx]->GetItemType());
2181-
if (!(keyBaseType->GetKind() == ETypeAnnotationKind::Data && keyBaseType->IsComparable() && keyBaseType->IsEquatable())) {
2220+
auto keyBaseType = GetBasePgOrDataType(RowType->GetItems()[*idx]->GetItemType());
2221+
if (!(keyBaseType && keyBaseType->IsComparable() && keyBaseType->IsEquatable())) {
21822222
idx = {};
21832223
}
21842224
}

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

+47-15
Original file line numberDiff line numberDiff line change
@@ -10917,12 +10917,18 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1091710917

1091810918
bool IsValidTypeForRanges(const TTypeAnnotationNode* type) {
1091910919
YQL_ENSURE(type);
10920+
// top level optional is always present and used for +- infinity value
1092010921
if (type->GetKind() != ETypeAnnotationKind::Optional) {
1092110922
return false;
1092210923
}
10924+
type = type->Cast<TOptionalExprType>()->GetItemType();
10925+
bool hasOptionals = type->GetKind() == ETypeAnnotationKind::Optional;
1092310926
type = RemoveAllOptionals(type);
1092410927
YQL_ENSURE(type);
10925-
return type->GetKind() == ETypeAnnotationKind::Data && type->IsComparable() && type->IsEquatable();
10928+
if (!type->IsComparable() || !type->IsEquatable()) {
10929+
return false;
10930+
}
10931+
return type->GetKind() == ETypeAnnotationKind::Data || (type->GetKind() == ETypeAnnotationKind::Pg && !hasOptionals);
1092610932
}
1092710933

1092810934
bool EnsureValidRangeBoundary(TPositionHandle pos, const TTypeAnnotationNode* type, TExprContext& ctx) {
@@ -10957,7 +10963,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1095710963
ctx.AddError(TIssue(ctx.GetPosition(pos),
1095810964
TStringBuilder() << "Expected " << i <<
1095910965
"th component of range boundary tuple to be (multi) optional of "
10960-
"comparable and equatable Data type, but got: " << *itemType));
10966+
"comparable and equatable Data or Pg type, but got: " << *itemType));
1096110967
return false;
1096210968
}
1096310969
}
@@ -11043,7 +11049,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1104311049
if (!IsValidTypeForRanges(items[i])) {
1104411050
ctx.AddError(TIssue(ctx.GetPosition(pos),
1104511051
TStringBuilder() << "Expected " << i << "th component of range boundary tuple to be (multi) optional of "
11046-
"comparable and equatable Data type, but got: " << *items[i]));
11052+
"comparable and equatable Data or Pg type, but got: " << *items[i]));
1104711053
return false;
1104811054
}
1104911055
resultBoundaryItems.push_back(int32Type);
@@ -11152,7 +11158,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1115211158
for (auto& optKeyType : optKeys) {
1115311159
if (!IsValidTypeForRanges(optKeyType)) {
1115411160
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Head().Pos()),
11155-
TStringBuilder() << "Expected (multi) optional of comparable and equatable Data type, but got: " << *optKeyType));
11161+
TStringBuilder() << "Expected (multi) optional of comparable and equatable Data or Pg type, but got: " << *optKeyType));
1115611162
return IGraphTransformer::TStatus::Error;
1115711163
}
1115811164
resultBoundaryItems.push_back(int32Type);
@@ -11198,8 +11204,13 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1119811204
if (op != "Exists" && op != "NotExists") {
1119911205
valueBaseType = RemoveAllOptionals(valueType);
1120011206
YQL_ENSURE(valueBaseType);
11201-
if (valueBaseType->GetKind() != ETypeAnnotationKind::Data &&
11202-
valueBaseType->GetKind() != ETypeAnnotationKind::Null)
11207+
const auto valueKind = valueBaseType->GetKind();
11208+
if (valueKind != ETypeAnnotationKind::Pg) {
11209+
valueBaseType = RemoveAllOptionals(valueType);
11210+
}
11211+
if (valueKind != ETypeAnnotationKind::Data &&
11212+
valueKind != ETypeAnnotationKind::Null &&
11213+
valueKind != ETypeAnnotationKind::Pg)
1120311214
{
1120411215
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Child(1)->Pos()),
1120511216
TStringBuilder() << "Expecting (optional) Data as second argument, but got: " << *valueType));
@@ -11218,14 +11229,29 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1121811229
auto optKeyType = ctx.Expr.MakeType<TOptionalExprType>(keyType);
1121911230
if (!IsValidTypeForRanges(optKeyType)) {
1122011231
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Head().Pos()),
11221-
TStringBuilder() << "Expected (optional) of comparable and equatable Data type, but got: " << *keyType));
11232+
TStringBuilder() << "Expected (optional) of comparable and equatable Data or Pg type, but got: " << *keyType));
1122211233
return IGraphTransformer::TStatus::Error;
1122311234
}
1122411235

11225-
if (valueBaseType && CanCompare<false>(RemoveAllOptionals(keyType), valueBaseType) == ECompareOptions::Uncomparable) {
11226-
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
11227-
TStringBuilder() << "Uncompatible key and value types: " << *keyType << " and " << *valueType));
11228-
return IGraphTransformer::TStatus::Error;
11236+
if (valueBaseType) {
11237+
if (keyType->GetKind() == ETypeAnnotationKind::Pg) {
11238+
if (!IsSameAnnotation(*keyType, *valueBaseType)) {
11239+
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
11240+
TStringBuilder() << "Unequal key/value types are not supported: " << *keyType << " and " << *valueType));
11241+
return IGraphTransformer::TStatus::Error;
11242+
}
11243+
if (op == "StartsWith" || op == "NotStartsWith") {
11244+
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
11245+
TStringBuilder() << "Operation " << op << "is unsupported for pg type " << *keyType));
11246+
return IGraphTransformer::TStatus::Error;
11247+
11248+
}
11249+
}
11250+
if (CanCompare<false>(RemoveAllOptionals(keyType), valueBaseType) == ECompareOptions::Uncomparable) {
11251+
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
11252+
TStringBuilder() << "Uncompatible key and value types: " << *keyType << " and " << *valueType));
11253+
return IGraphTransformer::TStatus::Error;
11254+
}
1122911255
}
1123011256

1123111257
auto int32Type = ctx.Expr.MakeType<TDataExprType>(EDataSlot::Int32);
@@ -11412,7 +11438,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1141211438
auto optKeyType = ctx.Expr.MakeType<TOptionalExprType>(keyType);
1141311439
if (!IsValidTypeForRanges(optKeyType)) {
1141411440
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(keyNode->Pos()),
11415-
TStringBuilder() << "Unsupported index column type: expecting Data or (multi) optional of Data, "
11441+
TStringBuilder() << "Unsupported index column type: expecting Data or (multi) optional of Data (or Pg), "
1141611442
<< "got: " << *keyType << " for column '" << key << "'"));
1141711443
return IGraphTransformer::TStatus::Error;
1141811444
}
@@ -11438,7 +11464,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1143811464
return IGraphTransformer::TStatus::Error;
1143911465
}
1144011466

11441-
if (!EnsureDataType(input->Head(), ctx.Expr)) {
11467+
if (!EnsureDataOrPgType(input->Head(), ctx.Expr)) {
1144211468
return IGraphTransformer::TStatus::Error;
1144311469
}
1144411470

@@ -11447,12 +11473,12 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1144711473
}
1144811474

1144911475
auto dstType = input->Tail().GetTypeAnn()->Cast<TTypeExprType>()->GetType();
11450-
if (!EnsureDataType(input->Tail().Pos(), *dstType, ctx.Expr)) {
11476+
if (!EnsureDataOrPgType(input->Tail().Pos(), *dstType, ctx.Expr)) {
1145111477
return IGraphTransformer::TStatus::Error;
1145211478
}
1145311479

1145411480
auto srcType = input->Head().GetTypeAnn();
11455-
if (CanCompare<false>(srcType, dstType) != ECompareOptions::Comparable) {
11481+
if (CanCompare<false>(srcType, dstType) == ECompareOptions::Uncomparable) {
1145611482
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
1145711483
TStringBuilder() << "Uncompatible types in rounding: " << *srcType << " " << input->Content() << " to " << *dstType));
1145811484
return IGraphTransformer::TStatus::Error;
@@ -11463,6 +11489,12 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1146311489
return IGraphTransformer::TStatus::Repeat;
1146411490
}
1146511491

11492+
if (dstType->GetKind() == ETypeAnnotationKind::Pg) {
11493+
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
11494+
TStringBuilder() << "Pg types in rounding are not supported: " << *srcType << " " << input->Content() << " to " << *dstType));
11495+
return IGraphTransformer::TStatus::Error;
11496+
}
11497+
1146611498
auto sSlot = srcType->Cast<TDataExprType>()->GetSlot();
1146711499
auto tSlot = dstType->Cast<TDataExprType>()->GetSlot();
1146811500
auto resultType = ctx.Expr.MakeType<TOptionalExprType>(dstType);

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

+54
Original file line numberDiff line numberDiff line change
@@ -2335,6 +2335,60 @@ bool EnsureVariantType(TPositionHandle position, const TTypeAnnotationNode& type
23352335
return true;
23362336
}
23372337

2338+
bool EnsureDataOrPgType(const TExprNode& node, TExprContext& ctx) {
2339+
if (HasError(node.GetTypeAnn(), ctx)) {
2340+
return false;
2341+
}
2342+
2343+
if (!node.GetTypeAnn()) {
2344+
YQL_ENSURE(node.Type() == TExprNode::Lambda);
2345+
ctx.AddError(TIssue(ctx.GetPosition(node.Pos()), TStringBuilder() << "Expected data or pg type, but got lambda"));
2346+
return false;
2347+
}
2348+
2349+
return EnsureDataOrPgType(node.Pos(), *node.GetTypeAnn(), ctx);
2350+
}
2351+
2352+
bool EnsureDataOrPgType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx) {
2353+
if (HasError(&type, ctx)) {
2354+
return false;
2355+
}
2356+
2357+
if (type.GetKind() != ETypeAnnotationKind::Data && type.GetKind() != ETypeAnnotationKind::Pg) {
2358+
ctx.AddError(TIssue(ctx.GetPosition(position), TStringBuilder() << "Expected data or pg type, but got: " << type));
2359+
return false;
2360+
}
2361+
2362+
return true;
2363+
}
2364+
2365+
bool EnsurePgType(const TExprNode& node, TExprContext& ctx) {
2366+
if (HasError(node.GetTypeAnn(), ctx)) {
2367+
return false;
2368+
}
2369+
2370+
if (!node.GetTypeAnn()) {
2371+
YQL_ENSURE(node.Type() == TExprNode::Lambda);
2372+
ctx.AddError(TIssue(ctx.GetPosition(node.Pos()), TStringBuilder() << "Expected pg type, but got lambda"));
2373+
return false;
2374+
}
2375+
2376+
return EnsurePgType(node.Pos(), *node.GetTypeAnn(), ctx);
2377+
}
2378+
2379+
bool EnsurePgType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx) {
2380+
if (HasError(&type, ctx)) {
2381+
return false;
2382+
}
2383+
2384+
if (type.GetKind() != ETypeAnnotationKind::Pg) {
2385+
ctx.AddError(TIssue(ctx.GetPosition(position), TStringBuilder() << "Expected pg type, but got: " << type));
2386+
return false;
2387+
}
2388+
2389+
return true;
2390+
}
2391+
23382392
bool EnsureDataType(const TExprNode& node, TExprContext& ctx) {
23392393
if (HasError(node.GetTypeAnn(), ctx)) {
23402394
return false;

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

+4
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,10 @@ bool EnsureMultiType(const TExprNode& node, TExprContext& ctx);
9292
bool EnsureMultiType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);
9393
bool EnsureVariantType(const TExprNode& node, TExprContext& ctx);
9494
bool EnsureVariantType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);
95+
bool EnsureDataOrPgType(const TExprNode& node, TExprContext& ctx);
96+
bool EnsureDataOrPgType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);
97+
bool EnsurePgType(const TExprNode& node, TExprContext& ctx);
98+
bool EnsurePgType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);
9599
bool EnsureDataType(const TExprNode& node, TExprContext& ctx);
96100
bool EnsureDataType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);
97101
bool EnsureSpecificDataType(const TExprNode& node, EDataSlot expectedDataSlot, TExprContext& ctx, bool allowOptional = false);

0 commit comments

Comments
 (0)