little topsort/keeptop optimization

1 files changed, 10 insertions, 18 deletions

diff --git a/ydb/library/yql/minikql/mkql_program_builder.cpp b/ydb/library/yql/minikql/mkql_program_builder.cpp
index 501e708c0af..6eb2f14802b 100644
--- a/ydb/library/yql/minikql/mkql_program_builder.cpp
+++ b/ydb/library/yql/minikql/mkql_program_builder.cpp
@@ -1663,18 +1663,14 @@ TRuntimeNode TProgramBuilder::Sort(TRuntimeNode list, TRuntimeNode ascending, co
 
 TRuntimeNode TProgramBuilder::Top(TRuntimeNode flow, TRuntimeNode count, TRuntimeNode ascending, const TUnaryLambda& keyExtractor) {
     if (const auto flowType = flow.GetStaticType(); flowType->IsFlow() || flowType->IsStream()) {
-        const auto itemType = flowType->IsFlow() ? AS_TYPE(TFlowType, flowType)->GetItemType() : AS_TYPE(TStreamType, flowType)->GetItemType();
-        const auto finalKeyExtractor = [&](TRuntimeNode item) { return Nth(item, 1U); };
 
-        return Map(FlatMap(Condense1(flow,
-                [&](TRuntimeNode item) { return AsList(NewTuple({Pickle(item), keyExtractor(item)})); },
+        return FlatMap(Condense1(flow,
+                [&](TRuntimeNode item) { return AsList(item); },
                 [this](TRuntimeNode, TRuntimeNode) { return NewDataLiteral<bool>(false); },
-                [&](TRuntimeNode item, TRuntimeNode state) { return KeepTop(count, state, NewTuple({Pickle(item), keyExtractor(item)}), ascending, finalKeyExtractor); }
+                [&](TRuntimeNode item, TRuntimeNode state) { return KeepTop(count, state, item, ascending, keyExtractor); }
             ),
-            [&](TRuntimeNode list) { return Top(list, count, ascending, finalKeyExtractor); }
-        ), [&](TRuntimeNode item) {
-            return Unpickle(itemType, Nth(item, 0U));
-        });
+            [&](TRuntimeNode list) { return Top(list, count, ascending, keyExtractor); }
+        );
     }
 
     return BuildListNth(__func__, flow, count, ascending, keyExtractor);
@@ -1682,19 +1678,15 @@ TRuntimeNode TProgramBuilder::Top(TRuntimeNode flow, TRuntimeNode count, TRuntim
 
 TRuntimeNode TProgramBuilder::TopSort(TRuntimeNode flow, TRuntimeNode count, TRuntimeNode ascending, const TUnaryLambda& keyExtractor) {
     if (const auto flowType = flow.GetStaticType(); flowType->IsFlow() || flowType->IsStream()) {
-        const auto itemType = flowType->IsFlow() ? AS_TYPE(TFlowType, flowType)->GetItemType() : AS_TYPE(TStreamType, flowType)->GetItemType();
-        const auto finalKeyExtractor = [&](TRuntimeNode item) { return Nth(item, 1U); };
-        return Map(FlatMap(Condense1(flow,
-                [&](TRuntimeNode item) { return AsList(NewTuple({Pickle(item), keyExtractor(item)})); },
+        return FlatMap(Condense1(flow,
+                [&](TRuntimeNode item) { return AsList(item); },
                 [this](TRuntimeNode, TRuntimeNode) { return NewDataLiteral<bool>(false); },
                 [&](TRuntimeNode item, TRuntimeNode state) {
-                    return KeepTop(count, state, NewTuple({Pickle(item), keyExtractor(item)}), ascending, finalKeyExtractor);
+                    return KeepTop(count, state, item, ascending, keyExtractor);
                 }
             ),
-            [&](TRuntimeNode list) { return TopSort(list, count, ascending, finalKeyExtractor); }
-        ), [&](TRuntimeNode item) {
-            return Unpickle(itemType, Nth(item, 0U));
-        });
+            [&](TRuntimeNode list) { return TopSort(list, count, ascending, keyExtractor); }
+        );
     }
 
     if constexpr (RuntimeVersion >= 25U)