support Take/TopSort rules for secondary indexes

feature(kqp_scan_query): support Take/TopSort rules for secondary index

3 files changed, 93 insertions, 44 deletions

diff --git a/ydb/core/kqp/opt/kqp_opt_kql.cpp b/ydb/core/kqp/opt/kqp_opt_kql.cpp
index 2081c4236a0..c9f1c846974 100644
--- a/ydb/core/kqp/opt/kqp_opt_kql.cpp
+++ b/ydb/core/kqp/opt/kqp_opt_kql.cpp
@@ -518,6 +518,12 @@ TExprNode::TPtr HandleReadTable(const TKiReadTable& read, TExprContext& ctx, con
             return nullptr;
         }
 
+        if (kqpCtx->IsScanQuery() && !kqpCtx->Config->FeatureFlags.GetEnableKqpScanQueryStreamLookup()) {
+            const TString err = "Secondary index is not supported for ScanQuery";
+            ctx.AddError(YqlIssue(ctx.GetPosition(read.Pos()), TIssuesIds::KIKIMR_BAD_REQUEST, err));
+            return nullptr;
+        }
+
         auto [metadata, state] = tableData.Metadata->GetIndexMetadata(indexName);
         YQL_ENSURE(metadata, "unable to find metadata for index: " << indexName);
         YQL_ENSURE(state == TIndexDescription::EIndexState::Ready
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_indexes.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_indexes.cpp
index 10e142a848b..ed3f4576acf 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_indexes.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_indexes.cpp
@@ -27,7 +27,7 @@ bool CheckIndexCovering(const TRead& read, const TIntrusivePtr<TKikimrTableMetad
 }
 
 TExprBase DoRewriteIndexRead(const TKqlReadTableIndex& read, TExprContext& ctx,
-    const TKikimrTableDescription& tableDesc, TIntrusivePtr<TKikimrTableMetadata> indexMeta,
+    const TKikimrTableDescription& tableDesc, TIntrusivePtr<TKikimrTableMetadata> indexMeta, bool useStreamLookup,
     const TVector<TString>& extraColumns, const std::function<TExprBase(const TExprBase&)>& middleFilter = {})
 {
     const bool needDataRead = CheckIndexCovering(read, indexMeta);
@@ -99,11 +99,19 @@ TExprBase DoRewriteIndexRead(const TKqlReadTableIndex& read, TExprContext& ctx,
             .Done();
     }
 
-    return Build<TKqlLookupTable>(ctx, read.Pos())
-        .Table(read.Table())
-        .LookupKeys(readIndexTable.Ptr())
-        .Columns(read.Columns())
-        .Done();
+    if (useStreamLookup) {
+        return Build<TKqlStreamLookupTable>(ctx, read.Pos())
+            .Table(read.Table())
+            .LookupKeys(readIndexTable.Ptr())
+            .Columns(read.Columns())
+            .Done();
+    } else {
+        return Build<TKqlLookupTable>(ctx, read.Pos())
+            .Table(read.Table())
+            .LookupKeys(readIndexTable.Ptr())
+            .Columns(read.Columns())
+            .Done();
+    }
 }
 
 } // namespace
@@ -115,34 +123,7 @@ TExprBase KqpRewriteIndexRead(const TExprBase& node, TExprContext& ctx, const TK
         const auto& tableDesc = GetTableData(*kqpCtx.Tables, kqpCtx.Cluster, indexRead.Table().Path());
         const auto& [indexMeta, _ ] = tableDesc.Metadata->GetIndexMetadata(TString(indexRead.Index().Value()));
 
-        if (kqpCtx.IsDataQuery()) {
-            return DoRewriteIndexRead(indexRead, ctx, tableDesc, indexMeta, {});
-        }
-
-        const bool needDataRead = CheckIndexCovering(indexRead, indexMeta);
-        if (!needDataRead) {
-            return Build<TKqlReadTable>(ctx, indexRead.Pos())
-                .Table(BuildTableMeta(*indexMeta, indexRead.Pos(), ctx))
-                .Range(indexRead.Range())
-                .Columns(indexRead.Columns())
-                .Settings(indexRead.Settings())
-                .Done();
-        }
-
-        auto keyColumnsList = BuildKeyColumnsList(tableDesc, indexRead.Pos(), ctx);
-
-        TExprBase readIndexTable = Build<TKqlReadTable>(ctx, indexRead.Pos())
-            .Table(BuildTableMeta(*indexMeta, indexRead.Pos(), ctx))
-            .Range(indexRead.Range())
-            .Columns(keyColumnsList)
-            .Settings(indexRead.Settings())
-            .Done();
-
-        return Build<TKqlStreamLookupTable>(ctx, indexRead.Pos())
-            .Table(indexRead.Table())
-            .LookupKeys(readIndexTable.Ptr())
-            .Columns(indexRead.Columns())
-            .Done();
+        return DoRewriteIndexRead(indexRead, ctx, tableDesc, indexMeta, kqpCtx.IsScanQuery(), {});
     }
 
     return node;
@@ -229,10 +210,6 @@ TExprBase KqpRewriteStreamLookupIndex(const TExprBase& node, TExprContext& ctx,
 // through TKqlLookupTable.
 // The simplest way is to match TopSort or Take over TKqlReadTableIndex.
 TExprBase KqpRewriteTopSortOverIndexRead(const TExprBase& node, TExprContext& ctx, const TKqpOptimizeContext& kqpCtx) {
-    if (!kqpCtx.IsDataQuery()) {
-        return node;
-    }
-
     if (!node.Maybe<TCoTopSort>()) {
         return node;
     }
@@ -262,7 +239,8 @@ TExprBase KqpRewriteTopSortOverIndexRead(const TExprBase& node, TExprContext& ct
             return TExprBase(newTopSort);
         };
 
-        auto lookup = DoRewriteIndexRead(readTableIndex, ctx, tableDesc, indexMeta, sortByColumns, filter);
+        auto lookup = DoRewriteIndexRead(readTableIndex, ctx, tableDesc, indexMeta,
+            kqpCtx.IsScanQuery(), sortByColumns, filter);
 
         return Build<TCoTopSort>(ctx, node.Pos())
             .Input(lookup)
@@ -276,10 +254,6 @@ TExprBase KqpRewriteTopSortOverIndexRead(const TExprBase& node, TExprContext& ct
 }
 
 TExprBase KqpRewriteTakeOverIndexRead(const TExprBase& node, TExprContext& ctx, const TKqpOptimizeContext& kqpCtx) {
-    if (!kqpCtx.IsDataQuery()) {
-        return node;
-    }
-
     if (!node.Maybe<TCoTake>()) {
         return node;
     }
@@ -297,7 +271,7 @@ TExprBase KqpRewriteTakeOverIndexRead(const TExprBase& node, TExprContext& ctx,
             return TExprBase(ctx.ChangeChild(*node.Ptr(), 0, in.Ptr()));
         };
 
-        return DoRewriteIndexRead(readTableIndex, ctx, tableDesc, indexMeta, {}, filter);
+        return DoRewriteIndexRead(readTableIndex, ctx, tableDesc, indexMeta, kqpCtx.IsScanQuery(), {}, filter);
     }
 
     return node;
diff --git a/ydb/core/kqp/ut/kqp_scan_ut.cpp b/ydb/core/kqp/ut/kqp_scan_ut.cpp
index 62c9e64cd0d..1e6e8b844ac 100644
--- a/ydb/core/kqp/ut/kqp_scan_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_scan_ut.cpp
@@ -2129,7 +2129,76 @@ Y_UNIT_TEST_SUITE(KqpScan) {
             PRAGMA kikimr.OptEnablePredicateExtract = "false";
             SELECT Value FROM `/Root/Table` WHERE Key IN AsList(1, 2, 3);
         )");
+    }
+
+    Y_UNIT_TEST_TWIN(LimitOverSecondaryIndexRead, UseSessionActor) {
+        auto kikimr = KikimrRunnerEnableSessionActor(UseSessionActor);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        CreateSampleTablesWithIndex(session);
+
+        TStreamExecScanQuerySettings querySettings;
+        querySettings.Explain(true);
+
+        auto itIndex = db.StreamExecuteScanQuery(R"(
+            SELECT *
+            FROM `/Root/SecondaryComplexKeys` VIEW Index
+            WHERE Fk1 == 1
+            LIMIT 2;
+        )", querySettings).GetValueSync();
+
+        UNIT_ASSERT_C(itIndex.IsSuccess(), itIndex.GetIssues().ToString());
+
+        auto res = CollectStreamResult(itIndex);
+        UNIT_ASSERT(res.PlanJson);
+
+        Cerr << *res.PlanJson;
+
+        NJson::TJsonValue plan;
+        NJson::ReadJsonTree(*res.PlanJson, &plan, true);
+
+        auto indexRead = FindPlanNodeByKv(plan, "Node Type", "Limit-TablePointLookup");
+        UNIT_ASSERT(indexRead.IsDefined());
+        auto indexTable = FindPlanNodeByKv(indexRead, "Table", "SecondaryComplexKeys/Index/indexImplTable");
+        UNIT_ASSERT(indexTable.IsDefined());
+        auto limit = FindPlanNodeByKv(indexRead, "Limit", "2");
+        UNIT_ASSERT(limit.IsDefined());
+    }
+
+    Y_UNIT_TEST_TWIN(TopSortOverSecondaryIndexRead, UseSessionActor) {
+        auto kikimr = KikimrRunnerEnableSessionActor(UseSessionActor);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        CreateSampleTablesWithIndex(session);
+
+        TStreamExecScanQuerySettings querySettings;
+        querySettings.Explain(true);
+
+        auto itIndex = db.StreamExecuteScanQuery(R"(
+            SELECT *
+            FROM `/Root/SecondaryComplexKeys` VIEW Index
+            WHERE Fk1 == 1
+            ORDER BY Fk1 LIMIT 2;
+        )", querySettings).GetValueSync();
+
+        UNIT_ASSERT_C(itIndex.IsSuccess(), itIndex.GetIssues().ToString());
+
+        auto res = CollectStreamResult(itIndex);
+        UNIT_ASSERT(res.PlanJson);
+
+        Cerr << *res.PlanJson;
+
+        NJson::TJsonValue plan;
+        NJson::ReadJsonTree(*res.PlanJson, &plan, true);
 
+        auto indexRead = FindPlanNodeByKv(plan, "Node Type", "Limit-TablePointLookup");
+        UNIT_ASSERT(indexRead.IsDefined());
+        auto indexTable = FindPlanNodeByKv(indexRead, "Table", "SecondaryComplexKeys/Index/indexImplTable");
+        UNIT_ASSERT(indexTable.IsDefined());
+        auto limit = FindPlanNodeByKv(indexRead, "Limit", "2");
+        UNIT_ASSERT(limit.IsDefined());
     }
 }