Skip to content

Commit ee2a708

Browse files
gridnevvvitgridnevvvit
authored
Merge bugfixes 24 4 (#16685)
1 parent 7083890 commit ee2a708

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
@@ -403,21 +403,13 @@ RewriteInputForConstraint(const TExprBase& inputRows, const THashSet<TStringBuf>
403403
const THashSet<TString> indexKeyColumns = CreateKeyColumnSetToRead(indexes);
404404
const THashSet<TString> indexDataColumns = CreateDataColumnSetToRead(indexes);
405405

406-
for (const auto& x : indexKeyColumns) {
407-
columns.push_back(Build<TCoAtom>(ctx, pos).Value(x).Done());
408-
}
406+
THashSet<TString> columnsToReadInPrecomputeLookupDict;
407+
columnsToReadInPrecomputeLookupDict.insert(indexKeyColumns.begin(), indexKeyColumns.end());
408+
columnsToReadInPrecomputeLookupDict.insert(indexDataColumns.begin(), indexDataColumns.end());
409+
columnsToReadInPrecomputeLookupDict.insert(mainPk.begin(), mainPk.end());
410+
columnsToReadInPrecomputeLookupDict.insert(checkDefaults.begin(), checkDefaults.end());
409411

410-
for (const auto& x : indexDataColumns) {
411-
// Handle the case of multiple indexes
412-
// one of them has 'foo' as data column but for another one foo is just indexed column
413-
if (indexKeyColumns.contains(x))
414-
continue;
415-
columns.push_back(Build<TCoAtom>(ctx, pos).Value(x).Done());
416-
}
417-
418-
for (const auto& x : mainPk) {
419-
if (indexKeyColumns.contains(x))
420-
continue;
412+
for (const auto& x : columnsToReadInPrecomputeLookupDict) {
421413
columns.push_back(Build<TCoAtom>(ctx, pos).Value(x).Done());
422414
}
423415

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

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,206 @@ Y_UNIT_TEST(ReturningSerial) {
241241
}
242242
}
243243

244+
TString ExecuteReturningQuery(TKikimrRunner& kikimr, bool queryService, TString query) {
245+
if (queryService) {
246+
auto qdb = kikimr.GetQueryClient();
247+
auto qSession = qdb.GetSession().GetValueSync().GetSession();
248+
auto result = qSession.ExecuteQuery(
249+
query, NYdb::NQuery::TTxControl::BeginTx().CommitTx()).ExtractValueSync();
250+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
251+
return FormatResultSetYson(result.GetResultSet(0));
252+
}
253+
254+
auto db = kikimr.GetTableClient();
255+
auto session = db.CreateSession().GetValueSync().GetSession();
256+
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).ExtractValueSync();
257+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
258+
return FormatResultSetYson(result.GetResultSet(0));
259+
}
260+
261+
Y_UNIT_TEST_TWIN(ReturningWorks, QueryService) {
262+
auto kikimr = DefaultKikimrRunner();
263+
auto db = kikimr.GetTableClient();
264+
auto session = db.CreateSession().GetValueSync().GetSession();
265+
CreateSampleTablesWithIndex(session, true);
266+
CompareYson(
267+
R"([[[101];[101];["Payload1"]];])",
268+
ExecuteReturningQuery(kikimr, QueryService, R"(
269+
UPSERT INTO `/Root/SecondaryKeys` (Key, Fk, Value) VALUES (101, 101, "Payload1") RETURNING *;
270+
)")
271+
);
272+
CompareYson(
273+
R"(
274+
[[#;#;["Payload8"]];
275+
[[1];[1];["Payload1"]];
276+
[[2];[2];["Payload2"]];
277+
[[5];[5];["Payload5"]];
278+
[#;[7];["Payload7"]];
279+
[[101];[101];["Payload1"]]
280+
])",
281+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
282+
);
283+
}
284+
285+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedUpsert, QueryService) {
286+
auto kikimr = DefaultKikimrRunner();
287+
auto db = kikimr.GetTableClient();
288+
auto session = db.CreateSession().GetValueSync().GetSession();
289+
CreateSampleTablesWithIndex(session, true);
290+
CompareYson(R"([
291+
[[110];[110];["Payload5"]];
292+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
293+
$v1 = (SELECT Key + 100 as Key, Fk + 100 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
294+
$v2 = (SELECT Key + 105 as Key, Fk + 105 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
295+
UPSERT INTO `/Root/SecondaryKeys`
296+
SELECT * FROM (SELECT * FROM $v1 UNION ALL SELECT * FROM $v2) WHERE Key > 107 RETURNING *;
297+
)"));
298+
CompareYson(
299+
R"(
300+
[[#;#;["Payload8"]];
301+
[[1];[1];["Payload1"]];
302+
[[2];[2];["Payload2"]];
303+
[[5];[5];["Payload5"]];
304+
[#;[7];["Payload7"]];
305+
[[110];[110];["Payload5"]]
306+
])",
307+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
308+
);
309+
}
310+
311+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedDelete, QueryService) {
312+
auto kikimr = DefaultKikimrRunner();
313+
auto db = kikimr.GetTableClient();
314+
auto session = db.CreateSession().GetValueSync().GetSession();
315+
CreateSampleTablesWithIndex(session, true);
316+
CompareYson(R"([
317+
[[5];[5];["Payload5"]];
318+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
319+
$v1 = (SELECT Key, Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL AND Key >= 1);
320+
$v2 = (SELECT Key, Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL AND Key <= 5);
321+
DELETE FROM `/Root/SecondaryKeys` ON
322+
SELECT * FROM (SELECT * FROM $v1 UNION ALL SELECT * FROM $v2) WHERE Key >= 5 RETURNING *;
323+
)"));
324+
CompareYson(
325+
R"(
326+
[[#;#;["Payload8"]];
327+
[[1];[1];["Payload1"]];
328+
[[2];[2];["Payload2"]];
329+
[#;[7];["Payload7"]];
330+
])",
331+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
332+
);
333+
}
334+
335+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedDeleteV2, QueryService) {
336+
auto kikimr = DefaultKikimrRunner();
337+
auto db = kikimr.GetTableClient();
338+
auto session = db.CreateSession().GetValueSync().GetSession();
339+
CreateSampleTablesWithIndex(session, true);
340+
CompareYson(R"([
341+
[[1];[1];["Payload1"]];
342+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
343+
DELETE FROM `/Root/SecondaryKeys` WHERE Key = 1 RETURNING *;
344+
)"));
345+
CompareYson(
346+
R"(
347+
[[#;#;["Payload8"]];
348+
[[2];[2];["Payload2"]];
349+
[[5];[5];["Payload5"]];
350+
[#;[7];["Payload7"]];
351+
])",
352+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
353+
);
354+
}
355+
356+
357+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedInsert, QueryService) {
358+
auto kikimr = DefaultKikimrRunner();
359+
auto db = kikimr.GetTableClient();
360+
auto session = db.CreateSession().GetValueSync().GetSession();
361+
CreateSampleTablesWithIndex(session, true);
362+
363+
CompareYson(R"([
364+
[[101];[101];["Payload1"]];
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 + 205 as Key, Fk + 205 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
368+
INSERT INTO `/Root/SecondaryKeys`
369+
SELECT * FROM (SELECT * FROM $v1 UNION ALL SELECT * FROM $v2 ) WHERE Key < 102 RETURNING *;
370+
)"));
371+
372+
CompareYson(
373+
R"(
374+
[[#;#;["Payload8"]];
375+
[[1];[1];["Payload1"]];
376+
[[2];[2];["Payload2"]];
377+
[[5];[5];["Payload5"]];
378+
[#;[7];["Payload7"]];
379+
[[101];[101];["Payload1"]]
380+
])",
381+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
382+
);
383+
}
384+
385+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedReplace, QueryService) {
386+
auto kikimr = DefaultKikimrRunner();
387+
auto db = kikimr.GetTableClient();
388+
auto session = db.CreateSession().GetValueSync().GetSession();
389+
CreateSampleTablesWithIndex(session, true);
390+
391+
CompareYson(R"([
392+
[[101];[101];["Payload1"]];
393+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
394+
$v1 = (SELECT Key + 100 as Key, Fk + 100 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
395+
$v2 = (SELECT Key + 205 as Key, Fk + 205 as Fk, Value FROM `/Root/SecondaryKeys` WHERE Key IS NOT NULL AND Fk IS NOT NULL);
396+
REPLACE INTO `/Root/SecondaryKeys`
397+
SELECT * FROM (SELECT * FROM $v1 UNION ALL SELECT * FROM $v2 ) WHERE Key < 102 RETURNING *;
398+
)"));
399+
400+
CompareYson(
401+
R"(
402+
[[#;#;["Payload8"]];
403+
[[1];[1];["Payload1"]];
404+
[[2];[2];["Payload2"]];
405+
[[5];[5];["Payload5"]];
406+
[#;[7];["Payload7"]];
407+
[[101];[101];["Payload1"]]
408+
])",
409+
ExecuteReturningQuery(kikimr, QueryService, "SELECT * FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
410+
);
411+
}
412+
413+
Y_UNIT_TEST_TWIN(ReturningWorksIndexedOperationsWithDefault, QueryService) {
414+
auto kikimr = DefaultKikimrRunner();
415+
auto db = kikimr.GetTableClient();
416+
auto session = db.CreateSession().GetValueSync().GetSession();
417+
{
418+
auto res = session.ExecuteSchemeQuery(R"(
419+
--!syntax_v1
420+
CREATE TABLE `/Root/SecondaryKeys` (
421+
Key Serial,
422+
Fk Int32,
423+
Value String,
424+
PRIMARY KEY (Key),
425+
INDEX Index GLOBAL ON (Fk)
426+
);
427+
)").GetValueSync();
428+
}
429+
430+
CompareYson(R"([
431+
[1;[1];["Payload"]];
432+
])", ExecuteReturningQuery(kikimr, QueryService, R"(
433+
REPLACE INTO `/Root/SecondaryKeys` (Fk, Value) VALUES (1, "Payload") RETURNING Key, Fk, Value;
434+
)"));
435+
436+
CompareYson(
437+
R"([
438+
[1;[1];["Payload"]];
439+
])",
440+
ExecuteReturningQuery(kikimr, QueryService, "SELECT Key, Fk, Value FROM `/Root/SecondaryKeys` ORDER BY Key, Fk;")
441+
);
442+
}
443+
244444
Y_UNIT_TEST(ReturningColumnsOrder) {
245445
auto kikimr = DefaultKikimrRunner();
246446

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

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

11921192
}
11931193

1194+
Y_UNIT_TEST(DefaultAndIndexesTestDefaultColumnNotIncludedInIndex) {
1195+
NKikimrConfig::TAppConfig appConfig;
1196+
TKikimrRunner kikimr(TKikimrSettings().SetPQConfig(DefaultPQConfig()).SetAppConfig(appConfig));
1197+
1198+
auto db = kikimr.GetTableClient();
1199+
auto session = db.CreateSession().GetValueSync().GetSession();
1200+
1201+
{
1202+
auto query = R"(
1203+
--!syntax_v1
1204+
CREATE TABLE test (
1205+
A Int64 NOT NULL,
1206+
B Int64,
1207+
Created Int32 DEFAULT 1,
1208+
Deleted Int32 DEFAULT 0,
1209+
PRIMARY KEY (A ),
1210+
INDEX testIndex GLOBAL ON (B, A)
1211+
)
1212+
)";
1213+
1214+
auto result = session.ExecuteSchemeQuery(query).GetValueSync();
1215+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS,
1216+
result.GetIssues().ToString());
1217+
}
1218+
1219+
auto fQuery = [&](TString query) -> TString {
1220+
NYdb::NTable::TExecDataQuerySettings execSettings;
1221+
execSettings.KeepInQueryCache(true);
1222+
execSettings.CollectQueryStats(ECollectQueryStatsMode::Basic);
1223+
1224+
auto result =
1225+
session
1226+
.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx(),
1227+
execSettings)
1228+
.ExtractValueSync();
1229+
1230+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS,
1231+
result.GetIssues().ToString());
1232+
if (result.GetResultSets().size() > 0)
1233+
return NYdb::FormatResultSetYson(result.GetResultSet(0));
1234+
return "";
1235+
};
1236+
1237+
fQuery(R"(
1238+
upsert into test (A, B, Created, Deleted) values (5, 15, 1, 0)
1239+
)");
1240+
1241+
fQuery(R"(
1242+
$to_upsert = (
1243+
select A from
1244+
`test`
1245+
where A = 5
1246+
);
1247+
1248+
upsert into `test` (A, Deleted)
1249+
select A, 10 as Deleted from $to_upsert;
1250+
)");
1251+
1252+
CompareYson(
1253+
R"(
1254+
[
1255+
[5;[15];[1];[10]]
1256+
]
1257+
)",
1258+
fQuery(R"(
1259+
SELECT A, B, Created, Deleted FROM `test` ORDER BY A;
1260+
)")
1261+
);
1262+
}
1263+
11941264
Y_UNIT_TEST(Utf8AndDefault) {
11951265

11961266
NKikimrConfig::TAppConfig appConfig;

0 commit comments

Comments
 (0)