From cd76b7649ccf66952a41cea2b68a46aef2159376 Mon Sep 17 00:00:00 2001 From: paragjain Date: Tue, 16 Jul 2024 16:40:48 +0000 Subject: [PATCH] Initial changes --- .../distributed/planner/merge_planner.c | 5 +- .../planner/query_colocation_checker.c | 8 ++- .../distributed/planner/recursive_planning.c | 7 ++- .../distributed/query_colocation_checker.h | 3 +- src/test/regress/sql/merge_vcore.sql | 57 +++++++++++++++++++ 5 files changed, 72 insertions(+), 8 deletions(-) diff --git a/src/backend/distributed/planner/merge_planner.c b/src/backend/distributed/planner/merge_planner.c index 1f9d17c43..92fea1c9d 100644 --- a/src/backend/distributed/planner/merge_planner.c +++ b/src/backend/distributed/planner/merge_planner.c @@ -857,8 +857,11 @@ ConvertRelationRTEIntoSubquery(Query *mergeQuery, RangeTblEntry *sourceRte, /* set the FROM expression to the subquery */ newRangeTableRef->rtindex = SINGLE_RTE_INDEX; sourceResultsQuery->jointree = makeFromExpr(list_make1(newRangeTableRef), NULL); + + bool isMergeQuery = true; sourceResultsQuery->targetList = - CreateAllTargetListForRelation(sourceRte->relid, requiredAttributes); + CreateAllTargetListForRelation(sourceRte->relid, requiredAttributes, + isMergeQuery); List *restrictionList = GetRestrictInfoListForRelation(sourceRte, plannerRestrictionContext); List *copyRestrictionList = copyObject(restrictionList); diff --git a/src/backend/distributed/planner/query_colocation_checker.c b/src/backend/distributed/planner/query_colocation_checker.c index bef91618e..6c34d01a8 100644 --- a/src/backend/distributed/planner/query_colocation_checker.c +++ b/src/backend/distributed/planner/query_colocation_checker.c @@ -325,13 +325,14 @@ WrapRteRelationIntoSubquery(RangeTblEntry *rteRelation, * as a NULL column. */ List * -CreateAllTargetListForRelation(Oid relationId, List *requiredAttributes) +CreateAllTargetListForRelation(Oid relationId, List *requiredAttributes, bool + isMergeQuery) { Relation relation = relation_open(relationId, AccessShareLock); int numberOfAttributes = RelationGetNumberOfAttributes(relation); List *targetList = NIL; - int varAttrNo = 1; + int colAppendIdx = 1; for (int attrNum = 1; attrNum <= numberOfAttributes; attrNum++) { @@ -361,8 +362,9 @@ CreateAllTargetListForRelation(Oid relationId, List *requiredAttributes) } else { + int varAttNum = isMergeQuery ? attrNum : colAppendIdx++; TargetEntry *targetEntry = - CreateTargetEntryForColumn(attributeTuple, SINGLE_RTE_INDEX, varAttrNo++, + CreateTargetEntryForColumn(attributeTuple, SINGLE_RTE_INDEX, varAttNum, resNo); targetList = lappend(targetList, targetEntry); } diff --git a/src/backend/distributed/planner/recursive_planning.c b/src/backend/distributed/planner/recursive_planning.c index 9f520fa5f..98f28efb9 100644 --- a/src/backend/distributed/planner/recursive_planning.c +++ b/src/backend/distributed/planner/recursive_planning.c @@ -1767,8 +1767,11 @@ ReplaceRTERelationWithRteSubquery(RangeTblEntry *rangeTableEntry, { Query *subquery = WrapRteRelationIntoSubquery(rangeTableEntry, requiredAttrNumbers, perminfo); + + bool isMergeQuery = false; List *outerQueryTargetList = CreateAllTargetListForRelation(rangeTableEntry->relid, - requiredAttrNumbers); + requiredAttrNumbers, + isMergeQuery); List *restrictionList = GetRestrictInfoListForRelation(rangeTableEntry, @@ -2109,7 +2112,6 @@ TransformFunctionRTE(RangeTblEntry *rangeTblEntry) subquery->targetList = lappend(subquery->targetList, targetEntry); } } - /* * If tupleDesc is NULL we have 2 different cases: * @@ -2159,7 +2161,6 @@ TransformFunctionRTE(RangeTblEntry *rangeTblEntry) columnType = list_nth_oid(rangeTblFunction->funccoltypes, targetColumnIndex); } - /* use the types in the function definition otherwise */ else { diff --git a/src/include/distributed/query_colocation_checker.h b/src/include/distributed/query_colocation_checker.h index 2a46d364c..ae6062e27 100644 --- a/src/include/distributed/query_colocation_checker.h +++ b/src/include/distributed/query_colocation_checker.h @@ -38,6 +38,7 @@ extern bool SubqueryColocated(Query *subquery, ColocatedJoinChecker *context); extern Query * WrapRteRelationIntoSubquery(RangeTblEntry *rteRelation, List *requiredAttributes, RTEPermissionInfo *perminfo); -extern List * CreateAllTargetListForRelation(Oid relationId, List *requiredAttributes); +extern List * CreateAllTargetListForRelation(Oid relationId, List *requiredAttributes, + bool isMergeQuery); #endif /* QUERY_COLOCATION_CHECKER_H */ diff --git a/src/test/regress/sql/merge_vcore.sql b/src/test/regress/sql/merge_vcore.sql index 472bbfe91..88f478ea7 100644 --- a/src/test/regress/sql/merge_vcore.sql +++ b/src/test/regress/sql/merge_vcore.sql @@ -312,6 +312,63 @@ WHEN MATCHED THEN DO NOTHING; DROP TABLE IF EXISTS source; DROP TABLE IF EXISTS target; + +-- Bug Fix Test as part of this PR +-- Test 1 +CREATE TABLE source ( + id int, + age int, + salary int +); + +CREATE TABLE target ( + id int, + age int, + salary int +); + +SELECT create_distributed_table('source', 'id', colocate_with=>'none'); +SELECT create_distributed_table('target', 'id', colocate_with=>'none'); + +INSERT INTO source (id, age, salary) VALUES (1,30, 100000); + +MERGE INTO ONLY target USING source ON (source.id = target.id) +WHEN NOT MATCHED THEN +INSERT (id, salary) VALUES (source.id, source.salary); + +SELECT * FROM TARGET; +DROP TABLE IF EXISTS source; +DROP TABLE IF EXISTS target; + + +-- Test 2 +CREATE TABLE source ( + id int, + age int, + salary int +); + +CREATE TABLE target ( + id int, + age int, + salary int +); + +SELECT create_distributed_table('source', 'id', colocate_with=>'none'); +SELECT create_distributed_table('target', 'id', colocate_with=>'none'); + +INSERT INTO source (id, age, salary) VALUES (1,30, 100000); + +MERGE INTO ONLY target USING source ON (source.id = target.id) +WHEN NOT MATCHED THEN +INSERT (salary, id) VALUES (source.salary, source.id); + +SELECT * FROM TARGET; +DROP TABLE IF EXISTS source; +DROP TABLE IF EXISTS target; + + + DROP SCHEMA IF EXISTS merge_vcore_schema CASCADE;