diff --git a/src/backend/distributed/planner/multi_logical_optimizer.c b/src/backend/distributed/planner/multi_logical_optimizer.c index 9db945cba..59c5c4ee0 100644 --- a/src/backend/distributed/planner/multi_logical_optimizer.c +++ b/src/backend/distributed/planner/multi_logical_optimizer.c @@ -3710,6 +3710,8 @@ PartitionColumnOpExpressionList(Query *query) Node *whereNode = (Node *) lfirst(whereClauseCell); Node *leftArgument = NULL; Node *rightArgument = NULL; + Node *strippedLeftArgument = NULL; + Node *strippedRightArgument = NULL; OpExpr *whereClause = NULL; List *argumentList = NIL; List *rangetableList = NIL; @@ -3744,14 +3746,16 @@ PartitionColumnOpExpressionList(Query *query) leftArgument = (Node *) linitial(argumentList); rightArgument = (Node *) lsecond(argumentList); + strippedLeftArgument = strip_implicit_coercions(leftArgument); + strippedRightArgument = strip_implicit_coercions(rightArgument); - if (IsA(leftArgument, Var) && IsA(rightArgument, Const)) + if (IsA(strippedLeftArgument, Var) && IsA(strippedRightArgument, Const)) { - candidatePartitionColumn = (Var *) leftArgument; + candidatePartitionColumn = (Var *) strippedLeftArgument; } - else if (IsA(leftArgument, Const) && IsA(leftArgument, Var)) + else if (IsA(strippedLeftArgument, Const) && IsA(strippedRightArgument, Var)) { - candidatePartitionColumn = (Var *) rightArgument; + candidatePartitionColumn = (Var *) strippedRightArgument; } else { @@ -3798,14 +3802,16 @@ ReplaceColumnsInOpExpressionList(List *opExpressionList, Var *newColumn) Node *leftArgument = (Node *) linitial(argumentList); Node *rightArgument = (Node *) lsecond(argumentList); + Node *strippedLeftArgument = strip_implicit_coercions(leftArgument); + Node *strippedRightArgument = strip_implicit_coercions(rightArgument); - if (IsA(leftArgument, Var)) + if (IsA(strippedLeftArgument, Var)) { - newArgumentList = list_make2(newColumn, rightArgument); + newArgumentList = list_make2(newColumn, strippedRightArgument); } - else if (IsA(leftArgument, Var)) + else if (IsA(strippedRightArgument, Var)) { - newArgumentList = list_make2(leftArgument, newColumn); + newArgumentList = list_make2(strippedLeftArgument, newColumn); } copyOpExpression->args = newArgumentList; diff --git a/src/test/regress/input/multi_subquery.source b/src/test/regress/input/multi_subquery.source index 0accf9135..cee05ac3d 100644 --- a/src/test/regress/input/multi_subquery.source +++ b/src/test/regress/input/multi_subquery.source @@ -294,3 +294,31 @@ FROM l_orderkey = o_orderkey GROUP BY l_orderkey) AS unit_prices; + +-- Check that we can prune shards in subqueries with VARCHAR partition columns + +CREATE TABLE subquery_pruning_varchar_test_table +( + a varchar, + b int +); + +SELECT master_create_distributed_table('subquery_pruning_varchar_test_table', 'a', 'hash'); +SELECT master_create_worker_shards('subquery_pruning_varchar_test_table', 4, 1); + +SET citus.subquery_pushdown TO TRUE; +SET client_min_messages TO DEBUG2; + +SELECT * FROM + (SELECT count(*) FROM subquery_pruning_varchar_test_table WHERE a = 'onder' GROUP BY a) +AS foo; + +SELECT * FROM + (SELECT count(*) FROM subquery_pruning_varchar_test_table WHERE 'eren' = a GROUP BY a) +AS foo; + +SET client_min_messages TO NOTICE; + +DROP TABLE subquery_pruning_varchar_test_table; + +ALTER SEQUENCE pg_catalog.pg_dist_shardid_seq RESTART 102026; diff --git a/src/test/regress/output/multi_subquery.source b/src/test/regress/output/multi_subquery.source index f0e6f0250..8f5f30c22 100644 --- a/src/test/regress/output/multi_subquery.source +++ b/src/test/regress/output/multi_subquery.source @@ -326,3 +326,46 @@ FROM l_orderkey) AS unit_prices; ERROR: cannot push down this subquery DETAIL: Shards of relations in subquery need to have 1-to-1 shard partitioning +-- Check that we can prune shards in subqueries with VARCHAR partition columns +CREATE TABLE subquery_pruning_varchar_test_table +( + a varchar, + b int +); +SELECT master_create_distributed_table('subquery_pruning_varchar_test_table', 'a', 'hash'); + master_create_distributed_table +--------------------------------- + +(1 row) + +SELECT master_create_worker_shards('subquery_pruning_varchar_test_table', 4, 1); + master_create_worker_shards +----------------------------- + +(1 row) + +SET citus.subquery_pushdown TO TRUE; +SET client_min_messages TO DEBUG2; +SELECT * FROM + (SELECT count(*) FROM subquery_pruning_varchar_test_table WHERE a = 'onder' GROUP BY a) +AS foo; +DEBUG: predicate pruning for shardId 102029 +DEBUG: predicate pruning for shardId 102027 +DEBUG: predicate pruning for shardId 102026 + count +------- +(0 rows) + +SELECT * FROM + (SELECT count(*) FROM subquery_pruning_varchar_test_table WHERE 'eren' = a GROUP BY a) +AS foo; +DEBUG: predicate pruning for shardId 102029 +DEBUG: predicate pruning for shardId 102028 +DEBUG: predicate pruning for shardId 102026 + count +------- +(0 rows) + +SET client_min_messages TO NOTICE; +DROP TABLE subquery_pruning_varchar_test_table; +ALTER SEQUENCE pg_catalog.pg_dist_shardid_seq RESTART 102026;