From 9faa564525a22d7622f3023ce2a93d4f87dac1a3 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 29 Jun 2026 11:49:54 +0300 Subject: [PATCH 01/25] Add citus.enable_shard_local_batching for colocated INSERT..SELECT batching pushdown Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../planner/insert_select_planner.c | 28 +++- .../distributed/planner/merge_planner.c | 2 +- .../planner/multi_router_planner.c | 2 +- .../planner/query_pushdown_planning.c | 154 ++++++++++-------- src/backend/distributed/shared_library_init.c | 16 ++ .../distributed/query_pushdown_planning.h | 8 +- 6 files changed, 134 insertions(+), 76 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 3191615d663..e6c2c46e342 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -752,7 +752,8 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, } } - if (FindNodeMatchingCheckFunction((Node *) queryTree, CitusIsVolatileFunction)) + if (!EnableShardLocalBatching && + FindNodeMatchingCheckFunction((Node *) queryTree, CitusIsVolatileFunction)) { return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, "volatile functions are not allowed in distributed " @@ -771,14 +772,15 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, /* first apply toplevel pushdown checks to SELECT query */ error = DeferErrorIfUnsupportedSubqueryPushdown(subquery, plannerRestrictionContext, - true); + true, EnableShardLocalBatching); if (error) { return error; } /* then apply subquery pushdown checks to SELECT query */ - error = DeferErrorIfCannotPushdownSubquery(subquery, false); + error = DeferErrorIfCannotPushdownSubquery(subquery, false, + EnableShardLocalBatching); if (error) { return error; @@ -823,12 +825,22 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, if (HasDistributionKey(targetRelationId)) { - /* ensure that INSERT's partition column comes from SELECT's partition column */ - error = InsertPartitionColumnMatchesSelect(queryTree, insertRte, subqueryRte, - &selectPartitionColumnTableId); - if (error) + /* + * ensure that INSERT's partition column comes from SELECT's partition + * column. With shard-local batching enabled the distribution value is + * typically derived (e.g. unnest(array_agg(text_id))), so we cannot + * match it to a plain SELECT partition column; colocation enforced + * below keeps batches shard-local. + */ + if (!EnableShardLocalBatching) { - return error; + error = InsertPartitionColumnMatchesSelect(queryTree, insertRte, + subqueryRte, + &selectPartitionColumnTableId); + if (error) + { + return error; + } } } } diff --git a/src/backend/distributed/planner/merge_planner.c b/src/backend/distributed/planner/merge_planner.c index a05eef493d8..95a5217ac29 100644 --- a/src/backend/distributed/planner/merge_planner.c +++ b/src/backend/distributed/planner/merge_planner.c @@ -1166,7 +1166,7 @@ DeferErrorIfRoutableMergeNotSupported(Query *query, List *rangeTableList, deferredError = DeferErrorIfUnsupportedSubqueryPushdown(query, plannerRestrictionContext, - true); + true, false); if (deferredError) { ereport(DEBUG1, (errmsg("Sub-query is not pushable, try repartitioning"))); diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index 313d2674ddc..693fd678610 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -1342,7 +1342,7 @@ MultiShardUpdateDeleteSupported(Query *originalQuery, errorMessage = DeferErrorIfUnsupportedSubqueryPushdown( originalQuery, plannerRestrictionContext, - true); + true, false); } return errorMessage; diff --git a/src/backend/distributed/planner/query_pushdown_planning.c b/src/backend/distributed/planner/query_pushdown_planning.c index 753643b1929..2ddcda1718e 100644 --- a/src/backend/distributed/planner/query_pushdown_planning.c +++ b/src/backend/distributed/planner/query_pushdown_planning.c @@ -81,6 +81,16 @@ typedef struct RelidsReferenceWalkerContext bool SubqueryPushdown = false; /* is subquery pushdown enabled */ int ValuesMaterializationThreshold = 100; +/* + * Allows shard-local batching to be pushed down for colocated INSERT ... SELECT. + * When enabled, GROUP BY / window / aggregate-without-group-by on non-distribution + * columns and volatile functions are permitted in colocated INSERT ... SELECT so + * that batching (and any batch UDF call) executes on the shards instead of pulling + * data to the coordinator. The user takes responsibility for keeping batches + * shard-local; colocation is still enforced. Off by default. + */ +bool EnableShardLocalBatching = false; + /* Local functions forward declarations */ static bool JoinTreeContainsSubqueryWalker(Node *joinTreeNode, void *context); static bool IsFunctionOrValuesRTE(Node *node); @@ -92,7 +102,8 @@ static DeferredErrorMessage * DeferredErrorIfUnsupportedRecurringTuplesJoin( static DeferredErrorMessage * DeferErrorIfUnsupportedTableCombination(Query *queryTree); static DeferredErrorMessage * DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, - char *referencedThing); + char *referencedThing, + bool shardLocalBatching); static bool ExtractSetOperationStatementWalker(Node *node, List **setOperationList); static RecurringTuplesType FetchFirstRecurType(PlannerInfo *plannerInfo, Relids relids); @@ -550,7 +561,7 @@ SubqueryMultiNodeTree(Query *originalQuery, Query *queryTree, DeferredErrorMessage *subqueryPushdownError = DeferErrorIfUnsupportedSubqueryPushdown( originalQuery, plannerRestrictionContext, - false); + false, false); if (subqueryPushdownError != NULL) { @@ -575,7 +586,7 @@ DeferredErrorMessage * DeferErrorIfUnsupportedSubqueryPushdown(Query *originalQuery, PlannerRestrictionContext * plannerRestrictionContext, - bool plannerPhase) + bool plannerPhase, bool shardLocalBatching) { bool outerMostQueryHasLimit = false; ListCell *subqueryCell = NULL; @@ -648,7 +659,8 @@ DeferErrorIfUnsupportedSubqueryPushdown(Query *originalQuery, { Query *subquery = lfirst(subqueryCell); error = DeferErrorIfCannotPushdownSubquery(subquery, - outerMostQueryHasLimit); + outerMostQueryHasLimit, + shardLocalBatching); if (error) { return error; @@ -970,8 +982,8 @@ DeferredErrorIfUnsupportedRecurringTuplesJoin(PlannerRestrictionContext * bool CanPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit) { - return DeferErrorIfCannotPushdownSubquery(subqueryTree, outerMostQueryHasLimit) == - NULL; + return DeferErrorIfCannotPushdownSubquery(subqueryTree, outerMostQueryHasLimit, + false) == NULL; } @@ -998,7 +1010,8 @@ CanPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit) * features of underlying tables. */ DeferredErrorMessage * -DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit) +DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit, + bool shardLocalBatching) { bool preconditionsSatisfied = true; char *errorDetail = NULL; @@ -1027,7 +1040,8 @@ DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLi if (!ContainsReferencesToOuterQuery(subqueryTree)) { deferredError = DeferErrorIfSubqueryRequiresMerge(subqueryTree, false, - "another query"); + "another query", + shardLocalBatching); if (deferredError) { return deferredError; @@ -1128,7 +1142,7 @@ FlattenGroupExprs(Query *queryTree) */ static DeferredErrorMessage * DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, - char *referencedThing) + char *referencedThing, bool shardLocalBatching) { bool preconditionsSatisfied = true; char *errorDetail = NULL; @@ -1152,68 +1166,80 @@ DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, referencedThing); } - /* group clause list must include partition column */ - if (subqueryTree->groupClause) + /* + * With shard-local batching the entire subquery executes on a single shard + * (colocation is enforced separately), so grouping / aggregation / window / + * distinct on non-distribution columns stays shard-local and does not require + * a merge step. Skip the partition-column requirements in that case. + */ + if (!shardLocalBatching) { - List *groupClauseList = subqueryTree->groupClause; - List *targetEntryList = subqueryTree->targetList; - List *groupTargetEntryList = GroupTargetEntryList(groupClauseList, - targetEntryList); - bool groupOnPartitionColumn = - TargetListOnPartitionColumn(subqueryTree, groupTargetEntryList); - if (!groupOnPartitionColumn) + /* group clause list must include partition column */ + if (subqueryTree->groupClause) { - preconditionsSatisfied = false; - errorDetail = psprintf("Group by list without partition column is currently " - "unsupported when a %ssubquery references a column " - "from %s", lateralString, referencedThing); + List *groupClauseList = subqueryTree->groupClause; + List *targetEntryList = subqueryTree->targetList; + List *groupTargetEntryList = GroupTargetEntryList(groupClauseList, + targetEntryList); + bool groupOnPartitionColumn = + TargetListOnPartitionColumn(subqueryTree, groupTargetEntryList); + if (!groupOnPartitionColumn) + { + preconditionsSatisfied = false; + errorDetail = psprintf( + "Group by list without partition column is currently " + "unsupported when a %ssubquery references a column " + "from %s", lateralString, referencedThing); + } } - } - /* we don't support aggregates without group by */ - if (subqueryTree->hasAggs && (subqueryTree->groupClause == NULL)) - { - preconditionsSatisfied = false; - errorDetail = psprintf("Aggregates without group by are currently unsupported " - "when a %ssubquery references a column from %s", - lateralString, referencedThing); - } - - /* having clause without group by on partition column is not supported */ - if (subqueryTree->havingQual && (subqueryTree->groupClause == NULL)) - { - preconditionsSatisfied = false; - errorDetail = psprintf("Having qual without group by on partition column is " - "currently unsupported when a %ssubquery references " - "a column from %s", lateralString, referencedThing); - } + /* we don't support aggregates without group by */ + if (subqueryTree->hasAggs && (subqueryTree->groupClause == NULL)) + { + preconditionsSatisfied = false; + errorDetail = psprintf( + "Aggregates without group by are currently unsupported " + "when a %ssubquery references a column from %s", + lateralString, referencedThing); + } - /* - * We support window functions when the window function - * is partitioned on distribution column. - */ - StringInfo errorInfo = NULL; - if (subqueryTree->hasWindowFuncs && !SafeToPushdownWindowFunction(subqueryTree, - &errorInfo)) - { - errorDetail = (char *) errorInfo->data; - preconditionsSatisfied = false; - } + /* having clause without group by on partition column is not supported */ + if (subqueryTree->havingQual && (subqueryTree->groupClause == NULL)) + { + preconditionsSatisfied = false; + errorDetail = psprintf( + "Having qual without group by on partition column is " + "currently unsupported when a %ssubquery references " + "a column from %s", lateralString, referencedThing); + } - /* distinct clause list must include partition column */ - if (subqueryTree->distinctClause) - { - List *distinctClauseList = subqueryTree->distinctClause; - List *targetEntryList = subqueryTree->targetList; - List *distinctTargetEntryList = GroupTargetEntryList(distinctClauseList, - targetEntryList); - bool distinctOnPartitionColumn = - TargetListOnPartitionColumn(subqueryTree, distinctTargetEntryList); - if (!distinctOnPartitionColumn) + /* + * We support window functions when the window function + * is partitioned on distribution column. + */ + StringInfo errorInfo = NULL; + if (subqueryTree->hasWindowFuncs && !SafeToPushdownWindowFunction(subqueryTree, + &errorInfo)) { + errorDetail = (char *) errorInfo->data; preconditionsSatisfied = false; - errorDetail = "Distinct on columns without partition column is " - "currently unsupported"; + } + + /* distinct clause list must include partition column */ + if (subqueryTree->distinctClause) + { + List *distinctClauseList = subqueryTree->distinctClause; + List *targetEntryList = subqueryTree->targetList; + List *distinctTargetEntryList = GroupTargetEntryList(distinctClauseList, + targetEntryList); + bool distinctOnPartitionColumn = + TargetListOnPartitionColumn(subqueryTree, distinctTargetEntryList); + if (!distinctOnPartitionColumn) + { + preconditionsSatisfied = false; + errorDetail = "Distinct on columns without partition column is " + "currently unsupported"; + } } } @@ -1763,7 +1789,7 @@ DeferredErrorIfUnsupportedLateralSubquery(PlannerInfo *plannerInfo, /* property number 3, has a merge step */ DeferredErrorMessage *deferredError = DeferErrorIfSubqueryRequiresMerge( - rangeTableEntry->subquery, true, recurTypeDescription); + rangeTableEntry->subquery, true, recurTypeDescription, false); if (deferredError) { return deferredError; diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c index 9ea35038f8e..d5094e015bb 100644 --- a/src/backend/distributed/shared_library_init.c +++ b/src/backend/distributed/shared_library_init.c @@ -987,6 +987,22 @@ RegisterCitusConfigVariables(void) GUC_STANDARD, NULL, NULL, NULL); + DefineCustomBoolVariable( + "citus.enable_shard_local_batching", + gettext_noop("Enables shard-local batching pushdown for colocated " + "INSERT ... SELECT."), + gettext_noop("When enabled, Citus allows GROUP BY / window / aggregate " + "constructs on non-distribution columns as well as volatile " + "functions in colocated INSERT ... SELECT, so that batching and " + "any batch UDF call run on the shards instead of pulling data to " + "the coordinator. Colocation is still enforced; the user takes " + "responsibility for keeping batches shard-local."), + &EnableShardLocalBatching, + false, + PGC_USERSET, + GUC_STANDARD, + NULL, NULL, NULL); + DefineCustomBoolVariable( "citus.allow_aggregate_worker_combine_on_internal_types", gettext_noop("Enables aggregate worker partial aggregates on aggregates that " diff --git a/src/include/distributed/query_pushdown_planning.h b/src/include/distributed/query_pushdown_planning.h index 0b69d36c75f..aaca0bcf729 100644 --- a/src/include/distributed/query_pushdown_planning.h +++ b/src/include/distributed/query_pushdown_planning.h @@ -22,6 +22,7 @@ /* Config variables managed via guc.c */ extern bool SubqueryPushdown; extern int ValuesMaterializationThreshold; +extern bool EnableShardLocalBatching; extern bool AllowAggregateWorkerCombineOnInternalTypes; @@ -44,10 +45,13 @@ extern DeferredErrorMessage * DeferErrorIfUnsupportedSubqueryPushdown(Query * PlannerRestrictionContext * plannerRestrictionContext, - bool plannerPhase); + bool plannerPhase, + bool shardLocalBatching); extern DeferredErrorMessage * DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool - outerMostQueryHasLimit); + outerMostQueryHasLimit, + bool + shardLocalBatching); extern DeferredErrorMessage * DeferErrorIfUnsupportedUnionQuery(Query *queryTree); extern bool IsJsonTableRTE(RangeTblEntry *rte); extern bool IsOuterJoinExpr(Node *node); From c889c53abf37bfd8342d28b243d3e4a7076b9d4d Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 29 Jun 2026 12:04:16 +0300 Subject: [PATCH 02/25] No-op partition NOT NULL filter for shard-local batching Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../distributed/planner/multi_router_planner.c | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index 693fd678610..07383de112f 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -369,6 +369,17 @@ AddPartitionKeyNotNullFilterToSelect(Query *subqery) } } + /* + * Normally the SELECT projects the distribution column as a plain Var. + * With shard-local batching the distribution column may be derived (e.g. + * unnest(array_agg(text_id))), so there is no Var to attach a NOT NULL + * filter to; the batch stays shard-local so the filter is unnecessary. + */ + if (targetPartitionColumnVar == NULL && EnableShardLocalBatching) + { + return; + } + /* we should have found target partition column */ Assert(targetPartitionColumnVar != NULL); From 0595921785b70eb1003392dfbd2ce84cac029f9c Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 29 Jun 2026 12:12:17 +0300 Subject: [PATCH 03/25] Add shard_local_batching regression test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/test/regress/multi_schedule | 1 + src/test/regress/sql/shard_local_batching.sql | 59 +++++++++++++++++++ 2 files changed, 60 insertions(+) create mode 100644 src/test/regress/sql/shard_local_batching.sql diff --git a/src/test/regress/multi_schedule b/src/test/regress/multi_schedule index 49891fcdb37..ef4a16f0c26 100644 --- a/src/test/regress/multi_schedule +++ b/src/test/regress/multi_schedule @@ -19,6 +19,7 @@ test: multi_behavioral_analytics_basics multi_behavioral_analytics_single_shard_ # We don't parallelize the following test with the ones above because they're # not idempotent and hence causing flaky test detection check to fail. test: multi_insert_select_non_pushable_queries multi_insert_select +test: shard_local_batching test: multi_shard_update_delete recursive_dml_with_different_planners_executors test: insert_select_repartition window_functions dml_recursive multi_insert_select_window diff --git a/src/test/regress/sql/shard_local_batching.sql b/src/test/regress/sql/shard_local_batching.sql new file mode 100644 index 00000000000..454bcf3a535 --- /dev/null +++ b/src/test/regress/sql/shard_local_batching.sql @@ -0,0 +1,59 @@ +-- +-- SHARD_LOCAL_BATCHING +-- +-- Tests citus.enable_shard_local_batching, which lets a colocated INSERT .. SELECT +-- push down GROUP BY / window batching and a volatile batch UDF to the shards +-- instead of pulling rows to the coordinator. +-- +CREATE SCHEMA shard_local_batching; +SET search_path = shard_local_batching; +SET citus.next_shard_id TO 14000000; +SET citus.shard_count = 4; +SET citus.shard_replication_factor = 1; + +CREATE TABLE dist(text_id int, text_col text); +CREATE TABLE emb(text_id int, embedding int); +SELECT create_distributed_table('dist', 'text_id'); +SELECT create_distributed_table('emb', 'text_id'); + +INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; + +-- a volatile batch UDF: returns one value per input, mimicking a batched API call +CREATE FUNCTION create_embeddings(t text[]) RETURNS int[] +LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; +SELECT create_distributed_function('create_embeddings(text[])'); + +-- default off: batching is done after pulling rows to the coordinator +EXPLAIN (COSTS OFF) INSERT INTO emb(text_id, embedding) +SELECT id, vec FROM ( + SELECT b, unnest(array_agg(text_id)) id, + unnest(create_embeddings(array_agg(text_col))) vec + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s; + +SET citus.enable_shard_local_batching TO on; + +-- now the batching and UDF call run on the shards +EXPLAIN (COSTS OFF) INSERT INTO emb(text_id, embedding) +SELECT id, vec FROM ( + SELECT b, unnest(array_agg(text_id)) id, + unnest(create_embeddings(array_agg(text_col))) vec + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s; + +INSERT INTO emb(text_id, embedding) +SELECT id, vec FROM ( + SELECT b, unnest(array_agg(text_id)) id, + unnest(create_embeddings(array_agg(text_col))) vec + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s; + +-- every text_id should be matched to the right vector +SELECT count(*), count(*) FILTER (WHERE embedding = length('t' || text_id)) AS ok +FROM emb JOIN dist USING (text_id); + +SET client_min_messages TO WARNING; +DROP SCHEMA shard_local_batching CASCADE; From b7537685bc8f94afdb6fe8200ec258d27b8e25ae Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 29 Jun 2026 12:26:37 +0300 Subject: [PATCH 04/25] Add expected output for shard_local_batching Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../regress/expected/shard_local_batching.out | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 src/test/regress/expected/shard_local_batching.out diff --git a/src/test/regress/expected/shard_local_batching.out b/src/test/regress/expected/shard_local_batching.out new file mode 100644 index 00000000000..ddb5ae15fb8 --- /dev/null +++ b/src/test/regress/expected/shard_local_batching.out @@ -0,0 +1,112 @@ +-- +-- SHARD_LOCAL_BATCHING +-- +-- Tests citus.enable_shard_local_batching, which lets a colocated INSERT .. SELECT +-- push down GROUP BY / window batching and a volatile batch UDF to the shards +-- instead of pulling rows to the coordinator. +-- +CREATE SCHEMA shard_local_batching; +SET search_path = shard_local_batching; +SET citus.next_shard_id TO 14000000; +SET citus.shard_count = 4; +SET citus.shard_replication_factor = 1; +CREATE TABLE dist(text_id int, text_col text); +CREATE TABLE emb(text_id int, embedding int); +SELECT create_distributed_table('dist', 'text_id'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +SELECT create_distributed_table('emb', 'text_id'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; +-- a volatile batch UDF: returns one value per input, mimicking a batched API call +CREATE FUNCTION create_embeddings(t text[]) RETURNS int[] +LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; +SELECT create_distributed_function('create_embeddings(text[])'); +NOTICE: procedure shard_local_batching.create_embeddings is already distributed +DETAIL: Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] commands + create_distributed_function +--------------------------------------------------------------------- + +(1 row) + +-- default off: batching is done after pulling rows to the coordinator +EXPLAIN (COSTS OFF) INSERT INTO emb(text_id, embedding) +SELECT id, vec FROM ( + SELECT b, unnest(array_agg(text_id)) id, + unnest(create_embeddings(array_agg(text_col))) vec + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s; + QUERY PLAN +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> WindowAgg + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Subquery Scan on s + -> ProjectSet + -> HashAggregate + Group Key: intermediate_result.b + -> Function Scan on read_intermediate_result intermediate_result +(20 rows) + +SET citus.enable_shard_local_batching TO on; +-- now the batching and UDF call run on the shards +EXPLAIN (COSTS OFF) INSERT INTO emb(text_id, embedding) +SELECT id, vec FROM ( + SELECT b, unnest(array_agg(text_id)) id, + unnest(create_embeddings(array_agg(text_col))) vec + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s; + QUERY PLAN +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on emb_14000004 + -> Subquery Scan on s + -> ProjectSet + -> HashAggregate + Group Key: ((row_number() OVER (?) - 1) / 100) + -> WindowAgg + -> Seq Scan on dist_14000000 dist +(12 rows) + +INSERT INTO emb(text_id, embedding) +SELECT id, vec FROM ( + SELECT b, unnest(array_agg(text_id)) id, + unnest(create_embeddings(array_agg(text_col))) vec + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s; +-- every text_id should be matched to the right vector +SELECT count(*), count(*) FILTER (WHERE embedding = length('t' || text_id)) AS ok +FROM emb JOIN dist USING (text_id); + count | ok +--------------------------------------------------------------------- + 500 | 500 +(1 row) + +SET client_min_messages TO WARNING; +DROP SCHEMA shard_local_batching CASCADE; From d9aa5320bc2aa9f34a3bec7a8544ebabab25c1ef Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 29 Jun 2026 14:50:45 +0300 Subject: [PATCH 05/25] Rename GUC to citus.allow_unsafe_insert_select_pushdown Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../planner/insert_select_planner.c | 17 ++++++----- .../planner/multi_router_planner.c | 9 +++--- .../planner/query_pushdown_planning.c | 29 ++++++++++--------- src/backend/distributed/shared_library_init.c | 22 +++++++------- .../distributed/query_pushdown_planning.h | 6 ++-- ...> allow_unsafe_insert_select_pushdown.out} | 18 ++++++------ src/test/regress/multi_schedule | 2 +- ...> allow_unsafe_insert_select_pushdown.sql} | 16 +++++----- 8 files changed, 62 insertions(+), 57 deletions(-) rename src/test/regress/expected/{shard_local_batching.out => allow_unsafe_insert_select_pushdown.out} (87%) rename src/test/regress/sql/{shard_local_batching.sql => allow_unsafe_insert_select_pushdown.sql} (79%) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index e6c2c46e342..5f615a3e5d9 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -752,7 +752,7 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, } } - if (!EnableShardLocalBatching && + if (!AllowUnsafeInsertSelectPushdown && FindNodeMatchingCheckFunction((Node *) queryTree, CitusIsVolatileFunction)) { return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, @@ -772,7 +772,7 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, /* first apply toplevel pushdown checks to SELECT query */ error = DeferErrorIfUnsupportedSubqueryPushdown(subquery, plannerRestrictionContext, - true, EnableShardLocalBatching); + true, AllowUnsafeInsertSelectPushdown); if (error) { return error; @@ -780,7 +780,7 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, /* then apply subquery pushdown checks to SELECT query */ error = DeferErrorIfCannotPushdownSubquery(subquery, false, - EnableShardLocalBatching); + AllowUnsafeInsertSelectPushdown); if (error) { return error; @@ -827,12 +827,13 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, { /* * ensure that INSERT's partition column comes from SELECT's partition - * column. With shard-local batching enabled the distribution value is - * typically derived (e.g. unnest(array_agg(text_id))), so we cannot - * match it to a plain SELECT partition column; colocation enforced - * below keeps batches shard-local. + * column. With unsafe INSERT ... SELECT pushdown enabled the + * distribution value is typically derived (e.g. + * unnest(array_agg(text_id))), so we cannot match it to a plain SELECT + * partition column; colocation enforced below keeps batches + * shard-local. */ - if (!EnableShardLocalBatching) + if (!AllowUnsafeInsertSelectPushdown) { error = InsertPartitionColumnMatchesSelect(queryTree, insertRte, subqueryRte, diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index 07383de112f..6c05edc9d1f 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -371,11 +371,12 @@ AddPartitionKeyNotNullFilterToSelect(Query *subqery) /* * Normally the SELECT projects the distribution column as a plain Var. - * With shard-local batching the distribution column may be derived (e.g. - * unnest(array_agg(text_id))), so there is no Var to attach a NOT NULL - * filter to; the batch stays shard-local so the filter is unnecessary. + * With unsafe INSERT ... SELECT pushdown the distribution column may be + * derived (e.g. unnest(array_agg(text_id))), so there is no Var to attach a + * NOT NULL filter to; the batch stays shard-local so the filter is + * unnecessary. */ - if (targetPartitionColumnVar == NULL && EnableShardLocalBatching) + if (targetPartitionColumnVar == NULL && AllowUnsafeInsertSelectPushdown) { return; } diff --git a/src/backend/distributed/planner/query_pushdown_planning.c b/src/backend/distributed/planner/query_pushdown_planning.c index 2ddcda1718e..c7efe9af4db 100644 --- a/src/backend/distributed/planner/query_pushdown_planning.c +++ b/src/backend/distributed/planner/query_pushdown_planning.c @@ -82,14 +82,14 @@ bool SubqueryPushdown = false; /* is subquery pushdown enabled */ int ValuesMaterializationThreshold = 100; /* - * Allows shard-local batching to be pushed down for colocated INSERT ... SELECT. + * Allows otherwise-unsafe colocated INSERT ... SELECT queries to be pushed down. * When enabled, GROUP BY / window / aggregate-without-group-by on non-distribution * columns and volatile functions are permitted in colocated INSERT ... SELECT so * that batching (and any batch UDF call) executes on the shards instead of pulling * data to the coordinator. The user takes responsibility for keeping batches - * shard-local; colocation is still enforced. Off by default. + * shard-local and order-preserving; colocation is still enforced. Off by default. */ -bool EnableShardLocalBatching = false; +bool AllowUnsafeInsertSelectPushdown = false; /* Local functions forward declarations */ static bool JoinTreeContainsSubqueryWalker(Node *joinTreeNode, void *context); @@ -103,7 +103,7 @@ static DeferredErrorMessage * DeferErrorIfUnsupportedTableCombination(Query *que static DeferredErrorMessage * DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, char *referencedThing, - bool shardLocalBatching); + bool allowUnsafePushdown); static bool ExtractSetOperationStatementWalker(Node *node, List **setOperationList); static RecurringTuplesType FetchFirstRecurType(PlannerInfo *plannerInfo, Relids relids); @@ -586,7 +586,7 @@ DeferredErrorMessage * DeferErrorIfUnsupportedSubqueryPushdown(Query *originalQuery, PlannerRestrictionContext * plannerRestrictionContext, - bool plannerPhase, bool shardLocalBatching) + bool plannerPhase, bool allowUnsafePushdown) { bool outerMostQueryHasLimit = false; ListCell *subqueryCell = NULL; @@ -660,7 +660,7 @@ DeferErrorIfUnsupportedSubqueryPushdown(Query *originalQuery, Query *subquery = lfirst(subqueryCell); error = DeferErrorIfCannotPushdownSubquery(subquery, outerMostQueryHasLimit, - shardLocalBatching); + allowUnsafePushdown); if (error) { return error; @@ -1011,7 +1011,7 @@ CanPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit) */ DeferredErrorMessage * DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit, - bool shardLocalBatching) + bool allowUnsafePushdown) { bool preconditionsSatisfied = true; char *errorDetail = NULL; @@ -1041,7 +1041,7 @@ DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLi { deferredError = DeferErrorIfSubqueryRequiresMerge(subqueryTree, false, "another query", - shardLocalBatching); + allowUnsafePushdown); if (deferredError) { return deferredError; @@ -1142,7 +1142,7 @@ FlattenGroupExprs(Query *queryTree) */ static DeferredErrorMessage * DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, - char *referencedThing, bool shardLocalBatching) + char *referencedThing, bool allowUnsafePushdown) { bool preconditionsSatisfied = true; char *errorDetail = NULL; @@ -1167,12 +1167,13 @@ DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, } /* - * With shard-local batching the entire subquery executes on a single shard - * (colocation is enforced separately), so grouping / aggregation / window / - * distinct on non-distribution columns stays shard-local and does not require - * a merge step. Skip the partition-column requirements in that case. + * With unsafe INSERT ... SELECT pushdown the entire subquery executes on a + * single shard (colocation is enforced separately), so grouping / aggregation + * / window / distinct on non-distribution columns stays shard-local and does + * not require a merge step. Skip the partition-column requirements in that + * case. */ - if (!shardLocalBatching) + if (!allowUnsafePushdown) { /* group clause list must include partition column */ if (subqueryTree->groupClause) diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c index d5094e015bb..d3d0288d49f 100644 --- a/src/backend/distributed/shared_library_init.c +++ b/src/backend/distributed/shared_library_init.c @@ -988,16 +988,18 @@ RegisterCitusConfigVariables(void) NULL, NULL, NULL); DefineCustomBoolVariable( - "citus.enable_shard_local_batching", - gettext_noop("Enables shard-local batching pushdown for colocated " - "INSERT ... SELECT."), - gettext_noop("When enabled, Citus allows GROUP BY / window / aggregate " - "constructs on non-distribution columns as well as volatile " - "functions in colocated INSERT ... SELECT, so that batching and " - "any batch UDF call run on the shards instead of pulling data to " - "the coordinator. Colocation is still enforced; the user takes " - "responsibility for keeping batches shard-local."), - &EnableShardLocalBatching, + "citus.allow_unsafe_insert_select_pushdown", + gettext_noop("Allows pushdown of otherwise-unsafe colocated " + "INSERT ... SELECT queries."), + gettext_noop("When enabled, Citus relaxes safety checks (GROUP BY / window / " + "aggregate constructs on non-distribution columns, volatile " + "functions, partition-column matching) for colocated " + "INSERT ... SELECT, so that batching and any batch UDF call run " + "on the shards instead of pulling data to the coordinator. " + "Colocation is still enforced, but the user takes responsibility " + "for keeping batches shard-local and order-preserving; otherwise " + "results may be silently incorrect."), + &AllowUnsafeInsertSelectPushdown, false, PGC_USERSET, GUC_STANDARD, diff --git a/src/include/distributed/query_pushdown_planning.h b/src/include/distributed/query_pushdown_planning.h index aaca0bcf729..bc32a2ea176 100644 --- a/src/include/distributed/query_pushdown_planning.h +++ b/src/include/distributed/query_pushdown_planning.h @@ -22,7 +22,7 @@ /* Config variables managed via guc.c */ extern bool SubqueryPushdown; extern int ValuesMaterializationThreshold; -extern bool EnableShardLocalBatching; +extern bool AllowUnsafeInsertSelectPushdown; extern bool AllowAggregateWorkerCombineOnInternalTypes; @@ -46,12 +46,12 @@ extern DeferredErrorMessage * DeferErrorIfUnsupportedSubqueryPushdown(Query * * plannerRestrictionContext, bool plannerPhase, - bool shardLocalBatching); + bool allowUnsafePushdown); extern DeferredErrorMessage * DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit, bool - shardLocalBatching); + allowUnsafePushdown); extern DeferredErrorMessage * DeferErrorIfUnsupportedUnionQuery(Query *queryTree); extern bool IsJsonTableRTE(RangeTblEntry *rte); extern bool IsOuterJoinExpr(Node *node); diff --git a/src/test/regress/expected/shard_local_batching.out b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out similarity index 87% rename from src/test/regress/expected/shard_local_batching.out rename to src/test/regress/expected/allow_unsafe_insert_select_pushdown.out index ddb5ae15fb8..a7c2608b077 100644 --- a/src/test/regress/expected/shard_local_batching.out +++ b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out @@ -1,12 +1,12 @@ -- --- SHARD_LOCAL_BATCHING +-- ALLOW_UNSAFE_INSERT_SELECT_PUSHDOWN -- --- Tests citus.enable_shard_local_batching, which lets a colocated INSERT .. SELECT --- push down GROUP BY / window batching and a volatile batch UDF to the shards --- instead of pulling rows to the coordinator. +-- Tests citus.allow_unsafe_insert_select_pushdown, which lets a colocated +-- INSERT .. SELECT push down GROUP BY / window batching and a volatile batch UDF +-- to the shards instead of pulling rows to the coordinator. -- -CREATE SCHEMA shard_local_batching; -SET search_path = shard_local_batching; +CREATE SCHEMA allow_unsafe_insert_select_pushdown; +SET search_path = allow_unsafe_insert_select_pushdown; SET citus.next_shard_id TO 14000000; SET citus.shard_count = 4; SET citus.shard_replication_factor = 1; @@ -29,7 +29,7 @@ INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; CREATE FUNCTION create_embeddings(t text[]) RETURNS int[] LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; SELECT create_distributed_function('create_embeddings(text[])'); -NOTICE: procedure shard_local_batching.create_embeddings is already distributed +NOTICE: procedure allow_unsafe_insert_select_pushdown.create_embeddings is already distributed DETAIL: Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] commands create_distributed_function --------------------------------------------------------------------- @@ -68,7 +68,7 @@ SELECT id, vec FROM ( -> Function Scan on read_intermediate_result intermediate_result (20 rows) -SET citus.enable_shard_local_batching TO on; +SET citus.allow_unsafe_insert_select_pushdown TO on; -- now the batching and UDF call run on the shards EXPLAIN (COSTS OFF) INSERT INTO emb(text_id, embedding) SELECT id, vec FROM ( @@ -109,4 +109,4 @@ FROM emb JOIN dist USING (text_id); (1 row) SET client_min_messages TO WARNING; -DROP SCHEMA shard_local_batching CASCADE; +DROP SCHEMA allow_unsafe_insert_select_pushdown CASCADE; diff --git a/src/test/regress/multi_schedule b/src/test/regress/multi_schedule index ef4a16f0c26..0960831df83 100644 --- a/src/test/regress/multi_schedule +++ b/src/test/regress/multi_schedule @@ -19,7 +19,7 @@ test: multi_behavioral_analytics_basics multi_behavioral_analytics_single_shard_ # We don't parallelize the following test with the ones above because they're # not idempotent and hence causing flaky test detection check to fail. test: multi_insert_select_non_pushable_queries multi_insert_select -test: shard_local_batching +test: allow_unsafe_insert_select_pushdown test: multi_shard_update_delete recursive_dml_with_different_planners_executors test: insert_select_repartition window_functions dml_recursive multi_insert_select_window diff --git a/src/test/regress/sql/shard_local_batching.sql b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql similarity index 79% rename from src/test/regress/sql/shard_local_batching.sql rename to src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql index 454bcf3a535..fa0acb185cc 100644 --- a/src/test/regress/sql/shard_local_batching.sql +++ b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql @@ -1,12 +1,12 @@ -- --- SHARD_LOCAL_BATCHING +-- ALLOW_UNSAFE_INSERT_SELECT_PUSHDOWN -- --- Tests citus.enable_shard_local_batching, which lets a colocated INSERT .. SELECT --- push down GROUP BY / window batching and a volatile batch UDF to the shards --- instead of pulling rows to the coordinator. +-- Tests citus.allow_unsafe_insert_select_pushdown, which lets a colocated +-- INSERT .. SELECT push down GROUP BY / window batching and a volatile batch UDF +-- to the shards instead of pulling rows to the coordinator. -- -CREATE SCHEMA shard_local_batching; -SET search_path = shard_local_batching; +CREATE SCHEMA allow_unsafe_insert_select_pushdown; +SET search_path = allow_unsafe_insert_select_pushdown; SET citus.next_shard_id TO 14000000; SET citus.shard_count = 4; SET citus.shard_replication_factor = 1; @@ -32,7 +32,7 @@ SELECT id, vec FROM ( GROUP BY b ) s; -SET citus.enable_shard_local_batching TO on; +SET citus.allow_unsafe_insert_select_pushdown TO on; -- now the batching and UDF call run on the shards EXPLAIN (COSTS OFF) INSERT INTO emb(text_id, embedding) @@ -56,4 +56,4 @@ SELECT count(*), count(*) FILTER (WHERE embedding = length('t' || text_id)) AS o FROM emb JOIN dist USING (text_id); SET client_min_messages TO WARNING; -DROP SCHEMA shard_local_batching CASCADE; +DROP SCHEMA allow_unsafe_insert_select_pushdown CASCADE; From c8dbe0245bdf9c4eb6c45558be009667a1251cee Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 29 Jun 2026 15:47:29 +0000 Subject: [PATCH 06/25] more --- .../planner/insert_select_planner.c | 31 +++++++++---------- .../planner/query_pushdown_planning.c | 24 +++++--------- .../distributed/query_pushdown_planning.h | 4 +-- 3 files changed, 24 insertions(+), 35 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 5f615a3e5d9..35b1637087d 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -823,25 +823,22 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, "table", NULL, NULL); } - if (HasDistributionKey(targetRelationId)) + /* + * Ensure that INSERT's partition column comes from SELECT's partition + * column. With unsafe INSERT ... SELECT pushdown enabled the + * distribution value is typically derived (e.g. + * unnest(array_agg(text_id))), so we cannot match it to a plain SELECT + * partition column; colocation enforced below keeps batches + * shard-local. + */ + if (!AllowUnsafeInsertSelectPushdown && HasDistributionKey(targetRelationId)) { - /* - * ensure that INSERT's partition column comes from SELECT's partition - * column. With unsafe INSERT ... SELECT pushdown enabled the - * distribution value is typically derived (e.g. - * unnest(array_agg(text_id))), so we cannot match it to a plain SELECT - * partition column; colocation enforced below keeps batches - * shard-local. - */ - if (!AllowUnsafeInsertSelectPushdown) + error = InsertPartitionColumnMatchesSelect(queryTree, insertRte, + subqueryRte, + &selectPartitionColumnTableId); + if (error) { - error = InsertPartitionColumnMatchesSelect(queryTree, insertRte, - subqueryRte, - &selectPartitionColumnTableId); - if (error) - { - return error; - } + return error; } } } diff --git a/src/backend/distributed/planner/query_pushdown_planning.c b/src/backend/distributed/planner/query_pushdown_planning.c index c7efe9af4db..14439854d34 100644 --- a/src/backend/distributed/planner/query_pushdown_planning.c +++ b/src/backend/distributed/planner/query_pushdown_planning.c @@ -80,15 +80,6 @@ typedef struct RelidsReferenceWalkerContext /* Config variable managed via guc.c */ bool SubqueryPushdown = false; /* is subquery pushdown enabled */ int ValuesMaterializationThreshold = 100; - -/* - * Allows otherwise-unsafe colocated INSERT ... SELECT queries to be pushed down. - * When enabled, GROUP BY / window / aggregate-without-group-by on non-distribution - * columns and volatile functions are permitted in colocated INSERT ... SELECT so - * that batching (and any batch UDF call) executes on the shards instead of pulling - * data to the coordinator. The user takes responsibility for keeping batches - * shard-local and order-preserving; colocation is still enforced. Off by default. - */ bool AllowUnsafeInsertSelectPushdown = false; /* Local functions forward declarations */ @@ -103,7 +94,7 @@ static DeferredErrorMessage * DeferErrorIfUnsupportedTableCombination(Query *que static DeferredErrorMessage * DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, char *referencedThing, - bool allowUnsafePushdown); + bool allowUnsafeShardLocalGrouping); static bool ExtractSetOperationStatementWalker(Node *node, List **setOperationList); static RecurringTuplesType FetchFirstRecurType(PlannerInfo *plannerInfo, Relids relids); @@ -586,7 +577,8 @@ DeferredErrorMessage * DeferErrorIfUnsupportedSubqueryPushdown(Query *originalQuery, PlannerRestrictionContext * plannerRestrictionContext, - bool plannerPhase, bool allowUnsafePushdown) + bool plannerPhase, + bool allowUnsafeShardLocalGroupingForSubqueries) { bool outerMostQueryHasLimit = false; ListCell *subqueryCell = NULL; @@ -660,7 +652,7 @@ DeferErrorIfUnsupportedSubqueryPushdown(Query *originalQuery, Query *subquery = lfirst(subqueryCell); error = DeferErrorIfCannotPushdownSubquery(subquery, outerMostQueryHasLimit, - allowUnsafePushdown); + allowUnsafeShardLocalGroupingForSubqueries); if (error) { return error; @@ -1011,7 +1003,7 @@ CanPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit) */ DeferredErrorMessage * DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit, - bool allowUnsafePushdown) + bool allowUnsafeShardLocalGrouping) { bool preconditionsSatisfied = true; char *errorDetail = NULL; @@ -1041,7 +1033,7 @@ DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLi { deferredError = DeferErrorIfSubqueryRequiresMerge(subqueryTree, false, "another query", - allowUnsafePushdown); + allowUnsafeShardLocalGrouping); if (deferredError) { return deferredError; @@ -1142,7 +1134,7 @@ FlattenGroupExprs(Query *queryTree) */ static DeferredErrorMessage * DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, - char *referencedThing, bool allowUnsafePushdown) + char *referencedThing, bool allowUnsafeShardLocalGrouping) { bool preconditionsSatisfied = true; char *errorDetail = NULL; @@ -1173,7 +1165,7 @@ DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, * not require a merge step. Skip the partition-column requirements in that * case. */ - if (!allowUnsafePushdown) + if (!allowUnsafeShardLocalGrouping) { /* group clause list must include partition column */ if (subqueryTree->groupClause) diff --git a/src/include/distributed/query_pushdown_planning.h b/src/include/distributed/query_pushdown_planning.h index bc32a2ea176..3172f35434d 100644 --- a/src/include/distributed/query_pushdown_planning.h +++ b/src/include/distributed/query_pushdown_planning.h @@ -46,12 +46,12 @@ extern DeferredErrorMessage * DeferErrorIfUnsupportedSubqueryPushdown(Query * * plannerRestrictionContext, bool plannerPhase, - bool allowUnsafePushdown); + bool allowUnsafeShardLocalGroupingForSubqueries); extern DeferredErrorMessage * DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit, bool - allowUnsafePushdown); + allowUnsafeShardLocalGrouping); extern DeferredErrorMessage * DeferErrorIfUnsupportedUnionQuery(Query *queryTree); extern bool IsJsonTableRTE(RangeTblEntry *rte); extern bool IsOuterJoinExpr(Node *node); From 48e6ffac77ddc8ef5a66d09a092c1b7a33d7984f Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 29 Jun 2026 18:52:29 +0300 Subject: [PATCH 07/25] Document allowUnsafeShardLocalGrouping params; fix comments/indent Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../planner/insert_select_planner.c | 4 ++-- .../planner/query_pushdown_planning.c | 24 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 35b1637087d..9e2c1ca0ab7 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -824,11 +824,11 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, } /* - * Ensure that INSERT's partition column comes from SELECT's partition + * Ensure that INSERT's partition column comes from SELECT's partition * column. With unsafe INSERT ... SELECT pushdown enabled the * distribution value is typically derived (e.g. * unnest(array_agg(text_id))), so we cannot match it to a plain SELECT - * partition column; colocation enforced below keeps batches + * partition column; the co-location check further below keeps batches * shard-local. */ if (!AllowUnsafeInsertSelectPushdown && HasDistributionKey(targetRelationId)) diff --git a/src/backend/distributed/planner/query_pushdown_planning.c b/src/backend/distributed/planner/query_pushdown_planning.c index 14439854d34..93a85e4a450 100644 --- a/src/backend/distributed/planner/query_pushdown_planning.c +++ b/src/backend/distributed/planner/query_pushdown_planning.c @@ -572,6 +572,15 @@ SubqueryMultiNodeTree(Query *originalQuery, Query *queryTree, * entry list and uses helper functions to check if we can push down subquery * to worker nodes. These helper functions returns a deferred error if we * cannot push down the subquery. + * + * allowUnsafeShardLocalGroupingForSubqueries is forwarded as-is to + * DeferErrorIfCannotPushdownSubquery for every subquery checked here. It must only + * be set for colocated INSERT ... SELECT under + * citus.allow_unsafe_insert_select_pushdown; when true, the GROUP BY / aggregate / + * window / DISTINCT merge-step requirements are skipped, trusting the caller that + * those grouping constructs stay shard-local. All other pushdown checks + * (co-location, joins on the distribution column, recurring tuples, LIMIT/OFFSET) + * remain enforced. */ DeferredErrorMessage * DeferErrorIfUnsupportedSubqueryPushdown(Query *originalQuery, @@ -1000,6 +1009,14 @@ CanPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit) * a subquery has a group by on another subquery which includes order by with * limit, we let this query to run, but results could be wrong depending on the * features of underlying tables. + * + * When allowUnsafeShardLocalGrouping is true, the caller asserts that any GROUP BY / + * aggregate / window / DISTINCT in the subquery stays within a single shard, so the + * partition-column requirements that would otherwise force a coordinator merge step + * are skipped (this flag is forwarded to DeferErrorIfSubqueryRequiresMerge). + * LIMIT/OFFSET handling is unaffected. It is only set for colocated + * INSERT ... SELECT under citus.allow_unsafe_insert_select_pushdown, where keeping + * batches shard-local becomes the user's responsibility. */ DeferredErrorMessage * DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit, @@ -1131,6 +1148,13 @@ FlattenGroupExprs(Query *queryTree) * DeferErrorIfSubqueryRequiresMerge returns a deferred error if the subquery * requires a merge step on the coordinator (e.g. limit, group by non-distribution * column, etc.). + * + * When allowUnsafeShardLocalGrouping is true, the partition-column requirements for + * GROUP BY / aggregate / window / DISTINCT are skipped: the caller guarantees these + * grouping constructs are shard-local, so they do not need a coordinator merge. + * LIMIT/OFFSET are still rejected, because those require a merge regardless of the + * distribution column. The flag is only ever true for colocated INSERT ... SELECT + * under citus.allow_unsafe_insert_select_pushdown. */ static DeferredErrorMessage * DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, From d82e27c1a15c1c30d272d9ecdcfb02d96391b44e Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 29 Jun 2026 15:54:52 +0000 Subject: [PATCH 08/25] Apply citus_indent Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/backend/distributed/planner/insert_select_planner.c | 3 ++- src/backend/distributed/planner/query_pushdown_planning.c | 6 ++++-- src/include/distributed/query_pushdown_planning.h | 3 ++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 9e2c1ca0ab7..254a4c4f9aa 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -772,7 +772,8 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, /* first apply toplevel pushdown checks to SELECT query */ error = DeferErrorIfUnsupportedSubqueryPushdown(subquery, plannerRestrictionContext, - true, AllowUnsafeInsertSelectPushdown); + true, AllowUnsafeInsertSelectPushdown) + ; if (error) { return error; diff --git a/src/backend/distributed/planner/query_pushdown_planning.c b/src/backend/distributed/planner/query_pushdown_planning.c index 93a85e4a450..152c3e45682 100644 --- a/src/backend/distributed/planner/query_pushdown_planning.c +++ b/src/backend/distributed/planner/query_pushdown_planning.c @@ -94,7 +94,8 @@ static DeferredErrorMessage * DeferErrorIfUnsupportedTableCombination(Query *que static DeferredErrorMessage * DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, char *referencedThing, - bool allowUnsafeShardLocalGrouping); + bool + allowUnsafeShardLocalGrouping); static bool ExtractSetOperationStatementWalker(Node *node, List **setOperationList); static RecurringTuplesType FetchFirstRecurType(PlannerInfo *plannerInfo, Relids relids); @@ -1158,7 +1159,8 @@ FlattenGroupExprs(Query *queryTree) */ static DeferredErrorMessage * DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral, - char *referencedThing, bool allowUnsafeShardLocalGrouping) + char *referencedThing, bool + allowUnsafeShardLocalGrouping) { bool preconditionsSatisfied = true; char *errorDetail = NULL; diff --git a/src/include/distributed/query_pushdown_planning.h b/src/include/distributed/query_pushdown_planning.h index 3172f35434d..fec709fbdc5 100644 --- a/src/include/distributed/query_pushdown_planning.h +++ b/src/include/distributed/query_pushdown_planning.h @@ -46,7 +46,8 @@ extern DeferredErrorMessage * DeferErrorIfUnsupportedSubqueryPushdown(Query * * plannerRestrictionContext, bool plannerPhase, - bool allowUnsafeShardLocalGroupingForSubqueries); + bool + allowUnsafeShardLocalGroupingForSubqueries); extern DeferredErrorMessage * DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLimit, From cd3b22dcfc1dd6fdc6654419702d08cef75d290c Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 29 Jun 2026 20:05:40 +0300 Subject: [PATCH 09/25] rename some stuff Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../allow_unsafe_insert_select_pushdown.out | 36 +++++++++---------- .../allow_unsafe_insert_select_pushdown.sql | 32 ++++++++--------- 2 files changed, 34 insertions(+), 34 deletions(-) diff --git a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out index a7c2608b077..d3ed378223a 100644 --- a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out +++ b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out @@ -11,14 +11,14 @@ SET citus.next_shard_id TO 14000000; SET citus.shard_count = 4; SET citus.shard_replication_factor = 1; CREATE TABLE dist(text_id int, text_col text); -CREATE TABLE emb(text_id int, embedding int); +CREATE TABLE res(text_id int, val int); SELECT create_distributed_table('dist', 'text_id'); create_distributed_table --------------------------------------------------------------------- (1 row) -SELECT create_distributed_table('emb', 'text_id'); +SELECT create_distributed_table('res', 'text_id'); create_distributed_table --------------------------------------------------------------------- @@ -26,10 +26,10 @@ SELECT create_distributed_table('emb', 'text_id'); INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; -- a volatile batch UDF: returns one value per input, mimicking a batched API call -CREATE FUNCTION create_embeddings(t text[]) RETURNS int[] +CREATE FUNCTION batch_transform(t text[]) RETURNS int[] LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; -SELECT create_distributed_function('create_embeddings(text[])'); -NOTICE: procedure allow_unsafe_insert_select_pushdown.create_embeddings is already distributed +SELECT create_distributed_function('batch_transform(text[])'); +NOTICE: procedure allow_unsafe_insert_select_pushdown.batch_transform is already distributed DETAIL: Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] commands create_distributed_function --------------------------------------------------------------------- @@ -37,10 +37,10 @@ DETAIL: Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] (1 row) -- default off: batching is done after pulling rows to the coordinator -EXPLAIN (COSTS OFF) INSERT INTO emb(text_id, embedding) -SELECT id, vec FROM ( +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( SELECT b, unnest(array_agg(text_id)) id, - unnest(create_embeddings(array_agg(text_col))) vec + unnest(batch_transform(array_agg(text_col))) val FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q GROUP BY b ) s; @@ -70,10 +70,10 @@ SELECT id, vec FROM ( SET citus.allow_unsafe_insert_select_pushdown TO on; -- now the batching and UDF call run on the shards -EXPLAIN (COSTS OFF) INSERT INTO emb(text_id, embedding) -SELECT id, vec FROM ( +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( SELECT b, unnest(array_agg(text_id)) id, - unnest(create_embeddings(array_agg(text_col))) vec + unnest(batch_transform(array_agg(text_col))) val FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q GROUP BY b ) s; @@ -84,7 +84,7 @@ SELECT id, vec FROM ( Tasks Shown: One of 4 -> Task Node: host=localhost port=xxxxx dbname=regression - -> Insert on emb_14000004 + -> Insert on res_14000004 -> Subquery Scan on s -> ProjectSet -> HashAggregate @@ -93,16 +93,16 @@ SELECT id, vec FROM ( -> Seq Scan on dist_14000000 dist (12 rows) -INSERT INTO emb(text_id, embedding) -SELECT id, vec FROM ( +INSERT INTO res(text_id, val) +SELECT id, val FROM ( SELECT b, unnest(array_agg(text_id)) id, - unnest(create_embeddings(array_agg(text_col))) vec + unnest(batch_transform(array_agg(text_col))) val FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q GROUP BY b ) s; --- every text_id should be matched to the right vector -SELECT count(*), count(*) FILTER (WHERE embedding = length('t' || text_id)) AS ok -FROM emb JOIN dist USING (text_id); +-- every text_id should be matched to the right value +SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok +FROM res JOIN dist USING (text_id); count | ok --------------------------------------------------------------------- 500 | 500 diff --git a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql index fa0acb185cc..a6de7b9bbc4 100644 --- a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql +++ b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql @@ -12,22 +12,22 @@ SET citus.shard_count = 4; SET citus.shard_replication_factor = 1; CREATE TABLE dist(text_id int, text_col text); -CREATE TABLE emb(text_id int, embedding int); +CREATE TABLE res(text_id int, val int); SELECT create_distributed_table('dist', 'text_id'); -SELECT create_distributed_table('emb', 'text_id'); +SELECT create_distributed_table('res', 'text_id'); INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; -- a volatile batch UDF: returns one value per input, mimicking a batched API call -CREATE FUNCTION create_embeddings(t text[]) RETURNS int[] +CREATE FUNCTION batch_transform(t text[]) RETURNS int[] LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; -SELECT create_distributed_function('create_embeddings(text[])'); +SELECT create_distributed_function('batch_transform(text[])'); -- default off: batching is done after pulling rows to the coordinator -EXPLAIN (COSTS OFF) INSERT INTO emb(text_id, embedding) -SELECT id, vec FROM ( +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( SELECT b, unnest(array_agg(text_id)) id, - unnest(create_embeddings(array_agg(text_col))) vec + unnest(batch_transform(array_agg(text_col))) val FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q GROUP BY b ) s; @@ -35,25 +35,25 @@ SELECT id, vec FROM ( SET citus.allow_unsafe_insert_select_pushdown TO on; -- now the batching and UDF call run on the shards -EXPLAIN (COSTS OFF) INSERT INTO emb(text_id, embedding) -SELECT id, vec FROM ( +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( SELECT b, unnest(array_agg(text_id)) id, - unnest(create_embeddings(array_agg(text_col))) vec + unnest(batch_transform(array_agg(text_col))) val FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q GROUP BY b ) s; -INSERT INTO emb(text_id, embedding) -SELECT id, vec FROM ( +INSERT INTO res(text_id, val) +SELECT id, val FROM ( SELECT b, unnest(array_agg(text_id)) id, - unnest(create_embeddings(array_agg(text_col))) vec + unnest(batch_transform(array_agg(text_col))) val FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q GROUP BY b ) s; --- every text_id should be matched to the right vector -SELECT count(*), count(*) FILTER (WHERE embedding = length('t' || text_id)) AS ok -FROM emb JOIN dist USING (text_id); +-- every text_id should be matched to the right value +SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok +FROM res JOIN dist USING (text_id); SET client_min_messages TO WARNING; DROP SCHEMA allow_unsafe_insert_select_pushdown CASCADE; From b238537b5999f9b5fd423395e014a65f5d9dd27d Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Tue, 30 Jun 2026 10:59:26 +0300 Subject: [PATCH 10/25] Sort allow_unsafe_insert_select_pushdown GUC alphabetically Fixes check-style ci/check_gucs_are_alphabetically_sorted.sh. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/backend/distributed/shared_library_init.c | 36 +++++++++---------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c index d3d0288d49f..226fa12ac57 100644 --- a/src/backend/distributed/shared_library_init.c +++ b/src/backend/distributed/shared_library_init.c @@ -987,24 +987,6 @@ RegisterCitusConfigVariables(void) GUC_STANDARD, NULL, NULL, NULL); - DefineCustomBoolVariable( - "citus.allow_unsafe_insert_select_pushdown", - gettext_noop("Allows pushdown of otherwise-unsafe colocated " - "INSERT ... SELECT queries."), - gettext_noop("When enabled, Citus relaxes safety checks (GROUP BY / window / " - "aggregate constructs on non-distribution columns, volatile " - "functions, partition-column matching) for colocated " - "INSERT ... SELECT, so that batching and any batch UDF call run " - "on the shards instead of pulling data to the coordinator. " - "Colocation is still enforced, but the user takes responsibility " - "for keeping batches shard-local and order-preserving; otherwise " - "results may be silently incorrect."), - &AllowUnsafeInsertSelectPushdown, - false, - PGC_USERSET, - GUC_STANDARD, - NULL, NULL, NULL); - DefineCustomBoolVariable( "citus.allow_aggregate_worker_combine_on_internal_types", gettext_noop("Enables aggregate worker partial aggregates on aggregates that " @@ -1075,6 +1057,24 @@ RegisterCitusConfigVariables(void) GUC_NO_SHOW_ALL | GUC_NOT_IN_SAMPLE, NULL, NULL, NULL); + DefineCustomBoolVariable( + "citus.allow_unsafe_insert_select_pushdown", + gettext_noop("Allows pushdown of otherwise-unsafe colocated " + "INSERT ... SELECT queries."), + gettext_noop("When enabled, Citus relaxes safety checks (GROUP BY / window / " + "aggregate constructs on non-distribution columns, volatile " + "functions, partition-column matching) for colocated " + "INSERT ... SELECT, so that batching and any batch UDF call run " + "on the shards instead of pulling data to the coordinator. " + "Colocation is still enforced, but the user takes responsibility " + "for keeping batches shard-local and order-preserving; otherwise " + "results may be silently incorrect."), + &AllowUnsafeInsertSelectPushdown, + false, + PGC_USERSET, + GUC_STANDARD, + NULL, NULL, NULL); + DefineCustomBoolVariable( "citus.allow_unsafe_locks_from_workers", gettext_noop("Enables acquiring a distributed lock from a worker " From f6f01e775d42df3d50cbb93b7cf79ed974b0df95 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Tue, 30 Jun 2026 08:11:21 +0000 Subject: [PATCH 11/25] Make EXPLAIN output stable across PG16/17/18 via explain_filter Wrap the two EXPLAIN statements in public.explain_filter(..., true) so the PG18-only "Window:" line is stripped and the plan footer row counts match across supported Postgres versions. Regenerate the expected output, which also re-syncs blank lines that had drifted from the .sql. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../allow_unsafe_insert_select_pushdown.out | 14 ++++++++++---- .../sql/allow_unsafe_insert_select_pushdown.sql | 10 ++++++++-- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out index d3ed378223a..b94d4ef057f 100644 --- a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out +++ b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out @@ -37,14 +37,18 @@ DETAIL: Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] (1 row) -- default off: batching is done after pulling rows to the coordinator +-- (explain_filter strips the PG18-only "Window:" line so the plan is +-- comparable across supported Postgres versions) +SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT id, val FROM ( SELECT b, unnest(array_agg(text_id)) id, unnest(batch_transform(array_agg(text_col))) val FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q GROUP BY b -) s; - QUERY PLAN +) s +$$, true); + explain_filter --------------------------------------------------------------------- Custom Scan (Citus INSERT ... SELECT) INSERT/SELECT method: pull to coordinator @@ -70,14 +74,16 @@ SELECT id, val FROM ( SET citus.allow_unsafe_insert_select_pushdown TO on; -- now the batching and UDF call run on the shards +SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT id, val FROM ( SELECT b, unnest(array_agg(text_id)) id, unnest(batch_transform(array_agg(text_col))) val FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q GROUP BY b -) s; - QUERY PLAN +) s +$$, true); + explain_filter --------------------------------------------------------------------- Custom Scan (Citus Adaptive) Task Count: 4 diff --git a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql index a6de7b9bbc4..48bc9b8aea3 100644 --- a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql +++ b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql @@ -24,24 +24,30 @@ LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; SELECT create_distributed_function('batch_transform(text[])'); -- default off: batching is done after pulling rows to the coordinator +-- (explain_filter strips the PG18-only "Window:" line so the plan is +-- comparable across supported Postgres versions) +SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT id, val FROM ( SELECT b, unnest(array_agg(text_id)) id, unnest(batch_transform(array_agg(text_col))) val FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q GROUP BY b -) s; +) s +$$, true); SET citus.allow_unsafe_insert_select_pushdown TO on; -- now the batching and UDF call run on the shards +SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT id, val FROM ( SELECT b, unnest(array_agg(text_id)) id, unnest(batch_transform(array_agg(text_col))) val FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q GROUP BY b -) s; +) s +$$, true); INSERT INTO res(text_id, val) SELECT id, val FROM ( From 3c3c048d5daf478131344fbce5988666d90f37a7 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Tue, 30 Jun 2026 08:57:45 +0000 Subject: [PATCH 12/25] Add raw PG18 EXPLAIN test for insert-select pushdown Validate the shard-local INSERT..SELECT batching pushdown plan directly on PG18 (no public.explain_filter), so the PG18-only WindowAgg "Window:" line is exercised in pg18.sql. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/test/regress/expected/pg18.out | 72 ++++++++++++++++++++++++++++++ src/test/regress/sql/pg18.sql | 42 +++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/src/test/regress/expected/pg18.out b/src/test/regress/expected/pg18.out index 861c12f4946..9fe61b4808d 100644 --- a/src/test/regress/expected/pg18.out +++ b/src/test/regress/expected/pg18.out @@ -3165,6 +3165,78 @@ drop cascades to type product_rating drop cascades to table product_ratings drop cascades to table record_arg_t -- END: PG18: MIN/MAX aggregate OID resolution for ANYARRAY and RECORD +-- PG18: shard-local INSERT .. SELECT batching pushdown +-- (citus.allow_unsafe_insert_select_pushdown). This mirrors a case from +-- allow_unsafe_insert_select_pushdown.sql, which wraps EXPLAIN in +-- public.explain_filter so the plan is comparable across supported Postgres +-- versions. Here we keep the raw EXPLAIN output because this file only runs on +-- PG18+, so we exercise the real plan including the PG18-only WindowAgg +-- "Window:" line. +CREATE SCHEMA pg18_insert_select_pushdown; +SET search_path TO pg18_insert_select_pushdown; +SET citus.next_shard_id TO 14100000; +SET citus.shard_count = 4; +SET citus.shard_replication_factor = 1; +CREATE TABLE dist(text_id int, text_col text); +CREATE TABLE res(text_id int, val int); +SELECT create_distributed_table('dist', 'text_id'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +SELECT create_distributed_table('res', 'text_id'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; +-- a volatile batch UDF: returns one value per input, mimicking a batched API call +CREATE FUNCTION batch_transform(t text[]) RETURNS int[] +LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; +SELECT create_distributed_function('batch_transform(text[])'); +NOTICE: procedure pg18_insert_select_pushdown.batch_transform is already distributed +DETAIL: Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] commands + create_distributed_function +--------------------------------------------------------------------- + +(1 row) + +SET citus.allow_unsafe_insert_select_pushdown TO on; +-- the batching and volatile batch UDF run on the shards; the raw PG18 plan +-- includes the WindowAgg "Window:" line +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT b, unnest(array_agg(text_id)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s; + QUERY PLAN +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14100004 + -> Subquery Scan on s + -> ProjectSet + -> HashAggregate + Group Key: ((row_number() OVER (?) - 1) / 100) + -> WindowAgg + Window: w1 AS (ROWS UNBOUNDED PRECEDING) + -> Seq Scan on dist_14100000 dist +(13 rows) + +RESET citus.allow_unsafe_insert_select_pushdown; +DROP SCHEMA pg18_insert_select_pushdown CASCADE; +NOTICE: drop cascades to 3 other objects +DETAIL: drop cascades to table dist +drop cascades to table res +drop cascades to function batch_transform(text[]) +SET search_path TO pg18_nn; -- cleanup with minimum verbosity SET client_min_messages TO ERROR; RESET search_path; diff --git a/src/test/regress/sql/pg18.sql b/src/test/regress/sql/pg18.sql index aedc786473d..d585abc20de 100644 --- a/src/test/regress/sql/pg18.sql +++ b/src/test/regress/sql/pg18.sql @@ -2007,6 +2007,48 @@ DROP SCHEMA pg18_minmax CASCADE; -- END: PG18: MIN/MAX aggregate OID resolution for ANYARRAY and RECORD +-- PG18: shard-local INSERT .. SELECT batching pushdown +-- (citus.allow_unsafe_insert_select_pushdown). This mirrors a case from +-- allow_unsafe_insert_select_pushdown.sql, which wraps EXPLAIN in +-- public.explain_filter so the plan is comparable across supported Postgres +-- versions. Here we keep the raw EXPLAIN output because this file only runs on +-- PG18+, so we exercise the real plan including the PG18-only WindowAgg +-- "Window:" line. +CREATE SCHEMA pg18_insert_select_pushdown; +SET search_path TO pg18_insert_select_pushdown; +SET citus.next_shard_id TO 14100000; +SET citus.shard_count = 4; +SET citus.shard_replication_factor = 1; + +CREATE TABLE dist(text_id int, text_col text); +CREATE TABLE res(text_id int, val int); +SELECT create_distributed_table('dist', 'text_id'); +SELECT create_distributed_table('res', 'text_id'); + +INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; + +-- a volatile batch UDF: returns one value per input, mimicking a batched API call +CREATE FUNCTION batch_transform(t text[]) RETURNS int[] +LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; +SELECT create_distributed_function('batch_transform(text[])'); + +SET citus.allow_unsafe_insert_select_pushdown TO on; + +-- the batching and volatile batch UDF run on the shards; the raw PG18 plan +-- includes the WindowAgg "Window:" line +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT b, unnest(array_agg(text_id)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s; + +RESET citus.allow_unsafe_insert_select_pushdown; +DROP SCHEMA pg18_insert_select_pushdown CASCADE; +SET search_path TO pg18_nn; + + -- cleanup with minimum verbosity SET client_min_messages TO ERROR; RESET search_path; From c70602523876dd038820060390af57af466801e3 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Tue, 30 Jun 2026 09:22:29 +0000 Subject: [PATCH 13/25] Add branch-coverage tests for insert-select pushdown GUC Cover each relaxed branch (volatile fn, derived insert dist col, GROUP BY on non-dist col, aggregates and HAVING without GROUP BY, window not on dist col, DISTINCT without dist col), a few combinations, and the same constructs nested in subqueries. Add negative tests showing the GUC does not apply to reference-table targets and never relaxes LIMIT/OFFSET. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../allow_unsafe_insert_select_pushdown.out | 366 ++++++++++++++++++ .../allow_unsafe_insert_select_pushdown.sql | 139 +++++++ 2 files changed, 505 insertions(+) diff --git a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out index b94d4ef057f..3df0a0a21c9 100644 --- a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out +++ b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out @@ -114,5 +114,371 @@ FROM res JOIN dist USING (text_id); 500 | 500 (1 row) +-- --------------------------------------------------------------------- +-- Per-branch coverage. Each construct below used to force a coordinator merge +-- (or was rejected outright); with the GUC enabled the whole colocated +-- INSERT .. SELECT is pushed to the shards. The pushed-down plan is a +-- Custom Scan (Citus Adaptive) whose task runs the INSERT on a shard, with no +-- Distributed Subplan / intermediate results. explain_filter keeps the plan +-- comparable across Postgres versions. +-- --------------------------------------------------------------------- +-- branch: volatile function in the SELECT +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id, (batch_transform(ARRAY[text_col]))[1] FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 + -> Seq Scan on dist_14000000 dist + Filter: (text_id IS NOT NULL) +(8 rows) + +-- branch: INSERT distribution column derived from SELECT (not a plain +-- partition-column Var, so InsertPartitionColumnMatchesSelect is skipped) +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id + 0, length(text_col) FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> Seq Scan on dist_14000000 dist +(7 rows) + +-- branch: GROUP BY on a non-distribution column +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) +FROM dist GROUP BY text_id % 10 +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> Subquery Scan on "*SELECT*" + -> ProjectSet + -> HashAggregate + Group Key: (dist.text_id % 10) + -> Seq Scan on dist_14000000 dist +(11 rows) + +-- branch: aggregates without GROUP BY +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> ProjectSet + -> Aggregate + -> Seq Scan on dist_14000000 dist +(9 rows) + +-- branch: HAVING without GROUP BY on the distribution column +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) +FROM dist HAVING count(*) > 0 +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> ProjectSet + -> Aggregate + Filter: (count(*) > 0) + -> Seq Scan on dist_14000000 dist +(10 rows) + +-- branch: window function not partitioned on the distribution column +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id, row_number() OVER (ORDER BY text_col) FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> Subquery Scan on "*SELECT*" + -> WindowAgg + -> Sort + Sort Key: dist.text_col + -> Seq Scan on dist_14000000 dist + Filter: (text_id IS NOT NULL) +(12 rows) + +-- branch: DISTINCT without the distribution column +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT DISTINCT length(text_col), text_id % 7 FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> HashAggregate + Group Key: length(dist.text_col), (dist.text_id % 7) + -> Seq Scan on dist_14000000 dist +(9 rows) + +-- combination: GROUP BY + DISTINCT +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT DISTINCT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) +FROM dist GROUP BY text_id % 10 +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> Subquery Scan on "*SELECT*" + -> HashAggregate + Group Key: unnest((array_agg(dist.text_id))), unnest((array_agg(length(dist.text_col)))) + -> ProjectSet + -> HashAggregate + Group Key: (dist.text_id % 10) + -> Seq Scan on dist_14000000 dist +(13 rows) + +-- combination: window + GROUP BY, relaxed constructs living in nested subqueries +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, v FROM ( + SELECT unnest(array_agg(text_id)) id, unnest(array_agg(length(text_col))) v + FROM (SELECT text_id, text_col, row_number() OVER (ORDER BY text_col) rn FROM dist) q + GROUP BY rn % 5 +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> Subquery Scan on s + -> ProjectSet + -> HashAggregate + Group Key: (q.rn % '5'::bigint) + -> Subquery Scan on q + -> WindowAgg + -> Sort + Sort Key: dist.text_col + -> Seq Scan on dist_14000000 dist +(15 rows) + +-- subquery: a relaxed construct (DISTINCT without the distribution column) +-- nested one level down still pushes down (checks recurse into subqueries) +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT k, k FROM ( + SELECT DISTINCT length(text_col) k FROM dist +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> Subquery Scan on s + -> HashAggregate + Group Key: length(dist.text_col) + -> Seq Scan on dist_14000000 dist +(10 rows) + +-- correctness for a GROUP BY batch: per-shard aggregation keeps each text_id +-- matched to its own value +TRUNCATE res; +INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) +FROM dist GROUP BY text_id % 10; +SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok +FROM res JOIN dist USING (text_id); + count | ok +--------------------------------------------------------------------- + 500 | 500 +(1 row) + +-- --------------------------------------------------------------------- +-- Paths where the GUC does NOT help. It only relaxes colocated +-- INSERT .. SELECT into a distributed table. +-- --------------------------------------------------------------------- +-- reference table target: the GUC does not apply. The plan is a coordinator +-- merge whether the GUC is disabled or enabled (identical), unlike the +-- distributed-target cases above which push down once the GUC is enabled. +CREATE TABLE ref(text_id int, val int); +SELECT create_reference_table('ref'); + create_reference_table +--------------------------------------------------------------------- + +(1 row) + +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO ref(text_id, val) +SELECT text_id % 10, count(*)::int FROM dist GROUP BY text_id % 10 +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> HashAggregate + Group Key: remote_scan.text_id + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> HashAggregate + Group Key: (text_id % 10) + -> Seq Scan on dist_14000000 dist +(12 rows) + +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO ref(text_id, val) +SELECT text_id % 10, count(*)::int FROM dist GROUP BY text_id % 10 +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> HashAggregate + Group Key: remote_scan.text_id + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> HashAggregate + Group Key: (text_id % 10) + -> Seq Scan on dist_14000000 dist +(12 rows) + +-- a GROUP BY on a non-distribution column into a distributed target pushes +-- down with the GUC enabled ... +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, v FROM ( + SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> HashAggregate + Group Key: (dist.text_id % 10) + -> Seq Scan on dist_14000000 dist +(9 rows) + +-- ... but a LIMIT in the subquery still forces a coordinator merge even with +-- the GUC enabled (LIMIT/OFFSET are never relaxed) +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, v FROM ( + SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 LIMIT 100 +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> Limit + -> HashAggregate + Group Key: remote_scan.id + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> HashAggregate + Group Key: (text_id % 10) + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Function Scan on read_intermediate_result intermediate_result +(20 rows) + +-- ... and likewise for OFFSET +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, v FROM ( + SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 OFFSET 5 +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> Limit + -> HashAggregate + Group Key: remote_scan.id + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> HashAggregate + Group Key: (text_id % 10) + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Function Scan on read_intermediate_result intermediate_result +(20 rows) + SET client_min_messages TO WARNING; DROP SCHEMA allow_unsafe_insert_select_pushdown CASCADE; diff --git a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql index 48bc9b8aea3..292f6ca38bf 100644 --- a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql +++ b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql @@ -61,5 +61,144 @@ SELECT id, val FROM ( SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok FROM res JOIN dist USING (text_id); +-- --------------------------------------------------------------------- +-- Per-branch coverage. Each construct below used to force a coordinator merge +-- (or was rejected outright); with the GUC enabled the whole colocated +-- INSERT .. SELECT is pushed to the shards. The pushed-down plan is a +-- Custom Scan (Citus Adaptive) whose task runs the INSERT on a shard, with no +-- Distributed Subplan / intermediate results. explain_filter keeps the plan +-- comparable across Postgres versions. +-- --------------------------------------------------------------------- + +-- branch: volatile function in the SELECT +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id, (batch_transform(ARRAY[text_col]))[1] FROM dist +$$, true); + +-- branch: INSERT distribution column derived from SELECT (not a plain +-- partition-column Var, so InsertPartitionColumnMatchesSelect is skipped) +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id + 0, length(text_col) FROM dist +$$, true); + +-- branch: GROUP BY on a non-distribution column +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) +FROM dist GROUP BY text_id % 10 +$$, true); + +-- branch: aggregates without GROUP BY +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) FROM dist +$$, true); + +-- branch: HAVING without GROUP BY on the distribution column +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) +FROM dist HAVING count(*) > 0 +$$, true); + +-- branch: window function not partitioned on the distribution column +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id, row_number() OVER (ORDER BY text_col) FROM dist +$$, true); + +-- branch: DISTINCT without the distribution column +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT DISTINCT length(text_col), text_id % 7 FROM dist +$$, true); + +-- combination: GROUP BY + DISTINCT +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT DISTINCT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) +FROM dist GROUP BY text_id % 10 +$$, true); + +-- combination: window + GROUP BY, relaxed constructs living in nested subqueries +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, v FROM ( + SELECT unnest(array_agg(text_id)) id, unnest(array_agg(length(text_col))) v + FROM (SELECT text_id, text_col, row_number() OVER (ORDER BY text_col) rn FROM dist) q + GROUP BY rn % 5 +) s +$$, true); + +-- subquery: a relaxed construct (DISTINCT without the distribution column) +-- nested one level down still pushes down (checks recurse into subqueries) +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT k, k FROM ( + SELECT DISTINCT length(text_col) k FROM dist +) s +$$, true); + +-- correctness for a GROUP BY batch: per-shard aggregation keeps each text_id +-- matched to its own value +TRUNCATE res; +INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) +FROM dist GROUP BY text_id % 10; + +SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok +FROM res JOIN dist USING (text_id); + +-- --------------------------------------------------------------------- +-- Paths where the GUC does NOT help. It only relaxes colocated +-- INSERT .. SELECT into a distributed table. +-- --------------------------------------------------------------------- + +-- reference table target: the GUC does not apply. The plan is a coordinator +-- merge whether the GUC is disabled or enabled (identical), unlike the +-- distributed-target cases above which push down once the GUC is enabled. +CREATE TABLE ref(text_id int, val int); +SELECT create_reference_table('ref'); + +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO ref(text_id, val) +SELECT text_id % 10, count(*)::int FROM dist GROUP BY text_id % 10 +$$, true); + +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO ref(text_id, val) +SELECT text_id % 10, count(*)::int FROM dist GROUP BY text_id % 10 +$$, true); + +-- a GROUP BY on a non-distribution column into a distributed target pushes +-- down with the GUC enabled ... +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, v FROM ( + SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 +) s +$$, true); + +-- ... but a LIMIT in the subquery still forces a coordinator merge even with +-- the GUC enabled (LIMIT/OFFSET are never relaxed) +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, v FROM ( + SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 LIMIT 100 +) s +$$, true); + +-- ... and likewise for OFFSET +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, v FROM ( + SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 OFFSET 5 +) s +$$, true); + SET client_min_messages TO WARNING; DROP SCHEMA allow_unsafe_insert_select_pushdown CASCADE; From 0ecb0064c33ff8beca60078c439710fdfe9cb0c6 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Tue, 30 Jun 2026 16:15:07 +0000 Subject: [PATCH 14/25] Restrict insert-select GUC to grouping/partition-col, keep volatile ban The unsafe insert-select pushdown GUC no longer relaxes the volatile function ban: a batched-embeddings style UDF is immutable, so the GUC only needs to relax grouping / window / aggregate / DISTINCT on non-distribution columns and partition-column matching. Volatile functions stay blocked from shard pushdown and fall back to the coordinator plan as before. - insert_select_planner.c: always defer on volatile functions - shared_library_init.c: drop "volatile functions" from the GUC description - tests: make batch_transform IMMUTABLE PARALLEL SAFE, replace the volatile EXPLAIN test with the batched-embeddings benchmark query shape (row_number()/batch_size bucketing + array_agg + batch UDF + unnest zip-back) as the SELECT of an INSERT .. SELECT, with an id<->val correctness check, and a negative test showing a volatile function is still not pushed down even with the GUC enabled. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../planner/insert_select_planner.c | 3 +- src/backend/distributed/shared_library_init.c | 4 +- .../allow_unsafe_insert_select_pushdown.out | 107 +++++++++++++++--- .../allow_unsafe_insert_select_pushdown.sql | 62 ++++++++-- 4 files changed, 144 insertions(+), 32 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 254a4c4f9aa..e89ea207442 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -752,8 +752,7 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, } } - if (!AllowUnsafeInsertSelectPushdown && - FindNodeMatchingCheckFunction((Node *) queryTree, CitusIsVolatileFunction)) + if (FindNodeMatchingCheckFunction((Node *) queryTree, CitusIsVolatileFunction)) { return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, "volatile functions are not allowed in distributed " diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c index 226fa12ac57..a75e0d4df41 100644 --- a/src/backend/distributed/shared_library_init.c +++ b/src/backend/distributed/shared_library_init.c @@ -1062,8 +1062,8 @@ RegisterCitusConfigVariables(void) gettext_noop("Allows pushdown of otherwise-unsafe colocated " "INSERT ... SELECT queries."), gettext_noop("When enabled, Citus relaxes safety checks (GROUP BY / window / " - "aggregate constructs on non-distribution columns, volatile " - "functions, partition-column matching) for colocated " + "aggregate / DISTINCT constructs on non-distribution columns, " + "partition-column matching) for colocated " "INSERT ... SELECT, so that batching and any batch UDF call run " "on the shards instead of pulling data to the coordinator. " "Colocation is still enforced, but the user takes responsibility " diff --git a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out index 3df0a0a21c9..6a9a6466012 100644 --- a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out +++ b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out @@ -2,7 +2,7 @@ -- ALLOW_UNSAFE_INSERT_SELECT_PUSHDOWN -- -- Tests citus.allow_unsafe_insert_select_pushdown, which lets a colocated --- INSERT .. SELECT push down GROUP BY / window batching and a volatile batch UDF +-- INSERT .. SELECT push down GROUP BY / window batching and a batch UDF -- to the shards instead of pulling rows to the coordinator. -- CREATE SCHEMA allow_unsafe_insert_select_pushdown; @@ -25,9 +25,10 @@ SELECT create_distributed_table('res', 'text_id'); (1 row) INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; --- a volatile batch UDF: returns one value per input, mimicking a batched API call +-- a batched UDF: returns one value per input, mimicking a batched API call. +-- immutable + parallel safe, like a real embeddings UDF. CREATE FUNCTION batch_transform(t text[]) RETURNS int[] -LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; +LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; SELECT create_distributed_function('batch_transform(text[])'); NOTICE: procedure allow_unsafe_insert_select_pushdown.batch_transform is already distributed DETAIL: Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] commands @@ -90,7 +91,7 @@ $$, true); Tasks Shown: One of 4 -> Task Node: host=localhost port=xxxxx dbname=regression - -> Insert on res_14000004 + -> Insert on res_14000004 citus_table_alias -> Subquery Scan on s -> ProjectSet -> HashAggregate @@ -122,22 +123,40 @@ FROM res JOIN dist USING (text_id); -- Distributed Subplan / intermediate results. explain_filter keeps the plan -- comparable across Postgres versions. -- --------------------------------------------------------------------- --- branch: volatile function in the SELECT +-- the batched-embeddings benchmark shape: bucket rows into fixed-size batches +-- with row_number()/batch_size, array_agg each batch (id and text in the same +-- order), call the batch UDF once per batch, then unnest back to one row per id. +-- This is the query the GUC is meant to push down to the shards. SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT text_id, (batch_transform(ARRAY[text_col]))[1] FROM dist +SELECT id, val FROM ( + SELECT + unnest(array_agg(text_id ORDER BY text_id)) id, + unnest(batch_transform(array_agg(text_col ORDER BY text_id))) val + FROM ( + SELECT text_id, text_col, (row_number() OVER () - 1) / 100 batch FROM dist + ) q + GROUP BY batch +) s $$, true); - explain_filter + explain_filter --------------------------------------------------------------------- Custom Scan (Citus Adaptive) Task Count: 4 Tasks Shown: One of 4 -> Task Node: host=localhost port=xxxxx dbname=regression - -> Insert on res_14000004 - -> Seq Scan on dist_14000000 dist - Filter: (text_id IS NOT NULL) -(8 rows) + -> Insert on res_14000004 citus_table_alias + -> Subquery Scan on s + -> ProjectSet + -> GroupAggregate + Group Key: q.batch + -> Sort + Sort Key: q.batch, q.text_id + -> Subquery Scan on q + -> WindowAgg + -> Seq Scan on dist_14000000 dist +(15 rows) -- branch: INSERT distribution column derived from SELECT (not a plain -- partition-column Var, so InsertPartitionColumnMatchesSelect is skipped) @@ -340,10 +359,26 @@ FROM res JOIN dist USING (text_id); 500 | 500 (1 row) --- --------------------------------------------------------------------- --- Paths where the GUC does NOT help. It only relaxes colocated --- INSERT .. SELECT into a distributed table. --- --------------------------------------------------------------------- +-- correctness for the batched-embeddings benchmark shape: every text_id keeps +-- its own value after batching, the UDF call, and the unnest zip-back +TRUNCATE res; +INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT + unnest(array_agg(text_id ORDER BY text_id)) id, + unnest(batch_transform(array_agg(text_col ORDER BY text_id))) val + FROM ( + SELECT text_id, text_col, (row_number() OVER () - 1) / 100 batch FROM dist + ) q + GROUP BY batch +) s; +SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok +FROM res JOIN dist USING (text_id); + count | ok +--------------------------------------------------------------------- + 500 | 500 +(1 row) + -- reference table target: the GUC does not apply. The plan is a coordinator -- merge whether the GUC is disabled or enabled (identical), unlike the -- distributed-target cases above which push down once the GUC is enabled. @@ -396,8 +431,46 @@ $$, true); -> Seq Scan on dist_14000000 dist (12 rows) --- a GROUP BY on a non-distribution column into a distributed target pushes --- down with the GUC enabled ... +-- volatile functions: the GUC relaxes only grouping / partition-column matching, +-- never volatility. A volatile function in the SELECT is still not pushed to the +-- shards with the GUC enabled -- it falls back to a coordinator plan, identical +-- to the GUC-disabled case. +CREATE FUNCTION volatile_transform(t text) RETURNS int +LANGUAGE sql VOLATILE AS $$ SELECT length(t) $$; +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id, volatile_transform(text_col) FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: repartition + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist +(8 rows) + +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id, volatile_transform(text_col) FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: repartition + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist +(8 rows) + SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT id, v FROM ( diff --git a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql index 292f6ca38bf..562456aff38 100644 --- a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql +++ b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql @@ -2,7 +2,7 @@ -- ALLOW_UNSAFE_INSERT_SELECT_PUSHDOWN -- -- Tests citus.allow_unsafe_insert_select_pushdown, which lets a colocated --- INSERT .. SELECT push down GROUP BY / window batching and a volatile batch UDF +-- INSERT .. SELECT push down GROUP BY / window batching and a batch UDF -- to the shards instead of pulling rows to the coordinator. -- CREATE SCHEMA allow_unsafe_insert_select_pushdown; @@ -18,9 +18,10 @@ SELECT create_distributed_table('res', 'text_id'); INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; --- a volatile batch UDF: returns one value per input, mimicking a batched API call +-- a batched UDF: returns one value per input, mimicking a batched API call. +-- immutable + parallel safe, like a real embeddings UDF. CREATE FUNCTION batch_transform(t text[]) RETURNS int[] -LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; +LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; SELECT create_distributed_function('batch_transform(text[])'); -- default off: batching is done after pulling rows to the coordinator @@ -70,10 +71,21 @@ FROM res JOIN dist USING (text_id); -- comparable across Postgres versions. -- --------------------------------------------------------------------- --- branch: volatile function in the SELECT +-- the batched-embeddings benchmark shape: bucket rows into fixed-size batches +-- with row_number()/batch_size, array_agg each batch (id and text in the same +-- order), call the batch UDF once per batch, then unnest back to one row per id. +-- This is the query the GUC is meant to push down to the shards. SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT text_id, (batch_transform(ARRAY[text_col]))[1] FROM dist +SELECT id, val FROM ( + SELECT + unnest(array_agg(text_id ORDER BY text_id)) id, + unnest(batch_transform(array_agg(text_col ORDER BY text_id))) val + FROM ( + SELECT text_id, text_col, (row_number() OVER () - 1) / 100 batch FROM dist + ) q + GROUP BY batch +) s $$, true); -- branch: INSERT distribution column derived from SELECT (not a plain @@ -151,10 +163,22 @@ FROM dist GROUP BY text_id % 10; SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok FROM res JOIN dist USING (text_id); --- --------------------------------------------------------------------- --- Paths where the GUC does NOT help. It only relaxes colocated --- INSERT .. SELECT into a distributed table. --- --------------------------------------------------------------------- +-- correctness for the batched-embeddings benchmark shape: every text_id keeps +-- its own value after batching, the UDF call, and the unnest zip-back +TRUNCATE res; +INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT + unnest(array_agg(text_id ORDER BY text_id)) id, + unnest(batch_transform(array_agg(text_col ORDER BY text_id))) val + FROM ( + SELECT text_id, text_col, (row_number() OVER () - 1) / 100 batch FROM dist + ) q + GROUP BY batch +) s; + +SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok +FROM res JOIN dist USING (text_id); -- reference table target: the GUC does not apply. The plan is a coordinator -- merge whether the GUC is disabled or enabled (identical), unlike the @@ -174,8 +198,24 @@ EXPLAIN (COSTS OFF) INSERT INTO ref(text_id, val) SELECT text_id % 10, count(*)::int FROM dist GROUP BY text_id % 10 $$, true); --- a GROUP BY on a non-distribution column into a distributed target pushes --- down with the GUC enabled ... +-- volatile functions: the GUC relaxes only grouping / partition-column matching, +-- never volatility. A volatile function in the SELECT is still not pushed to the +-- shards with the GUC enabled -- it falls back to a coordinator plan, identical +-- to the GUC-disabled case. +CREATE FUNCTION volatile_transform(t text) RETURNS int +LANGUAGE sql VOLATILE AS $$ SELECT length(t) $$; + +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id, volatile_transform(text_col) FROM dist +$$, true); + +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id, volatile_transform(text_col) FROM dist +$$, true); SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT id, v FROM ( From 051a842798ff91be1ee8cd7465731e67cba4c85b Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Wed, 1 Jul 2026 13:50:28 +0000 Subject: [PATCH 15/25] defensive --- .../planner/insert_select_planner.c | 75 ++- .../planner/multi_logical_optimizer.c | 135 ++++++ .../planner/multi_router_planner.c | 21 +- src/backend/distributed/shared_library_init.c | 15 +- .../distributed/multi_logical_optimizer.h | 1 + .../allow_unsafe_insert_select_pushdown.out | 450 +++++++++++++----- src/test/regress/expected/pg18.out | 7 +- .../allow_unsafe_insert_select_pushdown.sql | 175 +++++-- src/test/regress/sql/pg18.sql | 7 +- 9 files changed, 708 insertions(+), 178 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index e89ea207442..92a7869f883 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -89,6 +89,9 @@ static DeferredErrorMessage * InsertPartitionColumnMatchesSelect(Query *query, subqueryRte, Oid * selectPartitionColumnTableId); +static bool InsertPartitionColumnIsBatchPassThrough(Query *query, + RangeTblEntry *insertRte, + RangeTblEntry *subqueryRte); static DistributedPlan * CreateNonPushableInsertSelectPlan(uint64 planId, Query *parse, ParamListInfo boundParams); static DeferredErrorMessage * NonPushableInsertSelectSupported(Query *insertSelectQuery); @@ -825,17 +828,27 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, /* * Ensure that INSERT's partition column comes from SELECT's partition - * column. With unsafe INSERT ... SELECT pushdown enabled the - * distribution value is typically derived (e.g. - * unnest(array_agg(text_id))), so we cannot match it to a plain SELECT - * partition column; the co-location check further below keeps batches - * shard-local. + * column. Normally this requires a plain-Var partition-column match. + * With unsafe INSERT ... SELECT pushdown enabled we additionally accept + * a provably shard-local batch pass-through of the distribution column, + * i.e. unnest(array_agg(dist_key)); those values can only hash back into + * the current shard, so the co-location check further below keeps the + * batch shard-local. Any other derived distribution value is still + * rejected, because it could route rows that actually belong to a different + * shard. */ - if (!AllowUnsafeInsertSelectPushdown && HasDistributionKey(targetRelationId)) + if (HasDistributionKey(targetRelationId)) { error = InsertPartitionColumnMatchesSelect(queryTree, insertRte, subqueryRte, &selectPartitionColumnTableId); + if (error && AllowUnsafeInsertSelectPushdown && + InsertPartitionColumnIsBatchPassThrough(queryTree, insertRte, + subqueryRte)) + { + error = NULL; + } + if (error) { return error; @@ -1206,6 +1219,56 @@ ReorderInsertSelectTargetLists(Query *originalQuery, RangeTblEntry *insertRte, } +/* + * InsertPartitionColumnIsBatchPassThrough returns true if the SELECT target + * entry that feeds the INSERT's partition column is a provably shard-local + * unnest(array_agg()) batch pass-through. It mirrors the + * partition-column position mapping used by InsertPartitionColumnMatchesSelect + * and then defers the actual pattern check to IsBatchUnnestArrayAggPartitionColumn. + */ +static bool +InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, + RangeTblEntry *subqueryRte) +{ + Oid insertRelationId = insertRte->relid; + Var *insertPartitionColumn = PartitionColumn(insertRelationId, 1); + Query *subquery = subqueryRte->subquery; + + ListCell *targetEntryCell = NULL; + foreach(targetEntryCell, query->targetList) + { + TargetEntry *targetEntry = (TargetEntry *) lfirst(targetEntryCell); + List *insertTargetEntryColumnList = pull_var_clause_default((Node *) targetEntry); + + if (list_length(insertTargetEntryColumnList) != 1) + { + continue; + } + + Var *insertVar = (Var *) linitial(insertTargetEntryColumnList); + + /* skip processing of target table non-partition columns */ + if (targetEntry->resno != insertPartitionColumn->varattno) + { + continue; + } + + if (insertVar->varattno > list_length(subquery->targetList)) + { + return false; + } + + TargetEntry *subqueryTargetEntry = list_nth(subquery->targetList, + insertVar->varattno - 1); + + return IsBatchUnnestArrayAggPartitionColumn(subqueryTargetEntry->expr, + subquery); + } + + return false; +} + + /* * InsertPartitionColumnMatchesSelect returns NULL the partition column in the * table targeted by INSERTed matches with the any of the SELECTed table's diff --git a/src/backend/distributed/planner/multi_logical_optimizer.c b/src/backend/distributed/planner/multi_logical_optimizer.c index 3a4aef85042..2f8924d46f4 100644 --- a/src/backend/distributed/planner/multi_logical_optimizer.c +++ b/src/backend/distributed/planner/multi_logical_optimizer.c @@ -4763,6 +4763,141 @@ IsPartitionColumn(Expr *columnExpression, Query *query, bool skipOuterVars) } +/* + * IsBatchUnnestArrayAggPartitionColumn returns true if the given SELECT target + * expression is a provably shard-local "batch pass-through" of the distribution + * column, i.e. it has the shape + * + * unnest(array_agg()) + * + * (optionally projected through one or more plain-Var subquery indirections, and + * with array_agg allowed to carry ORDER BY / DISTINCT / FILTER modifiers). + * + * Such an expression only ever emits distribution-column values that were read + * from rows of the current shard, so - given source and target are colocated - + * every produced value hashes back into this shard's range. That makes it safe + * to push a colocated INSERT ... SELECT down to the shards even though the + * distribution value is technically a derived expression rather than a plain + * Var. Any intermediate transformation (e.g. unnest(array_agg(dist_key + 1)) or + * unnest(f(array_agg(dist_key)))) is rejected because it could produce values + * that would hash to a different shard. + */ +bool +IsBatchUnnestArrayAggPartitionColumn(Expr *expr, Query *query) +{ + Query *leafQuery = query; + + /* + * Peel plain-Var subquery projection indirection down to the underlying + * expression. We only follow the simple "SELECT FROM (subquery)" + * projection form; anything else makes us conservatively bail out. + */ + for (;;) + { + expr = (Expr *) strip_implicit_coercions((Node *) expr); + + if (!IsA(expr, Var)) + { + break; + } + + Var *var = (Var *) expr; + if (var->varlevelsup != 0 || var->varattno <= InvalidAttrNumber) + { + return false; + } + + if (var->varno <= 0 || var->varno > list_length(leafQuery->rtable)) + { + return false; + } + + RangeTblEntry *rte = rt_fetch(var->varno, leafQuery->rtable); + if (rte->rtekind != RTE_SUBQUERY) + { + return false; + } + + Query *subquery = rte->subquery; + if (var->varattno > list_length(subquery->targetList)) + { + return false; + } + + TargetEntry *targetEntry = list_nth(subquery->targetList, var->varattno - 1); + expr = targetEntry->expr; + leafQuery = subquery; + } + + /* the leaf expression must be unnest(...) over a single array argument */ + if (!IsA(expr, FuncExpr)) + { + return false; + } + + FuncExpr *unnestExpr = (FuncExpr *) expr; + if (unnestExpr->funcid != F_UNNEST_ANYARRAY || + list_length(unnestExpr->args) != 1) + { + return false; + } + + /* the unnest argument must be array_agg(...) with no wrapping transform */ + Expr *unnestArg = (Expr *) strip_implicit_coercions( + (Node *) linitial(unnestExpr->args)); + if (!IsA(unnestArg, Aggref)) + { + return false; + } + + Aggref *arrayAgg = (Aggref *) unnestArg; + if (arrayAgg->aggfnoid != F_ARRAY_AGG_ANYNONARRAY && + arrayAgg->aggfnoid != F_ARRAY_AGG_ANYARRAY) + { + return false; + } + + /* + * Locate the single aggregated value argument. ORDER BY keys that differ + * from the aggregated value appear as additional resjunk target entries; + * they only reorder the (shard-local) values and do not affect routing. + */ + TargetEntry *aggValueTargetEntry = NULL; + TargetEntry *aggArgTargetEntry = NULL; + foreach_declared_ptr(aggArgTargetEntry, arrayAgg->args) + { + if (aggArgTargetEntry->resjunk) + { + continue; + } + + if (aggValueTargetEntry != NULL) + { + /* + * array_agg (the OIDs we matched above) is a single-argument + * aggregate, so a second non-resjunk entry cannot occur for a + * well-formed tree. Assert to catch a broken invariant in debug + * builds, but still bail out gracefully in production: a false + * result only forgoes the optimization, it is never unsafe. + */ + Assert(false); + return false; + } + + aggValueTargetEntry = aggArgTargetEntry; + } + + if (aggValueTargetEntry == NULL) + { + return false; + } + + /* the aggregated value must be the untransformed source partition column */ + bool skipOuterVars = false; + return IsPartitionColumn(aggValueTargetEntry->expr, leafQuery, skipOuterVars); +} + + /* * FindReferencedTableColumn recursively traverses query tree to find actual relation * id, and column that columnExpression refers to. If columnExpression is a diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index 6c05edc9d1f..bea08f039b9 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -370,15 +370,24 @@ AddPartitionKeyNotNullFilterToSelect(Query *subqery) } /* - * Normally the SELECT projects the distribution column as a plain Var. - * With unsafe INSERT ... SELECT pushdown the distribution column may be - * derived (e.g. unnest(array_agg(text_id))), so there is no Var to attach a - * NOT NULL filter to; the batch stays shard-local so the filter is - * unnecessary. + * Normally the SELECT projects the distribution column as a plain Var. With + * unsafe INSERT ... SELECT pushdown the distribution column may instead be a + * provably shard-local batch pass-through, i.e. unnest(array_agg(dist_col)). + * In that case there is no plain Var to attach a NOT NULL filter to, and the + * batch stays shard-local, so the filter is unnecessary; skip it. Any other + * derived distribution value is rejected earlier during planning, so if we + * reach here without a plain-Var partition column it must be this pattern. */ if (targetPartitionColumnVar == NULL && AllowUnsafeInsertSelectPushdown) { - return; + TargetEntry *batchTargetEntry = NULL; + foreach_declared_ptr(batchTargetEntry, targetList) + { + if (IsBatchUnnestArrayAggPartitionColumn(batchTargetEntry->expr, subqery)) + { + return; + } + } } /* we should have found target partition column */ diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c index a75e0d4df41..058fe764be3 100644 --- a/src/backend/distributed/shared_library_init.c +++ b/src/backend/distributed/shared_library_init.c @@ -1062,13 +1062,16 @@ RegisterCitusConfigVariables(void) gettext_noop("Allows pushdown of otherwise-unsafe colocated " "INSERT ... SELECT queries."), gettext_noop("When enabled, Citus relaxes safety checks (GROUP BY / window / " - "aggregate / DISTINCT constructs on non-distribution columns, " - "partition-column matching) for colocated " - "INSERT ... SELECT, so that batching and any batch UDF call run " - "on the shards instead of pulling data to the coordinator. " + "aggregate / DISTINCT constructs on non-distribution columns) " + "for colocated INSERT ... SELECT, so that batching and any batch " + "UDF call run on the shards instead of pulling data to the " + "coordinator. The INSERT's distribution column may then be a " + "provably shard-local unnest(array_agg()); " + "any other derived distribution value is still rejected, since it " + "could route rows that actually belong to a different shard. " "Colocation is still enforced, but the user takes responsibility " - "for keeping batches shard-local and order-preserving; otherwise " - "results may be silently incorrect."), + "for keeping batches order-preserving; otherwise results may be " + "silently incorrect."), &AllowUnsafeInsertSelectPushdown, false, PGC_USERSET, diff --git a/src/include/distributed/multi_logical_optimizer.h b/src/include/distributed/multi_logical_optimizer.h index 940cbc12358..d62bd7b0e7f 100644 --- a/src/include/distributed/multi_logical_optimizer.h +++ b/src/include/distributed/multi_logical_optimizer.h @@ -176,6 +176,7 @@ extern List * SubqueryMultiTableList(MultiNode *multiNode); extern List * GroupTargetEntryList(List *groupClauseList, List *targetEntryList); extern bool ExtractQueryWalker(Node *node, List **queryList); extern bool IsPartitionColumn(Expr *columnExpression, Query *query, bool skipOuterVars); +extern bool IsBatchUnnestArrayAggPartitionColumn(Expr *expr, Query *query); extern void FindReferencedTableColumn(Expr *columnExpression, List *parentQueryList, Query *query, Var **column, RangeTblEntry **rteContainingReferencedColumn, diff --git a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out index 6a9a6466012..d79f779669a 100644 --- a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out +++ b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out @@ -2,8 +2,15 @@ -- ALLOW_UNSAFE_INSERT_SELECT_PUSHDOWN -- -- Tests citus.allow_unsafe_insert_select_pushdown, which lets a colocated --- INSERT .. SELECT push down GROUP BY / window batching and a batch UDF --- to the shards instead of pulling rows to the coordinator. +-- INSERT .. SELECT push GROUP BY / window / DISTINCT batching and a batch UDF +-- down to the shards instead of pulling rows to the coordinator. +-- +-- The distribution column of the target must still come from the source +-- distribution column unchanged: either as a plain Var, or as the provably +-- shard-local batch pass-through unnest(array_agg(dist_col)). Any other derived +-- distribution value (an arithmetic/function transform, or a transform wrapped +-- inside the array_agg) is rejected even with the GUC enabled, because it could +-- route a row to a different shard. -- CREATE SCHEMA allow_unsafe_insert_select_pushdown; SET search_path = allow_unsafe_insert_select_pushdown; @@ -26,7 +33,7 @@ SELECT create_distributed_table('res', 'text_id'); INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; -- a batched UDF: returns one value per input, mimicking a batched API call. --- immutable + parallel safe, like a real embeddings UDF. +-- immutable + parallel safe, like a real batch UDF. CREATE FUNCTION batch_transform(t text[]) RETURNS int[] LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; SELECT create_distributed_function('batch_transform(text[])'); @@ -116,15 +123,17 @@ FROM res JOIN dist USING (text_id); (1 row) -- --------------------------------------------------------------------- --- Per-branch coverage. Each construct below used to force a coordinator merge --- (or was rejected outright); with the GUC enabled the whole colocated --- INSERT .. SELECT is pushed to the shards. The pushed-down plan is a --- Custom Scan (Citus Adaptive) whose task runs the INSERT on a shard, with no --- Distributed Subplan / intermediate results. explain_filter keeps the plan --- comparable across Postgres versions. +-- Positive per-branch coverage. Each construct below used to force a +-- coordinator merge (or was rejected outright). With the GUC enabled the whole +-- colocated INSERT .. SELECT is pushed to the shards because the distribution +-- column is either a plain partition-column Var or the unnest(array_agg(text_id)) +-- batch pass-through. The pushed-down plan is a Custom Scan (Citus Adaptive) +-- whose task runs the INSERT on a shard, with no Distributed Subplan / +-- intermediate results. explain_filter keeps the plan comparable across +-- Postgres versions. -- --------------------------------------------------------------------- --- the batched-embeddings benchmark shape: bucket rows into fixed-size batches --- with row_number()/batch_size, array_agg each batch (id and text in the same +-- the batched benchmark shape: bucket rows into fixed-size batches with +-- row_number()/batch_size, array_agg each batch (id and text in the same -- order), call the batch UDF once per batch, then unnest back to one row per id. -- This is the query the GUC is meant to push down to the shards. SELECT public.explain_filter($$ @@ -158,24 +167,8 @@ $$, true); -> Seq Scan on dist_14000000 dist (15 rows) --- branch: INSERT distribution column derived from SELECT (not a plain --- partition-column Var, so InsertPartitionColumnMatchesSelect is skipped) -SELECT public.explain_filter($$ -EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT text_id + 0, length(text_col) FROM dist -$$, true); - explain_filter ---------------------------------------------------------------------- - Custom Scan (Citus Adaptive) - Task Count: 4 - Tasks Shown: One of 4 - -> Task - Node: host=localhost port=xxxxx dbname=regression - -> Insert on res_14000004 citus_table_alias - -> Seq Scan on dist_14000000 dist -(7 rows) - --- branch: GROUP BY on a non-distribution column +-- branch: GROUP BY on a non-distribution column, distribution column projected +-- as the unnest(array_agg(text_id)) batch pass-through SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) @@ -234,7 +227,8 @@ $$, true); -> Seq Scan on dist_14000000 dist (10 rows) --- branch: window function not partitioned on the distribution column +-- branch: window function not partitioned on the distribution column, with the +-- distribution column projected as a plain partition-column Var SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT text_id, row_number() OVER (ORDER BY text_col) FROM dist @@ -255,25 +249,7 @@ $$, true); Filter: (text_id IS NOT NULL) (12 rows) --- branch: DISTINCT without the distribution column -SELECT public.explain_filter($$ -EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT DISTINCT length(text_col), text_id % 7 FROM dist -$$, true); - explain_filter ---------------------------------------------------------------------- - Custom Scan (Citus Adaptive) - Task Count: 4 - Tasks Shown: One of 4 - -> Task - Node: host=localhost port=xxxxx dbname=regression - -> Insert on res_14000004 citus_table_alias - -> HashAggregate - Group Key: length(dist.text_col), (dist.text_id % 7) - -> Seq Scan on dist_14000000 dist -(9 rows) - --- combination: GROUP BY + DISTINCT +-- combination: GROUP BY + DISTINCT, distribution column as the batch pass-through SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT DISTINCT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) @@ -324,28 +300,298 @@ $$, true); -> Seq Scan on dist_14000000 dist (15 rows) --- subquery: a relaxed construct (DISTINCT without the distribution column) --- nested one level down still pushes down (checks recurse into subqueries) +-- --------------------------------------------------------------------- +-- Negative coverage: pattern requirement. With the GUC enabled the batching +-- relaxations still fire, but the distribution column is a *transformed* value +-- rather than the source partition column, so the query is not pushed down and +-- falls back to a coordinator merge -- identical to the GUC-disabled plan. +-- --------------------------------------------------------------------- +-- distribution column derived by arithmetic on the partition column +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id + 0, length(text_col) FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: repartition + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist +(8 rows) + +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id + 0, length(text_col) FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: repartition + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist +(8 rows) + +-- distribution column derived from a non-distribution column (DISTINCT) +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT DISTINCT length(text_col), text_id % 7 FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> HashAggregate + Group Key: remote_scan.text_id, remote_scan.val + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> HashAggregate + Group Key: length(text_col), (text_id % 7) + -> Seq Scan on dist_14000000 dist +(12 rows) + +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT DISTINCT length(text_col), text_id % 7 FROM dist +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> HashAggregate + Group Key: remote_scan.text_id, remote_scan.val + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> HashAggregate + Group Key: length(text_col), (text_id % 7) + -> Seq Scan on dist_14000000 dist +(12 rows) + +-- distribution column derived from a non-distribution column, one subquery down +SET citus.allow_unsafe_insert_select_pushdown TO off; SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT k, k FROM ( SELECT DISTINCT length(text_col) k FROM dist ) s $$, true); - explain_filter + explain_filter --------------------------------------------------------------------- - Custom Scan (Citus Adaptive) - Task Count: 4 - Tasks Shown: One of 4 - -> Task - Node: host=localhost port=xxxxx dbname=regression - -> Insert on res_14000004 citus_table_alias + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> HashAggregate + Group Key: remote_scan.k + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> HashAggregate + Group Key: length(text_col) + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Function Scan on read_intermediate_result intermediate_result +(19 rows) + +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT k, k FROM ( + SELECT DISTINCT length(text_col) k FROM dist +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> HashAggregate + Group Key: remote_scan.k + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> HashAggregate + Group Key: length(text_col) + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Function Scan on read_intermediate_result intermediate_result +(19 rows) + +-- the batched benchmark shape, but with the distribution column transformed +-- *inside* array_agg (unnest(array_agg(text_id + 1))): the batch pass-through no +-- longer carries the untransformed partition column, so the same shape that +-- pushes down above is now rejected and falls back to a coordinator merge +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id + 1)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> WindowAgg + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression -> Subquery Scan on s - -> HashAggregate - Group Key: length(dist.text_col) - -> Seq Scan on dist_14000000 dist -(10 rows) + -> ProjectSet + -> HashAggregate + Group Key: intermediate_result.b + -> Function Scan on read_intermediate_result intermediate_result +(20 rows) + +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id + 1)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> WindowAgg + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Subquery Scan on s + -> ProjectSet + -> HashAggregate + Group Key: intermediate_result.b + -> Function Scan on read_intermediate_result intermediate_result +(20 rows) + +-- the batched benchmark shape, but with a transform wrapping array_agg for the +-- distribution column (unnest(batch_transform(array_agg(text_col)))): the unnest +-- argument must be a bare array_agg of the partition column, so this is rejected +-- too and falls back to a coordinator merge +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(batch_transform(array_agg(text_col))) id, + unnest(array_agg(text_id)) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> WindowAgg + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Subquery Scan on s + -> ProjectSet + -> HashAggregate + Group Key: intermediate_result.b + -> Function Scan on read_intermediate_result intermediate_result +(20 rows) +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(batch_transform(array_agg(text_col))) id, + unnest(array_agg(text_id)) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> WindowAgg + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Subquery Scan on s + -> ProjectSet + -> HashAggregate + Group Key: intermediate_result.b + -> Function Scan on read_intermediate_result intermediate_result +(20 rows) + +-- --------------------------------------------------------------------- +-- Correctness of the pushed-down batches. +-- --------------------------------------------------------------------- -- correctness for a GROUP BY batch: per-shard aggregation keeps each text_id -- matched to its own value TRUNCATE res; @@ -359,8 +605,8 @@ FROM res JOIN dist USING (text_id); 500 | 500 (1 row) --- correctness for the batched-embeddings benchmark shape: every text_id keeps --- its own value after batching, the UDF call, and the unnest zip-back +-- correctness for the batched benchmark shape: every text_id keeps its own value +-- after batching, the UDF call, and the unnest zip-back TRUNCATE res; INSERT INTO res(text_id, val) SELECT id, val FROM ( @@ -379,6 +625,9 @@ FROM res JOIN dist USING (text_id); 500 | 500 (1 row) +-- --------------------------------------------------------------------- +-- Negative coverage: constructs the GUC never relaxes. +-- --------------------------------------------------------------------- -- reference table target: the GUC does not apply. The plan is a coordinator -- merge whether the GUC is disabled or enabled (identical), unlike the -- distributed-target cases above which push down once the GUC is enabled. @@ -471,87 +720,66 @@ $$, true); -> Seq Scan on dist_14000000 dist (8 rows) +-- a LIMIT forces a coordinator merge even with the GUC enabled and an otherwise +-- pushdown-eligible plan (grouped on the distribution column); LIMIT/OFFSET are +-- never relaxed SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT id, v FROM ( - SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 -) s -$$, true); - explain_filter ---------------------------------------------------------------------- - Custom Scan (Citus Adaptive) - Task Count: 4 - Tasks Shown: One of 4 - -> Task - Node: host=localhost port=xxxxx dbname=regression - -> Insert on res_14000004 citus_table_alias - -> HashAggregate - Group Key: (dist.text_id % 10) - -> Seq Scan on dist_14000000 dist -(9 rows) - --- ... but a LIMIT in the subquery still forces a coordinator merge even with --- the GUC enabled (LIMIT/OFFSET are never relaxed) -SELECT public.explain_filter($$ -EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT id, v FROM ( - SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 LIMIT 100 +SELECT id, val FROM ( + SELECT text_id id, count(*)::int val FROM dist GROUP BY text_id LIMIT 100 ) s $$, true); - explain_filter + explain_filter --------------------------------------------------------------------- Custom Scan (Citus INSERT ... SELECT) INSERT/SELECT method: pull to coordinator -> Custom Scan (Citus Adaptive) -> Distributed Subplan XXX_1 -> Limit - -> HashAggregate - Group Key: remote_scan.id - -> Custom Scan (Citus Adaptive) - Task Count: 4 - Tasks Shown: One of 4 - -> Task - Node: host=localhost port=xxxxx dbname=regression + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Limit -> HashAggregate - Group Key: (text_id % 10) + Group Key: text_id -> Seq Scan on dist_14000000 dist Task Count: 1 Tasks Shown: All -> Task Node: host=localhost port=xxxxx dbname=regression -> Function Scan on read_intermediate_result intermediate_result -(20 rows) +(19 rows) -- ... and likewise for OFFSET SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT id, v FROM ( - SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 OFFSET 5 +SELECT id, val FROM ( + SELECT text_id id, count(*)::int val FROM dist GROUP BY text_id OFFSET 5 ) s $$, true); - explain_filter + explain_filter --------------------------------------------------------------------- Custom Scan (Citus INSERT ... SELECT) INSERT/SELECT method: pull to coordinator -> Custom Scan (Citus Adaptive) -> Distributed Subplan XXX_1 -> Limit - -> HashAggregate - Group Key: remote_scan.id - -> Custom Scan (Citus Adaptive) - Task Count: 4 - Tasks Shown: One of 4 - -> Task - Node: host=localhost port=xxxxx dbname=regression - -> HashAggregate - Group Key: (text_id % 10) - -> Seq Scan on dist_14000000 dist + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> HashAggregate + Group Key: text_id + -> Seq Scan on dist_14000000 dist Task Count: 1 Tasks Shown: All -> Task Node: host=localhost port=xxxxx dbname=regression -> Function Scan on read_intermediate_result intermediate_result -(20 rows) +(18 rows) SET client_min_messages TO WARNING; DROP SCHEMA allow_unsafe_insert_select_pushdown CASCADE; diff --git a/src/test/regress/expected/pg18.out b/src/test/regress/expected/pg18.out index 9fe61b4808d..e53d916000c 100644 --- a/src/test/regress/expected/pg18.out +++ b/src/test/regress/expected/pg18.out @@ -3192,9 +3192,10 @@ SELECT create_distributed_table('res', 'text_id'); (1 row) INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; --- a volatile batch UDF: returns one value per input, mimicking a batched API call +-- a batch UDF: returns one value per input, mimicking a batched API call. +-- immutable + parallel safe, like a real batch UDF. CREATE FUNCTION batch_transform(t text[]) RETURNS int[] -LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; +LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; SELECT create_distributed_function('batch_transform(text[])'); NOTICE: procedure pg18_insert_select_pushdown.batch_transform is already distributed DETAIL: Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] commands @@ -3204,7 +3205,7 @@ DETAIL: Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] (1 row) SET citus.allow_unsafe_insert_select_pushdown TO on; --- the batching and volatile batch UDF run on the shards; the raw PG18 plan +-- the batching and batch UDF run on the shards; the raw PG18 plan -- includes the WindowAgg "Window:" line EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT id, val FROM ( diff --git a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql index 562456aff38..4f154da7568 100644 --- a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql +++ b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql @@ -2,8 +2,15 @@ -- ALLOW_UNSAFE_INSERT_SELECT_PUSHDOWN -- -- Tests citus.allow_unsafe_insert_select_pushdown, which lets a colocated --- INSERT .. SELECT push down GROUP BY / window batching and a batch UDF --- to the shards instead of pulling rows to the coordinator. +-- INSERT .. SELECT push GROUP BY / window / DISTINCT batching and a batch UDF +-- down to the shards instead of pulling rows to the coordinator. +-- +-- The distribution column of the target must still come from the source +-- distribution column unchanged: either as a plain Var, or as the provably +-- shard-local batch pass-through unnest(array_agg(dist_col)). Any other derived +-- distribution value (an arithmetic/function transform, or a transform wrapped +-- inside the array_agg) is rejected even with the GUC enabled, because it could +-- route a row to a different shard. -- CREATE SCHEMA allow_unsafe_insert_select_pushdown; SET search_path = allow_unsafe_insert_select_pushdown; @@ -19,7 +26,7 @@ SELECT create_distributed_table('res', 'text_id'); INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; -- a batched UDF: returns one value per input, mimicking a batched API call. --- immutable + parallel safe, like a real embeddings UDF. +-- immutable + parallel safe, like a real batch UDF. CREATE FUNCTION batch_transform(t text[]) RETURNS int[] LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; SELECT create_distributed_function('batch_transform(text[])'); @@ -63,16 +70,18 @@ SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok FROM res JOIN dist USING (text_id); -- --------------------------------------------------------------------- --- Per-branch coverage. Each construct below used to force a coordinator merge --- (or was rejected outright); with the GUC enabled the whole colocated --- INSERT .. SELECT is pushed to the shards. The pushed-down plan is a --- Custom Scan (Citus Adaptive) whose task runs the INSERT on a shard, with no --- Distributed Subplan / intermediate results. explain_filter keeps the plan --- comparable across Postgres versions. +-- Positive per-branch coverage. Each construct below used to force a +-- coordinator merge (or was rejected outright). With the GUC enabled the whole +-- colocated INSERT .. SELECT is pushed to the shards because the distribution +-- column is either a plain partition-column Var or the unnest(array_agg(text_id)) +-- batch pass-through. The pushed-down plan is a Custom Scan (Citus Adaptive) +-- whose task runs the INSERT on a shard, with no Distributed Subplan / +-- intermediate results. explain_filter keeps the plan comparable across +-- Postgres versions. -- --------------------------------------------------------------------- --- the batched-embeddings benchmark shape: bucket rows into fixed-size batches --- with row_number()/batch_size, array_agg each batch (id and text in the same +-- the batched benchmark shape: bucket rows into fixed-size batches with +-- row_number()/batch_size, array_agg each batch (id and text in the same -- order), call the batch UDF once per batch, then unnest back to one row per id. -- This is the query the GUC is meant to push down to the shards. SELECT public.explain_filter($$ @@ -88,14 +97,8 @@ SELECT id, val FROM ( ) s $$, true); --- branch: INSERT distribution column derived from SELECT (not a plain --- partition-column Var, so InsertPartitionColumnMatchesSelect is skipped) -SELECT public.explain_filter($$ -EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT text_id + 0, length(text_col) FROM dist -$$, true); - --- branch: GROUP BY on a non-distribution column +-- branch: GROUP BY on a non-distribution column, distribution column projected +-- as the unnest(array_agg(text_id)) batch pass-through SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) @@ -115,19 +118,14 @@ SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) FROM dist HAVING count(*) > 0 $$, true); --- branch: window function not partitioned on the distribution column +-- branch: window function not partitioned on the distribution column, with the +-- distribution column projected as a plain partition-column Var SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT text_id, row_number() OVER (ORDER BY text_col) FROM dist $$, true); --- branch: DISTINCT without the distribution column -SELECT public.explain_filter($$ -EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT DISTINCT length(text_col), text_id % 7 FROM dist -$$, true); - --- combination: GROUP BY + DISTINCT +-- combination: GROUP BY + DISTINCT, distribution column as the batch pass-through SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT DISTINCT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) @@ -144,8 +142,46 @@ SELECT id, v FROM ( ) s $$, true); --- subquery: a relaxed construct (DISTINCT without the distribution column) --- nested one level down still pushes down (checks recurse into subqueries) +-- --------------------------------------------------------------------- +-- Negative coverage: pattern requirement. With the GUC enabled the batching +-- relaxations still fire, but the distribution column is a *transformed* value +-- rather than the source partition column, so the query is not pushed down and +-- falls back to a coordinator merge -- identical to the GUC-disabled plan. +-- --------------------------------------------------------------------- + +-- distribution column derived by arithmetic on the partition column +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id + 0, length(text_col) FROM dist +$$, true); +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT text_id + 0, length(text_col) FROM dist +$$, true); + +-- distribution column derived from a non-distribution column (DISTINCT) +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT DISTINCT length(text_col), text_id % 7 FROM dist +$$, true); +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT DISTINCT length(text_col), text_id % 7 FROM dist +$$, true); + +-- distribution column derived from a non-distribution column, one subquery down +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT k, k FROM ( + SELECT DISTINCT length(text_col) k FROM dist +) s +$$, true); +SET citus.allow_unsafe_insert_select_pushdown TO on; SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT k, k FROM ( @@ -153,6 +189,60 @@ SELECT k, k FROM ( ) s $$, true); +-- the batched benchmark shape, but with the distribution column transformed +-- *inside* array_agg (unnest(array_agg(text_id + 1))): the batch pass-through no +-- longer carries the untransformed partition column, so the same shape that +-- pushes down above is now rejected and falls back to a coordinator merge +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id + 1)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id + 1)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + +-- the batched benchmark shape, but with a transform wrapping array_agg for the +-- distribution column (unnest(batch_transform(array_agg(text_col)))): the unnest +-- argument must be a bare array_agg of the partition column, so this is rejected +-- too and falls back to a coordinator merge +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(batch_transform(array_agg(text_col))) id, + unnest(array_agg(text_id)) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(batch_transform(array_agg(text_col))) id, + unnest(array_agg(text_id)) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + +-- --------------------------------------------------------------------- +-- Correctness of the pushed-down batches. +-- --------------------------------------------------------------------- + -- correctness for a GROUP BY batch: per-shard aggregation keeps each text_id -- matched to its own value TRUNCATE res; @@ -163,8 +253,8 @@ FROM dist GROUP BY text_id % 10; SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok FROM res JOIN dist USING (text_id); --- correctness for the batched-embeddings benchmark shape: every text_id keeps --- its own value after batching, the UDF call, and the unnest zip-back +-- correctness for the batched benchmark shape: every text_id keeps its own value +-- after batching, the UDF call, and the unnest zip-back TRUNCATE res; INSERT INTO res(text_id, val) SELECT id, val FROM ( @@ -180,6 +270,10 @@ SELECT id, val FROM ( SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok FROM res JOIN dist USING (text_id); +-- --------------------------------------------------------------------- +-- Negative coverage: constructs the GUC never relaxes. +-- --------------------------------------------------------------------- + -- reference table target: the GUC does not apply. The plan is a coordinator -- merge whether the GUC is disabled or enabled (identical), unlike the -- distributed-target cases above which push down once the GUC is enabled. @@ -216,27 +310,22 @@ SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT text_id, volatile_transform(text_col) FROM dist $$, true); -SELECT public.explain_filter($$ -EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT id, v FROM ( - SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 -) s -$$, true); --- ... but a LIMIT in the subquery still forces a coordinator merge even with --- the GUC enabled (LIMIT/OFFSET are never relaxed) +-- a LIMIT forces a coordinator merge even with the GUC enabled and an otherwise +-- pushdown-eligible plan (grouped on the distribution column); LIMIT/OFFSET are +-- never relaxed SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT id, v FROM ( - SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 LIMIT 100 +SELECT id, val FROM ( + SELECT text_id id, count(*)::int val FROM dist GROUP BY text_id LIMIT 100 ) s $$, true); -- ... and likewise for OFFSET SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) -SELECT id, v FROM ( - SELECT text_id % 10 id, count(*)::int v FROM dist GROUP BY text_id % 10 OFFSET 5 +SELECT id, val FROM ( + SELECT text_id id, count(*)::int val FROM dist GROUP BY text_id OFFSET 5 ) s $$, true); diff --git a/src/test/regress/sql/pg18.sql b/src/test/regress/sql/pg18.sql index d585abc20de..8e6e5b505e0 100644 --- a/src/test/regress/sql/pg18.sql +++ b/src/test/regress/sql/pg18.sql @@ -2027,14 +2027,15 @@ SELECT create_distributed_table('res', 'text_id'); INSERT INTO dist SELECT g, 't' || g FROM generate_series(1, 500) g; --- a volatile batch UDF: returns one value per input, mimicking a batched API call +-- a batch UDF: returns one value per input, mimicking a batched API call. +-- immutable + parallel safe, like a real batch UDF. CREATE FUNCTION batch_transform(t text[]) RETURNS int[] -LANGUAGE sql VOLATILE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; +LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$ SELECT array_agg(length(x)) FROM unnest(t) x $$; SELECT create_distributed_function('batch_transform(text[])'); SET citus.allow_unsafe_insert_select_pushdown TO on; --- the batching and volatile batch UDF run on the shards; the raw PG18 plan +-- the batching and batch UDF run on the shards; the raw PG18 plan -- includes the WindowAgg "Window:" line EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT id, val FROM ( From aa9725c6004a2df41af9ad773942804afd5a8c53 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Wed, 1 Jul 2026 17:16:32 +0300 Subject: [PATCH 16/25] Add citus_table_alias to pg18 pushdown expected plan PG18 emits 'Insert on res_... citus_table_alias' for the shard-local INSERT..SELECT pushdown plan; the committed pg18.out was missing the alias, causing the pg18 test (and all Test flakyness shards) to fail on CI. Match CI output. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/test/regress/expected/pg18.out | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/test/regress/expected/pg18.out b/src/test/regress/expected/pg18.out index e53d916000c..419342506c5 100644 --- a/src/test/regress/expected/pg18.out +++ b/src/test/regress/expected/pg18.out @@ -3221,7 +3221,7 @@ SELECT id, val FROM ( Tasks Shown: One of 4 -> Task Node: host=localhost port=xxxxx dbname=regression - -> Insert on res_14100004 + -> Insert on res_14100004 citus_table_alias -> Subquery Scan on s -> ProjectSet -> HashAggregate From 5612f7db8f0bc52e2ac41e5f855fe4b0349721d5 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Thu, 2 Jul 2026 11:25:19 +0300 Subject: [PATCH 17/25] Dedup INSERT distribution-column lookup in insert-select planner Extract the shared INSERT target-list walk that maps the target table's distribution column to the SELECT (subquery) target entry into a single helper, SelectTargetEntryForInsertPartitionColumn(). Both InsertPartitionColumnMatchesSelect() and InsertPartitionColumnIsBatchPassThrough() previously duplicated this lookup; they now call the helper. The helper optionally reports the matching INSERT target entry so InsertPartitionColumnMatchesSelect() can still inspect casting on the distribution column. Pure refactor; no behavior change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../planner/insert_select_planner.c | 387 +++++++++--------- 1 file changed, 197 insertions(+), 190 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 92a7869f883..098318dc5f9 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -92,6 +92,11 @@ static DeferredErrorMessage * InsertPartitionColumnMatchesSelect(Query *query, static bool InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, RangeTblEntry *subqueryRte); +static TargetEntry * SelectTargetEntryForInsertPartitionColumn(Query *query, + RangeTblEntry *insertRte, + RangeTblEntry *subqueryRte, + TargetEntry ** + insertTargetEntry); static DistributedPlan * CreateNonPushableInsertSelectPlan(uint64 planId, Query *parse, ParamListInfo boundParams); static DeferredErrorMessage * NonPushableInsertSelectSupported(Query *insertSelectQuery); @@ -1220,26 +1225,44 @@ ReorderInsertSelectTargetLists(Query *originalQuery, RangeTblEntry *insertRte, /* - * InsertPartitionColumnIsBatchPassThrough returns true if the SELECT target - * entry that feeds the INSERT's partition column is a provably shard-local - * unnest(array_agg()) batch pass-through. It mirrors the - * partition-column position mapping used by InsertPartitionColumnMatchesSelect - * and then defers the actual pattern check to IsBatchUnnestArrayAggPartitionColumn. + * SelectTargetEntryForInsertPartitionColumn walks the INSERT target list to find + * the entry that feeds the target table's distribution column and returns the + * corresponding SELECT (subquery) target entry that produces its value. Returns + * NULL when the INSERT does not project the distribution column via a single Var. + * + * When insertTargetEntry is not NULL, it is set to the matching INSERT target + * entry so that callers can inspect casting on the distribution column. */ -static bool -InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, - RangeTblEntry *subqueryRte) +static TargetEntry * +SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte, + RangeTblEntry *subqueryRte, + TargetEntry **insertTargetEntry) { Oid insertRelationId = insertRte->relid; Var *insertPartitionColumn = PartitionColumn(insertRelationId, 1); Query *subquery = subqueryRte->subquery; + if (insertTargetEntry != NULL) + { + *insertTargetEntry = NULL; + } + ListCell *targetEntryCell = NULL; foreach(targetEntryCell, query->targetList) { TargetEntry *targetEntry = (TargetEntry *) lfirst(targetEntryCell); - List *insertTargetEntryColumnList = pull_var_clause_default((Node *) targetEntry); + List *insertTargetEntryColumnList = + pull_var_clause_default((Node *) targetEntry); + /* + * We only consider target entries that include a single column. Note that + * this is slightly different than directly checking whether the + * targetEntry->expr is a Var since the Var could be wrapped into an + * implicit/explicit casting. + * + * Also note that we skip the target entry if it does not contain a Var, + * which corresponds to columns with DEFAULT values on the target list. + */ if (list_length(insertTargetEntryColumnList) != 1) { continue; @@ -1255,17 +1278,39 @@ InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, if (insertVar->varattno > list_length(subquery->targetList)) { - return false; + return NULL; } - TargetEntry *subqueryTargetEntry = list_nth(subquery->targetList, - insertVar->varattno - 1); + if (insertTargetEntry != NULL) + { + *insertTargetEntry = targetEntry; + } - return IsBatchUnnestArrayAggPartitionColumn(subqueryTargetEntry->expr, - subquery); + return list_nth(subquery->targetList, insertVar->varattno - 1); } - return false; + return NULL; +} + + +/* + * InsertPartitionColumnIsBatchPassThrough returns true if the SELECT target + * entry that feeds the INSERT's distribution column is a provably shard-local + * unnest(array_agg()) batch pass-through. + */ +static bool +InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, + RangeTblEntry *subqueryRte) +{ + TargetEntry *subqueryTargetEntry = + SelectTargetEntryForInsertPartitionColumn(query, insertRte, subqueryRte, NULL); + if (subqueryTargetEntry == NULL) + { + return false; + } + + return IsBatchUnnestArrayAggPartitionColumn(subqueryTargetEntry->expr, + subqueryRte->subquery); } @@ -1282,221 +1327,183 @@ InsertPartitionColumnMatchesSelect(Query *query, RangeTblEntry *insertRte, RangeTblEntry *subqueryRte, Oid *selectPartitionColumnTableId) { - ListCell *targetEntryCell = NULL; - uint32 rangeTableId = 1; - Oid insertRelationId = insertRte->relid; - Var *insertPartitionColumn = PartitionColumn(insertRelationId, rangeTableId); Query *subquery = subqueryRte->subquery; - bool targetTableHasPartitionColumn = false; - foreach(targetEntryCell, query->targetList) + TargetEntry *insertTargetEntry = NULL; + TargetEntry *subqueryTargetEntry = + SelectTargetEntryForInsertPartitionColumn(query, insertRte, subqueryRte, + &insertTargetEntry); + if (subqueryTargetEntry == NULL) { - TargetEntry *targetEntry = (TargetEntry *) lfirst(targetEntryCell); - List *insertTargetEntryColumnList = pull_var_clause_default((Node *) targetEntry); - Var *subqueryPartitionColumn = NULL; - - /* - * We only consider target entries that include a single column. Note that this - * is slightly different than directly checking the whether the targetEntry->expr - * is a var since the var could be wrapped into an implicit/explicit casting. - * - * Also note that we skip the target entry if it does not contain a Var, which - * corresponds to columns with DEFAULT values on the target list. - */ - if (list_length(insertTargetEntryColumnList) != 1) - { - continue; - } + return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, + "cannot perform distributed INSERT INTO ... SELECT " + "because the partition columns in the source table " + "and subquery do not match", + "the query doesn't include the target table's " + "partition column", + NULL); + } - Var *insertVar = (Var *) linitial(insertTargetEntryColumnList); - AttrNumber originalAttrNo = targetEntry->resno; + Expr *selectTargetExpr = subqueryTargetEntry->expr; + + Var *subqueryPartitionColumn = NULL; + RangeTblEntry *subqueryPartitionColumnRelationIdRTE = NULL; + List *parentQueryList = list_make2(query, subquery); + bool skipOuterVars = false; + FindReferencedTableColumn(selectTargetExpr, + parentQueryList, subquery, + &subqueryPartitionColumn, + &subqueryPartitionColumnRelationIdRTE, + skipOuterVars); + Oid subqueryPartitionColumnRelationId = subqueryPartitionColumnRelationIdRTE ? + subqueryPartitionColumnRelationIdRTE->relid : + InvalidOid; - /* skip processing of target table non-partition columns */ - if (originalAttrNo != insertPartitionColumn->varattno) - { - continue; - } + /* + * Corresponding (i.e., in the same ordinal position as the target table's + * partition column) select target entry does not directly belong a table. + * Evaluate its expression type and error out properly. + */ + if (subqueryPartitionColumnRelationId == InvalidOid) + { + char *errorDetailTemplate = "Subquery contains %s in the " + "same position as the target table's " + "partition column."; - /* INSERT query includes the partition column */ - targetTableHasPartitionColumn = true; - - TargetEntry *subqueryTargetEntry = list_nth(subquery->targetList, - insertVar->varattno - 1); - Expr *selectTargetExpr = subqueryTargetEntry->expr; - - RangeTblEntry *subqueryPartitionColumnRelationIdRTE = NULL; - List *parentQueryList = list_make2(query, subquery); - bool skipOuterVars = false; - FindReferencedTableColumn(selectTargetExpr, - parentQueryList, subquery, - &subqueryPartitionColumn, - &subqueryPartitionColumnRelationIdRTE, - skipOuterVars); - Oid subqueryPartitionColumnRelationId = subqueryPartitionColumnRelationIdRTE ? - subqueryPartitionColumnRelationIdRTE-> - relid : - InvalidOid; + char *exprDescription = ""; - /* - * Corresponding (i.e., in the same ordinal position as the target table's - * partition column) select target entry does not directly belong a table. - * Evaluate its expression type and error out properly. - */ - if (subqueryPartitionColumnRelationId == InvalidOid) + switch (selectTargetExpr->type) { - char *errorDetailTemplate = "Subquery contains %s in the " - "same position as the target table's " - "partition column."; - - char *exprDescription = ""; + case T_Const: + { + exprDescription = "a constant value"; + break; + } - switch (selectTargetExpr->type) + case T_OpExpr: { - case T_Const: - { - exprDescription = "a constant value"; - break; - } + exprDescription = "an operator"; + break; + } - case T_OpExpr: - { - exprDescription = "an operator"; - break; - } + case T_FuncExpr: + { + FuncExpr *subqueryFunctionExpr = (FuncExpr *) selectTargetExpr; - case T_FuncExpr: + switch (subqueryFunctionExpr->funcformat) { - FuncExpr *subqueryFunctionExpr = (FuncExpr *) selectTargetExpr; - - switch (subqueryFunctionExpr->funcformat) + case COERCE_EXPLICIT_CALL: { - case COERCE_EXPLICIT_CALL: - { - exprDescription = "a function call"; - break; - } - - case COERCE_EXPLICIT_CAST: - { - exprDescription = "an explicit cast"; - break; - } - - case COERCE_IMPLICIT_CAST: - { - exprDescription = "an implicit cast"; - break; - } - - default: - { - exprDescription = "a function call"; - break; - } + exprDescription = "a function call"; + break; } - break; - } - case T_Aggref: - { - exprDescription = "an aggregation"; - break; - } + case COERCE_EXPLICIT_CAST: + { + exprDescription = "an explicit cast"; + break; + } - case T_CaseExpr: - { - exprDescription = "a case expression"; - break; - } + case COERCE_IMPLICIT_CAST: + { + exprDescription = "an implicit cast"; + break; + } - case T_CoalesceExpr: - { - exprDescription = "a coalesce expression"; - break; + default: + { + exprDescription = "a function call"; + break; + } } + break; + } - case T_RowExpr: - { - exprDescription = "a row expression"; - break; - } + case T_Aggref: + { + exprDescription = "an aggregation"; + break; + } - case T_MinMaxExpr: - { - exprDescription = "a min/max expression"; - break; - } + case T_CaseExpr: + { + exprDescription = "a case expression"; + break; + } - case T_CoerceViaIO: - { - exprDescription = "an explicit coercion"; - break; - } + case T_CoalesceExpr: + { + exprDescription = "a coalesce expression"; + break; + } - default: - { - exprDescription = - "an expression that is not a simple column reference"; - break; - } + case T_RowExpr: + { + exprDescription = "a row expression"; + break; } - return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, - "cannot perform distributed INSERT INTO ... SELECT " - "because the partition columns in the source table " - "and subquery do not match", - psprintf(errorDetailTemplate, exprDescription), - "Ensure the target table's partition column has a " - "corresponding simple column reference to a distributed " - "table's partition column in the subquery."); - } + case T_MinMaxExpr: + { + exprDescription = "a min/max expression"; + break; + } - /* - * Insert target expression could only be non-var if the select target - * entry does not have the same type (i.e., target column requires casting). - */ - if (!IsA(targetEntry->expr, Var)) - { - return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, - "cannot perform distributed INSERT INTO ... SELECT " - "because the partition columns in the source table " - "and subquery do not match", - "The data type of the target table's partition column " - "should exactly match the data type of the " - "corresponding simple column reference in the subquery.", - NULL); - } + case T_CoerceViaIO: + { + exprDescription = "an explicit coercion"; + break; + } - /* finally, check that the select target column is a partition column */ - if (!IsPartitionColumn(selectTargetExpr, subquery, skipOuterVars)) - { - return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, - "cannot perform distributed INSERT INTO ... SELECT " - "because the partition columns in the source table " - "and subquery do not match", - "The target table's partition column should correspond " - "to a partition column in the subquery.", - NULL); + default: + { + exprDescription = + "an expression that is not a simple column reference"; + break; + } } - /* finally, check that the select target column is a partition column */ - /* we can set the select relation id */ - *selectPartitionColumnTableId = subqueryPartitionColumnRelationId; + return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, + "cannot perform distributed INSERT INTO ... SELECT " + "because the partition columns in the source table " + "and subquery do not match", + psprintf(errorDetailTemplate, exprDescription), + "Ensure the target table's partition column has a " + "corresponding simple column reference to a distributed " + "table's partition column in the subquery."); + } - break; + /* + * Insert target expression could only be non-var if the select target + * entry does not have the same type (i.e., target column requires casting). + */ + if (!IsA(insertTargetEntry->expr, Var)) + { + return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, + "cannot perform distributed INSERT INTO ... SELECT " + "because the partition columns in the source table " + "and subquery do not match", + "The data type of the target table's partition column " + "should exactly match the data type of the " + "corresponding simple column reference in the subquery.", + NULL); } - if (!targetTableHasPartitionColumn) + /* finally, check that the select target column is a partition column */ + if (!IsPartitionColumn(selectTargetExpr, subquery, skipOuterVars)) { return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, "cannot perform distributed INSERT INTO ... SELECT " "because the partition columns in the source table " "and subquery do not match", - "the query doesn't include the target table's " - "partition column", + "The target table's partition column should correspond " + "to a partition column in the subquery.", NULL); } + /* we can set the select relation id */ + *selectPartitionColumnTableId = subqueryPartitionColumnRelationId; + return NULL; } From b3740cb2d08d1ed9909f4119716352c87de6116a Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Thu, 2 Jul 2026 11:32:43 +0300 Subject: [PATCH 18/25] Compute batch pass-through at call site instead of reading GUC AddPartitionKeyNotNullFilterToSelect() previously consulted the global GUC AllowUnsafeInsertSelectPushdown (and re-walked the target list) to decide whether to skip the NOT NULL filter for a batch pass-through distribution column. That coupled a low-level deparse/router helper to a planner-level GUC and re-derived a fact the caller already knows. Instead, take a bool distributionColumnIsBatchPassThrough parameter and skip the filter when it is set and there is no plain-Var partition column. Both callers (RouterModifyTaskForShardInterval and deparse_shard_query) compute the flag via the now-exported InsertPartitionColumnIsBatchPassThrough(), which inspects the query shape directly. A batch-shape query with the GUC off is rejected during planning and never reaches these call sites, so the query-shape predicate is the correct signal here. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../distributed/planner/deparse_shard_query.c | 7 +++++- .../planner/insert_select_planner.c | 11 ++++---- .../planner/multi_router_planner.c | 25 +++++++++---------- .../distributed/insert_select_planner.h | 3 +++ .../distributed/multi_router_planner.h | 4 ++- 5 files changed, 30 insertions(+), 20 deletions(-) diff --git a/src/backend/distributed/planner/deparse_shard_query.c b/src/backend/distributed/planner/deparse_shard_query.c index 4b3d3664e87..9d6f080a26b 100644 --- a/src/backend/distributed/planner/deparse_shard_query.c +++ b/src/backend/distributed/planner/deparse_shard_query.c @@ -109,7 +109,12 @@ RebuildQueryStrings(Job *workerJob) /* there are no restrictions to add for reference and citus local tables */ if (IsCitusTableType(shardInterval->relationId, DISTRIBUTED_TABLE)) { - AddPartitionKeyNotNullFilterToSelect(copiedSubquery); + bool distributionColumnIsBatchPassThrough = + InsertPartitionColumnIsBatchPassThrough(query, copiedInsertRte, + copiedSubqueryRte); + AddPartitionKeyNotNullFilterToSelect(copiedSubquery, + distributionColumnIsBatchPassThrough) + ; } ReorderInsertSelectTargetLists(query, copiedInsertRte, copiedSubqueryRte); diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 098318dc5f9..245fdb08315 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -89,9 +89,6 @@ static DeferredErrorMessage * InsertPartitionColumnMatchesSelect(Query *query, subqueryRte, Oid * selectPartitionColumnTableId); -static bool InsertPartitionColumnIsBatchPassThrough(Query *query, - RangeTblEntry *insertRte, - RangeTblEntry *subqueryRte); static TargetEntry * SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte, RangeTblEntry *subqueryRte, @@ -983,7 +980,11 @@ RouterModifyTaskForShardInterval(Query *originalQuery, copiedSubquery); if (subqueryRteListProperties->hasDistTableWithShardKey) { - AddPartitionKeyNotNullFilterToSelect(copiedSubquery); + bool distributionColumnIsBatchPassThrough = + InsertPartitionColumnIsBatchPassThrough(copiedQuery, copiedInsertRte, + copiedSubqueryRte); + AddPartitionKeyNotNullFilterToSelect(copiedSubquery, + distributionColumnIsBatchPassThrough); } /* mark that we don't want the router planner to generate dummy hosts/queries */ @@ -1298,7 +1299,7 @@ SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte * entry that feeds the INSERT's distribution column is a provably shard-local * unnest(array_agg()) batch pass-through. */ -static bool +bool InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, RangeTblEntry *subqueryRte) { diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index bea08f039b9..1ad9e7f35d6 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -347,9 +347,15 @@ ShardIntervalOpExpressions(ShardInterval *shardInterval, Index rteIndex) * * The function expects and asserts that subquery's target list contains a partition * column value. Thus, this function should never be called with reference tables. + * + * The exception is unsafe INSERT ... SELECT pushdown, where the distribution + * column may be a provably shard-local batch pass-through with no plain Var to + * filter on; callers signal that case via distributionColumnIsBatchPassThrough so + * the filter is skipped. */ void -AddPartitionKeyNotNullFilterToSelect(Query *subqery) +AddPartitionKeyNotNullFilterToSelect(Query *subqery, bool + distributionColumnIsBatchPassThrough) { List *targetList = subqery->targetList; ListCell *targetEntryCell = NULL; @@ -374,20 +380,13 @@ AddPartitionKeyNotNullFilterToSelect(Query *subqery) * unsafe INSERT ... SELECT pushdown the distribution column may instead be a * provably shard-local batch pass-through, i.e. unnest(array_agg(dist_col)). * In that case there is no plain Var to attach a NOT NULL filter to, and the - * batch stays shard-local, so the filter is unnecessary; skip it. Any other - * derived distribution value is rejected earlier during planning, so if we - * reach here without a plain-Var partition column it must be this pattern. + * batch stays shard-local, so the filter is unnecessary; skip it. The caller + * determines whether the query has this shape and passes the result in + * distributionColumnIsBatchPassThrough. */ - if (targetPartitionColumnVar == NULL && AllowUnsafeInsertSelectPushdown) + if (targetPartitionColumnVar == NULL && distributionColumnIsBatchPassThrough) { - TargetEntry *batchTargetEntry = NULL; - foreach_declared_ptr(batchTargetEntry, targetList) - { - if (IsBatchUnnestArrayAggPartitionColumn(batchTargetEntry->expr, subqery)) - { - return; - } - } + return; } /* we should have found target partition column */ diff --git a/src/include/distributed/insert_select_planner.h b/src/include/distributed/insert_select_planner.h index a9100b02dfd..32984401018 100644 --- a/src/include/distributed/insert_select_planner.h +++ b/src/include/distributed/insert_select_planner.h @@ -46,6 +46,9 @@ extern DistributedPlan * CreateInsertSelectIntoLocalTablePlan(uint64 planId, extern char * InsertSelectResultIdPrefix(uint64 planId); extern bool PlanningInsertSelect(void); extern Query * WrapSubquery(Query *subquery); +extern bool InsertPartitionColumnIsBatchPassThrough(Query *query, + RangeTblEntry *insertRte, + RangeTblEntry *subqueryRte); #endif /* INSERT_SELECT_PLANNER_H */ diff --git a/src/include/distributed/multi_router_planner.h b/src/include/distributed/multi_router_planner.h index af3421e965c..22763f97053 100644 --- a/src/include/distributed/multi_router_planner.h +++ b/src/include/distributed/multi_router_planner.h @@ -79,7 +79,9 @@ extern RangeTblEntry * ExtractResultRelationRTE(Query *query); extern RangeTblEntry * ExtractResultRelationRTEOrError(Query *query); extern RangeTblEntry * ExtractDistributedInsertValuesRTE(Query *query); extern bool IsMultiRowInsert(Query *query); -extern void AddPartitionKeyNotNullFilterToSelect(Query *subqery); +extern void AddPartitionKeyNotNullFilterToSelect(Query *subqery, + bool distributionColumnIsBatchPassThrough + ); extern bool UpdateOrDeleteOrMergeQuery(Query *query); extern bool IsMergeQuery(Query *query); From 0fd26d26d67fafd3ae26d9ac6c02b316ad9a36f7 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Thu, 2 Jul 2026 11:36:51 +0300 Subject: [PATCH 19/25] Move batch pass-through predicate into insert-select planner IsBatchUnnestArrayAggPartitionColumn() lived in multi_logical_optimizer.c and was exported, but after the previous commit its only caller is InsertPartitionColumnIsBatchPassThrough() in insert_select_planner.c. Move it next to its caller as a file-local static and rename it to IsInsertSelectBatchPassThroughDistributionColumn(), which better describes what it checks and where it is used. Drop the header export and add utils/fmgroids.h to insert_select_planner.c for the array_agg/unnest function OIDs. No behavior change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../planner/insert_select_planner.c | 141 +++++++++++++++++- .../planner/multi_logical_optimizer.c | 135 ----------------- .../distributed/multi_logical_optimizer.h | 1 - 3 files changed, 139 insertions(+), 138 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 245fdb08315..91442408d89 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -26,6 +26,7 @@ #include "parser/parsetree.h" #include "tcop/tcopprot.h" #include "utils/builtins.h" +#include "utils/fmgroids.h" #include "utils/lsyscache.h" #include "utils/rel.h" @@ -94,6 +95,7 @@ static TargetEntry * SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *subqueryRte, TargetEntry ** insertTargetEntry); +static bool IsInsertSelectBatchPassThroughDistributionColumn(Expr *expr, Query *query); static DistributedPlan * CreateNonPushableInsertSelectPlan(uint64 planId, Query *parse, ParamListInfo boundParams); static DeferredErrorMessage * NonPushableInsertSelectSupported(Query *insertSelectQuery); @@ -1294,6 +1296,141 @@ SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte } +/* + * IsInsertSelectBatchPassThroughDistributionColumn returns true if the given SELECT target + * expression is a provably shard-local "batch pass-through" of the distribution + * column, i.e. it has the shape + * + * unnest(array_agg()) + * + * (optionally projected through one or more plain-Var subquery indirections, and + * with array_agg allowed to carry ORDER BY / DISTINCT / FILTER modifiers). + * + * Such an expression only ever emits distribution-column values that were read + * from rows of the current shard, so - given source and target are colocated - + * every produced value hashes back into this shard's range. That makes it safe + * to push a colocated INSERT ... SELECT down to the shards even though the + * distribution value is technically a derived expression rather than a plain + * Var. Any intermediate transformation (e.g. unnest(array_agg(dist_key + 1)) or + * unnest(f(array_agg(dist_key)))) is rejected because it could produce values + * that would hash to a different shard. + */ +bool +IsInsertSelectBatchPassThroughDistributionColumn(Expr *expr, Query *query) +{ + Query *leafQuery = query; + + /* + * Peel plain-Var subquery projection indirection down to the underlying + * expression. We only follow the simple "SELECT FROM (subquery)" + * projection form; anything else makes us conservatively bail out. + */ + for (;;) + { + expr = (Expr *) strip_implicit_coercions((Node *) expr); + + if (!IsA(expr, Var)) + { + break; + } + + Var *var = (Var *) expr; + if (var->varlevelsup != 0 || var->varattno <= InvalidAttrNumber) + { + return false; + } + + if (var->varno <= 0 || var->varno > list_length(leafQuery->rtable)) + { + return false; + } + + RangeTblEntry *rte = rt_fetch(var->varno, leafQuery->rtable); + if (rte->rtekind != RTE_SUBQUERY) + { + return false; + } + + Query *subquery = rte->subquery; + if (var->varattno > list_length(subquery->targetList)) + { + return false; + } + + TargetEntry *targetEntry = list_nth(subquery->targetList, var->varattno - 1); + expr = targetEntry->expr; + leafQuery = subquery; + } + + /* the leaf expression must be unnest(...) over a single array argument */ + if (!IsA(expr, FuncExpr)) + { + return false; + } + + FuncExpr *unnestExpr = (FuncExpr *) expr; + if (unnestExpr->funcid != F_UNNEST_ANYARRAY || + list_length(unnestExpr->args) != 1) + { + return false; + } + + /* the unnest argument must be array_agg(...) with no wrapping transform */ + Expr *unnestArg = (Expr *) strip_implicit_coercions( + (Node *) linitial(unnestExpr->args)); + if (!IsA(unnestArg, Aggref)) + { + return false; + } + + Aggref *arrayAgg = (Aggref *) unnestArg; + if (arrayAgg->aggfnoid != F_ARRAY_AGG_ANYNONARRAY && + arrayAgg->aggfnoid != F_ARRAY_AGG_ANYARRAY) + { + return false; + } + + /* + * Locate the single aggregated value argument. ORDER BY keys that differ + * from the aggregated value appear as additional resjunk target entries; + * they only reorder the (shard-local) values and do not affect routing. + */ + TargetEntry *aggValueTargetEntry = NULL; + TargetEntry *aggArgTargetEntry = NULL; + foreach_declared_ptr(aggArgTargetEntry, arrayAgg->args) + { + if (aggArgTargetEntry->resjunk) + { + continue; + } + + if (aggValueTargetEntry != NULL) + { + /* + * array_agg (the OIDs we matched above) is a single-argument + * aggregate, so a second non-resjunk entry cannot occur for a + * well-formed tree. Assert to catch a broken invariant in debug + * builds, but still bail out gracefully in production: a false + * result only forgoes the optimization, it is never unsafe. + */ + Assert(false); + return false; + } + + aggValueTargetEntry = aggArgTargetEntry; + } + + if (aggValueTargetEntry == NULL) + { + return false; + } + + /* the aggregated value must be the untransformed source partition column */ + bool skipOuterVars = false; + return IsPartitionColumn(aggValueTargetEntry->expr, leafQuery, skipOuterVars); +} + + /* * InsertPartitionColumnIsBatchPassThrough returns true if the SELECT target * entry that feeds the INSERT's distribution column is a provably shard-local @@ -1310,8 +1447,8 @@ InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, return false; } - return IsBatchUnnestArrayAggPartitionColumn(subqueryTargetEntry->expr, - subqueryRte->subquery); + return IsInsertSelectBatchPassThroughDistributionColumn(subqueryTargetEntry->expr, + subqueryRte->subquery); } diff --git a/src/backend/distributed/planner/multi_logical_optimizer.c b/src/backend/distributed/planner/multi_logical_optimizer.c index 2f8924d46f4..3a4aef85042 100644 --- a/src/backend/distributed/planner/multi_logical_optimizer.c +++ b/src/backend/distributed/planner/multi_logical_optimizer.c @@ -4763,141 +4763,6 @@ IsPartitionColumn(Expr *columnExpression, Query *query, bool skipOuterVars) } -/* - * IsBatchUnnestArrayAggPartitionColumn returns true if the given SELECT target - * expression is a provably shard-local "batch pass-through" of the distribution - * column, i.e. it has the shape - * - * unnest(array_agg()) - * - * (optionally projected through one or more plain-Var subquery indirections, and - * with array_agg allowed to carry ORDER BY / DISTINCT / FILTER modifiers). - * - * Such an expression only ever emits distribution-column values that were read - * from rows of the current shard, so - given source and target are colocated - - * every produced value hashes back into this shard's range. That makes it safe - * to push a colocated INSERT ... SELECT down to the shards even though the - * distribution value is technically a derived expression rather than a plain - * Var. Any intermediate transformation (e.g. unnest(array_agg(dist_key + 1)) or - * unnest(f(array_agg(dist_key)))) is rejected because it could produce values - * that would hash to a different shard. - */ -bool -IsBatchUnnestArrayAggPartitionColumn(Expr *expr, Query *query) -{ - Query *leafQuery = query; - - /* - * Peel plain-Var subquery projection indirection down to the underlying - * expression. We only follow the simple "SELECT FROM (subquery)" - * projection form; anything else makes us conservatively bail out. - */ - for (;;) - { - expr = (Expr *) strip_implicit_coercions((Node *) expr); - - if (!IsA(expr, Var)) - { - break; - } - - Var *var = (Var *) expr; - if (var->varlevelsup != 0 || var->varattno <= InvalidAttrNumber) - { - return false; - } - - if (var->varno <= 0 || var->varno > list_length(leafQuery->rtable)) - { - return false; - } - - RangeTblEntry *rte = rt_fetch(var->varno, leafQuery->rtable); - if (rte->rtekind != RTE_SUBQUERY) - { - return false; - } - - Query *subquery = rte->subquery; - if (var->varattno > list_length(subquery->targetList)) - { - return false; - } - - TargetEntry *targetEntry = list_nth(subquery->targetList, var->varattno - 1); - expr = targetEntry->expr; - leafQuery = subquery; - } - - /* the leaf expression must be unnest(...) over a single array argument */ - if (!IsA(expr, FuncExpr)) - { - return false; - } - - FuncExpr *unnestExpr = (FuncExpr *) expr; - if (unnestExpr->funcid != F_UNNEST_ANYARRAY || - list_length(unnestExpr->args) != 1) - { - return false; - } - - /* the unnest argument must be array_agg(...) with no wrapping transform */ - Expr *unnestArg = (Expr *) strip_implicit_coercions( - (Node *) linitial(unnestExpr->args)); - if (!IsA(unnestArg, Aggref)) - { - return false; - } - - Aggref *arrayAgg = (Aggref *) unnestArg; - if (arrayAgg->aggfnoid != F_ARRAY_AGG_ANYNONARRAY && - arrayAgg->aggfnoid != F_ARRAY_AGG_ANYARRAY) - { - return false; - } - - /* - * Locate the single aggregated value argument. ORDER BY keys that differ - * from the aggregated value appear as additional resjunk target entries; - * they only reorder the (shard-local) values and do not affect routing. - */ - TargetEntry *aggValueTargetEntry = NULL; - TargetEntry *aggArgTargetEntry = NULL; - foreach_declared_ptr(aggArgTargetEntry, arrayAgg->args) - { - if (aggArgTargetEntry->resjunk) - { - continue; - } - - if (aggValueTargetEntry != NULL) - { - /* - * array_agg (the OIDs we matched above) is a single-argument - * aggregate, so a second non-resjunk entry cannot occur for a - * well-formed tree. Assert to catch a broken invariant in debug - * builds, but still bail out gracefully in production: a false - * result only forgoes the optimization, it is never unsafe. - */ - Assert(false); - return false; - } - - aggValueTargetEntry = aggArgTargetEntry; - } - - if (aggValueTargetEntry == NULL) - { - return false; - } - - /* the aggregated value must be the untransformed source partition column */ - bool skipOuterVars = false; - return IsPartitionColumn(aggValueTargetEntry->expr, leafQuery, skipOuterVars); -} - - /* * FindReferencedTableColumn recursively traverses query tree to find actual relation * id, and column that columnExpression refers to. If columnExpression is a diff --git a/src/include/distributed/multi_logical_optimizer.h b/src/include/distributed/multi_logical_optimizer.h index d62bd7b0e7f..940cbc12358 100644 --- a/src/include/distributed/multi_logical_optimizer.h +++ b/src/include/distributed/multi_logical_optimizer.h @@ -176,7 +176,6 @@ extern List * SubqueryMultiTableList(MultiNode *multiNode); extern List * GroupTargetEntryList(List *groupClauseList, List *targetEntryList); extern bool ExtractQueryWalker(Node *node, List **queryList); extern bool IsPartitionColumn(Expr *columnExpression, Query *query, bool skipOuterVars); -extern bool IsBatchUnnestArrayAggPartitionColumn(Expr *expr, Query *query); extern void FindReferencedTableColumn(Expr *columnExpression, List *parentQueryList, Query *query, Var **column, RangeTblEntry **rteContainingReferencedColumn, From 71bc03df0bef56192e8ae29ea02401b096a7edf8 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Thu, 2 Jul 2026 11:41:16 +0300 Subject: [PATCH 20/25] Tighten batch pass-through safety wording in comments Several comments described unnest(array_agg(dist_key)) as "provably shard-local". That overstates what the shape check proves: it only validates the distribution-column expression, and its safety argument holds only if the row reaching the INSERT actually carries one of the original source-row distribution values rather than a targetlist/SRF artifact (e.g. a NULL-padded row from a set-returning function). Reword the acceptance-site comment, the moved predicate's doc, the InsertPartitionColumnIsBatchPassThrough doc, and the router NOT NULL filter comment to state the precise invariant and explicitly call out the unguarded SRF NULL-padding gap as the reason the GUC is unsafe. Comments only; no behavior change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../planner/insert_select_planner.c | 46 +++++++++++-------- .../planner/multi_router_planner.c | 4 +- 2 files changed, 30 insertions(+), 20 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 91442408d89..ae80e424e05 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -833,13 +833,16 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, /* * Ensure that INSERT's partition column comes from SELECT's partition * column. Normally this requires a plain-Var partition-column match. - * With unsafe INSERT ... SELECT pushdown enabled we additionally accept - * a provably shard-local batch pass-through of the distribution column, - * i.e. unnest(array_agg(dist_key)); those values can only hash back into - * the current shard, so the co-location check further below keeps the - * batch shard-local. Any other derived distribution value is still - * rejected, because it could route rows that actually belong to a different - * shard. + * With unsafe INSERT ... SELECT pushdown enabled we additionally accept a + * batch pass-through of the distribution column, i.e. + * unnest(array_agg(dist_key)), which re-emits distribution values verbatim + * from the current shard's rows; combined with the co-location check below, + * those values route back into this shard. Any other derived distribution + * value is still rejected, because it could route rows that belong to a + * different shard. This only constrains the distribution-column expression's + * shape; it does not by itself prove every projected row carries an original + * value (an SRF NULL-padded row, for example, is not guarded), which is part + * of why the GUC is explicitly unsafe. */ if (HasDistributionKey(targetRelationId)) { @@ -1298,7 +1301,7 @@ SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte /* * IsInsertSelectBatchPassThroughDistributionColumn returns true if the given SELECT target - * expression is a provably shard-local "batch pass-through" of the distribution + * expression is a batch pass-through of the distribution * column, i.e. it has the shape * * unnest(array_agg()) @@ -1306,14 +1309,19 @@ SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte * (optionally projected through one or more plain-Var subquery indirections, and * with array_agg allowed to carry ORDER BY / DISTINCT / FILTER modifiers). * - * Such an expression only ever emits distribution-column values that were read - * from rows of the current shard, so - given source and target are colocated - - * every produced value hashes back into this shard's range. That makes it safe - * to push a colocated INSERT ... SELECT down to the shards even though the - * distribution value is technically a derived expression rather than a plain - * Var. Any intermediate transformation (e.g. unnest(array_agg(dist_key + 1)) or - * unnest(f(array_agg(dist_key)))) is rejected because it could produce values - * that would hash to a different shard. + * The invariant is narrow: unnest(array_agg()) re-emits + * distribution values verbatim from source rows of the current shard, so - + * given source and target are colocated - each value routes back into this + * shard's range. Any intermediate transformation (e.g. unnest(array_agg(dist_key + * + 1)) or unnest(f(array_agg(dist_key)))) is rejected because it could produce + * values that route to a different shard. + * + * This shape check alone does not prove the whole INSERT is shard-local. It + * assumes the row reaching the INSERT actually carries one of those original + * distribution values, and not a targetlist/SRF artifact such as a NULL-padded + * row emitted when set-returning functions in the same target list return + * differing row counts. That case is not guarded here, which is part of why the + * enabling GUC is explicitly unsafe and opt-in. */ bool IsInsertSelectBatchPassThroughDistributionColumn(Expr *expr, Query *query) @@ -1433,8 +1441,10 @@ IsInsertSelectBatchPassThroughDistributionColumn(Expr *expr, Query *query) /* * InsertPartitionColumnIsBatchPassThrough returns true if the SELECT target - * entry that feeds the INSERT's distribution column is a provably shard-local - * unnest(array_agg()) batch pass-through. + * entry that feeds the INSERT's distribution column is a batch pass-through, + * i.e. unnest(array_agg()). See + * IsInsertSelectBatchPassThroughDistributionColumn for the exact invariant and + * its known limitations. */ bool InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index 1ad9e7f35d6..efaa1e0583d 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -349,7 +349,7 @@ ShardIntervalOpExpressions(ShardInterval *shardInterval, Index rteIndex) * column value. Thus, this function should never be called with reference tables. * * The exception is unsafe INSERT ... SELECT pushdown, where the distribution - * column may be a provably shard-local batch pass-through with no plain Var to + * column may be a batch pass-through with no plain Var to * filter on; callers signal that case via distributionColumnIsBatchPassThrough so * the filter is skipped. */ @@ -378,7 +378,7 @@ AddPartitionKeyNotNullFilterToSelect(Query *subqery, bool /* * Normally the SELECT projects the distribution column as a plain Var. With * unsafe INSERT ... SELECT pushdown the distribution column may instead be a - * provably shard-local batch pass-through, i.e. unnest(array_agg(dist_col)). + * batch pass-through, i.e. unnest(array_agg(dist_col)). * In that case there is no plain Var to attach a NOT NULL filter to, and the * batch stays shard-local, so the filter is unnecessary; skip it. The caller * determines whether the query has this shape and passes the result in From 8ede5f5a65f22f4d9c8acfb0fe835b1f649d0e37 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Thu, 2 Jul 2026 15:40:14 +0300 Subject: [PATCH 21/25] Guard batch pass-through lookup against null-distribution-key targets SelectTargetEntryForInsertPartitionColumn() (added when de-duplicating the INSERT distribution-column lookup) unconditionally dereferences the target table's distribution column via PartitionColumn(insertRelationId, 1). For single-shard (null distribution key) target tables that helper returns NULL, so the subsequent insertPartitionColumn->varattno access segfaulted the backend. This path is reached from RouterModifyTaskForShardInterval when a shard-key source is inserted into a single-shard target (e.g. INSERT INTO nullkey_c1_t1 SELECT * FROM range_table), because InsertPartitionColumnIsBatchPassThrough() is now computed unconditionally at the call site to feed AddPartitionKeyNotNullFilterToSelect(). The earlier InsertPartitionColumnMatchesSelect() caller was guarded by HasDistributionKey(targetRelationId), which is why the crash only surfaced on the batch pass-through path. Return NULL early when the target has no distribution column. This restores the pre-refactor behavior: the batch pass-through check reports false and the plain-Var NOT NULL filter is still added as before. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../distributed/planner/insert_select_planner.c | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index ae80e424e05..bee6cd8d71d 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -1253,6 +1253,16 @@ SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte *insertTargetEntry = NULL; } + /* + * Single-shard (null distribution key) target tables have no + * distribution column, so there is no INSERT partition column to look + * up. Bail out before dereferencing insertPartitionColumn below. + */ + if (insertPartitionColumn == NULL) + { + return NULL; + } + ListCell *targetEntryCell = NULL; foreach(targetEntryCell, query->targetList) { From 731195afef3f7ad1dc656ce87e0b1c1ceaf37791 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Fri, 3 Jul 2026 16:28:04 +0000 Subject: [PATCH 22/25] Trim batch pass-through comment caveats --- .../distributed/planner/insert_select_planner.c | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index bee6cd8d71d..83bbfe1c4ce 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -839,10 +839,7 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte, * from the current shard's rows; combined with the co-location check below, * those values route back into this shard. Any other derived distribution * value is still rejected, because it could route rows that belong to a - * different shard. This only constrains the distribution-column expression's - * shape; it does not by itself prove every projected row carries an original - * value (an SRF NULL-padded row, for example, is not guarded), which is part - * of why the GUC is explicitly unsafe. + * different shard. */ if (HasDistributionKey(targetRelationId)) { @@ -1237,7 +1234,7 @@ ReorderInsertSelectTargetLists(Query *originalQuery, RangeTblEntry *insertRte, * NULL when the INSERT does not project the distribution column via a single Var. * * When insertTargetEntry is not NULL, it is set to the matching INSERT target - * entry so that callers can inspect casting on the distribution column. + * entry so that callers can inspect it if needed. */ static TargetEntry * SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte, @@ -1256,7 +1253,7 @@ SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte /* * Single-shard (null distribution key) target tables have no * distribution column, so there is no INSERT partition column to look - * up. Bail out before dereferencing insertPartitionColumn below. + * up. */ if (insertPartitionColumn == NULL) { @@ -1267,8 +1264,7 @@ SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte foreach(targetEntryCell, query->targetList) { TargetEntry *targetEntry = (TargetEntry *) lfirst(targetEntryCell); - List *insertTargetEntryColumnList = - pull_var_clause_default((Node *) targetEntry); + List *insertTargetEntryColumnList = pull_var_clause_default((Node *) targetEntry); /* * We only consider target entries that include a single column. Note that From bdf309b25e6b06bd063604a066feb90ee63a2219 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Fri, 3 Jul 2026 16:35:56 +0000 Subject: [PATCH 23/25] minor style: order insert-select batch helpers after their callers Our convention is to define file-local helper functions after the functions that call them. SelectTargetEntryForInsertPartitionColumn and IsInsertSelectBatchPassThroughDistributionColumn were defined before their callers InsertPartitionColumnIsBatchPassThrough and InsertPartitionColumnMatchesSelect. Move both helpers after the callers. Pure reordering; no behavior change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../planner/insert_select_planner.c | 434 +++++++++--------- 1 file changed, 217 insertions(+), 217 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 83bbfe1c4ce..9a28248e5f4 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -1227,6 +1227,223 @@ ReorderInsertSelectTargetLists(Query *originalQuery, RangeTblEntry *insertRte, } +/* + * InsertPartitionColumnIsBatchPassThrough returns true if the SELECT target + * entry that feeds the INSERT's distribution column is a batch pass-through, + * i.e. unnest(array_agg()). See + * IsInsertSelectBatchPassThroughDistributionColumn for the exact invariant and + * its known limitations. + */ +bool +InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, + RangeTblEntry *subqueryRte) +{ + TargetEntry *subqueryTargetEntry = + SelectTargetEntryForInsertPartitionColumn(query, insertRte, subqueryRte, NULL); + if (subqueryTargetEntry == NULL) + { + return false; + } + + return IsInsertSelectBatchPassThroughDistributionColumn(subqueryTargetEntry->expr, + subqueryRte->subquery); +} + + +/* + * InsertPartitionColumnMatchesSelect returns NULL the partition column in the + * table targeted by INSERTed matches with the any of the SELECTed table's + * partition column. Returns the error description if there's no match. + * + * On return without error (i.e., if partition columns match), the function + * also sets selectPartitionColumnTableId. + */ +static DeferredErrorMessage * +InsertPartitionColumnMatchesSelect(Query *query, RangeTblEntry *insertRte, + RangeTblEntry *subqueryRte, + Oid *selectPartitionColumnTableId) +{ + Query *subquery = subqueryRte->subquery; + + TargetEntry *insertTargetEntry = NULL; + TargetEntry *subqueryTargetEntry = + SelectTargetEntryForInsertPartitionColumn(query, insertRte, subqueryRte, + &insertTargetEntry); + if (subqueryTargetEntry == NULL) + { + return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, + "cannot perform distributed INSERT INTO ... SELECT " + "because the partition columns in the source table " + "and subquery do not match", + "the query doesn't include the target table's " + "partition column", + NULL); + } + + Expr *selectTargetExpr = subqueryTargetEntry->expr; + + Var *subqueryPartitionColumn = NULL; + RangeTblEntry *subqueryPartitionColumnRelationIdRTE = NULL; + List *parentQueryList = list_make2(query, subquery); + bool skipOuterVars = false; + FindReferencedTableColumn(selectTargetExpr, + parentQueryList, subquery, + &subqueryPartitionColumn, + &subqueryPartitionColumnRelationIdRTE, + skipOuterVars); + Oid subqueryPartitionColumnRelationId = subqueryPartitionColumnRelationIdRTE ? + subqueryPartitionColumnRelationIdRTE->relid : + InvalidOid; + + /* + * Corresponding (i.e., in the same ordinal position as the target table's + * partition column) select target entry does not directly belong a table. + * Evaluate its expression type and error out properly. + */ + if (subqueryPartitionColumnRelationId == InvalidOid) + { + char *errorDetailTemplate = "Subquery contains %s in the " + "same position as the target table's " + "partition column."; + + char *exprDescription = ""; + + switch (selectTargetExpr->type) + { + case T_Const: + { + exprDescription = "a constant value"; + break; + } + + case T_OpExpr: + { + exprDescription = "an operator"; + break; + } + + case T_FuncExpr: + { + FuncExpr *subqueryFunctionExpr = (FuncExpr *) selectTargetExpr; + + switch (subqueryFunctionExpr->funcformat) + { + case COERCE_EXPLICIT_CALL: + { + exprDescription = "a function call"; + break; + } + + case COERCE_EXPLICIT_CAST: + { + exprDescription = "an explicit cast"; + break; + } + + case COERCE_IMPLICIT_CAST: + { + exprDescription = "an implicit cast"; + break; + } + + default: + { + exprDescription = "a function call"; + break; + } + } + break; + } + + case T_Aggref: + { + exprDescription = "an aggregation"; + break; + } + + case T_CaseExpr: + { + exprDescription = "a case expression"; + break; + } + + case T_CoalesceExpr: + { + exprDescription = "a coalesce expression"; + break; + } + + case T_RowExpr: + { + exprDescription = "a row expression"; + break; + } + + case T_MinMaxExpr: + { + exprDescription = "a min/max expression"; + break; + } + + case T_CoerceViaIO: + { + exprDescription = "an explicit coercion"; + break; + } + + default: + { + exprDescription = + "an expression that is not a simple column reference"; + break; + } + } + + return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, + "cannot perform distributed INSERT INTO ... SELECT " + "because the partition columns in the source table " + "and subquery do not match", + psprintf(errorDetailTemplate, exprDescription), + "Ensure the target table's partition column has a " + "corresponding simple column reference to a distributed " + "table's partition column in the subquery."); + } + + /* + * Insert target expression could only be non-var if the select target + * entry does not have the same type (i.e., target column requires casting). + */ + if (!IsA(insertTargetEntry->expr, Var)) + { + return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, + "cannot perform distributed INSERT INTO ... SELECT " + "because the partition columns in the source table " + "and subquery do not match", + "The data type of the target table's partition column " + "should exactly match the data type of the " + "corresponding simple column reference in the subquery.", + NULL); + } + + /* finally, check that the select target column is a partition column */ + if (!IsPartitionColumn(selectTargetExpr, subquery, skipOuterVars)) + { + return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, + "cannot perform distributed INSERT INTO ... SELECT " + "because the partition columns in the source table " + "and subquery do not match", + "The target table's partition column should correspond " + "to a partition column in the subquery.", + NULL); + } + + /* we can set the select relation id */ + *selectPartitionColumnTableId = subqueryPartitionColumnRelationId; + + return NULL; +} + + /* * SelectTargetEntryForInsertPartitionColumn walks the INSERT target list to find * the entry that feeds the target table's distribution column and returns the @@ -1445,223 +1662,6 @@ IsInsertSelectBatchPassThroughDistributionColumn(Expr *expr, Query *query) } -/* - * InsertPartitionColumnIsBatchPassThrough returns true if the SELECT target - * entry that feeds the INSERT's distribution column is a batch pass-through, - * i.e. unnest(array_agg()). See - * IsInsertSelectBatchPassThroughDistributionColumn for the exact invariant and - * its known limitations. - */ -bool -InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, - RangeTblEntry *subqueryRte) -{ - TargetEntry *subqueryTargetEntry = - SelectTargetEntryForInsertPartitionColumn(query, insertRte, subqueryRte, NULL); - if (subqueryTargetEntry == NULL) - { - return false; - } - - return IsInsertSelectBatchPassThroughDistributionColumn(subqueryTargetEntry->expr, - subqueryRte->subquery); -} - - -/* - * InsertPartitionColumnMatchesSelect returns NULL the partition column in the - * table targeted by INSERTed matches with the any of the SELECTed table's - * partition column. Returns the error description if there's no match. - * - * On return without error (i.e., if partition columns match), the function - * also sets selectPartitionColumnTableId. - */ -static DeferredErrorMessage * -InsertPartitionColumnMatchesSelect(Query *query, RangeTblEntry *insertRte, - RangeTblEntry *subqueryRte, - Oid *selectPartitionColumnTableId) -{ - Query *subquery = subqueryRte->subquery; - - TargetEntry *insertTargetEntry = NULL; - TargetEntry *subqueryTargetEntry = - SelectTargetEntryForInsertPartitionColumn(query, insertRte, subqueryRte, - &insertTargetEntry); - if (subqueryTargetEntry == NULL) - { - return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, - "cannot perform distributed INSERT INTO ... SELECT " - "because the partition columns in the source table " - "and subquery do not match", - "the query doesn't include the target table's " - "partition column", - NULL); - } - - Expr *selectTargetExpr = subqueryTargetEntry->expr; - - Var *subqueryPartitionColumn = NULL; - RangeTblEntry *subqueryPartitionColumnRelationIdRTE = NULL; - List *parentQueryList = list_make2(query, subquery); - bool skipOuterVars = false; - FindReferencedTableColumn(selectTargetExpr, - parentQueryList, subquery, - &subqueryPartitionColumn, - &subqueryPartitionColumnRelationIdRTE, - skipOuterVars); - Oid subqueryPartitionColumnRelationId = subqueryPartitionColumnRelationIdRTE ? - subqueryPartitionColumnRelationIdRTE->relid : - InvalidOid; - - /* - * Corresponding (i.e., in the same ordinal position as the target table's - * partition column) select target entry does not directly belong a table. - * Evaluate its expression type and error out properly. - */ - if (subqueryPartitionColumnRelationId == InvalidOid) - { - char *errorDetailTemplate = "Subquery contains %s in the " - "same position as the target table's " - "partition column."; - - char *exprDescription = ""; - - switch (selectTargetExpr->type) - { - case T_Const: - { - exprDescription = "a constant value"; - break; - } - - case T_OpExpr: - { - exprDescription = "an operator"; - break; - } - - case T_FuncExpr: - { - FuncExpr *subqueryFunctionExpr = (FuncExpr *) selectTargetExpr; - - switch (subqueryFunctionExpr->funcformat) - { - case COERCE_EXPLICIT_CALL: - { - exprDescription = "a function call"; - break; - } - - case COERCE_EXPLICIT_CAST: - { - exprDescription = "an explicit cast"; - break; - } - - case COERCE_IMPLICIT_CAST: - { - exprDescription = "an implicit cast"; - break; - } - - default: - { - exprDescription = "a function call"; - break; - } - } - break; - } - - case T_Aggref: - { - exprDescription = "an aggregation"; - break; - } - - case T_CaseExpr: - { - exprDescription = "a case expression"; - break; - } - - case T_CoalesceExpr: - { - exprDescription = "a coalesce expression"; - break; - } - - case T_RowExpr: - { - exprDescription = "a row expression"; - break; - } - - case T_MinMaxExpr: - { - exprDescription = "a min/max expression"; - break; - } - - case T_CoerceViaIO: - { - exprDescription = "an explicit coercion"; - break; - } - - default: - { - exprDescription = - "an expression that is not a simple column reference"; - break; - } - } - - return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, - "cannot perform distributed INSERT INTO ... SELECT " - "because the partition columns in the source table " - "and subquery do not match", - psprintf(errorDetailTemplate, exprDescription), - "Ensure the target table's partition column has a " - "corresponding simple column reference to a distributed " - "table's partition column in the subquery."); - } - - /* - * Insert target expression could only be non-var if the select target - * entry does not have the same type (i.e., target column requires casting). - */ - if (!IsA(insertTargetEntry->expr, Var)) - { - return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, - "cannot perform distributed INSERT INTO ... SELECT " - "because the partition columns in the source table " - "and subquery do not match", - "The data type of the target table's partition column " - "should exactly match the data type of the " - "corresponding simple column reference in the subquery.", - NULL); - } - - /* finally, check that the select target column is a partition column */ - if (!IsPartitionColumn(selectTargetExpr, subquery, skipOuterVars)) - { - return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, - "cannot perform distributed INSERT INTO ... SELECT " - "because the partition columns in the source table " - "and subquery do not match", - "The target table's partition column should correspond " - "to a partition column in the subquery.", - NULL); - } - - /* we can set the select relation id */ - *selectPartitionColumnTableId = subqueryPartitionColumnRelationId; - - return NULL; -} - - /* * CreateNonPushableInsertSelectPlan creates a query plan for a SELECT into a * distributed table. The query plan can also be executed on a worker in MX. From 1e80db8be5e6d57c15b6f26949d4fe1e45a67ecf Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Fri, 3 Jul 2026 16:56:26 +0000 Subject: [PATCH 24/25] Reject DISTINCT/FILTER on the batch pass-through distribution column The shard-local batch pass-through unnest(array_agg()) is only correct when the array_agg emits exactly one value per group row. A DISTINCT or FILTER clause on that aggregate can emit fewer values, so when a sibling set-returning target in the same projection is longer PostgreSQL ProjectSet NULL-pads the shorter distribution-column set. Because the batch path skips the NOT NULL partition-column filter, that NULL (mis-routed) distribution key was inserted silently. Reject aggdistinct/aggfilter in IsInsertSelectBatchPassThroughDistributionColumn so such a shape is no longer recognized as a pass-through: the INSERT..SELECT falls back to the coordinator merge, which re-routes each row and raises the usual not-NULL partition-column error instead of corrupting data. ORDER BY inside array_agg stays allowed (it only reorders, never drops, values). Add regression coverage: FILTER and DISTINCT on the distribution-column array_agg now fall back to a coordinator merge under the GUC, and the FILTER shape run for real raises the not-NULL partition-column error. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../planner/insert_select_planner.c | 27 ++- .../allow_unsafe_insert_select_pushdown.out | 173 ++++++++++++++++++ .../allow_unsafe_insert_select_pushdown.sql | 71 +++++++ 3 files changed, 267 insertions(+), 4 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index 9a28248e5f4..d559c196a79 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -1530,7 +1530,10 @@ SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte * unnest(array_agg()) * * (optionally projected through one or more plain-Var subquery indirections, and - * with array_agg allowed to carry ORDER BY / DISTINCT / FILTER modifiers). + * with array_agg allowed to carry an ORDER BY modifier). DISTINCT and FILTER on + * the aggregate are rejected: they make the pass-through emit fewer values than + * the group has rows, so ProjectSet would NULL-pad the distribution column when a + * sibling set-returning target is longer (see the DISTINCT/FILTER guard below). * * The invariant is narrow: unnest(array_agg()) re-emits * distribution values verbatim from source rows of the current shard, so - @@ -1542,9 +1545,11 @@ SelectTargetEntryForInsertPartitionColumn(Query *query, RangeTblEntry *insertRte * This shape check alone does not prove the whole INSERT is shard-local. It * assumes the row reaching the INSERT actually carries one of those original * distribution values, and not a targetlist/SRF artifact such as a NULL-padded - * row emitted when set-returning functions in the same target list return - * differing row counts. That case is not guarded here, which is part of why the - * enabling GUC is explicitly unsafe and opt-in. + * row emitted when a sibling set-returning target in the same projection emits + * more rows than this pass-through. We cannot rule that out here without also + * rejecting the supported batched shape, e.g. unnest(batch_udf(array_agg(col))), + * whose row count is a runtime property of the (user) batch function. That + * residual is exactly why the enabling GUC is explicitly unsafe and opt-in. */ bool IsInsertSelectBatchPassThroughDistributionColumn(Expr *expr, Query *query) @@ -1621,6 +1626,20 @@ IsInsertSelectBatchPassThroughDistributionColumn(Expr *expr, Query *query) return false; } + /* + * DISTINCT or FILTER make array_agg emit fewer values than the group has + * rows, so this pass-through set-returning expression becomes shorter than a + * sibling set-returning target in the same projection. ProjectSet then + * NULL-pads the distribution column, which would insert a NULL (mis-routed) + * distribution key. A plain array_agg re-emits exactly one value per source + * row, so reject DISTINCT/FILTER here (ORDER BY is fine: it only reorders the + * shard-local values, it does not change their count). + */ + if (arrayAgg->aggdistinct != NIL || arrayAgg->aggfilter != NULL) + { + return false; + } + /* * Locate the single aggregated value argument. ORDER BY keys that differ * from the aggregated value appear as additional resjunk target entries; diff --git a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out index d79f779669a..e8e032270f5 100644 --- a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out +++ b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out @@ -589,6 +589,179 @@ $$, true); -> Function Scan on read_intermediate_result intermediate_result (20 rows) +-- --------------------------------------------------------------------- +-- Negative coverage: FILTER / DISTINCT on the batch pass-through array_agg. +-- The distribution-column pass-through unnest(array_agg(text_id)) is only +-- safe when the array_agg emits exactly one element per group row. A FILTER +-- or DISTINCT clause on that array_agg can emit fewer elements, so when a +-- sibling set-returning column is longer PostgreSQL's ProjectSet NULL-pads +-- the shorter distribution-column set -- silently producing NULL (mis-routed) +-- partition keys. Such an array_agg is therefore rejected even with the GUC +-- enabled and falls back to a coordinator merge, which re-routes each row and +-- enforces the not-NULL partition-column invariant. ORDER BY inside array_agg +-- only reorders (never drops) elements and stays pushed down (covered above). +-- --------------------------------------------------------------------- +-- FILTER on the distribution-column array_agg: not pushed down +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id) FILTER (WHERE text_id % 2 = 0)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> WindowAgg + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Subquery Scan on s + -> ProjectSet + -> HashAggregate + Group Key: intermediate_result.b + -> Function Scan on read_intermediate_result intermediate_result +(20 rows) + +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id) FILTER (WHERE text_id % 2 = 0)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> WindowAgg + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Subquery Scan on s + -> ProjectSet + -> HashAggregate + Group Key: intermediate_result.b + -> Function Scan on read_intermediate_result intermediate_result +(20 rows) + +-- DISTINCT on the distribution-column array_agg: not pushed down +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(DISTINCT text_id)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> WindowAgg + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Subquery Scan on s + -> ProjectSet + -> GroupAggregate + Group Key: intermediate_result.b + -> Sort + Sort Key: intermediate_result.b, intermediate_result.text_id + -> Function Scan on read_intermediate_result intermediate_result +(22 rows) + +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(DISTINCT text_id)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus INSERT ... SELECT) + INSERT/SELECT method: pull to coordinator + -> Custom Scan (Citus Adaptive) + -> Distributed Subplan XXX_1 + -> WindowAgg + -> Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Seq Scan on dist_14000000 dist + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Subquery Scan on s + -> ProjectSet + -> GroupAggregate + Group Key: intermediate_result.b + -> Sort + Sort Key: intermediate_result.b, intermediate_result.text_id + -> Function Scan on read_intermediate_result intermediate_result +(22 rows) + +-- the FILTER shape rejected above, run for real under the GUC: the coordinator +-- fallback re-routes each row and raises the not-NULL partition-column error +-- instead of silently inserting NULL (mis-routed) partition keys +SET citus.allow_unsafe_insert_select_pushdown TO on; +TRUNCATE res; +INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id) FILTER (WHERE text_id % 2 = 0)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s; +ERROR: the partition column of table allow_unsafe_insert_select_pushdown.res cannot be NULL +SELECT count(*) AS rows_with_null_key FROM res WHERE text_id IS NULL; + rows_with_null_key +--------------------------------------------------------------------- + 0 +(1 row) + -- --------------------------------------------------------------------- -- Correctness of the pushed-down batches. -- --------------------------------------------------------------------- diff --git a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql index 4f154da7568..30184b30d28 100644 --- a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql +++ b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql @@ -239,6 +239,77 @@ SELECT id, val FROM ( ) s $$, true); +-- --------------------------------------------------------------------- +-- Negative coverage: FILTER / DISTINCT on the batch pass-through array_agg. +-- The distribution-column pass-through unnest(array_agg(text_id)) is only +-- safe when the array_agg emits exactly one element per group row. A FILTER +-- or DISTINCT clause on that array_agg can emit fewer elements, so when a +-- sibling set-returning column is longer PostgreSQL's ProjectSet NULL-pads +-- the shorter distribution-column set -- silently producing NULL (mis-routed) +-- partition keys. Such an array_agg is therefore rejected even with the GUC +-- enabled and falls back to a coordinator merge, which re-routes each row and +-- enforces the not-NULL partition-column invariant. ORDER BY inside array_agg +-- only reorders (never drops) elements and stays pushed down (covered above). +-- --------------------------------------------------------------------- + +-- FILTER on the distribution-column array_agg: not pushed down +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id) FILTER (WHERE text_id % 2 = 0)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id) FILTER (WHERE text_id % 2 = 0)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + +-- DISTINCT on the distribution-column array_agg: not pushed down +SET citus.allow_unsafe_insert_select_pushdown TO off; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(DISTINCT text_id)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); +SET citus.allow_unsafe_insert_select_pushdown TO on; +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(DISTINCT text_id)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s +$$, true); + +-- the FILTER shape rejected above, run for real under the GUC: the coordinator +-- fallback re-routes each row and raises the not-NULL partition-column error +-- instead of silently inserting NULL (mis-routed) partition keys +SET citus.allow_unsafe_insert_select_pushdown TO on; +TRUNCATE res; +INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id) FILTER (WHERE text_id % 2 = 0)) id, + unnest(batch_transform(array_agg(text_col))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 b FROM dist) q + GROUP BY b +) s; +SELECT count(*) AS rows_with_null_key FROM res WHERE text_id IS NULL; + -- --------------------------------------------------------------------- -- Correctness of the pushed-down batches. -- --------------------------------------------------------------------- From e544103df8311499899128f568c66fe5f9fb76fa Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Fri, 3 Jul 2026 18:08:05 +0000 Subject: [PATCH 25/25] Drop NULL (mis-routed) batch pass-through partition keys at runtime When a sibling set-returning target in the batch projection (e.g. an over-emitting batch UDF) is longer than the distribution-column unnest(array_agg(...)), PostgreSQL ProjectSet NULL-pads the shorter distribution-column set. The batch pushdown path skips AddPartitionKeyNotNullFilterToSelect, so such a NULL (mis-routed) partition key would otherwise be inserted silently. AddBatchPassThroughNotNullFilter now appends a distribution-column IS NOT NULL qual to the pushed-down SELECT for both shapes: the outer-subquery shape (distribution column is a plain Var, filtered in place) and the flat shape (distribution column is the bare unnest set-returning expression, wrapped in a pass-through subquery so it becomes a filterable Var). Adds regression coverage asserting no NULL key is inserted while every real key keeps its value. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../distributed/planner/deparse_shard_query.c | 5 + .../planner/insert_select_planner.c | 170 ++++++++++++++++++ .../distributed/insert_select_planner.h | 3 + .../allow_unsafe_insert_select_pushdown.out | 162 +++++++++++++++-- src/test/regress/expected/pg18.out | 3 +- .../allow_unsafe_insert_select_pushdown.sql | 67 +++++++ 6 files changed, 390 insertions(+), 20 deletions(-) diff --git a/src/backend/distributed/planner/deparse_shard_query.c b/src/backend/distributed/planner/deparse_shard_query.c index 9d6f080a26b..42ee5e41c80 100644 --- a/src/backend/distributed/planner/deparse_shard_query.c +++ b/src/backend/distributed/planner/deparse_shard_query.c @@ -115,6 +115,11 @@ RebuildQueryStrings(Job *workerJob) AddPartitionKeyNotNullFilterToSelect(copiedSubquery, distributionColumnIsBatchPassThrough) ; + if (distributionColumnIsBatchPassThrough) + { + AddBatchPassThroughNotNullFilter(query, copiedInsertRte, + copiedSubqueryRte); + } } ReorderInsertSelectTargetLists(query, copiedInsertRte, copiedSubqueryRte); diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index d559c196a79..505362c16fc 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -96,6 +96,9 @@ static TargetEntry * SelectTargetEntryForInsertPartitionColumn(Query *query, TargetEntry ** insertTargetEntry); static bool IsInsertSelectBatchPassThroughDistributionColumn(Expr *expr, Query *query); +static Query * WrapBatchSelectWithNotNullFilter(Query *selectQuery, + TargetEntry *distTargetEntry); +static Node * AppendNotNullTest(Node *existingQuals, Expr *arg); static DistributedPlan * CreateNonPushableInsertSelectPlan(uint64 planId, Query *parse, ParamListInfo boundParams); static DeferredErrorMessage * NonPushableInsertSelectSupported(Query *insertSelectQuery); @@ -987,6 +990,11 @@ RouterModifyTaskForShardInterval(Query *originalQuery, copiedSubqueryRte); AddPartitionKeyNotNullFilterToSelect(copiedSubquery, distributionColumnIsBatchPassThrough); + if (distributionColumnIsBatchPassThrough) + { + AddBatchPassThroughNotNullFilter(copiedQuery, copiedInsertRte, + copiedSubqueryRte); + } } /* mark that we don't want the router planner to generate dummy hosts/queries */ @@ -1250,6 +1258,168 @@ InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, } +/* + * AddBatchPassThroughNotNullFilter drops rows whose batch pass-through + * distribution value came out NULL, by adding a distribution-column IS NOT NULL + * qual to the pushed-down SELECT. + * + * The batch pass-through unnest(array_agg()) emits one value per group + * row, but a sibling set-returning target in the same projection - e.g. an + * over-emitting batch UDF unnest(f(array_agg())) - can be longer, so + * PostgreSQL's ProjectSet NULL-pads the shorter distribution-column set. Because + * the batch path skips AddPartitionKeyNotNullFilterToSelect, without this filter + * such a NULL (mis-routed) distribution key would be inserted silently. Dropping + * the NULL row matches how Citus handles a NULL distribution value in a normal + * INSERT ... SELECT (AddPartitionKeyNotNullFilterToSelect). + * + * The distribution column can surface in two shapes: + * - as a plain Var, when the unnest lives in an inner subquery + * (SELECT dist_var, ... FROM ( SELECT unnest(array_agg(dist_col)) dist_var, + * ... ) s); the qual is attached to the pushed-down SELECT directly, or + * - as a bare set-returning expression projected in the SELECT list + * (SELECT unnest(array_agg(dist_col)), ... GROUP BY ...); the SELECT is + * wrapped in a pass-through subquery so the column becomes a Var we can + * filter. + * + * Returns true if a filter was added. + */ +bool +AddBatchPassThroughNotNullFilter(Query *query, RangeTblEntry *insertRte, + RangeTblEntry *subqueryRte) +{ + Query *selectQuery = subqueryRte->subquery; + TargetEntry *distTargetEntry = + SelectTargetEntryForInsertPartitionColumn(query, insertRte, subqueryRte, NULL); + if (distTargetEntry == NULL) + { + return false; + } + + Expr *distExpr = (Expr *) strip_implicit_coercions((Node *) distTargetEntry->expr); + if (IsA(distExpr, Var)) + { + /* + * Outer-subquery shape: the distribution column already surfaces as a + * plain Var (the unnest lives in an inner subquery), so we can attach + * the qual to the pushed-down SELECT directly. + */ + selectQuery->jointree->quals = + AppendNotNullTest(selectQuery->jointree->quals, distExpr); + return true; + } + + /* + * Flat shape: the distribution column is projected as a bare set-returning + * expression (e.g. unnest(array_agg()) directly in the SELECT + * list), so there is no Var handle a WHERE clause can reference. Wrap the + * SELECT in a pass-through subquery, which turns the set-returning column + * into an ordinary output column, then filter that column for NULL: + * + * SELECT dist_col, ... + * FROM ( ) citus_insert_select_subquery + * WHERE dist_col IS NOT NULL + */ + subqueryRte->subquery = WrapBatchSelectWithNotNullFilter(selectQuery, + distTargetEntry); + return true; +} + + +/* + * WrapBatchSelectWithNotNullFilter wraps the given batch pass-through SELECT in + * a pass-through subquery and filters out rows whose distribution column (given + * by distTargetEntry) came out NULL: + * + * SELECT FROM ( ) citus_insert_select_subquery + * WHERE IS NOT NULL + * + * It is used when the distribution column is projected as a bare set-returning + * expression, so it has no Var handle a WHERE clause could reference in the + * original SELECT. Unlike WrapSubquery, this wrapper is a pure pass-through: it + * references every non-junk output column as an ordinary Var and never lifts + * expressions to the outer query, so the batch call keeps running on the shard. + */ +static Query * +WrapBatchSelectWithNotNullFilter(Query *selectQuery, TargetEntry *distTargetEntry) +{ + ParseState *pstate = make_parsestate(NULL); + Query *wrapperQuery = makeNode(Query); + wrapperQuery->commandType = CMD_SELECT; + + Alias *alias = makeAlias("citus_insert_select_subquery", NIL); + RangeTblEntry *subqueryRte = + RangeTableEntryFromNSItem( + addRangeTableEntryForSubquery(pstate, selectQuery, alias, + false /* not LATERAL */, + true /* in FROM clause */)); + wrapperQuery->rtable = list_make1(subqueryRte); + wrapperQuery->rteperminfos = NIL; + + RangeTblRef *rangeTableRef = makeNode(RangeTblRef); + rangeTableRef->rtindex = 1; + + List *wrapperTargetList = NIL; + Var *distColumnVar = NULL; + TargetEntry *innerTargetEntry = NULL; + foreach_declared_ptr(innerTargetEntry, selectQuery->targetList) + { + if (innerTargetEntry->resjunk) + { + continue; + } + + Var *outerVar = makeVar(1 /* single subquery RTE */, + innerTargetEntry->resno, + exprType((Node *) innerTargetEntry->expr), + exprTypmod((Node *) innerTargetEntry->expr), + exprCollation((Node *) innerTargetEntry->expr), + 0); + + TargetEntry *outerTargetEntry = + makeTargetEntry((Expr *) outerVar, + list_length(wrapperTargetList) + 1, + innerTargetEntry->resname, + false); + wrapperTargetList = lappend(wrapperTargetList, outerTargetEntry); + + if (innerTargetEntry == distTargetEntry) + { + distColumnVar = outerVar; + } + } + + Assert(distColumnVar != NULL); + wrapperQuery->targetList = wrapperTargetList; + + Node *notNullQual = AppendNotNullTest(NULL, (Expr *) distColumnVar); + wrapperQuery->jointree = makeFromExpr(list_make1(rangeTableRef), notNullQual); + + return wrapperQuery; +} + + +/* + * AppendNotNullTest returns existingQuals AND (arg IS NOT NULL). When + * existingQuals is NULL the bare NOT NULL test is returned. + */ +static Node * +AppendNotNullTest(Node *existingQuals, Expr *arg) +{ + NullTest *nullTest = makeNode(NullTest); + nullTest->nulltesttype = IS_NOT_NULL; + nullTest->arg = (Expr *) copyObject(arg); + nullTest->argisrow = false; + nullTest->location = -1; + + if (existingQuals == NULL) + { + return (Node *) nullTest; + } + + return make_and_qual(existingQuals, (Node *) nullTest); +} + + /* * InsertPartitionColumnMatchesSelect returns NULL the partition column in the * table targeted by INSERTed matches with the any of the SELECTed table's diff --git a/src/include/distributed/insert_select_planner.h b/src/include/distributed/insert_select_planner.h index 32984401018..5ba4ab75c55 100644 --- a/src/include/distributed/insert_select_planner.h +++ b/src/include/distributed/insert_select_planner.h @@ -49,6 +49,9 @@ extern Query * WrapSubquery(Query *subquery); extern bool InsertPartitionColumnIsBatchPassThrough(Query *query, RangeTblEntry *insertRte, RangeTblEntry *subqueryRte); +extern bool AddBatchPassThroughNotNullFilter(Query *query, + RangeTblEntry *insertRte, + RangeTblEntry *subqueryRte); #endif /* INSERT_SELECT_PLANNER_H */ diff --git a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out index e8e032270f5..41b1bd6e859 100644 --- a/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out +++ b/src/test/regress/expected/allow_unsafe_insert_select_pushdown.out @@ -100,12 +100,13 @@ $$, true); Node: host=localhost port=xxxxx dbname=regression -> Insert on res_14000004 citus_table_alias -> Subquery Scan on s + Filter: (s.id IS NOT NULL) -> ProjectSet -> HashAggregate Group Key: ((row_number() OVER (?) - 1) / 100) -> WindowAgg -> Seq Scan on dist_14000000 dist -(12 rows) +(13 rows) INSERT INTO res(text_id, val) SELECT id, val FROM ( @@ -157,6 +158,7 @@ $$, true); Node: host=localhost port=xxxxx dbname=regression -> Insert on res_14000004 citus_table_alias -> Subquery Scan on s + Filter: (s.id IS NOT NULL) -> ProjectSet -> GroupAggregate Group Key: q.batch @@ -165,7 +167,7 @@ $$, true); -> Subquery Scan on q -> WindowAgg -> Seq Scan on dist_14000000 dist -(15 rows) +(16 rows) -- branch: GROUP BY on a non-distribution column, distribution column projected -- as the unnest(array_agg(text_id)) batch pass-through @@ -174,7 +176,7 @@ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) FROM dist GROUP BY text_id % 10 $$, true); - explain_filter + explain_filter --------------------------------------------------------------------- Custom Scan (Citus Adaptive) Task Count: 4 @@ -182,19 +184,20 @@ $$, true); -> Task Node: host=localhost port=xxxxx dbname=regression -> Insert on res_14000004 citus_table_alias - -> Subquery Scan on "*SELECT*" + -> Subquery Scan on citus_insert_select_subquery + Filter: (citus_insert_select_subquery.unnest IS NOT NULL) -> ProjectSet -> HashAggregate Group Key: (dist.text_id % 10) -> Seq Scan on dist_14000000 dist -(11 rows) +(12 rows) -- branch: aggregates without GROUP BY SELECT public.explain_filter($$ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) FROM dist $$, true); - explain_filter + explain_filter --------------------------------------------------------------------- Custom Scan (Citus Adaptive) Task Count: 4 @@ -202,10 +205,12 @@ $$, true); -> Task Node: host=localhost port=xxxxx dbname=regression -> Insert on res_14000004 citus_table_alias - -> ProjectSet - -> Aggregate - -> Seq Scan on dist_14000000 dist -(9 rows) + -> Subquery Scan on citus_insert_select_subquery + Filter: (citus_insert_select_subquery.unnest IS NOT NULL) + -> ProjectSet + -> Aggregate + -> Seq Scan on dist_14000000 dist +(11 rows) -- branch: HAVING without GROUP BY on the distribution column SELECT public.explain_filter($$ @@ -213,7 +218,7 @@ EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) SELECT unnest(array_agg(text_id)), unnest(array_agg(length(text_col))) FROM dist HAVING count(*) > 0 $$, true); - explain_filter + explain_filter --------------------------------------------------------------------- Custom Scan (Citus Adaptive) Task Count: 4 @@ -221,11 +226,13 @@ $$, true); -> Task Node: host=localhost port=xxxxx dbname=regression -> Insert on res_14000004 citus_table_alias - -> ProjectSet - -> Aggregate - Filter: (count(*) > 0) - -> Seq Scan on dist_14000000 dist -(10 rows) + -> Subquery Scan on citus_insert_select_subquery + Filter: (citus_insert_select_subquery.unnest IS NOT NULL) + -> ProjectSet + -> Aggregate + Filter: (count(*) > 0) + -> Seq Scan on dist_14000000 dist +(12 rows) -- branch: window function not partitioned on the distribution column, with the -- distribution column projected as a plain partition-column Var @@ -263,14 +270,15 @@ $$, true); -> Task Node: host=localhost port=xxxxx dbname=regression -> Insert on res_14000004 citus_table_alias - -> Subquery Scan on "*SELECT*" + -> Subquery Scan on citus_insert_select_subquery + Filter: (citus_insert_select_subquery.unnest IS NOT NULL) -> HashAggregate Group Key: unnest((array_agg(dist.text_id))), unnest((array_agg(length(dist.text_col)))) -> ProjectSet -> HashAggregate Group Key: (dist.text_id % 10) -> Seq Scan on dist_14000000 dist -(13 rows) +(14 rows) -- combination: window + GROUP BY, relaxed constructs living in nested subqueries SELECT public.explain_filter($$ @@ -290,6 +298,7 @@ $$, true); Node: host=localhost port=xxxxx dbname=regression -> Insert on res_14000004 citus_table_alias -> Subquery Scan on s + Filter: (s.id IS NOT NULL) -> ProjectSet -> HashAggregate Group Key: (q.rn % '5'::bigint) @@ -298,7 +307,7 @@ $$, true); -> Sort Sort Key: dist.text_col -> Seq Scan on dist_14000000 dist -(15 rows) +(16 rows) -- --------------------------------------------------------------------- -- Negative coverage: pattern requirement. With the GUC enabled the batching @@ -798,6 +807,121 @@ FROM res JOIN dist USING (text_id); 500 | 500 (1 row) +-- --------------------------------------------------------------------- +-- Runtime NULL-key drop: an over-emitting batch UDF returns more elements than +-- the distribution-column array_agg, so PostgreSQL's ProjectSet NULL-pads the +-- (shorter) distribution column. The pushed-down plan filters those rows out +-- (distribution column IS NOT NULL) so no NULL (mis-routed) partition key is +-- inserted, while every real text_id keeps its own value. Both shapes are +-- covered: the outer-subquery shape (distribution column is a plain Var, +-- filtered in place) and the flat shape (distribution column is the unnest +-- set-returning expression, filtered via a pass-through subquery wrapper). +-- --------------------------------------------------------------------- +-- a batch UDF that emits one extra element per batch (appends a trailing 0), +-- mimicking a batched API call whose output is longer than its input +CREATE FUNCTION batch_transform_over(t text[]) RETURNS int[] +LANGUAGE sql IMMUTABLE PARALLEL SAFE AS +$$ SELECT (SELECT array_agg(length(x)) FROM unnest(t) x) || 0 $$; +SELECT create_distributed_function('batch_transform_over(text[])'); +NOTICE: procedure allow_unsafe_insert_select_pushdown.batch_transform_over is already distributed +DETAIL: Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] commands + create_distributed_function +--------------------------------------------------------------------- + +(1 row) + +SET citus.allow_unsafe_insert_select_pushdown TO on; +-- outer-subquery shape: the distribution column is a plain Var, so the +-- IS NOT NULL filter is attached to the pushed-down SELECT directly +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id ORDER BY text_id)) id, + unnest(batch_transform_over(array_agg(text_col ORDER BY text_id))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 batch FROM dist) q + GROUP BY batch +) s +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> Subquery Scan on s + Filter: (s.id IS NOT NULL) + -> ProjectSet + -> GroupAggregate + Group Key: q.batch + -> Sort + Sort Key: q.batch, q.text_id + -> Subquery Scan on q + -> WindowAgg + -> Seq Scan on dist_14000000 dist +(16 rows) + +TRUNCATE res; +INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id ORDER BY text_id)) id, + unnest(batch_transform_over(array_agg(text_col ORDER BY text_id))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 batch FROM dist) q + GROUP BY batch +) s; +-- no NULL (mis-routed) key inserted; all 500 rows kept their value +SELECT count(*) FILTER (WHERE text_id IS NULL) AS rows_with_null_key, + count(*) AS rows, + count(*) FILTER (WHERE val = length('t' || text_id)) AS ok +FROM res; + rows_with_null_key | rows | ok +--------------------------------------------------------------------- + 0 | 500 | 500 +(1 row) + +-- flat shape: the distribution column is the bare unnest set-returning +-- expression, so the SELECT is wrapped in a pass-through subquery whose +-- IS NOT NULL filter drops the NULL-padded rows +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id ORDER BY text_id)), + unnest(batch_transform_over(array_agg(text_col ORDER BY text_id))) +FROM dist GROUP BY text_id % 10 +$$, true); + explain_filter +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 4 + Tasks Shown: One of 4 + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on res_14000004 citus_table_alias + -> Subquery Scan on citus_insert_select_subquery + Filter: (citus_insert_select_subquery.unnest IS NOT NULL) + -> ProjectSet + -> GroupAggregate + Group Key: ((dist.text_id % 10)) + -> Sort + Sort Key: ((dist.text_id % 10)), dist.text_id + -> Seq Scan on dist_14000000 dist +(14 rows) + +TRUNCATE res; +INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id ORDER BY text_id)), + unnest(batch_transform_over(array_agg(text_col ORDER BY text_id))) +FROM dist GROUP BY text_id % 10; +-- no NULL (mis-routed) key inserted; all 500 rows kept their value +SELECT count(*) FILTER (WHERE text_id IS NULL) AS rows_with_null_key, + count(*) AS rows, + count(*) FILTER (WHERE val = length('t' || text_id)) AS ok +FROM res; + rows_with_null_key | rows | ok +--------------------------------------------------------------------- + 0 | 500 | 500 +(1 row) + -- --------------------------------------------------------------------- -- Negative coverage: constructs the GUC never relaxes. -- --------------------------------------------------------------------- diff --git a/src/test/regress/expected/pg18.out b/src/test/regress/expected/pg18.out index 419342506c5..ba89b45f4aa 100644 --- a/src/test/regress/expected/pg18.out +++ b/src/test/regress/expected/pg18.out @@ -3223,13 +3223,14 @@ SELECT id, val FROM ( Node: host=localhost port=xxxxx dbname=regression -> Insert on res_14100004 citus_table_alias -> Subquery Scan on s + Filter: (s.id IS NOT NULL) -> ProjectSet -> HashAggregate Group Key: ((row_number() OVER (?) - 1) / 100) -> WindowAgg Window: w1 AS (ROWS UNBOUNDED PRECEDING) -> Seq Scan on dist_14100000 dist -(13 rows) +(14 rows) RESET citus.allow_unsafe_insert_select_pushdown; DROP SCHEMA pg18_insert_select_pushdown CASCADE; diff --git a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql index 30184b30d28..a963abcebea 100644 --- a/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql +++ b/src/test/regress/sql/allow_unsafe_insert_select_pushdown.sql @@ -341,6 +341,73 @@ SELECT id, val FROM ( SELECT count(*), count(*) FILTER (WHERE val = length('t' || text_id)) AS ok FROM res JOIN dist USING (text_id); +-- --------------------------------------------------------------------- +-- Runtime NULL-key drop: an over-emitting batch UDF returns more elements than +-- the distribution-column array_agg, so PostgreSQL's ProjectSet NULL-pads the +-- (shorter) distribution column. The pushed-down plan filters those rows out +-- (distribution column IS NOT NULL) so no NULL (mis-routed) partition key is +-- inserted, while every real text_id keeps its own value. Both shapes are +-- covered: the outer-subquery shape (distribution column is a plain Var, +-- filtered in place) and the flat shape (distribution column is the unnest +-- set-returning expression, filtered via a pass-through subquery wrapper). +-- --------------------------------------------------------------------- + +-- a batch UDF that emits one extra element per batch (appends a trailing 0), +-- mimicking a batched API call whose output is longer than its input +CREATE FUNCTION batch_transform_over(t text[]) RETURNS int[] +LANGUAGE sql IMMUTABLE PARALLEL SAFE AS +$$ SELECT (SELECT array_agg(length(x)) FROM unnest(t) x) || 0 $$; +SELECT create_distributed_function('batch_transform_over(text[])'); + +SET citus.allow_unsafe_insert_select_pushdown TO on; + +-- outer-subquery shape: the distribution column is a plain Var, so the +-- IS NOT NULL filter is attached to the pushed-down SELECT directly +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id ORDER BY text_id)) id, + unnest(batch_transform_over(array_agg(text_col ORDER BY text_id))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 batch FROM dist) q + GROUP BY batch +) s +$$, true); + +TRUNCATE res; +INSERT INTO res(text_id, val) +SELECT id, val FROM ( + SELECT unnest(array_agg(text_id ORDER BY text_id)) id, + unnest(batch_transform_over(array_agg(text_col ORDER BY text_id))) val + FROM (SELECT text_id, text_col, (row_number() OVER () - 1) / 100 batch FROM dist) q + GROUP BY batch +) s; +-- no NULL (mis-routed) key inserted; all 500 rows kept their value +SELECT count(*) FILTER (WHERE text_id IS NULL) AS rows_with_null_key, + count(*) AS rows, + count(*) FILTER (WHERE val = length('t' || text_id)) AS ok +FROM res; + +-- flat shape: the distribution column is the bare unnest set-returning +-- expression, so the SELECT is wrapped in a pass-through subquery whose +-- IS NOT NULL filter drops the NULL-padded rows +SELECT public.explain_filter($$ +EXPLAIN (COSTS OFF) INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id ORDER BY text_id)), + unnest(batch_transform_over(array_agg(text_col ORDER BY text_id))) +FROM dist GROUP BY text_id % 10 +$$, true); + +TRUNCATE res; +INSERT INTO res(text_id, val) +SELECT unnest(array_agg(text_id ORDER BY text_id)), + unnest(batch_transform_over(array_agg(text_col ORDER BY text_id))) +FROM dist GROUP BY text_id % 10; +-- no NULL (mis-routed) key inserted; all 500 rows kept their value +SELECT count(*) FILTER (WHERE text_id IS NULL) AS rows_with_null_key, + count(*) AS rows, + count(*) FILTER (WHERE val = length('t' || text_id)) AS ok +FROM res; + -- --------------------------------------------------------------------- -- Negative coverage: constructs the GUC never relaxes. -- ---------------------------------------------------------------------