From 2e9a5794658f94ab9cd399cffbb7a92d87c2e599 Mon Sep 17 00:00:00 2001
From: Mikhail Surin <ssmike@ydb.tech>
Date: Sat, 28 Dec 2024 14:29:02 +0300
Subject: [PATCH] Don't omit precomputes in presense of a returning statement
 (#12965)

---
 ydb/core/kqp/opt/kqp_opt_phy_finalize.cpp     | 24 ++++++++++++++++++
 .../effects/kqp_opt_phy_returning.cpp         |  8 +++---
 ydb/core/kqp/ut/opt/kqp_returning_ut.cpp      | 25 +++++++++++++++++++
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/ydb/core/kqp/opt/kqp_opt_phy_finalize.cpp b/ydb/core/kqp/opt/kqp_opt_phy_finalize.cpp
index 139b5eb61a19..76b744bcf437 100644
--- a/ydb/core/kqp/opt/kqp_opt_phy_finalize.cpp
+++ b/ydb/core/kqp/opt/kqp_opt_phy_finalize.cpp
@@ -69,11 +69,35 @@ TStatus KqpBuildPureExprStagesResult(const TExprNode::TPtr& input, TExprNode::TP
             break;
         }
     }
+
+    std::function<bool(TExprBase)> hasReturning = [&] (TExprBase node) {
+        if (auto effect = node.Maybe<TKqlUpsertRows>()) {
+            return !effect.ReturningColumns().Cast().Empty();
+        }
+        if (auto effect = node.Maybe<TKqlDeleteRows>()) {
+            return !effect.ReturningColumns().Cast().Empty();
+        }
+        if (auto effect = node.Maybe<TExprList>()) {
+            for (auto item : effect.Cast()) {
+                if (hasReturning(item)) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    };
+
     for (const auto& effect : query.Effects()) {
         if (!hasPrecomputes(effect)) {
             omitResultPrecomputes = false;
             break;
         }
+        // returning works by forcing materialization of modified rows via precompute
+        // so omitting precomputes here breaks returning logic
+        if (hasReturning(effect)) {
+            omitResultPrecomputes = false;
+            break;
+        }
     }
 
     TNodeOnNodeOwnedMap replaces;
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 4ae0aaf7e3b4..ea208ecf461d 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
@@ -156,9 +156,9 @@ TExprBase KqpBuildReturning(TExprBase node, TExprContext& ctx, const TKqpOptimiz
         auto inputExpr = Build<TCoExtractMembers>(ctx, pos)
             .Input(input.Cast())
             .Members(returning.Columns())
-            .Done().Ptr();
+            .Done();
 
-        return TExprBase(ctx.ChangeChild(*returning.Raw(), TKqlReturningList::idx_Update, std::move(inputExpr)));
+        return TExprBase(ctx.ChangeChild(*returning.Raw(), TKqlReturningList::idx_Update, inputExpr.Ptr()));
     };
 
     if (auto maybeList = returning.Update().Maybe<TExprList>()) {
@@ -216,7 +216,7 @@ TExprBase KqpRewriteReturningUpsert(TExprBase node, TExprContext& ctx, const TKq
             .Table(upsert.Table())
             .Columns(upsert.Columns())
             .Settings(upsert.Settings())
-            .ReturningColumns<TCoAtomList>().Build()
+            .ReturningColumns(upsert.ReturningColumns())
             .Done();
 }
 
@@ -236,7 +236,7 @@ TExprBase KqpRewriteReturningDelete(TExprBase node, TExprContext& ctx, const TKq
                 .Input(del.Input())
                 .Build()
             .Table(del.Table())
-            .ReturningColumns<TCoAtomList>().Build()
+            .ReturningColumns(del.ReturningColumns())
             .Done();
 }
 
diff --git a/ydb/core/kqp/ut/opt/kqp_returning_ut.cpp b/ydb/core/kqp/ut/opt/kqp_returning_ut.cpp
index 4b123726e4fc..6a37e7b12abf 100644
--- a/ydb/core/kqp/ut/opt/kqp_returning_ut.cpp
+++ b/ydb/core/kqp/ut/opt/kqp_returning_ut.cpp
@@ -301,6 +301,31 @@ Y_UNIT_TEST(ReturningColumnsOrder) {
     
 }
 
+Y_UNIT_TEST(Random) {
+    auto kikimr = DefaultKikimrRunner();
+
+    auto client = kikimr.GetQueryClient();
+    auto settings = NYdb::NQuery::TExecuteQuerySettings()
+        .Syntax(NYdb::NQuery::ESyntax::YqlV1)
+        .ConcurrentResultSets(false);
+
+    {
+        auto result = client.ExecuteQuery("CREATE TABLE example (key Uint64, value String, PRIMARY KEY (key));",
+            NYdb::NQuery::TTxControl::NoTx(), settings).ExtractValueSync();
+        UNIT_ASSERT(result.IsSuccess());
+    }
+
+    {
+        auto result = client.ExecuteQuery(
+            R"(
+            UPSERT INTO example (key, value) VALUES (1, CAST(RandomUuid(1) AS String)) RETURNING *;
+            SELECT * FROM example;
+            )",
+            NYdb::NQuery::TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
+        CompareYson(FormatResultSetYson(result.GetResultSet(0)), FormatResultSetYson(result.GetResultSet(1)));
+    }
+}
+
 Y_UNIT_TEST(ReturningTypes) {
     auto kikimr = DefaultKikimrRunner();