Fix TSwitchFlowWrapper implementation. (KIKIMR-15086)

ref:068418fef58265db0b307851062e15495442a620

2 files changed, 87 insertions, 1 deletions

diff --git a/ydb/core/kqp/ut/kqp_sort_ut.cpp b/ydb/core/kqp/ut/kqp_sort_ut.cpp
index 0b7c65a3d5..da3a4e4699 100644
--- a/ydb/core/kqp/ut/kqp_sort_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_sort_ut.cpp
@@ -1148,6 +1148,87 @@ Y_UNIT_TEST_SUITE(KqpSort) {
             UNIT_ASSERT(result.GetAst().Contains("(Sort (KiPartialTake (Filter"));
         }
     }
+
+    Y_UNIT_TEST_NEW_ENGINE(Offset) {
+        TKikimrRunner kikimr;
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        auto result = session.ExecuteDataQuery(Q1_(R"(
+            $data = SELECT * FROM EightShard WHERE Text = "Value1" LIMIT 7;
+
+            SELECT * FROM $data LIMIT 3 OFFSET 0;
+            SELECT * FROM $data LIMIT 3 OFFSET 3;
+            SELECT * FROM $data LIMIT 3 OFFSET 6;
+        )"), TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        UNIT_ASSERT_VALUES_EQUAL(result.GetResultSet(0).RowsCount(), 3);
+        UNIT_ASSERT_VALUES_EQUAL(result.GetResultSet(1).RowsCount(), 3);
+        UNIT_ASSERT_VALUES_EQUAL(result.GetResultSet(2).RowsCount(), 1);
+    }
+
+    Y_UNIT_TEST_NEW_ENGINE(OffsetPk) {
+        TKikimrRunner kikimr;
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        auto result = session.ExecuteDataQuery(Q1_(R"(
+            $data = SELECT * FROM EightShard WHERE Text = "Value1" ORDER BY Key LIMIT 7;
+
+            SELECT * FROM $data ORDER BY Key LIMIT 3 OFFSET 0;
+            SELECT * FROM $data ORDER BY Key LIMIT 3 OFFSET 3;
+            SELECT * FROM $data ORDER BY Key LIMIT 3 OFFSET 6;
+        )"), TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        CompareYson(R"([
+            [[1];[101u];["Value1"]];
+            [[2];[201u];["Value1"]];
+            [[3];[301u];["Value1"]]
+        ])", FormatResultSetYson(result.GetResultSet(0)));
+
+        CompareYson(R"([
+            [[1];[401u];["Value1"]];
+            [[2];[501u];["Value1"]];
+            [[3];[601u];["Value1"]]
+        ])", FormatResultSetYson(result.GetResultSet(1)));
+
+        CompareYson(R"([
+            [[1];[701u];["Value1"]]
+        ])", FormatResultSetYson(result.GetResultSet(2)));
+    }
+
+    Y_UNIT_TEST_NEW_ENGINE(OffsetTopSort) {
+        TKikimrRunner kikimr;
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        auto result = session.ExecuteDataQuery(Q1_(R"(
+            $data = SELECT * FROM EightShard WHERE Text = "Value1" ORDER BY Data LIMIT 7;
+
+            SELECT * FROM $data ORDER BY Data LIMIT 3 OFFSET 0;
+            SELECT * FROM $data ORDER BY Data LIMIT 3 OFFSET 3;
+            SELECT * FROM $data ORDER BY Data LIMIT 3 OFFSET 6;
+        )"), TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        CompareYson(R"([
+            [[1];[101u];["Value1"]];
+            [[1];[401u];["Value1"]];
+            [[1];[701u];["Value1"]]
+        ])", FormatResultSetYson(result.GetResultSet(0)));
+
+        CompareYson(R"([
+            [[2];[201u];["Value1"]];
+            [[2];[501u];["Value1"]];
+            [[2];[801u];["Value1"]]
+        ])", FormatResultSetYson(result.GetResultSet(1)));
+
+        CompareYson(R"([
+            [[3];[301u];["Value1"]]
+        ])", FormatResultSetYson(result.GetResultSet(2)));
+    }
 }
 
 } // namespace NKqp
diff --git a/ydb/library/yql/minikql/comp_nodes/mkql_switch.cpp b/ydb/library/yql/minikql/comp_nodes/mkql_switch.cpp
index 4571141d80..75dd396cd6 100644
--- a/ydb/library/yql/minikql/comp_nodes/mkql_switch.cpp
+++ b/ydb/library/yql/minikql/comp_nodes/mkql_switch.cpp
@@ -75,11 +75,15 @@ private:
             Buffer.Clear();
         }
 
+        void ResetPosition() {
+            Position = 0U;
+        }
+
         NUdf::TUnboxedValuePod Handler(ui32 index, const TSwitchHandler& handler, TComputationContext& ctx) {
             while (true) {
                 auto current = Get(Position++);
                 if (current.IsSpecial()) {
-                    Position = 0U;
+                    ResetPosition();
                     return current;
                 }
 
@@ -150,6 +154,7 @@ public:
             const auto& handler = Handlers[ptr->ChildReadIndex];
             auto childRes = handler.NewItem->GetValue(ctx);
             if (childRes.IsSpecial()) {
+                ptr->ResetPosition();
                 if (++ptr->ChildReadIndex == Handlers.size()) {
                     ptr->Clear();
                 }