From 6e0d92edfb6c3688e93e9ff6e6ccbad4cd785da3 Mon Sep 17 00:00:00 2001
From: Mikhail Surin <ssmike@ydb.tech>
Date: Tue, 8 Oct 2024 22:45:42 +0300
Subject: [PATCH 1/2] Force precompute for returning effect input

---
 .../kqp/opt/physical/effects/kqp_opt_phy_returning.cpp |  4 ++--
 ydb/library/yql/dq/opt/dq_opt_phy.cpp                  | 10 ++++++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_returning.cpp b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_returning.cpp
index d4e97c698697..4ae0aaf7e3b4 100644
--- a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_returning.cpp
+++ b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_returning.cpp
@@ -204,7 +204,7 @@ TExprBase KqpRewriteReturningUpsert(TExprBase node, TExprContext& ctx, const TKq
         return node;
     }
 
-    if (!upsert.Input().Maybe<TDqPrecompute>() && !upsert.Input().Maybe<TDqPhyPrecompute>()) {
+    if (upsert.Input().Maybe<TDqPrecompute>() || upsert.Input().Maybe<TDqPhyPrecompute>()) {
         return node;
     }
 
@@ -226,7 +226,7 @@ TExprBase KqpRewriteReturningDelete(TExprBase node, TExprContext& ctx, const TKq
         return node;
     }
 
-    if (!del.Input().Maybe<TDqPrecompute>() && !del.Input().Maybe<TDqPhyPrecompute>()) {
+    if (del.Input().Maybe<TDqPrecompute>() || del.Input().Maybe<TDqPhyPrecompute>()) {
         return node;
     }
 
diff --git a/ydb/library/yql/dq/opt/dq_opt_phy.cpp b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
index 8ff1b8154de9..fb31fd8bebfb 100644
--- a/ydb/library/yql/dq/opt/dq_opt_phy.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
@@ -2733,6 +2733,16 @@ TExprBase DqPropagatePrecomuteTake(TExprBase node, TExprContext& ctx, IOptimizat
         return node;
     }
 
+    auto* typeAnn = precompute.Connection().Raw()->GetTypeAnn();
+    if (!typeAnn || typeAnn->GetKind() != ETypeAnnotationKind::Stream) {
+        return node;
+    }
+
+    typeAnn = typeAnn->Cast<TStreamExprType>()->GetItemType();
+    if (typeAnn->GetKind() != ETypeAnnotationKind::List && typeAnn->GetKind() != ETypeAnnotationKind::Optional) {
+        return node;
+    }
+
     auto takeLambda = Build<TCoLambda>(ctx, node.Pos())
         .Args({"list_stream"})
         .Body<TCoMap>()

From a83c9d68e1925a9bc91ef2100a5f2e64a0decf06 Mon Sep 17 00:00:00 2001
From: Mikhail Surin <ssmike@ydb.tech>
Date: Mon, 21 Oct 2024 11:32:53 +0300
Subject: [PATCH 2/2] make guard more specific

---
 ydb/library/yql/dq/opt/dq_opt_phy.cpp | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/ydb/library/yql/dq/opt/dq_opt_phy.cpp b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
index fb31fd8bebfb..e2c89c6537be 100644
--- a/ydb/library/yql/dq/opt/dq_opt_phy.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
@@ -2734,12 +2734,10 @@ TExprBase DqPropagatePrecomuteTake(TExprBase node, TExprContext& ctx, IOptimizat
     }
 
     auto* typeAnn = precompute.Connection().Raw()->GetTypeAnn();
-    if (!typeAnn || typeAnn->GetKind() != ETypeAnnotationKind::Stream) {
-        return node;
-    }
 
-    typeAnn = typeAnn->Cast<TStreamExprType>()->GetItemType();
-    if (typeAnn->GetKind() != ETypeAnnotationKind::List && typeAnn->GetKind() != ETypeAnnotationKind::Optional) {
+    YQL_ENSURE(typeAnn && typeAnn->GetKind() == ETypeAnnotationKind::List);
+    typeAnn = typeAnn->Cast<TListExprType>()->GetItemType();
+    if (typeAnn->GetKind() == ETypeAnnotationKind::Struct) {
         return node;
     }