Skip to content

Commit bc9e922

Browse files
ulya-sidorinaulya-sidorina
authored
fix(kqp): pass right table filter through stream index lookup join (#1034)
1 parent 4737601 commit bc9e922

File tree

10 files changed

+307
-207
lines changed

10 files changed

+307
-207
lines changed

‎ydb/core/kqp/expr_nodes/kqp_expr_nodes.json‎

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -180,21 +180,10 @@
180180
{
181181
"Name": "TKqlStreamLookupTable",
182182
"Base": "TKqlLookupTableBase",
183-
"Match": {"Type": "Callable", "Name": "KqlStreamLookupTable"}
184-
},
185-
{
186-
"Name": "TKqlStreamIdxLookupJoin",
187-
"Base": "TCallable",
188-
"Match": {"Type": "Callable", "Name": "KqlStreamIdxLookupJoin"},
183+
"Match": {"Type": "Callable", "Name": "KqlStreamLookupTable"},
189184
"Children": [
190-
{"Index": 0, "Name": "LeftInput", "Type": "TExprBase"},
191-
{"Index": 1, "Name": "LeftLabel", "Type": "TCoAtom"},
192-
{"Index": 2, "Name": "RightTable", "Type": "TKqpTable"},
193-
{"Index": 3, "Name": "RightColumns", "Type": "TCoAtomList"},
194-
{"Index": 4, "Name": "RightLabel", "Type": "TCoAtom"},
195-
{"Index": 5, "Name": "JoinType", "Type": "TCoAtom"}
185+
{"Index": 3, "Name": "LookupStrategy", "Type": "TCoAtom"}
196186
]
197-
198187
},
199188
{
200189
"Name": "TKqlSequencer",
@@ -469,16 +458,26 @@
469458
]
470459
},
471460
{
472-
"Name": "TKqpIndexLookupJoin",
461+
"Name": "TKqlIndexLookupJoinBase",
473462
"Base": "TCallable",
474-
"Match": {"Type": "Callable", "Name": "KqpIndexLookupJoin"},
463+
"Match": {"Type": "CallableBase"},
475464
"Children": [
476465
{"Index": 0, "Name": "Input", "Type": "TExprBase"},
477466
{"Index": 1, "Name": "JoinType", "Type": "TCoAtom"},
478467
{"Index": 2, "Name": "LeftLabel", "Type": "TCoAtom"},
479468
{"Index": 3, "Name": "RightLabel", "Type": "TCoAtom"}
480469
]
481470
},
471+
{
472+
"Name": "TKqlIndexLookupJoin",
473+
"Base": "TKqlIndexLookupJoinBase",
474+
"Match": {"Type": "Callable", "Name": "KqlIndexLookupJoin"}
475+
},
476+
{
477+
"Name": "TKqpIndexLookupJoin",
478+
"Base": "TKqlIndexLookupJoinBase",
479+
"Match": {"Type": "Callable", "Name": "KqpIndexLookupJoin"}
480+
},
482481
{
483482
"Name": "TKqpCnSequencer",
484483
"Base": "TKqpConnection",

‎ydb/core/kqp/host/kqp_type_ann.cpp‎

Lines changed: 68 additions & 118 deletions
Original file line numberDiff line numberDiff line change
@@ -425,7 +425,7 @@ TStatus AnnotateReadTableRanges(const TExprNode::TPtr& node, TExprContext& ctx,
425425
TStatus AnnotateLookupTable(const TExprNode::TPtr& node, TExprContext& ctx, const TString& cluster,
426426
const TKikimrTablesData& tablesData, bool withSystemColumns)
427427
{
428-
if (!EnsureArgsCount(*node, TKqlLookupIndexBase::Match(node.Get()) ? 4 : 3, ctx)) {
428+
if (!EnsureArgsCount(*node, TKqlLookupIndexBase::Match(node.Get()) || TKqlStreamLookupTable::Match(node.Get()) ? 4 : 3, ctx)) {
429429
return TStatus::Error;
430430
}
431431

@@ -467,10 +467,57 @@ TStatus AnnotateLookupTable(const TExprNode::TPtr& node, TExprContext& ctx, cons
467467
}
468468

469469
YQL_ENSURE(lookupType);
470-
if (!EnsureStructType(node->Pos(), *lookupType, ctx)) {
471-
return TStatus::Error;
470+
471+
const TStructExprType* structType = nullptr;
472+
bool isStreamLookup = TKqlStreamLookupTable::Match(node.Get());
473+
if (isStreamLookup) {
474+
auto lookupStrategy = node->Child(TKqlStreamLookupTable::idx_LookupStrategy);
475+
if (!EnsureAtom(*lookupStrategy, ctx)) {
476+
return TStatus::Error;
477+
}
478+
479+
if (lookupStrategy->Content() == TKqpStreamLookupJoinStrategyName) {
480+
if (!EnsureTupleType(node->Pos(), *lookupType, ctx)) {
481+
return TStatus::Error;
482+
}
483+
484+
if (!EnsureTupleTypeSize(node->Pos(), lookupType, 2, ctx)) {
485+
return TStatus::Error;
486+
}
487+
488+
auto tupleType = lookupType->Cast<TTupleExprType>();
489+
if (!EnsureStructType(node->Pos(), *tupleType->GetItems()[0], ctx)) {
490+
return TStatus::Error;
491+
}
492+
493+
if (!EnsureStructType(node->Pos(), *tupleType->GetItems()[1], ctx)) {
494+
return TStatus::Error;
495+
}
496+
497+
structType = tupleType->GetItems()[0]->Cast<TStructExprType>();
498+
auto leftRowType = tupleType->GetItems()[1]->Cast<TStructExprType>();
499+
500+
TVector<const TTypeAnnotationNode*> outputTypes;
501+
outputTypes.push_back(leftRowType);
502+
outputTypes.push_back(ctx.MakeType<TOptionalExprType>(rowType));
503+
504+
rowType = ctx.MakeType<TTupleExprType>(outputTypes);
505+
} else {
506+
if (!EnsureStructType(node->Pos(), *lookupType, ctx)) {
507+
return TStatus::Error;
508+
}
509+
510+
structType = lookupType->Cast<TStructExprType>();
511+
}
512+
} else {
513+
if (!EnsureStructType(node->Pos(), *lookupType, ctx)) {
514+
return TStatus::Error;
515+
}
516+
517+
structType = lookupType->Cast<TStructExprType>();
472518
}
473-
auto structType = lookupType->Cast<TStructExprType>();
519+
520+
YQL_ENSURE(structType);
474521

475522
ui32 keyColumnsCount = 0;
476523
if (TKqlLookupIndexBase::Match(node.Get())) {
@@ -1338,105 +1385,6 @@ TStatus AnnotateSequencer(const TExprNode::TPtr& node, TExprContext& ctx, const
13381385
return TStatus::Ok;
13391386
}
13401387

1341-
TStatus AnnotateStreamIdxLookupJoin(const TExprNode::TPtr& node, TExprContext& ctx, const TString& cluster,
1342-
const TKikimrTablesData& tablesData, bool withSystemColumns)
1343-
{
1344-
if (!EnsureArgsCount(*node, 6, ctx)) {
1345-
return TStatus::Error;
1346-
}
1347-
1348-
auto leftInputType = node->Child(TKqlStreamIdxLookupJoin::idx_LeftInput)->GetTypeAnn();
1349-
const TTypeAnnotationNode* leftInputItemType;
1350-
if (!EnsureNewSeqType<false>(node->Pos(), *leftInputType, ctx, &leftInputItemType)) {
1351-
return TStatus::Error;
1352-
}
1353-
1354-
YQL_ENSURE(leftInputItemType);
1355-
if (!EnsureTupleType(node->Pos(), *leftInputItemType, ctx)) {
1356-
return TStatus::Error;
1357-
}
1358-
1359-
if (!EnsureTupleTypeSize(node->Pos(), leftInputItemType, 2, ctx)) {
1360-
return TStatus::Error;
1361-
}
1362-
1363-
auto leftInputTupleType = leftInputItemType->Cast<TTupleExprType>();
1364-
if (!EnsureStructType(node->Pos(), *leftInputTupleType->GetItems()[0], ctx)) {
1365-
return TStatus::Error;
1366-
}
1367-
1368-
if (!EnsureStructType(node->Pos(), *leftInputTupleType->GetItems()[1], ctx)) {
1369-
return TStatus::Error;
1370-
}
1371-
1372-
if (!EnsureAtom(*node->Child(TKqlStreamIdxLookupJoin::idx_LeftLabel), ctx)) {
1373-
return TStatus::Error;
1374-
}
1375-
1376-
TCoAtom leftLabel(node->Child(TKqlStreamIdxLookupJoin::idx_LeftLabel));
1377-
1378-
auto rightTable = ResolveTable(node->Child(TKqlStreamIdxLookupJoin::idx_RightTable), ctx, cluster, tablesData);
1379-
if (!rightTable.second) {
1380-
return TStatus::Error;
1381-
}
1382-
1383-
const TStructExprType* inputKeysType = leftInputTupleType->GetItems()[0]->Cast<TStructExprType>();
1384-
for (const auto& inputKey : inputKeysType->GetItems()) {
1385-
if (!rightTable.second->GetKeyColumnIndex(TString(inputKey->GetName()))) {
1386-
return TStatus::Error;
1387-
}
1388-
}
1389-
1390-
if (!EnsureTupleOfAtoms(*node->Child(TKqlStreamIdxLookupJoin::idx_RightColumns), ctx)) {
1391-
return TStatus::Error;
1392-
}
1393-
1394-
TCoAtomList rightColumns{node->ChildPtr(TKqlStreamIdxLookupJoin::idx_RightColumns)};
1395-
for (const auto& rightColumn : rightColumns) {
1396-
if (!rightTable.second->GetColumnType(TString(rightColumn.Value()))) {
1397-
return TStatus::Error;
1398-
}
1399-
}
1400-
1401-
auto rightDataType = GetReadTableRowType(ctx, tablesData, cluster, rightTable.first, rightColumns, withSystemColumns);
1402-
if (!rightDataType) {
1403-
return TStatus::Error;
1404-
}
1405-
1406-
if (!EnsureAtom(*node->Child(TKqlStreamIdxLookupJoin::idx_RightLabel), ctx)) {
1407-
return TStatus::Error;
1408-
}
1409-
1410-
TCoAtom rightLabel(node->Child(TKqlStreamIdxLookupJoin::idx_RightLabel));
1411-
TCoAtom joinType(node->Child(TKqlStreamIdxLookupJoin::idx_JoinType));
1412-
1413-
const TStructExprType* leftDataType = leftInputTupleType->GetItems()[1]->Cast<TStructExprType>();
1414-
TVector<const TItemExprType*> resultStructItems;
1415-
for (const auto& member : leftDataType->GetItems()) {
1416-
resultStructItems.emplace_back(
1417-
ctx.MakeType<TItemExprType>(TString::Join(leftLabel.Value(), ".", member->GetName()), member->GetItemType())
1418-
);
1419-
}
1420-
1421-
if (RightJoinSideAllowed(joinType.Value())) {
1422-
for (const auto& member : rightDataType->Cast<TStructExprType>()->GetItems()) {
1423-
const bool makeOptional = RightJoinSideOptional(joinType.Value()) && !member->GetItemType()->IsOptionalOrNull();
1424-
1425-
const TTypeAnnotationNode* memberType = makeOptional
1426-
? ctx.MakeType<TOptionalExprType>(member->GetItemType())
1427-
: member->GetItemType();
1428-
1429-
resultStructItems.emplace_back(
1430-
ctx.MakeType<TItemExprType>(TString::Join(rightLabel.Value(), ".", member->GetName()), memberType)
1431-
);
1432-
}
1433-
}
1434-
1435-
auto rowType = ctx.MakeType<TStructExprType>(resultStructItems);
1436-
node->SetTypeAnn(ctx.MakeType<TListExprType>(rowType));
1437-
return TStatus::Ok;
1438-
}
1439-
14401388
TStatus AnnotateKqpProgram(const TExprNode::TPtr& node, TExprContext& ctx) {
14411389
if (!EnsureArgsCount(*node, 2, ctx)) {
14421390
return TStatus::Error;
@@ -1621,7 +1569,7 @@ TStatus AnnotateStreamLookupConnection(const TExprNode::TPtr& node, TExprContext
16211569

16221570
YQL_ENSURE(inputItemType);
16231571

1624-
if (lookupStrategy.Value() == "LookupRows") {
1572+
if (lookupStrategy.Value() == TKqpStreamLookupStrategyName) {
16251573
if (!EnsureStructType(node->Pos(), *inputItemType, ctx)) {
16261574
return TStatus::Error;
16271575
}
@@ -1640,7 +1588,7 @@ TStatus AnnotateStreamLookupConnection(const TExprNode::TPtr& node, TExprContext
16401588

16411589
node->SetTypeAnn(ctx.MakeType<TStreamExprType>(rowType));
16421590

1643-
} else if (lookupStrategy.Value() == "LookupJoinRows") {
1591+
} else if (lookupStrategy.Value() == TKqpStreamLookupJoinStrategyName) {
16441592
if (!EnsureTupleType(node->Pos(), *inputItemType, ctx)) {
16451593
return TStatus::Error;
16461594
}
@@ -1694,7 +1642,7 @@ TStatus AnnotateIndexLookupJoin(const TExprNode::TPtr& node, TExprContext& ctx)
16941642
return TStatus::Error;
16951643
}
16961644

1697-
auto inputType = node->Child(TKqpIndexLookupJoin::idx_Input)->GetTypeAnn();
1645+
auto inputType = node->Child(TKqlIndexLookupJoinBase::idx_Input)->GetTypeAnn();
16981646
const TTypeAnnotationNode* inputItemType;
16991647
if (!EnsureNewSeqType<false>(node->Pos(), *inputType, ctx, &inputItemType)) {
17001648
return TStatus::Error;
@@ -1725,22 +1673,22 @@ TStatus AnnotateIndexLookupJoin(const TExprNode::TPtr& node, TExprContext& ctx)
17251673
return TStatus::Error;
17261674
}
17271675

1728-
if (!EnsureAtom(*node->Child(TKqpIndexLookupJoin::idx_JoinType), ctx)) {
1676+
if (!EnsureAtom(*node->Child(TKqlIndexLookupJoinBase::idx_JoinType), ctx)) {
17291677
return TStatus::Error;
17301678
}
17311679

1732-
if (!EnsureAtom(*node->Child(TKqpIndexLookupJoin::idx_LeftLabel), ctx)) {
1680+
if (!EnsureAtom(*node->Child(TKqlIndexLookupJoinBase::idx_LeftLabel), ctx)) {
17331681
return TStatus::Error;
17341682
}
17351683

1736-
TCoAtom leftLabel(node->Child(TKqpIndexLookupJoin::idx_LeftLabel));
1684+
TCoAtom leftLabel(node->Child(TKqlIndexLookupJoinBase::idx_LeftLabel));
17371685

1738-
if (!EnsureAtom(*node->Child(TKqpIndexLookupJoin::idx_RightLabel), ctx)) {
1686+
if (!EnsureAtom(*node->Child(TKqlIndexLookupJoinBase::idx_RightLabel), ctx)) {
17391687
return TStatus::Error;
17401688
}
17411689

1742-
TCoAtom rightLabel(node->Child(TKqpIndexLookupJoin::idx_RightLabel));
1743-
TCoAtom joinType(node->Child(TKqpIndexLookupJoin::idx_JoinType));
1690+
TCoAtom rightLabel(node->Child(TKqlIndexLookupJoinBase::idx_RightLabel));
1691+
TCoAtom joinType(node->Child(TKqlIndexLookupJoinBase::idx_JoinType));
17441692

17451693
TVector<const TItemExprType*> resultStructItems;
17461694
for (const auto& item : leftRowType->GetItems()) {
@@ -1764,7 +1712,13 @@ TStatus AnnotateIndexLookupJoin(const TExprNode::TPtr& node, TExprContext& ctx)
17641712
}
17651713

17661714
auto outputRowType = ctx.MakeType<TStructExprType>(resultStructItems);
1767-
node->SetTypeAnn(ctx.MakeType<TStreamExprType>(outputRowType));
1715+
const bool isPhysical = TKqpIndexLookupJoin::Match(node.Get());
1716+
if (isPhysical) {
1717+
node->SetTypeAnn(ctx.MakeType<TStreamExprType>(outputRowType));
1718+
} else {
1719+
node->SetTypeAnn(ctx.MakeType<TListExprType>(outputRowType));
1720+
}
1721+
17681722
return TStatus::Ok;
17691723
}
17701724

@@ -1892,7 +1846,7 @@ TAutoPtr<IGraphTransformer> CreateKqpTypeAnnotationTransformer(const TString& cl
18921846
return AnnotateStreamLookupConnection(input, ctx, cluster, *tablesData, config->SystemColumnsEnabled());
18931847
}
18941848

1895-
if (TKqpIndexLookupJoin::Match(input.Get())) {
1849+
if (TKqlIndexLookupJoinBase::Match(input.Get())) {
18961850
return AnnotateIndexLookupJoin(input, ctx);
18971851
}
18981852

@@ -1924,10 +1878,6 @@ TAutoPtr<IGraphTransformer> CreateKqpTypeAnnotationTransformer(const TString& cl
19241878
return AnnotateSequencer(input, ctx, cluster, *tablesData);
19251879
}
19261880

1927-
if (TKqlStreamIdxLookupJoin::Match(input.Get())) {
1928-
return AnnotateStreamIdxLookupJoin(input, ctx, cluster, *tablesData, config->SystemColumnsEnabled());
1929-
}
1930-
19311881
if (TKqpProgram::Match(input.Get())) {
19321882
return AnnotateKqpProgram(input, ctx);
19331883
}

‎ydb/core/kqp/opt/logical/kqp_opt_log_extract.cpp‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,17 @@ TExprBase KqpApplyExtractMembersToLookupTable(TExprBase node, TExprContext& ctx,
214214
.Done();
215215
}
216216

217+
if (auto maybeStreamLookup = lookup.Maybe<TKqlStreamLookupTable>()) {
218+
auto streamLookup = maybeStreamLookup.Cast();
219+
220+
return Build<TKqlStreamLookupTable>(ctx, lookup.Pos())
221+
.Table(streamLookup.Table())
222+
.LookupKeys(streamLookup.LookupKeys())
223+
.Columns(usedColumns.Cast())
224+
.LookupStrategy(streamLookup.LookupStrategy())
225+
.Done();
226+
}
227+
217228
return Build<TKqlLookupTableBase>(ctx, lookup.Pos())
218229
.CallableName(lookup.CallableName())
219230
.Table(lookup.Table())

‎ydb/core/kqp/opt/logical/kqp_opt_log_indexes.cpp‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,7 @@ TExprBase DoRewriteIndexRead(const TReadMatch& read, TExprContext& ctx,
321321
.Table(read.Table())
322322
.LookupKeys(readIndexTable.Ptr())
323323
.Columns(read.Columns())
324+
.LookupStrategy().Build(TKqpStreamLookupStrategyName)
324325
.Done();
325326
} else {
326327
return Build<TKqlLookupTable>(ctx, read.Pos())
@@ -365,6 +366,7 @@ TExprBase KqpRewriteLookupIndex(const TExprBase& node, TExprContext& ctx, const
365366
.Table(BuildTableMeta(*indexMeta, node.Pos(), ctx))
366367
.LookupKeys(lookupIndex.LookupKeys())
367368
.Columns(lookupIndex.Columns())
369+
.LookupStrategy().Build(TKqpStreamLookupStrategyName)
368370
.Done();
369371
}
370372

@@ -382,12 +384,14 @@ TExprBase KqpRewriteLookupIndex(const TExprBase& node, TExprContext& ctx, const
382384
.Table(BuildTableMeta(*indexMeta, node.Pos(), ctx))
383385
.LookupKeys(lookupIndex.LookupKeys())
384386
.Columns(keyColumnsList)
387+
.LookupStrategy().Build(TKqpStreamLookupStrategyName)
385388
.Done();
386389

387390
return Build<TKqlStreamLookupTable>(ctx, node.Pos())
388391
.Table(lookupIndex.Table())
389392
.LookupKeys(lookupIndexTable.Ptr())
390393
.Columns(lookupIndex.Columns())
394+
.LookupStrategy().Build(TKqpStreamLookupStrategyName)
391395
.Done();
392396
}
393397

@@ -424,6 +428,7 @@ TExprBase KqpRewriteStreamLookupIndex(const TExprBase& node, TExprContext& ctx,
424428
.Table(BuildTableMeta(*indexMeta, node.Pos(), ctx))
425429
.LookupKeys(streamLookupIndex.LookupKeys())
426430
.Columns(streamLookupIndex.Columns())
431+
.LookupStrategy().Build(TKqpStreamLookupStrategyName)
427432
.Done();
428433
}
429434

@@ -433,12 +438,14 @@ TExprBase KqpRewriteStreamLookupIndex(const TExprBase& node, TExprContext& ctx,
433438
.Table(BuildTableMeta(*indexMeta, node.Pos(), ctx))
434439
.LookupKeys(streamLookupIndex.LookupKeys())
435440
.Columns(keyColumnsList)
441+
.LookupStrategy().Build(TKqpStreamLookupStrategyName)
436442
.Done();
437443

438444
return Build<TKqlStreamLookupTable>(ctx, node.Pos())
439445
.Table(streamLookupIndex.Table())
440446
.LookupKeys(lookupIndexTable.Ptr())
441447
.Columns(streamLookupIndex.Columns())
448+
.LookupStrategy().Build(TKqpStreamLookupStrategyName)
442449
.Done();
443450
}
444451

0 commit comments

Comments
 (0)