Fix rules order for lookupjoin

ydb-platform · Sep 25, 2023 · f7f3847 · f7f3847
1 parent 2ee6ad0
commit f7f3847
Show file tree

Hide file tree

Showing 2 changed files with 92 additions and 14 deletions.
diff --git 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<true>));
-        AddHandler(3, &TKqlReadTableRangesBase::Match, HNDL(ApplyExtractMembersToReadTableRanges<true>));
-        AddHandler(3, &TKqpReadOlapTableRangesBase::Match, HNDL(ApplyExtractMembersToReadOlapTable<true>));
-        AddHandler(3, &TKqlLookupTableBase::Match, HNDL(ApplyExtractMembersToLookupTable<true>));
+        AddHandler(3, &TKqlLookupTable::Match, HNDL(RewriteLookupTable));
+
+        AddHandler(4, &TKqlReadTableBase::Match, HNDL(ApplyExtractMembersToReadTable<true>));
+        AddHandler(4, &TKqlReadTableRangesBase::Match, HNDL(ApplyExtractMembersToReadTableRanges<true>));
+        AddHandler(4, &TKqpReadOlapTableRangesBase::Match, HNDL(ApplyExtractMembersToReadOlapTable<true>));
+        AddHandler(4, &TKqlLookupTableBase::Match, HNDL(ApplyExtractMembersToLookupTable<true>));
 
 #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
@@ -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<Uint64>;
+            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