Skip to content
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.

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

WIP Fix issue #7674 about UPDATE SET(..), with indirection #7675

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
153 changes: 153 additions & 0 deletions src/backend/distributed/deparser/ruleutils_16.c
Original file line number Diff line number Diff line change
Expand Up @@ -469,6 +469,9 @@ static void get_tablesample_def(TableSampleClause *tablesample,
deparse_context *context);
static void get_opclass_name(Oid opclass, Oid actual_datatype,
StringInfo buf);
static bool is_update_set_with_multiple_columns(List *targetList);
static List *processTargetsIndirection(List *targetList);
static AttrNumber extract_paramid_from_funcexpr(FuncExpr *func);
static Node *processIndirection(Node *node, deparse_context *context);
static void printSubscripts(SubscriptingRef *aref, deparse_context *context);
static char *get_relation_name(Oid relid);
Expand Down Expand Up @@ -3545,6 +3548,9 @@ get_update_query_targetlist_def(Query *query, List *targetList,
}
}
}
if (is_update_set_with_multiple_columns(targetList))
targetList = processTargetsIndirection(targetList);

next_ma_cell = list_head(ma_sublinks);
cur_ma_sublink = NULL;
remaining_ma_columns = 0;
Expand Down Expand Up @@ -8607,6 +8613,153 @@ get_opclass_name(Oid opclass, Oid actual_datatype,
ReleaseSysCache(ht_opc);
}

/*
* helper function to evaluate if we are in an SET (...)
* Caller is responsible to check the command type (UPDATE)
*/
static bool is_update_set_with_multiple_columns(List *targetList)
{
ListCell *lc;
foreach(lc, targetList) {
TargetEntry *tle = (TargetEntry *) lfirst(lc);
Node *expr;

if (tle->resjunk)
continue;

expr = strip_implicit_coercions((Node *) tle->expr);

if (expr && IsA(expr, Param) &&
((Param *) expr)->paramkind == PARAM_MULTIEXPR)
{
return true;
}
}

// No multi-column set expression found
return false;
}

/*
* processTargetsIndirection - reorder targets list (from indirection)
*
* We don't change anything but the order the target list.
* The purpose here is to be able to deparse a query tree as if it was
* provided by the PostgreSQL parser, not the rewriter (which is the one
* received by the planner hook).
*
* It's required only for UPDATE SET (MULTIEXPR) queries, other candidates
* are not supported by Citus.
*
* Returns the new target list, reordered.
*/
static List *processTargetsIndirection(List *targetList)
{
int nAssignableCols;
int targetListPosition;
bool sawJunk = false;
List *newTargetList = NIL;
ListCell *lc;

/* Count non-junk columns and ensure they precede junk columns */
nAssignableCols = 0;
foreach(lc, targetList)
{
TargetEntry *tle = lfirst_node(TargetEntry, lc);

if (tle->resjunk)
{
sawJunk = true;
}
else
{
if (sawJunk)
elog(ERROR, "Subplan target list is out of order");

nAssignableCols++;
}
}

/* If no assignable columns, return the original target list */
if (nAssignableCols == 0)
return targetList;

/* Reorder the target list */
/* we start from 1 */
targetListPosition = 1;
while (nAssignableCols > 0)
{
nAssignableCols--;

foreach(lc, targetList)
{
TargetEntry *tle = lfirst_node(TargetEntry, lc);

if (IsA(tle->expr, FuncExpr))
{
FuncExpr *funcexpr = (FuncExpr *) tle->expr;
AttrNumber attnum = extract_paramid_from_funcexpr(funcexpr);

if (attnum == targetListPosition)
{
ereport(DEBUG1, (errmsg("Adding FuncExpr resno: %d", tle->resno)));
newTargetList = lappend(newTargetList, tle);
targetListPosition++;
break;
}
}
else if (IsA(tle->expr, Param))
{
Param *param = (Param *) tle->expr;
AttrNumber attnum = param->paramid;

if (attnum == targetListPosition)
{
newTargetList = lappend(newTargetList, tle);
targetListPosition++;
break;
}
}
}
}

// TODO add check about what we did here ?

/* Append any remaining junk columns */
foreach(lc, targetList)
{
TargetEntry *tle = lfirst_node(TargetEntry, lc);
if (tle->resjunk)
newTargetList = lappend(newTargetList, tle);
}

return newTargetList;
}

/* Function to extract paramid from a FuncExpr node */
static AttrNumber extract_paramid_from_funcexpr(FuncExpr *func)
{
AttrNumber targetAttnum = InvalidAttrNumber;
ListCell *lc;

/* Iterate through the arguments of the FuncExpr */
foreach(lc, func->args)
{
Node *arg = (Node *) lfirst(lc);

/* Check if the argument is a PARAM node */
if (IsA(arg, Param))
{
Param *param = (Param *) arg;
targetAttnum = param->paramid;

break; // Exit loop once we find the PARAM node
}
}

return targetAttnum;
}

/*
* processIndirection - take care of array and subfield assignment
*
Expand Down
191 changes: 191 additions & 0 deletions src/test/regress/expected/multi_modifications.out
Original file line number Diff line number Diff line change
Expand Up @@ -890,6 +890,16 @@ SELECT * FROM summary_table ORDER BY id;
(2 rows)

-- try different syntax
UPDATE summary_table SET (average_value, min_value) =
(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
WHERE id = 2;
SELECT * FROM summary_table ORDER BY id;
id | min_value | average_value | count | uniques
---------------------------------------------------------------------
1 | | 200.0000000000000000 | |
2 | 400 | 450.0000000000000000 | |
(2 rows)

UPDATE summary_table SET (min_value, average_value) =
(SELECT min(value), avg(value) FROM raw_table WHERE id = 2)
WHERE id = 2;
Expand Down Expand Up @@ -1297,6 +1307,187 @@ CREATE TABLE multi_modifications.local (a int default 1, b int);
INSERT INTO multi_modifications.local VALUES (default, (SELECT min(id) FROM summary_table));
ERROR: subqueries are not supported within INSERT queries
HINT: Try rewriting your queries with 'INSERT INTO ... SELECT' syntax.
-- test advanced UPDATE SET () with indirection and physical reordering.
CREATE TABLE updateset (
id bigint primary key
, col_0 integer
, col_1 integer
, col_2 integer
, col_3 integer
);
select create_reference_table('updateset');
create_reference_table
---------------------------------------------------------------------

(1 row)

insert into updateset values (1, 0, 0, 0, 0);
-- default physical ordering
update updateset
SET (col_0, col_1, col_2, col_3)
= (SELECT 100, 111, 222, 333)
returning *;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
1 | 100 | 111 | 222 | 333
(1 row)

select * from updateset;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
1 | 100 | 111 | 222 | 333
(1 row)

-- check indirection
update updateset
SET (col_0, col_1, col_3, col_2)
= (SELECT 10, 11, 33, 22)
returning *;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
1 | 10 | 11 | 22 | 33
(1 row)

select * from updateset;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
1 | 10 | 11 | 22 | 33
(1 row)

update updateset
SET (col_0, col_3, col_1, col_2)
= (SELECT 100, 333, 111, 222)
returning *;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
1 | 100 | 111 | 222 | 333
(1 row)

select * from updateset;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
1 | 100 | 111 | 222 | 333
(1 row)

update updateset
SET (col_3, col_1)
= (SELECT 3, 1)
returning *;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
1 | 100 | 1 | 222 | 3
(1 row)

select * from updateset;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
1 | 100 | 1 | 222 | 3
(1 row)

-- check more complex queries with indirection
insert into updateset values (2, 0, 0, 0, 0);
update updateset
SET (col_0, col_1, col_3, col_2)
= (SELECT 10, 11, 33, 22)
where id = 2
returning *;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
2 | 10 | 11 | 22 | 33
(1 row)

select * from updateset where id = 2;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
2 | 10 | 11 | 22 | 33
(1 row)

update updateset
SET (col_0, col_3, col_1, col_2)
= (SELECT 100, 333, 111, 222)
where id = 2
returning *;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
2 | 100 | 111 | 222 | 333
(1 row)

select * from updateset where id = 2;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
2 | 100 | 111 | 222 | 333
(1 row)

update updateset
SET (col_3, col_1)
= (SELECT 3, 1)
where id = 2
returning *;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
2 | 100 | 1 | 222 | 3
(1 row)

select * from updateset where id = 2;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
2 | 100 | 1 | 222 | 3
(1 row)

-- the single row update is expected behavior
insert into updateset values (3, 0, 1, 2, 3);
insert into updateset values (4, 0, 0, 0, 0);
with qq as (
update updateset
SET (col_1, col_3)
= (SELECT 11, 33)
where id = 4
returning *
)
update updateset
set (col_2, col_1, col_3)
= (select col_2, col_1, col_3 from qq)
where id in (3, 4)
returning *;
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
3 | 0 | 11 | 0 | 33
(1 row)

select * from updateset where id in (3, 4);
id | col_0 | col_1 | col_2 | col_3
---------------------------------------------------------------------
3 | 0 | 11 | 0 | 33
4 | 0 | 11 | 0 | 33
(2 rows)

-- add more advanced queries ?
-- we want to ensure the reordering the targetlist
-- from indirection is not run when it should not
truncate updateset;
-- change physical ordering
alter table updateset drop col_2;
alter table updateset add col_2 integer;
insert into updateset values (1, 0, 0, 0, 0);
insert into updateset values (2, 0, 0, 0, 0);
update updateset
SET (col_0, col_1, col_2, col_3)
= (SELECT 10, 11, 22, 33)
returning *;
id | col_0 | col_1 | col_3 | col_2
---------------------------------------------------------------------
1 | 10 | 11 | 33 | 22
2 | 10 | 11 | 33 | 22
(2 rows)

select * from updateset;
id | col_0 | col_1 | col_3 | col_2
---------------------------------------------------------------------
1 | 10 | 11 | 33 | 22
2 | 10 | 11 | 33 | 22
(2 rows)

DROP TABLE updateset;
DROP TABLE insufficient_shards;
DROP TABLE raw_table;
DROP TABLE summary_table;
Expand Down
Loading
Loading