diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp index 9c07d36e05bc..13d19f971159 100644 --- a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp +++ b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp @@ -38,11 +38,7 @@ class TKqpLogicalOptTransformer : public TOptimizeTransformerBase { AddHandler(0, &TDqJoin::Match, HNDL(JoinToIndexLookup)); AddHandler(0, &TCoCalcOverWindowBase::Match, HNDL(ExpandWindowFunctions)); AddHandler(0, &TCoCalcOverWindowGroup::Match, HNDL(ExpandWindowFunctions)); - AddHandler(0, &TCoFlatMap::Match, HNDL(LatePushExtractedPredicateToReadTable)); AddHandler(0, &TCoTopSort::Match, HNDL(RewriteTopSortOverFlatMap)); - AddHandler(0, &TCoTop::Match, HNDL(RewriteTopSortOverIndexRead)); - AddHandler(0, &TCoTopSort::Match, HNDL(RewriteTopSortOverIndexRead)); - AddHandler(0, &TCoTake::Match, HNDL(RewriteTakeOverIndexRead)); AddHandler(0, &TCoFlatMapBase::Match, HNDL(RewriteFlatMapOverExtend)); AddHandler(0, &TKqlDeleteRows::Match, HNDL(DeleteOverLookup)); AddHandler(0, &TKqlUpsertRowsBase::Match, HNDL(ExcessUpsertInputColumns)); @@ -54,21 +50,26 @@ class TKqpLogicalOptTransformer : public TOptimizeTransformerBase { AddHandler(0, &TCoTop::Match, HNDL(TopSortOverExtend)); AddHandler(0, &TCoTopSort::Match, HNDL(TopSortOverExtend)); - AddHandler(1, &TKqlReadTableIndex::Match, HNDL(RewriteIndexRead)); - AddHandler(1, &TKqlLookupIndex::Match, HNDL(RewriteLookupIndex)); - AddHandler(1, &TKqlStreamLookupIndex::Match, HNDL(RewriteStreamLookupIndex)); - AddHandler(1, &TKqlReadTableIndexRanges::Match, HNDL(RewriteIndexRead)); + AddHandler(1, &TCoFlatMap::Match, HNDL(LatePushExtractedPredicateToReadTable)); + AddHandler(1, &TCoTop::Match, HNDL(RewriteTopSortOverIndexRead)); + AddHandler(1, &TCoTopSort::Match, HNDL(RewriteTopSortOverIndexRead)); + AddHandler(1, &TCoTake::Match, HNDL(RewriteTakeOverIndexRead)); - AddHandler(2, &TKqlLookupTable::Match, HNDL(RewriteLookupTable)); + AddHandler(2, &TKqlReadTableIndex::Match, HNDL(RewriteIndexRead)); + AddHandler(2, &TKqlLookupIndex::Match, HNDL(RewriteLookupIndex)); + AddHandler(2, &TKqlStreamLookupIndex::Match, HNDL(RewriteStreamLookupIndex)); + AddHandler(2, &TKqlReadTableIndexRanges::Match, HNDL(RewriteIndexRead)); - AddHandler(3, &TKqlReadTableBase::Match, HNDL(ApplyExtractMembersToReadTable)); - AddHandler(3, &TKqlReadTableRangesBase::Match, HNDL(ApplyExtractMembersToReadTableRanges)); - AddHandler(3, &TKqpReadOlapTableRangesBase::Match, HNDL(ApplyExtractMembersToReadOlapTable)); - AddHandler(3, &TKqlLookupTableBase::Match, HNDL(ApplyExtractMembersToLookupTable)); + AddHandler(3, &TKqlLookupTable::Match, HNDL(RewriteLookupTable)); + + AddHandler(4, &TKqlReadTableBase::Match, HNDL(ApplyExtractMembersToReadTable)); + AddHandler(4, &TKqlReadTableRangesBase::Match, HNDL(ApplyExtractMembersToReadTableRanges)); + AddHandler(4, &TKqpReadOlapTableRangesBase::Match, HNDL(ApplyExtractMembersToReadOlapTable)); + AddHandler(4, &TKqlLookupTableBase::Match, HNDL(ApplyExtractMembersToLookupTable)); #undef HNDL - SetGlobal(3u); + SetGlobal(4u); } protected: diff --git a/ydb/core/kqp/ut/opt/kqp_ne_ut.cpp b/ydb/core/kqp/ut/opt/kqp_ne_ut.cpp index 1fe0006bcae8..302324cad7fa 100644 --- a/ydb/core/kqp/ut/opt/kqp_ne_ut.cpp +++ b/ydb/core/kqp/ut/opt/kqp_ne_ut.cpp @@ -3912,6 +3912,83 @@ Y_UNIT_TEST_SUITE(KqpNewEngine) { AssertTableReads(result, "/Root/SecondaryKeys/Index/indexImplTable", 1); } + + Y_UNIT_TEST_TWIN(ComplexLookupLimit, NewPredicateExtract) { + if (NewPredicateExtract) { + return; + } + + TKikimrSettings settings; + NKikimrConfig::TAppConfig appConfig; + appConfig.MutableTableServiceConfig()->SetPredicateExtract20(NewPredicateExtract); + settings.SetAppConfig(appConfig); + + TKikimrRunner kikimr(settings); + auto db = kikimr.GetTableClient(); + + { + auto session = db.CreateSession().GetValueSync().GetSession(); + AssertSuccessResult(session.ExecuteSchemeQuery(R"( + --!syntax_v1 + + CREATE TABLE `/Root/Sample` ( + A Uint64, + B Uint64, + C Uint64, + D Uint64, + E Uint64, + PRIMARY KEY (A, B, C) + ); + + )").GetValueSync()); + + AssertSuccessResult(session.ExecuteDataQuery(R"( + REPLACE INTO `/Root/Sample` (A, B, C, D, E) VALUES + (1, 1, 1, 1, 1), + (1, 2, 2, 2, 2), + (2, 2, 2, 2, 2), + (3, 3, 3, 3, 3), + (4, 4, 4, 4, 4), + (4, 4, 4, 4, 4), + (5, 5, 5, 5, 5); + )", TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx()).GetValueSync()); + } + + auto params = + TParamsBuilder() + .AddParam("$lastCounterId").Uint64(1).Build() + .AddParam("$lastId").Uint64(1).Build() + .AddParam("$counterIds") + .BeginList() + .AddListItem().Uint64(1) + .AddListItem().Uint64(2) + .AddListItem().Uint64(3) + .EndList() + .Build() + .Build(); + + auto session = db.CreateSession().GetValueSync().GetSession(); + + NYdb::NTable::TExecDataQuerySettings querySettings; + querySettings.CollectQueryStats(ECollectQueryStatsMode::Profile); + auto result = session.ExecuteDataQuery(R"( + DECLARE $counterIds AS List; + DECLARE $lastCounterId AS Uint64; + DECLARE $lastId AS Uint64; + SELECT A, B FROM + `/Root/Sample` + WHERE + A in $counterIds and + (A,B) > ($lastCounterId, $lastId) + ORDER BY A,B + LIMIT 2; + )", TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), params, querySettings).GetValueSync(); + + AssertSuccessResult(result); + AssertTableReads(result, "/Root/Sample", NewPredicateExtract ? 2 : 4); + CompareYson(R"([[[1u];[2u]];[[2u];[2u]]])", FormatResultSetYson(result.GetResultSet(0))); + } + } } // namespace NKikimr::NKqp