Skip to content

Commit 7f54b71

Browse files
ssmikessmike
authored
Support returning list in delete statements (#1684)
1 parent 762916f commit 7f54b71

9 files changed

+130
-33
lines changed

‎ydb/core/kqp/expr_nodes/kqp_expr_nodes.json

+8-2
Original file line numberDiff line numberDiff line change
@@ -509,12 +509,18 @@
509509
{
510510
"Name": "TKqlDeleteRows",
511511
"Base": "TKqlDeleteRowsBase",
512-
"Match": {"Type": "Callable", "Name": "KqlDeleteRows"}
512+
"Match": {"Type": "Callable", "Name": "KqlDeleteRows"},
513+
"Children": [
514+
{"Index": 2, "Name": "ReturningColumns", "Type": "TCoAtomList"}
515+
]
513516
},
514517
{
515518
"Name": "TKqlDeleteRowsIndex",
516519
"Base": "TKqlDeleteRowsBase",
517-
"Match": {"Type": "Callable", "Name": "KqlDeleteRowsIndex"}
520+
"Match": {"Type": "Callable", "Name": "KqlDeleteRowsIndex"},
521+
"Children": [
522+
{"Index": 2, "Name": "ReturningColumns", "Type": "TCoAtomList"}
523+
]
518524
},
519525
{
520526
"Name": "TKqpDeleteRows",

‎ydb/core/kqp/host/kqp_type_ann.cpp

+1-1
Original file line numberDiff line numberDiff line change
@@ -830,7 +830,7 @@ TStatus AnnotateUpdateRows(const TExprNode::TPtr& node, TExprContext& ctx, const
830830
TStatus AnnotateDeleteRows(const TExprNode::TPtr& node, TExprContext& ctx, const TString& cluster,
831831
const TKikimrTablesData& tablesData)
832832
{
833-
if (!EnsureArgsCount(*node, 2, ctx)) {
833+
if (!EnsureMaxArgsCount(*node, 3, ctx) && !EnsureMinArgsCount(*node, 2, ctx)) {
834834
return TStatus::Error;
835835
}
836836

‎ydb/core/kqp/opt/kqp_opt_kql.cpp

+6
Original file line numberDiff line numberDiff line change
@@ -429,6 +429,7 @@ TExprBase BuildDeleteTable(const TKiWriteTable& write, const TKikimrTableDescrip
429429
return Build<TKqlDeleteRows>(ctx, write.Pos())
430430
.Table(BuildTableMeta(tableData, write.Pos(), ctx))
431431
.Input(keysToDelete)
432+
.ReturningColumns(write.ReturningColumns())
432433
.Done();
433434
}
434435

@@ -438,6 +439,7 @@ TExprBase BuildDeleteTableWithIndex(const TKiWriteTable& write, const TKikimrTab
438439
return Build<TKqlDeleteRowsIndex>(ctx, write.Pos())
439440
.Table(BuildTableMeta(tableData, write.Pos(), ctx))
440441
.Input(keysToDelete)
442+
.ReturningColumns(write.ReturningColumns())
441443
.Done();
442444
}
443445

@@ -464,6 +466,7 @@ TExprBase BuildDeleteTable(const TKiDeleteTable& del, const TKikimrTableDescript
464466
return Build<TKqlDeleteRows>(ctx, del.Pos())
465467
.Table(BuildTableMeta(tableData, del.Pos(), ctx))
466468
.Input(keysToDelete)
469+
.ReturningColumns<TCoAtomList>().Build()
467470
.Done();
468471
}
469472

@@ -483,6 +486,7 @@ TExprBase BuildDeleteTableWithIndex(const TKiDeleteTable& del, const TKikimrTabl
483486
auto tableDelete = Build<TKqlDeleteRows>(ctx, del.Pos())
484487
.Table(BuildTableMeta(tableData, del.Pos(), ctx))
485488
.Input(ProjectColumns(rowsToDelete, pk, ctx))
489+
.ReturningColumns<TCoAtomList>().Build()
486490
.Done();
487491

488492
TVector<TExprBase> effects;
@@ -506,6 +510,7 @@ TExprBase BuildDeleteTableWithIndex(const TKiDeleteTable& del, const TKikimrTabl
506510
auto indexDelete = Build<TKqlDeleteRows>(ctx, del.Pos())
507511
.Table(indexMeta)
508512
.Input(ProjectColumns(rowsToDelete, indexTableColumns, ctx))
513+
.ReturningColumns<TCoAtomList>().Build()
509514
.Done();
510515

511516
effects.push_back(indexDelete);
@@ -699,6 +704,7 @@ TExprBase BuildUpdateTableWithIndex(const TKiUpdateTable& update, const TKikimrT
699704
auto indexDelete = Build<TKqlDeleteRows>(ctx, update.Pos())
700705
.Table(indexMeta)
701706
.Input(ProjectColumns(rowsToUpdate, indexTableColumns, ctx))
707+
.ReturningColumns<TCoAtomList>().Build()
702708
.Done();
703709

704710
effects.push_back(indexDelete);

‎ydb/core/kqp/opt/logical/kqp_opt_log_effects.cpp

+1
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ TExprBase KqpDeleteOverLookup(const TExprBase& node, TExprContext& ctx, const TK
9797
return Build<TKqlDeleteRows>(ctx, deleteRows.Pos())
9898
.Table(deleteRows.Table())
9999
.Input(deleteInput.Cast())
100+
.ReturningColumns(deleteRows.ReturningColumns())
100101
.Done();
101102
}
102103

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

+2
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ TExprBase KqpBuildDeleteIndexStages(TExprBase node, TExprContext& ctx, const TKq
8181
auto tableDelete = Build<TKqlDeleteRows>(ctx, del.Pos())
8282
.Table(del.Table())
8383
.Input(lookupKeys)
84+
.ReturningColumns(del.ReturningColumns())
8485
.Done();
8586

8687
TVector<TExprBase> effects;
@@ -102,6 +103,7 @@ TExprBase KqpBuildDeleteIndexStages(TExprBase node, TExprContext& ctx, const TKq
102103
auto indexDelete = Build<TKqlDeleteRows>(ctx, del.Pos())
103104
.Table(tableNode)
104105
.Input(deleteIndexKeys)
106+
.ReturningColumns<TCoAtomList>().Build()
105107
.Done();
106108

107109
effects.emplace_back(std::move(indexDelete));

‎ydb/core/kqp/opt/physical/effects/kqp_opt_phy_effects_rules.h

+3
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@ NYql::NNodes::TExprBase KqpBuildReturning(NYql::NNodes::TExprBase node, NYql::TE
1212
NYql::NNodes::TExprBase KqpRewriteReturningUpsert(NYql::NNodes::TExprBase node, NYql::TExprContext& ctx,
1313
const TKqpOptimizeContext& kqpCtx);
1414

15+
NYql::NNodes::TExprBase KqpRewriteReturningDelete(NYql::NNodes::TExprBase node, NYql::TExprContext& ctx,
16+
const TKqpOptimizeContext& kqpCtx);
17+
1518
NYql::NNodes::TExprBase KqpRewriteGenerateIfInsert(NYql::NNodes::TExprBase node, NYql::TExprContext& ctx,
1619
const TKqpOptimizeContext& kqpCtx);
1720

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

+88-26
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,33 @@ TCoAtomList MakeColumnsList(Container rows, TExprContext& ctx, TPositionHandle p
1515
return Build<TCoAtomList>(ctx, pos).Add(columnsVector).Done();
1616
}
1717

18+
template<typename Container>
19+
TExprBase SelectFields(TExprBase node, Container fields, TExprContext& ctx, TPositionHandle pos) {
20+
TVector<TExprBase> items;
21+
for (auto&& field : fields) {
22+
TString name;
23+
24+
if constexpr (std::is_same_v<NYql::NNodes::TCoAtom&&, decltype(field)>) {
25+
name = field.Value();
26+
} else {
27+
name = field;
28+
}
29+
30+
auto tuple = Build<TCoNameValueTuple>(ctx, pos)
31+
.Name().Build(field)
32+
.template Value<TCoMember>()
33+
.Struct(node)
34+
.Name().Build(name)
35+
.Build()
36+
.Done();
37+
38+
items.emplace_back(tuple);
39+
}
40+
return Build<TCoAsStruct>(ctx, pos)
41+
.Add(items)
42+
.Done();
43+
}
44+
1845
TExprBase KqpBuildReturning(TExprBase node, TExprContext& ctx, const TKqpOptimizeContext& kqpCtx) {
1946
auto maybeReturning = node.Maybe<TKqlReturningList>();
2047
if (!maybeReturning) {
@@ -24,46 +51,60 @@ TExprBase KqpBuildReturning(TExprBase node, TExprContext& ctx, const TKqpOptimiz
2451
auto returning = maybeReturning.Cast();
2552
const auto& tableDesc = kqpCtx.Tables->ExistingTable(kqpCtx.Cluster, returning.Table().Path());
2653

27-
auto buildFromUpsert = [&](TMaybeNode<TKqlUpsertRows> upsert) -> TExprBase {
28-
auto rows = upsert.Cast().Input();
29-
auto pos = upsert.Input().Cast().Pos();
54+
auto buildReturningRows = [&](TExprBase rows, TCoAtomList columns, TCoAtomList returningColumns) -> TExprBase {
55+
auto pos = rows.Pos();
3056

3157
TSet<TString> inputColumns;
3258
TSet<TString> columnsToReadSet;
33-
34-
for (auto&& column : upsert.Columns().Cast()) {
59+
for (auto&& column : columns) {
3560
inputColumns.insert(TString(column.Value()));
3661
}
37-
for (auto&& column : upsert.ReturningColumns().Cast()) {
62+
for (auto&& column : returningColumns) {
3863
if (!inputColumns.contains(column) && !tableDesc.GetKeyColumnIndex(TString(column))) {
3964
columnsToReadSet.insert(TString(column));
4065
}
4166
}
42-
43-
TMaybeNode<TExprBase> input = upsert.Input();
67+
TMaybeNode<TExprBase> input = rows;
4468

4569
if (!columnsToReadSet.empty()) {
46-
TString upsertInputName = "upsertInput";
47-
TString tableInputName = "table";
70+
auto payloadSelectorArg = TCoArgument(ctx.NewArgument(pos, "payload_selector_row"));
71+
TVector<TExprBase> payloadTuples;
72+
for (const auto& column : columns) {
73+
payloadTuples.emplace_back(
74+
Build<TCoNameValueTuple>(ctx, pos)
75+
.Name(column)
76+
.Value<TCoMember>()
77+
.Struct(payloadSelectorArg)
78+
.Name(column)
79+
.Build()
80+
.Done());
81+
}
82+
83+
auto payloadSelector = Build<TCoLambda>(ctx, pos)
84+
.Args({payloadSelectorArg})
85+
.Body<TCoAsStruct>()
86+
.Add(payloadTuples)
87+
.Build()
88+
.Done();
4889

49-
auto payloadSelector = MakeRowsPayloadSelector(upsert.Columns().Cast(), tableDesc, pos, ctx);
5090
auto condenseResult = CondenseInputToDictByPk(input.Cast(), tableDesc, payloadSelector, ctx);
5191
if (!condenseResult) {
5292
return node;
5393
}
5494

5595
auto inputDictAndKeys = PrecomputeDictAndKeys(*condenseResult, pos, ctx);
56-
57-
TSet<TString> columnsToLookup = columnsToReadSet;
5896
for (auto&& column : tableDesc.Metadata->KeyColumnNames) {
5997
columnsToReadSet.insert(column);
6098
}
61-
99+
TSet<TString> columnsToLookup = columnsToReadSet;
62100
for (auto&& column : tableDesc.Metadata->KeyColumnNames) {
63101
columnsToReadSet.erase(column);
64102
}
65103
TCoAtomList additionalColumnsToRead = MakeColumnsList(columnsToReadSet, ctx, pos);
66104

105+
TCoArgument existingRow = Build<TCoArgument>(ctx, node.Pos())
106+
.Name("existing_row")
107+
.Done();
67108
auto prepareUpdateStage = Build<TDqStage>(ctx, pos)
68109
.Inputs()
69110
.Add(inputDictAndKeys.KeysPrecompute)
@@ -80,27 +121,21 @@ TExprBase KqpBuildReturning(TExprBase node, TExprContext& ctx, const TKqpOptimiz
80121
.Columns(MakeColumnsList(columnsToLookup, ctx, pos))
81122
.Build()
82123
.Lambda()
83-
.Args({"existingRow"})
124+
.Args({existingRow})
84125
.Body<TCoJust>()
85126
.Input<TCoFlattenMembers>()
86127
.Add()
87128
.Name().Build("")
88129
.Value<TCoUnwrap>() // Key should always exist in the dict
89130
.Optional<TCoLookup>()
90131
.Collection("dict")
91-
.Lookup<TCoExtractMembers>()
92-
.Input("existingRow")
93-
.Members(MakeColumnsList(tableDesc.Metadata->KeyColumnNames, ctx, pos))
94-
.Build()
132+
.Lookup(SelectFields(existingRow, tableDesc.Metadata->KeyColumnNames, ctx, pos))
95133
.Build()
96134
.Build()
97135
.Build()
98136
.Add()
99137
.Name().Build("")
100-
.Value<TCoExtractMembers>()
101-
.Input("existingRow")
102-
.Members(additionalColumnsToRead)
103-
.Build()
138+
.Value(SelectFields(existingRow, additionalColumnsToRead, ctx, pos))
104139
.Build()
105140
.Build()
106141
.Build()
@@ -128,20 +163,27 @@ TExprBase KqpBuildReturning(TExprBase node, TExprContext& ctx, const TKqpOptimiz
128163
for (auto item : maybeList.Cast()) {
129164
if (auto upsert = item.Maybe<TKqlUpsertRows>()) {
130165
if (upsert.Cast().Table().Raw() == returning.Table().Raw()) {
131-
return buildFromUpsert(upsert);
166+
return buildReturningRows(upsert.Input().Cast(), upsert.Columns().Cast(), returning.Columns());
167+
}
168+
}
169+
if (auto del = item.Maybe<TKqlDeleteRows>()) {
170+
if (del.Cast().Table().Raw() == returning.Table().Raw()) {
171+
return buildReturningRows(del.Input().Cast(), MakeColumnsList(tableDesc.Metadata->KeyColumnNames, ctx, node.Pos()), returning.Columns());
132172
}
133173
}
134174
}
135175
}
136176

137177
if (auto upsert = returning.Update().Maybe<TKqlUpsertRows>()) {
138-
return buildFromUpsert(upsert);
178+
return buildReturningRows(upsert.Input().Cast(), upsert.Columns().Cast(), returning.Columns());
179+
}
180+
if (auto del = returning.Update().Maybe<TKqlDeleteRows>()) {
181+
return buildReturningRows(del.Input().Cast(), MakeColumnsList(tableDesc.Metadata->KeyColumnNames, ctx, node.Pos()), returning.Columns());
139182
}
140183

141184
return node;
142185
}
143186

144-
145187
TExprBase KqpRewriteReturningUpsert(TExprBase node, TExprContext& ctx, const TKqpOptimizeContext&) {
146188
auto upsert = node.Cast<TKqlUpsertRows>();
147189
if (upsert.ReturningColumns().Empty()) {
@@ -164,4 +206,24 @@ TExprBase KqpRewriteReturningUpsert(TExprBase node, TExprContext& ctx, const TKq
164206
.Done();
165207
}
166208

209+
TExprBase KqpRewriteReturningDelete(TExprBase node, TExprContext& ctx, const TKqpOptimizeContext&) {
210+
auto del = node.Cast<TKqlDeleteRows>();
211+
if (del.ReturningColumns().Empty()) {
212+
return node;
213+
}
214+
215+
if (!del.Input().Maybe<TDqPrecompute>() && !del.Input().Maybe<TDqPhyPrecompute>()) {
216+
return node;
217+
}
218+
219+
return
220+
Build<TKqlDeleteRows>(ctx, del.Pos())
221+
.Input<TDqPrecompute>()
222+
.Input(del.Input())
223+
.Build()
224+
.Table(del.Table())
225+
.ReturningColumns<TCoAtomList>().Build()
226+
.Done();
227+
}
228+
167229
} // namespace NKikimr::NKqp::NOpt

‎ydb/core/kqp/opt/physical/kqp_opt_phy.cpp

+9-3
Original file line numberDiff line numberDiff line change
@@ -123,10 +123,10 @@ class TKqpPhysicalOptTransformer : public TOptimizeTransformerBase {
123123

124124
AddHandler(2, &TDqStage::Match, HNDL(RewriteKqpReadTable));
125125
AddHandler(2, &TDqStage::Match, HNDL(RewriteKqpLookupTable));
126+
AddHandler(2, &TKqlUpsertRows::Match, HNDL(RewriteReturningUpsert));
127+
AddHandler(2, &TKqlDeleteRows::Match, HNDL(RewriteReturningDelete));
126128

127-
AddHandler(3, &TKqlUpsertRows::Match, HNDL(RewriteReturningUpsert));
128-
129-
AddHandler(4, &TKqlReturningList::Match, HNDL(BuildReturning));
129+
AddHandler(3, &TKqlReturningList::Match, HNDL(BuildReturning));
130130
#undef HNDL
131131

132132
SetGlobal(1u);
@@ -145,6 +145,12 @@ class TKqpPhysicalOptTransformer : public TOptimizeTransformerBase {
145145
return output;
146146
}
147147

148+
TMaybeNode<TExprBase> RewriteReturningDelete(TExprBase node, TExprContext& ctx) {
149+
TExprBase output = KqpRewriteReturningDelete(node, ctx, KqpCtx);
150+
DumpAppliedRule("RewriteReturningDelete", node.Ptr(), output.Ptr(), ctx);
151+
return output;
152+
}
153+
148154
TMaybeNode<TExprBase> RewriteGenerateIfInsert(TExprBase node, TExprContext& ctx) {
149155
TExprBase output = KqpRewriteGenerateIfInsert(node, ctx, KqpCtx);
150156
DumpAppliedRule("RewriteGenerateIfInsert", node.Ptr(), output.Ptr(), ctx);

‎ydb/core/kqp/ut/pg/kqp_pg_ut.cpp

+12-1
Original file line numberDiff line numberDiff line change
@@ -1534,13 +1534,24 @@ Y_UNIT_TEST_SUITE(KqpPg) {
15341534
{
15351535
const auto query = Q_(R"(
15361536
--!syntax_pg
1537-
UPDATE ReturningTableExtraValue SET value2 = 3 where key = 2 RETURNING *;
1537+
DELETE FROM ReturningTableExtraValue WHERE key = 2 RETURNING key, value, value2;
15381538
)");
15391539

15401540
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).GetValueSync();
15411541
UNIT_ASSERT(result.IsSuccess());
15421542
CompareYson(R"([["2";"4";"3"]])", FormatResultSetYson(result.GetResultSet(0)));
15431543
}
1544+
1545+
{
1546+
const auto query = Q_(R"(
1547+
--!syntax_pg
1548+
DELETE FROM ReturningTable WHERE key <= 3 RETURNING key, value;
1549+
)");
1550+
1551+
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).GetValueSync();
1552+
UNIT_ASSERT(result.IsSuccess());
1553+
CompareYson(R"([["2";"2"];["3";"2"];["1";"3"]])", FormatResultSetYson(result.GetResultSet(0)));
1554+
}
15441555
}
15451556

15461557
Y_UNIT_TEST(DropTablePg) {

0 commit comments

Comments
 (0)