Fix KqpRewriteTopSortOverIndexRead KQP optimizer rule. (KIKIMR-15237)22.2.46

REVIEW: 2697057

REVIEW: 2697157
x-ydb-stable-ref: 5b63d35f55fa5bebbd00814d08a847018bfce1e1

3 files changed, 94 insertions, 5 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 4bb48475a7..955abe4e98 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 59b3b823c9..26daed51d7 100644
--- a/ydb/core/kqp/ut/kqp_indexes_multishard_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_indexes_multishard_ut.cpp
@@ -15,12 +15,20 @@ using namespace NYdb::NScripting;
 
 namespace {
 
+NYdb::NTable::TDataQueryResult ExecuteDataQuery(TSession& session, const TString& query) {
+    const auto txSettings = TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx();
+    return session.ExecuteDataQuery(query, txSettings).ExtractValueSync();
+}
+
 void CreateTableWithMultishardIndex(Tests::TClient& client) {
     const TString scheme =  R"(Name: "MultiShardIndexed"
         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);
@@ -30,6 +38,20 @@ void CreateTableWithMultishardIndex(Tests::TClient& client) {
     UNIT_ASSERT_VALUES_EQUAL(status, NMsgBusProxy::MSTATUS_OK);
 }
 
+template<bool UseNewEngine>
+void FillTable(NYdb::NTable::TSession& session) {
+    const TString query(Q_(R"(
+        UPSERT INTO `/Root/MultiShardIndexed` (key, fk, value) VALUES
+        (1, 1000000000, "v1"),
+        (2, 2000000000, "v2"),
+        (3, 3000000000, "v3"),
+        (4, 4294967295, "v4");
+    )"));
+
+    auto result = ExecuteDataQuery(session, query);
+    UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+}
+
 }
 
 Y_UNIT_TEST_SUITE(KqpMultishardIndex) {
@@ -190,6 +212,65 @@ Y_UNIT_TEST_SUITE(KqpMultishardIndex) {
         }
 
     }
+
+    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 39be9e8e0d..c3da54e201 100644
--- a/ydb/core/kqp/ut/kqp_indexes_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_indexes_ut.cpp
@@ -3074,7 +3074,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;
             )"));
 
            {