Skip to content

Fix returning if input is a parameter #28414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Nov 7, 2025
Merged

Fix returning if input is a parameter #28414

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0924d7e
implement the regression tests
gridnevvvit Nov 7, 2025
4032e74
implement a fix and fix the test
gridnevvvit Nov 7, 2025
ee8c231
add more tests
gridnevvvit Nov 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions ydb/core/kqp/opt/physical/effects/kqp_opt_phy_returning.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,7 @@ TExprBase KqpBuildReturning(TExprBase node, TExprContext& ctx, const TTypeAnnota
return TExprBase(ctx.ChangeChild(*returning.Raw(), TKqlReturningList::idx_Update, inputExpr.Ptr()));
};


if (auto maybeList = returning.Update().Maybe<TExprList>()) {
for (auto item : maybeList.Cast()) {
if (auto upsert = item.Maybe<TKqlUpsertRows>()) {
Expand Down Expand Up @@ -224,7 +225,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>() || upsert.Input().Maybe<TCoParameter>()) {
return node;
}

Expand All @@ -247,7 +248,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>() || del.Input().Maybe<TCoParameter>()) {
return node;
}

Expand Down
304 changes: 304 additions & 0 deletions ydb/core/kqp/ut/opt/kqp_returning_ut.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -377,6 +377,25 @@ TString ExecuteReturningQuery(TKikimrRunner& kikimr, bool queryService, TString
return FormatResultSetYson(result.GetResultSet(0));
}

TString ExecuteReturningQueryWithParams(TKikimrRunner& kikimr, bool queryService, TString query, const TParams& params) {
if (queryService) {
auto qdb = kikimr.GetQueryClient();
auto qSession = qdb.GetSession().GetValueSync().GetSession();
auto settings = NYdb::NQuery::TExecuteQuerySettings()
.Syntax(NYdb::NQuery::ESyntax::YqlV1);
auto result = qSession.ExecuteQuery(
query, NYdb::NQuery::TTxControl::BeginTx().CommitTx(), params, settings).ExtractValueSync();
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
return FormatResultSetYson(result.GetResultSet(0));
}

auto db = kikimr.GetTableClient();
auto session = db.CreateSession().GetValueSync().GetSession();
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx(), params).ExtractValueSync();
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
return FormatResultSetYson(result.GetResultSet(0));
}

Y_UNIT_TEST_TWIN(ReturningWorks, QueryService) {
auto kikimr = DefaultKikimrRunner();
auto db = kikimr.GetTableClient();
Expand Down Expand Up @@ -663,6 +682,291 @@ Y_UNIT_TEST(ReturningTypes) {
}
}

Y_UNIT_TEST_TWIN(ReturningUpsertAsTableListNotNullOnly, QueryService) {
// Test for issue #27021: Query fails when using RETURNING CLAUSE with UPSERT
// to table with only NOT NULL fields and query parameters of type List
auto kikimr = DefaultKikimrRunner();

auto client = kikimr.GetTableClient();
auto session = client.CreateSession().GetValueSync().GetSession();

// Create table with only NOT NULL fields
const auto queryCreate = Q_(R"(
--!syntax_v1
CREATE TABLE test_table (
id Uint64 NOT NULL,
value Utf8 NOT NULL,
PRIMARY KEY (id)
);
)");

auto resultCreate = session.ExecuteSchemeQuery(queryCreate).GetValueSync();
UNIT_ASSERT_C(resultCreate.IsSuccess(), resultCreate.GetIssues().ToString());

{
// Test case from issue #27021
const auto query = Q_(R"(
--!syntax_v1
DECLARE $data AS List<Struct<id: UInt64, value: Utf8>>;

UPSERT INTO test_table
SELECT * FROM AS_TABLE($data)
RETURNING *;
)");

auto paramsBuilder = TParamsBuilder();
auto& dataParam = paramsBuilder.AddParam("$data");

dataParam.BeginList();
dataParam.AddListItem()
.BeginStruct()
.AddMember("id")
.Uint64(1)
.AddMember("value")
.Utf8("test1")
.EndStruct();
dataParam.AddListItem()
.BeginStruct()
.AddMember("id")
.Uint64(2)
.AddMember("value")
.Utf8("test2")
.EndStruct();
dataParam.EndList();
dataParam.Build();

auto params = paramsBuilder.Build();

// This should succeed, but currently fails with infinite loop error
Copy link

Copilot AI Nov 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comment is misleading. Since this PR adds a fix for the issue (by handling TCoParameter in the optimization code), the test should now succeed. Consider updating the comment to reflect that this test validates the fix works correctly, e.g., "This should succeed with the fix for issue #27021".

Suggested change
// This should succeed, but currently fails with infinite loop error
// This should succeed with the fix for issue #27021

Copilot uses AI. Check for mistakes.
CompareYson(R"([[1u;"test1"];[2u;"test2"]])",
ExecuteReturningQueryWithParams(kikimr, QueryService, query, params));
}

{
// Test with explicit field names in SELECT clause
const auto query = Q_(R"(
--!syntax_v1
DECLARE $data AS List<Struct<id: UInt64, value: Utf8>>;

UPSERT INTO test_table
SELECT id, value FROM AS_TABLE($data)
RETURNING *;
)");

auto paramsBuilder = TParamsBuilder();
auto& dataParam = paramsBuilder.AddParam("$data");

dataParam.BeginList();
dataParam.AddListItem()
.BeginStruct()
.AddMember("id")
.Uint64(3)
.AddMember("value")
.Utf8("test3")
.EndStruct();
dataParam.EndList();
dataParam.Build();

auto params = paramsBuilder.Build();

CompareYson(R"([[3u;"test3"]])",
ExecuteReturningQueryWithParams(kikimr, QueryService, query, params));
}

{
// Test DELETE with RETURNING using AS_TABLE with List parameter (same issue #27021)
// First insert some data without RETURNING
const auto insertQuery = Q_(R"(
--!syntax_v1
DECLARE $data AS List<Struct<id: UInt64, value: Utf8>>;

UPSERT INTO test_table
SELECT * FROM AS_TABLE($data);
)");

auto insertParamsBuilder = TParamsBuilder();
auto& insertDataParam = insertParamsBuilder.AddParam("$data");

insertDataParam.BeginList();
insertDataParam.AddListItem()
.BeginStruct()
.AddMember("id")
.Uint64(10)
.AddMember("value")
.Utf8("delete1")
.EndStruct();
insertDataParam.EndList();
insertDataParam.Build();

auto insertParams = insertParamsBuilder.Build();

// Execute insert without RETURNING
if (QueryService) {
auto qdb = kikimr.GetQueryClient();
auto qSession = qdb.GetSession().GetValueSync().GetSession();
auto settings = NYdb::NQuery::TExecuteQuerySettings()
.Syntax(NYdb::NQuery::ESyntax::YqlV1);
auto insertResult = qSession.ExecuteQuery(
insertQuery, NYdb::NQuery::TTxControl::BeginTx().CommitTx(), insertParams, settings).ExtractValueSync();
UNIT_ASSERT_VALUES_EQUAL_C(insertResult.GetStatus(), EStatus::SUCCESS, insertResult.GetIssues().ToString());
} else {
auto db = kikimr.GetTableClient();
auto session = db.CreateSession().GetValueSync().GetSession();
auto insertResult = session.ExecuteDataQuery(insertQuery, TTxControl::BeginTx().CommitTx(), insertParams).ExtractValueSync();
UNIT_ASSERT_VALUES_EQUAL_C(insertResult.GetStatus(), EStatus::SUCCESS, insertResult.GetIssues().ToString());
}
Comment on lines +804 to +817
Copy link

Copilot AI Nov 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This code block is duplicated in both test functions (lines 804-817 and 926-939). Consider extracting this logic into a helper function similar to ExecuteReturningQueryWithParams to reduce code duplication and improve maintainability.

Copilot uses AI. Check for mistakes.

// Now DELETE with RETURNING using AS_TABLE
const auto deleteQuery = Q_(R"(
--!syntax_v1
DECLARE $data AS List<Struct<id: UInt64>>;

DELETE FROM test_table ON
SELECT * FROM AS_TABLE($data)
RETURNING *;
)");

auto deleteParamsBuilder = TParamsBuilder();
auto& deleteDataParam = deleteParamsBuilder.AddParam("$data");

deleteDataParam.BeginList();
deleteDataParam.AddListItem()
.BeginStruct()
.AddMember("id")
.Uint64(10)
.EndStruct();
deleteDataParam.EndList();
deleteDataParam.Build();

auto deleteParams = deleteParamsBuilder.Build();

// This should succeed, but currently fails with infinite loop error
Copy link

Copilot AI Nov 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comment is misleading. Since this PR adds a fix for the issue (by handling TCoParameter in the optimization code), the test should now succeed. Consider updating the comment to reflect that this test validates the fix works correctly, e.g., "This should succeed with the fix for issue #27021".

Suggested change
// This should succeed, but currently fails with infinite loop error
// This should succeed with the fix for issue #27021

Copilot uses AI. Check for mistakes.
CompareYson(R"([[10u;"delete1"]])",
ExecuteReturningQueryWithParams(kikimr, QueryService, deleteQuery, deleteParams));
}
}

Y_UNIT_TEST_TWIN(ReturningUpsertAsTableListWithNullable, QueryService) {
// Test that nullable columns work correctly (this should work even before the fix)
auto kikimr = DefaultKikimrRunner();

auto client = kikimr.GetTableClient();
auto session = client.CreateSession().GetValueSync().GetSession();

const auto queryCreate = Q_(R"(
--!syntax_v1
CREATE TABLE test_table_nullable (
id Uint64 NOT NULL,
value Utf8,
PRIMARY KEY (id)
);
)");

auto resultCreate = session.ExecuteSchemeQuery(queryCreate).GetValueSync();
UNIT_ASSERT_C(resultCreate.IsSuccess(), resultCreate.GetIssues().ToString());

{
const auto query = Q_(R"(
--!syntax_v1
DECLARE $data AS List<Struct<id: UInt64, value: Utf8?>>;

UPSERT INTO test_table_nullable
SELECT * FROM AS_TABLE($data)
RETURNING *;
)");

auto paramsBuilder = TParamsBuilder();
auto& dataParam = paramsBuilder.AddParam("$data");

dataParam.BeginList();
dataParam.AddListItem()
.BeginStruct()
.AddMember("id")
.Uint64(1)
.AddMember("value")
.OptionalUtf8("test1")
.EndStruct();
dataParam.EndList();
dataParam.Build();

auto params = paramsBuilder.Build();

CompareYson(R"([[1u;["test1"]]])",
ExecuteReturningQueryWithParams(kikimr, QueryService, query, params));
}

{
// Test DELETE with RETURNING using AS_TABLE with List parameter (with nullable columns)
// First insert some data without RETURNING
const auto insertQuery = Q_(R"(
--!syntax_v1
DECLARE $data AS List<Struct<id: UInt64, value: Utf8?>>;

UPSERT INTO test_table_nullable
SELECT * FROM AS_TABLE($data);
)");

auto insertParamsBuilder = TParamsBuilder();
auto& insertDataParam = insertParamsBuilder.AddParam("$data");

insertDataParam.BeginList();
insertDataParam.AddListItem()
.BeginStruct()
.AddMember("id")
.Uint64(20)
.AddMember("value")
.OptionalUtf8("delete_nullable1")
.EndStruct();
insertDataParam.EndList();
insertDataParam.Build();

auto insertParams = insertParamsBuilder.Build();

// Execute insert without RETURNING
if (QueryService) {
auto qdb = kikimr.GetQueryClient();
auto qSession = qdb.GetSession().GetValueSync().GetSession();
auto settings = NYdb::NQuery::TExecuteQuerySettings()
.Syntax(NYdb::NQuery::ESyntax::YqlV1);
auto insertResult = qSession.ExecuteQuery(
insertQuery, NYdb::NQuery::TTxControl::BeginTx().CommitTx(), insertParams, settings).ExtractValueSync();
UNIT_ASSERT_VALUES_EQUAL_C(insertResult.GetStatus(), EStatus::SUCCESS, insertResult.GetIssues().ToString());
} else {
auto db = kikimr.GetTableClient();
auto session = db.CreateSession().GetValueSync().GetSession();
auto insertResult = session.ExecuteDataQuery(insertQuery, TTxControl::BeginTx().CommitTx(), insertParams).ExtractValueSync();
UNIT_ASSERT_VALUES_EQUAL_C(insertResult.GetStatus(), EStatus::SUCCESS, insertResult.GetIssues().ToString());
}
Comment on lines +926 to +939
Copy link

Copilot AI Nov 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This code block duplicates the logic from lines 804-817. Consider extracting this logic into a helper function similar to ExecuteReturningQueryWithParams to reduce code duplication and improve maintainability.

Copilot uses AI. Check for mistakes.

// Now DELETE with RETURNING using AS_TABLE
const auto deleteQuery = Q_(R"(
--!syntax_v1
DECLARE $data AS List<Struct<id: UInt64>>;

DELETE FROM test_table_nullable ON
SELECT * FROM AS_TABLE($data)
RETURNING *;
)");

auto deleteParamsBuilder = TParamsBuilder();
auto& deleteDataParam = deleteParamsBuilder.AddParam("$data");

deleteDataParam.BeginList();
deleteDataParam.AddListItem()
.BeginStruct()
.AddMember("id")
.Uint64(20)
.EndStruct();
deleteDataParam.EndList();
deleteDataParam.Build();

auto deleteParams = deleteParamsBuilder.Build();

CompareYson(R"([[20u;["delete_nullable1"]]])",
ExecuteReturningQueryWithParams(kikimr, QueryService, deleteQuery, deleteParams));
}
}

}

} // namespace NKikimr::NKqp
Loading