fixed keys partitioning

1 files changed, 12 insertions, 2 deletions

diff --git a/ydb/library/yql/minikql/comp_nodes/mkql_wide_top_sort.cpp b/ydb/library/yql/minikql/comp_nodes/mkql_wide_top_sort.cpp
index a7e5f4b0ca..c1b9876635 100644
--- a/ydb/library/yql/minikql/comp_nodes/mkql_wide_top_sort.cpp
+++ b/ydb/library/yql/minikql/comp_nodes/mkql_wide_top_sort.cpp
@@ -516,15 +516,25 @@ IComputationNode* WrapWideTopT(TCallable& callable, const TComputationNodeFactor
     const auto keyWidth = (inputsWithCount >> 1U) - 1U;
     const auto inputType = AS_TYPE(TTupleType, AS_TYPE(TFlowType, callable.GetType()->GetReturnType())->GetItemType());
     std::vector<ui32> indexes(inputType->GetElementsCount());
-    std::iota(indexes.begin(), indexes.end(), 0U);
 
     TKeyTypes keyTypes(keyWidth);
+    std::unordered_set<ui32> keyIndexes;
     for (auto i = 0U; i < keyTypes.size(); ++i) {
         const auto keyIndex = AS_VALUE(TDataLiteral, callable.GetInput(((i + 1U) << 1U) - offset))->AsValue().Get<ui32>();
-        std::swap(indexes[i], indexes[indexes[keyIndex]]);
+        indexes[i] = keyIndex;
+        keyIndexes.emplace(keyIndex);
         keyTypes[i].first = *UnpackOptionalData(inputType->GetElementType(keyIndex), keyTypes[i].second)->GetDataSlot();
     }
 
+    size_t payloadPos = keyTypes.size();
+    for (auto i = 0U; i < indexes.size(); ++i) {
+        if (keyIndexes.contains(i)) {
+            continue;
+        }
+
+        indexes[payloadPos++] = i;
+    }
+
     std::vector<EValueRepresentation> representations(inputType->GetElementsCount());
     for (auto i = 0U; i < representations.size(); ++i)
         representations[i] = GetValueRepresentation(inputType->GetElementType(indexes[i]));