Skip to content

Commit 3f8e741

Browse files
gridnevvvitgridnevvvit
authored
merge bugfixes 25-1 (#15966)
1 parent 433d592 commit 3f8e741

File tree

4 files changed

+300
-14
lines changed

4 files changed

+300
-14
lines changed

‎ydb/core/kqp/opt/physical/effects/kqp_opt_phy_returning.cpp

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,22 @@ TExprBase KqpBuildReturning(TExprBase node, TExprContext& ctx, const TKqpOptimiz
151151
.Index().Build("0")
152152
.Build()
153153
.Done();
154+
} else if (NDq::IsDqPureExpr(input.Cast())) {
155+
input = Build<TDqCnUnionAll>(ctx, pos)
156+
.Output()
157+
.Stage<TDqStage>()
158+
.Inputs().Build()
159+
.Program()
160+
.Args({})
161+
.Body<TCoToFlow>()
162+
.Input(input.Cast())
163+
.Build()
164+
.Build()
165+
.Settings().Build()
166+
.Build()
167+
.Index().Build("0")
168+
.Build()
169+
.Done();
154170
}
155171

156172
auto inputExpr = Build<TCoExtractMembers>(ctx, pos)
@@ -173,6 +189,10 @@ TExprBase KqpBuildReturning(TExprBase node, TExprContext& ctx, const TKqpOptimiz
173189
return buildReturningRows(del.Input().Cast(), MakeColumnsList(tableDesc.Metadata->KeyColumnNames, ctx, node.Pos()), returning.Columns());
174190
}
175191
}
192+
193+
if (item.Maybe<TKqlTableEffect>()) {
194+
return node;
195+
}
176196
}
177197
}
178198

@@ -183,6 +203,10 @@ TExprBase KqpBuildReturning(TExprBase node, TExprContext& ctx, const TKqpOptimiz
183203
return buildReturningRows(del.Input().Cast(), MakeColumnsList(tableDesc.Metadata->KeyColumnNames, ctx, node.Pos()), returning.Columns());
184204
}
185205

206+
if (returning.Update().Maybe<TKqlTableEffect>()) {
207+
return node;
208+
}
209+
186210
TExprNode::TPtr result = returning.Update().Ptr();
187211
auto status = TryConvertTo(result, *result->GetTypeAnn(), *returning.Raw()->GetTypeAnn(), ctx);
188212
YQL_ENSURE(status.Level != IGraphTransformer::TStatus::Error, "wrong returning expr type");

‎ydb/core/kqp/opt/physical/effects/kqp_opt_phy_upsert_index.cpp

Lines changed: 6 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -391,21 +391,13 @@ RewriteInputForConstraint(const TExprBase& inputRows, const THashSet<TStringBuf>
391391
const THashSet<TString> indexKeyColumns = CreateKeyColumnSetToRead(indexes);
392392
const THashSet<TString> indexDataColumns = CreateDataColumnSetToRead(indexes);
393393

394-
for (const auto& x : indexKeyColumns) {
395-
columns.push_back(Build<TCoAtom>(ctx, pos).Value(x).Done());
396-
}
394+
THashSet<TString> columnsToReadInPrecomputeLookupDict;
395+
columnsToReadInPrecomputeLookupDict.insert(indexKeyColumns.begin(), indexKeyColumns.end());
396+
columnsToReadInPrecomputeLookupDict.insert(indexDataColumns.begin(), indexDataColumns.end());
397+
columnsToReadInPrecomputeLookupDict.insert(mainPk.begin(), mainPk.end());
398+
columnsToReadInPrecomputeLookupDict.insert(checkDefaults.begin(), checkDefaults.end());
397399

398-
for (const auto& x : indexDataColumns) {
399-
// Handle the case of multiple indexes
400-
// one of them has 'foo' as data column but for another one foo is just indexed column
401-
if (indexKeyColumns.contains(x))
402-
continue;
403-
columns.push_back(Build<TCoAtom>(ctx, pos).Value(x).Done());
404-
}
405-
406-
for (const auto& x : mainPk) {
407-
if (indexKeyColumns.contains(x))
408-
continue;
400+
for (const auto& x : columnsToReadInPrecomputeLookupDict) {
409401
columns.push_back(Build<TCoAtom>(ctx, pos).Value(x).Done());
410402
}
411403

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

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,206 @@ Y_UNIT_TEST(ReturningSerial) {
314314
}
315315
}
316316

317+
TString ExecuteReturningQuery(TKikimrRunner& kikimr, bool queryService, TString query) {
318+
if (queryService) {
319+
auto qdb = kikimr.GetQueryClient();
320+
auto qSession = qdb.GetSession().GetValueSync().GetSession();
321+
auto result = qSession.ExecuteQuery(
322+
query, NYdb::NQuery::TTxControl::BeginTx().CommitTx()).ExtractValueSync();
323+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
324+
return FormatResultSetYson(result.GetResultSet(0));
325+
}
326+
327+
auto db = kikimr.GetTableClient();
328+
auto session = db.CreateSession().GetValueSync().GetSession();
329+
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).ExtractValueSync();
330+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
331+
return FormatResultSetYson(result.GetResultSet(0));
332+
}
333+
334+
Y_UNIT_TEST_TWIN(ReturningWorks, QueryService) {
335+
auto kikimr = DefaultKikimrRunner();
336+
auto db = kikimr.GetTableClient();
337+
auto session = db.CreateSession().GetValueSync().GetSession();
338+
CreateSampleTablesWithIndex(session, true);
339+
CompareYson(
340+
R"([[[101];[101];["Payload1"]];])",
341+
ExecuteReturningQuery(kikimr, QueryService, R"(
342+
UPSERT INTO `/Root/SecondaryKeys` (Key, Fk, Value) VALUES (101, 101, "Payload1") RETURNING *;
343+
)")
344+
);
345+
CompareYson(
346+
R"(
347+
[[#;#;["Payload8"]];
348+
[[1];[1];["Payload1"]];
349+
[[2];[2];["Payload2"]];
350+
[[5];[5];["Payload5"]];
351+
[#;[7];["Payload7"]];
352+
[[101];[101];["Payload1"]]
353+
])",
354+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
355+
);
356+
}
357+
358+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedUpsert, QueryService) {
359+
auto kikimr = DefaultKikimrRunner();
360+
auto db = kikimr.GetTableClient();
361+
auto session = db.CreateSession().GetValueSync().GetSession();
362+
CreateSampleTablesWithIndex(session, true);
363+
CompareYson(R"([
364+
[[110];[110];["Payload5"]];
365+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
366+
$v1 = (SELECT Key + 100 as Key, Fk + 100 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
367+
$v2 = (SELECT Key + 105 as Key, Fk + 105 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
368+
UPSERT INTO `/Root/SecondaryKeys`
369+
SELECT * FROM (SELECT * FROM $v1 UNION ALL SELECT * FROM $v2) WHERE Key > 107 RETURNING *;
370+
)"));
371+
CompareYson(
372+
R"(
373+
[[#;#;["Payload8"]];
374+
[[1];[1];["Payload1"]];
375+
[[2];[2];["Payload2"]];
376+
[[5];[5];["Payload5"]];
377+
[#;[7];["Payload7"]];
378+
[[110];[110];["Payload5"]]
379+
])",
380+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
381+
);
382+
}
383+
384+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedDelete, QueryService) {
385+
auto kikimr = DefaultKikimrRunner();
386+
auto db = kikimr.GetTableClient();
387+
auto session = db.CreateSession().GetValueSync().GetSession();
388+
CreateSampleTablesWithIndex(session, true);
389+
CompareYson(R"([
390+
[[5];[5];["Payload5"]];
391+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
392+
$v1 = (SELECT Key, Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL AND Key >= 1);
393+
$v2 = (SELECT Key, Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL AND Key <= 5);
394+
DELETE FROM `/Root/SecondaryKeys` ON
395+
SELECT * FROM (SELECT * FROM $v1 UNION ALL SELECT * FROM $v2) WHERE Key >= 5 RETURNING *;
396+
)"));
397+
CompareYson(
398+
R"(
399+
[[#;#;["Payload8"]];
400+
[[1];[1];["Payload1"]];
401+
[[2];[2];["Payload2"]];
402+
[#;[7];["Payload7"]];
403+
])",
404+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
405+
);
406+
}
407+
408+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedDeleteV2, QueryService) {
409+
auto kikimr = DefaultKikimrRunner();
410+
auto db = kikimr.GetTableClient();
411+
auto session = db.CreateSession().GetValueSync().GetSession();
412+
CreateSampleTablesWithIndex(session, true);
413+
CompareYson(R"([
414+
[[1];[1];["Payload1"]];
415+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
416+
DELETE FROM `/Root/SecondaryKeys` WHERE Key = 1 RETURNING *;
417+
)"));
418+
CompareYson(
419+
R"(
420+
[[#;#;["Payload8"]];
421+
[[2];[2];["Payload2"]];
422+
[[5];[5];["Payload5"]];
423+
[#;[7];["Payload7"]];
424+
])",
425+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
426+
);
427+
}
428+
429+
430+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedInsert, QueryService) {
431+
auto kikimr = DefaultKikimrRunner();
432+
auto db = kikimr.GetTableClient();
433+
auto session = db.CreateSession().GetValueSync().GetSession();
434+
CreateSampleTablesWithIndex(session, true);
435+
436+
CompareYson(R"([
437+
[[101];[101];["Payload1"]];
438+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
439+
$v1 = (SELECT Key + 100 as Key, Fk + 100 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
440+
$v2 = (SELECT Key + 205 as Key, Fk + 205 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
441+
INSERT INTO `/Root/SecondaryKeys`
442+
SELECT * FROM (SELECT * FROM $v1 UNION ALL SELECT * FROM $v2 ) WHERE Key < 102 RETURNING *;
443+
)"));
444+
445+
CompareYson(
446+
R"(
447+
[[#;#;["Payload8"]];
448+
[[1];[1];["Payload1"]];
449+
[[2];[2];["Payload2"]];
450+
[[5];[5];["Payload5"]];
451+
[#;[7];["Payload7"]];
452+
[[101];[101];["Payload1"]]
453+
])",
454+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
455+
);
456+
}
457+
458+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedReplace, QueryService) {
459+
auto kikimr = DefaultKikimrRunner();
460+
auto db = kikimr.GetTableClient();
461+
auto session = db.CreateSession().GetValueSync().GetSession();
462+
CreateSampleTablesWithIndex(session, true);
463+
464+
CompareYson(R"([
465+
[[101];[101];["Payload1"]];
466+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
467+
$v1 = (SELECT Key + 100 as Key, Fk + 100 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
468+
$v2 = (SELECT Key + 205 as Key, Fk + 205 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
469+
REPLACE INTO `/Root/SecondaryKeys`
470+
SELECT * FROM (SELECT * FROM $v1 UNION ALL SELECT * FROM $v2 ) WHERE Key < 102 RETURNING *;
471+
)"));
472+
473+
CompareYson(
474+
R"(
475+
[[#;#;["Payload8"]];
476+
[[1];[1];["Payload1"]];
477+
[[2];[2];["Payload2"]];
478+
[[5];[5];["Payload5"]];
479+
[#;[7];["Payload7"]];
480+
[[101];[101];["Payload1"]]
481+
])",
482+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
483+
);
484+
}
485+
486+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedOperationsWithDefault, QueryService) {
487+
auto kikimr = DefaultKikimrRunner();
488+
auto db = kikimr.GetTableClient();
489+
auto session = db.CreateSession().GetValueSync().GetSession();
490+
{
491+
auto res = session.ExecuteSchemeQuery(R"(
492+
--!syntax_v1
493+
CREATE TABLE `/Root/SecondaryKeys` (
494+
Key Serial,
495+
Fk Int32,
496+
Value String,
497+
PRIMARY KEY (Key),
498+
INDEX Index GLOBAL ON (Fk)
499+
);
500+
)").GetValueSync();
501+
}
502+
503+
CompareYson(R"([
504+
[1;[1];["Payload"]];
505+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
506+
REPLACE INTO `/Root/SecondaryKeys` (Fk, Value) VALUES (1, "Payload") RETURNING Key, Fk, Value;
507+
)"));
508+
509+
CompareYson(
510+
R"([
511+
[1;[1];["Payload"]];
512+
])",
513+
ExecuteReturningQuery(kikimr, QueryService, "SELECT Key, Fk, Value FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
514+
);
515+
}
516+
317517
Y_UNIT_TEST(ReturningColumnsOrder) {
318518
auto kikimr = DefaultKikimrRunner();
319519

‎ydb/core/kqp/ut/scheme/kqp_constraints_ut.cpp

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1290,6 +1290,76 @@ Y_UNIT_TEST_SUITE(KqpConstraints) {
12901290

12911291
}
12921292

1293+
Y_UNIT_TEST(DefaultAndIndexesTestDefaultColumnNotIncludedInIndex) {
1294+
NKikimrConfig::TAppConfig appConfig;
1295+
TKikimrRunner kikimr(TKikimrSettings().SetPQConfig(DefaultPQConfig()).SetAppConfig(appConfig));
1296+
1297+
auto db = kikimr.GetTableClient();
1298+
auto session = db.CreateSession().GetValueSync().GetSession();
1299+
1300+
{
1301+
auto query = R"(
1302+
--!syntax_v1
1303+
CREATE TABLE test (
1304+
A Int64 NOT NULL,
1305+
B Int64,
1306+
Created Int32 DEFAULT 1,
1307+
Deleted Int32 DEFAULT 0,
1308+
PRIMARY KEY (A ),
1309+
INDEX testIndex GLOBAL ON (B, A)
1310+
)
1311+
)";
1312+
1313+
auto result = session.ExecuteSchemeQuery(query).GetValueSync();
1314+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS,
1315+
result.GetIssues().ToString());
1316+
}
1317+
1318+
auto fQuery = [&](TString query) -> TString {
1319+
NYdb::NTable::TExecDataQuerySettings execSettings;
1320+
execSettings.KeepInQueryCache(true);
1321+
execSettings.CollectQueryStats(ECollectQueryStatsMode::Basic);
1322+
1323+
auto result =
1324+
session
1325+
.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx(),
1326+
execSettings)
1327+
.ExtractValueSync();
1328+
1329+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS,
1330+
result.GetIssues().ToString());
1331+
if (result.GetResultSets().size() > 0)
1332+
return NYdb::FormatResultSetYson(result.GetResultSet(0));
1333+
return "";
1334+
};
1335+
1336+
fQuery(R"(
1337+
upsert into test (A, B, Created, Deleted) values (5, 15, 1, 0)
1338+
)");
1339+
1340+
fQuery(R"(
1341+
$to_upsert = (
1342+
select A from
1343+
`test`
1344+
where A = 5
1345+
);
1346+
1347+
upsert into `test` (A, Deleted)
1348+
select A, 10 as Deleted from $to_upsert;
1349+
)");
1350+
1351+
CompareYson(
1352+
R"(
1353+
[
1354+
[5;[15];[1];[10]]
1355+
]
1356+
)",
1357+
fQuery(R"(
1358+
SELECT A, B, Created, Deleted FROM `test` ORDER BY A;
1359+
)")
1360+
);
1361+
}
1362+
12931363
Y_UNIT_TEST(Utf8AndDefault) {
12941364

12951365
NKikimrConfig::TAppConfig appConfig;

0 commit comments

Comments
 (0)