Skip to content

Commit 75b77be

Browse files
author
Velikhov Pavel
committed
Working implementation
1 parent 5b212db commit 75b77be

File tree

2 files changed

+112
-11
lines changed

2 files changed

+112
-11
lines changed

‎ydb/core/kqp/opt/rbo/kqp_rbo_rules.cpp‎

Lines changed: 95 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,10 @@ std::shared_ptr<IOperator> TInlineSimpleInExistsSubplanRule::SimpleTestAndApply(
451451
return input;
452452
}
453453

454+
if (lambdaBody->IsCallable("ToPg")) {
455+
lambdaBody = lambdaBody->ChildPtr(0);
456+
}
457+
454458
// Decompose the conjunction into individual conjuncts
455459
TVector<TExprNode::TPtr> conjuncts;
456460
if (TCoAnd::Match(lambdaBody.Get())) {
@@ -470,14 +474,14 @@ std::shared_ptr<IOperator> TInlineSimpleInExistsSubplanRule::SimpleTestAndApply(
470474
for (conjunctIdx=0; conjunctIdx<conjuncts.size(); conjunctIdx++) {
471475
auto maybeSubplan = conjuncts[conjunctIdx];
472476

473-
if (maybeSubplan->IsCallable("ToPg")) {
477+
while (maybeSubplan->IsCallable("ToPg") || maybeSubplan->IsCallable("FromPg")) {
474478
maybeSubplan = maybeSubplan->ChildPtr(0);
475479
}
476480
if (TCoNot::Match(maybeSubplan.Get())){
477481
maybeSubplan = maybeSubplan->ChildPtr(0);
478482
negated = true;
479483
}
480-
if (maybeSubplan->IsCallable("FromPg")) {
484+
while (maybeSubplan->IsCallable("ToPg") || maybeSubplan->IsCallable("FromPg")) {
481485
maybeSubplan = maybeSubplan->ChildPtr(0);
482486
}
483487
if (TCoMember::Match(maybeSubplan.Get())) {
@@ -497,11 +501,21 @@ std::shared_ptr<IOperator> TInlineSimpleInExistsSubplanRule::SimpleTestAndApply(
497501
return input;
498502
}
499503

500-
// We build a semi-join or a left-only join
501-
auto leftJoinInput = filter->GetInput();
502-
TString joinKind = negated ? "LeftOnly" : "LeftSemi";
503-
TVector<std::pair<TInfoUnit, TInfoUnit>> joinKeys;
504+
std::shared_ptr<TOpJoin> join;
505+
506+
// We build a semi-join or a left-only join when processing IN subplan
504507
if (subplan.Type == ESubplanType::IN) {
508+
509+
auto leftJoinInput = filter->GetInput();
510+
TString joinKind;
511+
if (subplan.Type == ESubplanType::IN) {
512+
joinKind = negated ? "LeftOnly" : "LeftSemi";
513+
} else {
514+
joinKind = "Cross";
515+
}
516+
517+
TVector<std::pair<TInfoUnit, TInfoUnit>> joinKeys;
518+
505519
auto planIUs = subplan.Plan->GetOutputIUs();
506520
Y_ENSURE(subplan.Tuple.size() == planIUs.size());
507521

@@ -511,17 +525,21 @@ std::shared_ptr<IOperator> TInlineSimpleInExistsSubplanRule::SimpleTestAndApply(
511525

512526
auto arg = Build<TCoArgument>(ctx.ExprCtx, filter->Pos).Name("lambda_arg").Done();
513527

528+
// clang-format off
514529
auto toPgMember = Build<TCoMember>(ctx.ExprCtx, filter->Pos)
515530
.Struct(arg)
516531
.Name<TCoAtom>().Value(subplan.Tuple[i].GetFullName()).Build()
517532
.Done().Ptr();
533+
// clang-format on
518534

519535
toPgMember = ctx.ExprCtx.NewCallable(filter->Pos, "ToPg", {toPgMember});
520536

537+
// clang-format off
521538
auto toPgLambda = Build<TCoLambda>(ctx.ExprCtx, filter->Pos)
522539
.Args({arg})
523540
.Body(toPgMember)
524541
.Done().Ptr();
542+
// clang-format on
525543

526544
auto toPgVar = TInfoUnit("_rbo_arg_" + std::to_string(props.InternalVarIdx++), true);
527545
std::variant<TInfoUnit, TExprNode::TPtr> toPgElement = toPgLambda;
@@ -537,11 +555,77 @@ std::shared_ptr<IOperator> TInlineSimpleInExistsSubplanRule::SimpleTestAndApply(
537555
joinKeys.push_back(std::make_pair(subplan.Tuple[i], planIUs[i]));
538556
}
539557
}
558+
559+
join = std::make_shared<TOpJoin>(leftJoinInput, subplan.Plan, input->Pos, joinKind, joinKeys);
560+
conjuncts.erase(conjuncts.begin() + conjunctIdx);
561+
}
562+
563+
// EXISTS and NOT EXISTS
564+
else {
565+
566+
auto countResult = TInfoUnit("_rbo_arg_" + std::to_string(props.InternalVarIdx++), true);
567+
568+
TVector<std::pair<TInfoUnit, std::variant<TInfoUnit, TExprNode::TPtr>>> countMapElements;
569+
570+
auto zero = ctx.ExprCtx.NewCallable(filter->Pos, "Uint64", {ctx.ExprCtx.NewAtom(filter->Pos, "0")});
571+
572+
// clang-format off
573+
auto zeroLambda = Build<TCoLambda>(ctx.ExprCtx, filter->Pos)
574+
.Args({"arg"})
575+
.Body(zero)
576+
.Done().Ptr();
577+
// clang-format on
578+
579+
countMapElements.push_back(std::make_pair(countResult, zeroLambda));
580+
auto countMap = std::make_shared<TOpMap>(subplan.Plan, filter->Pos, countMapElements, true);
581+
582+
583+
TOpAggregationTraits aggFunction(countResult, "count");
584+
TVector<TOpAggregationTraits> aggs = {aggFunction};
585+
TVector<TInfoUnit> keyColumns;
586+
587+
auto agg = std::make_shared<TOpAggregate>(countMap, aggs, keyColumns, EAggregationPhase::Final, false, filter->Pos);
588+
589+
TVector<std::pair<TInfoUnit, std::variant<TInfoUnit, TExprNode::TPtr>>> mapElements;
590+
auto arg = Build<TCoArgument>(ctx.ExprCtx, filter->Pos).Name("lambda_arg").Done();
591+
592+
// FIXME: Apparently count(*) on empty result return null
593+
TString compareCallable = negated ? "==" : "!=";
594+
595+
// clang-format off
596+
auto member = Build<TCoMember>(ctx.ExprCtx, filter->Pos)
597+
.Struct(arg)
598+
.Name().Value(countResult.GetFullName()).Build()
599+
.Done().Ptr();
600+
// clang-format on
601+
602+
auto body = ctx.ExprCtx.NewCallable(filter->Pos, compareCallable, {member, zero});
603+
604+
// clang-format off
605+
auto lambda = Build<TCoLambda>(ctx.ExprCtx, filter->Pos)
606+
.Args({arg})
607+
.Body(body)
608+
.Done().Ptr();
609+
// clang-format on
610+
611+
auto compareResult = TInfoUnit("_rbo_arg_" + std::to_string(props.InternalVarIdx++), true);
612+
mapElements.push_back(std::make_pair(compareResult, lambda));
613+
auto map = std::make_shared<TOpMap>(agg, filter->Pos, mapElements, true);
614+
615+
TVector<std::pair<TInfoUnit, TInfoUnit>> joinKeys;
616+
join = std::make_shared<TOpJoin>(filter->GetInput(), map, filter->Pos, "Cross", joinKeys);
617+
618+
// clang-format off
619+
auto resultMember = Build<TCoMember>(ctx.ExprCtx, filter->Pos)
620+
.Struct(arg)
621+
.Name().Value(compareResult.GetFullName()).Build()
622+
.Done().Ptr();
623+
// clang-format on
624+
625+
conjuncts[conjunctIdx] = resultMember;
540626
}
541627

542-
auto join = std::make_shared<TOpJoin>(leftJoinInput, subplan.Plan, input->Pos, joinKind, joinKeys);
543628
props.Subplans.Remove(iu);
544-
conjuncts.erase(conjuncts.begin() + conjunctIdx);
545629

546630
// If there was a single conjunct, we can get rid of the filter completely
547631
if (conjuncts.empty()) {
@@ -556,6 +640,9 @@ std::shared_ptr<IOperator> TInlineSimpleInExistsSubplanRule::SimpleTestAndApply(
556640
newLambdaBody = conjuncts[0];
557641
} else {
558642
newLambdaBody = Build<TCoAnd>(ctx.ExprCtx, input->Pos).Add(conjuncts).Done().Ptr();
643+
if (props.PgSyntax) {
644+
newLambdaBody = ctx.ExprCtx.NewCallable(filter->Pos, "ToPg", {newLambdaBody});
645+
}
559646
}
560647
newLambdaBody = ReplaceArg(newLambdaBody, arg, ctx.ExprCtx);
561648

‎ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,7 @@ Y_UNIT_TEST_SUITE(KqpRboPg) {
237237
}
238238
}
239239

240-
Y_UNIT_TEST(ScalarSubquery) {
240+
Y_UNIT_TEST(ExpressionSubquery) {
241241
NKikimrConfig::TAppConfig appConfig;
242242
appConfig.MutableTableServiceConfig()->SetEnableNewRBO(true);
243243
appConfig.MutableTableServiceConfig()->SetEnableFallbackToYqlOptimizer(false);
@@ -310,13 +310,25 @@ Y_UNIT_TEST_SUITE(KqpRboPg) {
310310
R"(
311311
--!syntax_pg
312312
SET TablePathPrefix = "/Root/";
313-
SELECT bar.id FROM bar where EXISTS (SELECT id from foo);
313+
SELECT foo.id FROM foo where foo.id = 0 AND foo.id NOT IN (SELECT id from bar where id <> 0);
314314
)",
315315
R"(
316316
--!syntax_pg
317317
SET TablePathPrefix = "/Root/";
318-
SELECT bar.id FROM bar where NOT EXISTS (SELECT id from foo where id < 0);
318+
SELECT foo.id FROM foo where EXISTS (SELECT id from bar);
319+
)",
320+
R"(
321+
--!syntax_pg
322+
SET TablePathPrefix = "/Root/";
323+
SELECT foo.id FROM foo where foo.id = 0 AND EXISTS (SELECT id from bar);
324+
)",
325+
/*
326+
R"(
327+
--!syntax_pg
328+
SET TablePathPrefix = "/Root/";
329+
SELECT foo.id FROM foo where NOT EXISTS (SELECT id from bar where id = 100);
319330
)"
331+
*/
320332
};
321333

322334
// TODO: The order of result is not defined, we need order by to add more interesting tests.
@@ -325,6 +337,8 @@ Y_UNIT_TEST_SUITE(KqpRboPg) {
325337
R"([["0"]])",
326338
R"([["0"]])",
327339
R"([["0"]])",
340+
R"([["0"]])",
341+
R"([["0"]])",
328342
R"([["0"]])"
329343
};
330344

0 commit comments

Comments
 (0)