Ensure returning is in the same block (#6980)

Author: ssmike

ydb/core/kqp/provider/yql_kikimr_opt_build.cpp

@@ -247,14 +247,18 @@ struct TKiExploreTxResults {
         }
     }
 
-    void AddResult(const TExprBase& result) {
+    void PrepareForResult() {
         if (QueryBlocks.empty()) {
             AddQueryBlock();
         }
 
         if (!ConcurrentResults && QueryBlocks.back().Results.size() > 0) {
             AddQueryBlock();
         }
+    }
+
+    void AddResult(const TExprBase& result) {
+        PrepareForResult();
 
         auto& curBlock = QueryBlocks.back();
         curBlock.Results.push_back(result);
@@ -422,6 +426,10 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
         const auto& tableData = tablesData->ExistingTable(cluster, table);
         YQL_ENSURE(tableData.Metadata);
 
+        if (!write.ReturningColumns().Empty()) {
+            txRes.PrepareForResult();
+        }
+
         if (tableOp == TYdbOperation::UpdateOn) {
             auto inputColumnsSetting = GetSetting(write.Settings().Ref(), "input_columns");
             YQL_ENSURE(inputColumnsSetting);
@@ -482,6 +490,11 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
         for (const auto& item : updateStructType->GetItems()) {
             updateColumns.emplace(item->GetName());
         }
+
+        if (!update.ReturningColumns().Empty()) {
+            txRes.PrepareForResult();
+        }
+
         txRes.AddUpdateOpToQueryBlock(node, tableData.Metadata, updateColumns);
         if (!update.ReturningColumns().Empty()) {
             txRes.AddResult(
@@ -517,6 +530,10 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
         YQL_ENSURE(tablesData);
         const auto& tableData = tablesData->ExistingTable(cluster, table);
         YQL_ENSURE(tableData.Metadata);
+        if (!del.ReturningColumns().Empty()) {
+            txRes.PrepareForResult();
+        }
+
         txRes.AddWriteOpToQueryBlock(node, tableData.Metadata, tableOp & KikimrReadOps());
         if (!del.ReturningColumns().Empty()) {
             txRes.AddResult(

ydb/core/kqp/ut/opt/kqp_returning_ut.cpp

@@ -161,7 +161,8 @@ Y_UNIT_TEST(ReturningColumnsOrder) {
     }
 
     auto settings = NYdb::NQuery::TExecuteQuerySettings()
-        .Syntax(NYdb::NQuery::ESyntax::YqlV1);
+        .Syntax(NYdb::NQuery::ESyntax::YqlV1)
+        .ConcurrentResultSets(false);
     {
         auto result = db.ExecuteQuery(R"(
             UPSERT INTO test1 (id, v) VALUES (1, '321') RETURNING id, v;
@@ -171,6 +172,14 @@ Y_UNIT_TEST(ReturningColumnsOrder) {
         CompareYson(R"([[[1];["321"]]])", FormatResultSetYson(result.GetResultSet(0)));
         CompareYson(R"([[["111"];[1]]])", FormatResultSetYson(result.GetResultSet(1)));
     }
+    {
+        auto it = db.StreamExecuteQuery(R"(
+            UPSERT INTO test1 (id, v) VALUES (2, '321') RETURNING id, v;
+            REPLACE INTO test1 (id, v) VALUES (2, '111') RETURNING v, id;
+        )", NYdb::NQuery::TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(it.GetStatus(), EStatus::SUCCESS, it.GetIssues().ToString());
+        Cerr << StreamResultToYson(it);
+    }
 
 }