Fix KqpRewriteTopSortOverIndexRead KQP optimizer rule. (KIKIMR-15237)

ref:ca716939bd9f2359a1df160fb23f192c01c8190f

3 files changed, 79 insertions, 6 deletions

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 4bb48475a7a..955abe4e987 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_indexes.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_indexes.cpp
@@ -164,9 +164,9 @@ TExprBase KqpRewriteLookupIndex(const TExprBase& node, TExprContext& ctx, const
     return node;
 }
 
-// The index and main table have same number of rows, so we can push TCoTopSort or TCoTake through TKqlLookupTable.
+// The index and main table have same number of rows, so we can push a copy of TCoTopSort or TCoTake
+// 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;
@@ -201,7 +201,14 @@ TExprBase KqpRewriteTopSortOverIndexRead(const TExprBase& node, TExprContext& ct
             return TExprBase(newTopSort);
         };
 
-        return DoRewriteIndexRead(readTableIndex, ctx, tableDesc, indexMeta, sortByColumns, filter);
+        auto lookup = DoRewriteIndexRead(readTableIndex, ctx, tableDesc, indexMeta, sortByColumns, filter);
+
+        return Build<TCoTopSort>(ctx, node.Pos())
+            .Input(lookup)
+            .KeySelectorLambda(ctx.DeepCopyLambda(topSort.KeySelectorLambda().Ref()))
+            .SortDirections(topSort.SortDirections())
+            .Count(topSort.Count())
+            .Done();
     }
 
     return node;
diff --git a/ydb/core/kqp/ut/kqp_indexes_multishard_ut.cpp b/ydb/core/kqp/ut/kqp_indexes_multishard_ut.cpp
index 60b3f4d1678..7bee2f2914d 100644
--- a/ydb/core/kqp/ut/kqp_indexes_multishard_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_indexes_multishard_ut.cpp
@@ -32,7 +32,10 @@ void CreateTableWithMultishardIndex(Tests::TClient& client) {
         Columns { Name: "key"    Type: "Uint64" }
         Columns { Name: "fk"    Type: "Uint32" }
         Columns { Name: "value"  Type: "Utf8" }
-        KeyColumnNames: ["key"])";
+        KeyColumnNames: ["key"]
+        SplitBoundary { KeyPrefix { Tuple { Optional { Uint64: 3 } } } }
+        SplitBoundary { KeyPrefix { Tuple { Optional { Uint64: 100 } } } }
+    )";
 
     NKikimrSchemeOp::TTableDescription desc;
     bool parseOk = ::google::protobuf::TextFormat::ParseFromString(scheme, &desc);
@@ -48,7 +51,10 @@ void CreateTableWithMultishardIndexAndDataColumn(Tests::TClient& client) {
         Columns { Name: "fk"    Type: "Uint32" }
         Columns { Name: "value"  Type: "Utf8" }
         Columns { Name: "ext_value"  Type: "Utf8" }
-        KeyColumnNames: ["key"])";
+        KeyColumnNames: ["key"]
+        SplitBoundary { KeyPrefix { Tuple { Optional { Uint64: 3 } } } }
+        SplitBoundary { KeyPrefix { Tuple { Optional { Uint64: 100 } } } }
+    )";
 
     NKikimrSchemeOp::TTableDescription desc;
     bool parseOk = ::google::protobuf::TextFormat::ParseFromString(scheme, &desc);
@@ -1030,6 +1036,65 @@ Y_UNIT_TEST_SUITE(KqpMultishardIndex) {
             UNIT_ASSERT_VALUES_EQUAL(yson, expected);
         }
     }
+
+    Y_UNIT_TEST_QUAD(SortByPk, WithMvcc, UseNewEngine) {
+        auto serverSettings = TKikimrSettings()
+            .SetEnableMvcc(WithMvcc)
+            .SetEnableMvccSnapshotReads(WithMvcc);
+        TKikimrRunner kikimr(serverSettings);
+
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+        CreateTableWithMultishardIndex(kikimr.GetTestClient());
+        FillTable<UseNewEngine>(session);
+
+        AssertSuccessResult(session.ExecuteDataQuery(Q1_(R"(
+            UPSERT INTO `/Root/MultiShardIndexed` (key, fk, value) VALUES
+                (10u, 1000, "NewValue1"),
+                (11u, 1001, "NewValue2"),
+                (12u, 1002, "NewValue3"),
+                (13u, 1003, "NewValue4"),
+                (14u, 1004, "NewValue5"),
+                (15u, 1005, "NewValue6"),
+                (101u, 1011, "NewValue7");
+        )"), TTxControl::BeginTx().CommitTx()).ExtractValueSync());
+
+        auto query = Q1_(R"(
+            SELECT * FROM MultiShardIndexed VIEW index
+            WHERE fk > 100
+            ORDER BY fk, key
+            LIMIT 100;
+        )");
+
+        auto explainResult = session.ExplainDataQuery(query).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(explainResult.GetStatus(), EStatus::SUCCESS, explainResult.GetIssues().ToString());
+
+        // Cerr << explainResult.GetPlan() << Endl;
+
+        if (UseNewEngine) {
+            NJson::TJsonValue plan;
+            NJson::ReadJsonTree(explainResult.GetPlan(), &plan, true);
+            auto node = FindPlanNodeByKv(plan, "Name", "TopSort");
+            UNIT_ASSERT(node.IsDefined());
+        }
+
+        auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        CompareYson(R"([
+            [[1000u];[10u];["NewValue1"]];
+            [[1001u];[11u];["NewValue2"]];
+            [[1002u];[12u];["NewValue3"]];
+            [[1003u];[13u];["NewValue4"]];
+            [[1004u];[14u];["NewValue5"]];
+            [[1005u];[15u];["NewValue6"]];
+            [[1011u];[101u];["NewValue7"]];
+            [[1000000000u];[1u];["v1"]];
+            [[2000000000u];[2u];["v2"]];
+            [[3000000000u];[3u];["v3"]];
+            [[4294967295u];[4u];["v4"]]
+        ])", FormatResultSetYson(result.GetResultSet(0)));
+    }
 }
 
 }
diff --git a/ydb/core/kqp/ut/kqp_indexes_ut.cpp b/ydb/core/kqp/ut/kqp_indexes_ut.cpp
index 61327295c7b..c7d9d8e37ae 100644
--- a/ydb/core/kqp/ut/kqp_indexes_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_indexes_ut.cpp
@@ -3115,7 +3115,8 @@ Y_UNIT_TEST_SUITE(KqpIndexes) {
 
         {
             const TString query(Q1_(R"(
-                SELECT * FROM `/Root/TestTable` VIEW ix_cust as t WHERE t.customer = "Vasya" ORDER BY t.customer DESC;
+                SELECT * FROM `/Root/TestTable` VIEW ix_cust as t WHERE t.customer = "Vasya"
+                ORDER BY t.customer DESC, t.id DESC;
             )"));
 
            {