Don't omit precomputes in presense of a returning statement (#12965) (#…

…13117)
ydb-platform · Dec 28, 2024 · 43a1fb0 · 43a1fb0
1 parent b30bd4d
commit 43a1fb0
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 4 deletions.
diff --git 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
@@ -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
@@ -290,6 +290,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();