KIKIMR-14379: drop redundant Take over KqlLookupTable

ref:ff25109eccd63131cd1ee2e967d438908d98596f

5 files changed, 54 insertions, 1 deletions

diff --git a/ydb/core/kqp/executer/kqp_data_executer.cpp b/ydb/core/kqp/executer/kqp_data_executer.cpp
index e668d90684..134a324507 100644
--- a/ydb/core/kqp/executer/kqp_data_executer.cpp
+++ b/ydb/core/kqp/executer/kqp_data_executer.cpp
@@ -1569,7 +1569,10 @@ private:
                         auto* protoChannel = protoOutput->MutableChannels(outputChannelIndex);
 
                         ui64 dstTaskId = TasksGraph.GetChannel(outputChannelId).DstTask;
-                        YQL_ENSURE(dstTaskId); // result channel not allowed here
+
+                        if (dstTaskId == 0) {
+                            continue;
+                        }
 
                         auto& dstTask = TasksGraph.GetTask(dstTaskId);
                         if (dstTask.ComputeActorId) {
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
index 793b8a9343..4a2d137926 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
@@ -44,6 +44,7 @@ public:
         AddHandler(0, &TCoFlatMapBase::Match, HNDL(RewriteFlatMapOverExtend));
         AddHandler(0, &TKqlDeleteRows::Match, HNDL(DeleteOverLookup));
         AddHandler(0, &TKqlUpsertRowsBase::Match, HNDL(ExcessUpsertInputColumns));
+        AddHandler(0, &TCoTake::Match, HNDL(DropTakeOverLookupTable));
 
         AddHandler(1, &TKqlReadTableIndex::Match, HNDL(RewriteIndexRead));
         AddHandler(1, &TKqlLookupIndex::Match, HNDL(RewriteLookupIndex));
@@ -172,6 +173,12 @@ protected:
         return output;
     }
 
+    TMaybeNode<TExprBase> DropTakeOverLookupTable(TExprBase node, TExprContext& ctx) {
+        TExprBase output = KqpDropTakeOverLookupTable(node, ctx, KqpCtx);
+        DumpAppliedRule("DropTakeOverLookupTable", node.Ptr(), output.Ptr(), ctx);
+        return output;
+    }
+
 private:
     TTypeAnnotationContext& TypesCtx;
     const TKqpOptimizeContext& KqpCtx;
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_ranges.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_ranges.cpp
index 38933b19fb..acf6a707be 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_ranges.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_ranges.cpp
@@ -293,5 +293,44 @@ TExprBase KqpPushPredicateToReadTable(TExprBase node, TExprContext& ctx, const T
         .Done();
 }
 
+TExprBase KqpDropTakeOverLookupTable(const TExprBase& node, TExprContext&, const TKqpOptimizeContext& kqpCtx) {
+    if (!node.Maybe<TCoTake>().Input().Maybe<TKqlLookupTableBase>()) {
+        return node;
+    }
+
+    auto take = node.Cast<TCoTake>();
+    auto lookupTable = take.Input().Cast<TKqlLookupTableBase>();
+
+    if (!take.Count().Maybe<TCoUint64>()) {
+        return node;
+    }
+
+    const ui64 count = FromString<ui64>(take.Count().Cast<TCoUint64>().Literal().Value());
+    YQL_ENSURE(count > 0);
+
+    auto maybeAsList = lookupTable.LookupKeys().Maybe<TCoAsList>();
+    if (!maybeAsList) {
+        maybeAsList = lookupTable.LookupKeys().Maybe<TCoIterator>().List().Maybe<TCoAsList>();
+    }
+
+    if (!maybeAsList) {
+        return node;
+    }
+
+    if (maybeAsList.Cast().ArgCount() > count) {
+        return node;
+    }
+
+    const auto tablePath = lookupTable.Table().Path().Value();
+    const auto& table = kqpCtx.Tables->ExistingTable(kqpCtx.Cluster, tablePath);
+
+    const auto& lookupKeys = GetSeqItemType(lookupTable.LookupKeys().Ref().GetTypeAnn())->Cast<TStructExprType>()->GetItems();
+    if (table.Metadata->KeyColumnNames.size() != lookupKeys.size()) {
+        return node;
+    }
+
+    return lookupTable;
+}
+
 } // namespace NKikimr::NKqp::NOpt
 
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_rules.h b/ydb/core/kqp/opt/logical/kqp_opt_log_rules.h
index 75420f5cfe..ae001c4d9c 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_rules.h
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_rules.h
@@ -50,4 +50,7 @@ NYql::NNodes::TExprBase KqpDeleteOverLookup(const NYql::NNodes::TExprBase& node,
 
 NYql::NNodes::TExprBase KqpExcessUpsertInputColumns(const NYql::NNodes::TExprBase& node, NYql::TExprContext& ctx);
 
+NYql::NNodes::TExprBase KqpDropTakeOverLookupTable(const NYql::NNodes::TExprBase& node, NYql::TExprContext& ctx,
+    const TKqpOptimizeContext& kqpCtx);
+
 } // namespace NKikimr::NKqp::NOpt
diff --git a/ydb/core/kqp/ut/kqp_newengine_ut.cpp b/ydb/core/kqp/ut/kqp_newengine_ut.cpp
index 9e2e136ccc..4e99bea83d 100644
--- a/ydb/core/kqp/ut/kqp_newengine_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_newengine_ut.cpp
@@ -66,6 +66,7 @@ Y_UNIT_TEST_SUITE(KqpNewEngine) {
         auto explainResult = session.ExplainDataQuery(query).GetValueSync();
         UNIT_ASSERT_VALUES_EQUAL_C(explainResult.GetStatus(), EStatus::SUCCESS, explainResult.GetIssues().ToString());
         UNIT_ASSERT_C(explainResult.GetAst().Contains("KqpLookupTable"), explainResult.GetAst());
+        UNIT_ASSERT_C(!explainResult.GetAst().Contains("Take"), explainResult.GetAst());
 
         auto params = kikimr.GetTableClient().GetParamsBuilder()
             .AddParam("$key").Uint64(302).Build()