diff --git a/ydb/core/kqp/opt/kqp_opt_phy_finalize.cpp b/ydb/core/kqp/opt/kqp_opt_phy_finalize.cpp index bfd9839cdf62..f5ab68444be9 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 hasReturning = [&] (TExprBase node) { + if (auto effect = node.Maybe()) { + return !effect.ReturningColumns().Cast().Empty(); + } + if (auto effect = node.Maybe()) { + return !effect.ReturningColumns().Cast().Empty(); + } + if (auto effect = node.Maybe()) { + 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(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()) { @@ -216,7 +216,7 @@ TExprBase KqpRewriteReturningUpsert(TExprBase node, TExprContext& ctx, const TKq .Table(upsert.Table()) .Columns(upsert.Columns()) .Settings(upsert.Settings()) - .ReturningColumns().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().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 b23d9d38ace9..a06e07faf04a 100644 --- 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();