Better handle pseudotypes as partition keys

We fail to handle polymorphic types properly when they are used as
partition keys: we were unnecessarily adding a RelabelType node on top,
which confuses code examining the nodes.  In particular, this makes
predtest.c-based partition pruning not to work, and ruleutils.c to emit
expressions that are uglier than needed.  Fix it by not adding RelabelType
when not needed.

In master/11 the new pruning code is separate so it doesn't suffer from
this problem, since we already fixed it (in essentially the same way) in
e5dcbb88a15d, which also added a few tests; back-patch those tests to
pg10 also.  But since UPDATE/DELETE still uses predtest.c in pg11, this
change improves partitioning for those cases too.  Add tests for this.
The ruleutils.c behavior change is relevant in pg11/master too.

Co-authored-by: Amit Langote <Langote_Amit_f8@lab.ntt.co.jp>
Co-authored-by: Álvaro Herrera <alvherre@alvh.no-ip.org>
Reviewed-by: Álvaro Herrera <alvherre@alvh.no-ip.org>
Reviewed-by: Robert Haas <robertmhaas@gmail.com>
Discussion: https://postgr.es/m/54745d13-7ed4-54ac-97d8-ea1eec95ae25@lab.ntt.co.jp

4 files changed, 41 insertions, 31 deletions

diff --git a/src/backend/partitioning/partbounds.c b/src/backend/partitioning/partbounds.c
index b19c76acc8a..9015a05d323 100644
--- a/src/backend/partitioning/partbounds.c
+++ b/src/backend/partitioning/partbounds.c
@@ -1144,8 +1144,12 @@ get_partition_bound_num_indexes(PartitionBoundInfo bound)
 /*
  * get_partition_operator
  *
- * Return oid of the operator of given strategy for a given partition key
- * column.
+ * Return oid of the operator of the given strategy for the given partition
+ * key column.  It is assumed that the partitioning key is of the same type as
+ * the chosen partitioning opclass, or at least binary-compatible.  In the
+ * latter case, *need_relabel is set to true if the opclass is not of a
+ * polymorphic type (indicating a RelabelType node needed on top), otherwise
+ * false.
  */
 static Oid
 get_partition_operator(PartitionKey key, int col, StrategyNumber strategy,
@@ -1154,40 +1158,26 @@ get_partition_operator(PartitionKey key, int col, StrategyNumber strategy,
 	Oid			operoid;
 
 	/*
-	 * First check if there exists an operator of the given strategy, with
-	 * this column's type as both its lefttype and righttype, in the
-	 * partitioning operator family specified for the column.
+	 * Get the operator in the partitioning opfamily using the opclass'
+	 * declared input type as both left- and righttype.
 	 */
 	operoid = get_opfamily_member(key->partopfamily[col],
-								  key->parttypid[col],
-								  key->parttypid[col],
+								  key->partopcintype[col],
+								  key->partopcintype[col],
 								  strategy);
+	if (!OidIsValid(operoid))
+		elog(ERROR, "missing operator %d(%u,%u) in partition opfamily %u",
+			 strategy, key->partopcintype[col], key->partopcintype[col],
+			 key->partopfamily[col]);
 
 	/*
-	 * If one doesn't exist, we must resort to using an operator in the same
-	 * operator family but with the operator class declared input type.  It is
-	 * OK to do so, because the column's type is known to be binary-coercible
-	 * with the operator class input type (otherwise, the operator class in
-	 * question would not have been accepted as the partitioning operator
-	 * class).  We must however inform the caller to wrap the non-Const
-	 * expression with a RelabelType node to denote the implicit coercion. It
-	 * ensures that the resulting expression structurally matches similarly
-	 * processed expressions within the optimizer.
+	 * If the partition key column is not of the same type as the operator
+	 * class and not polymorphic, tell caller to wrap the non-Const expression
+	 * in a RelabelType.  This matches what parse_coerce.c does.
 	 */
-	if (!OidIsValid(operoid))
-	{
-		operoid = get_opfamily_member(key->partopfamily[col],
-									  key->partopcintype[col],
-									  key->partopcintype[col],
-									  strategy);
-		if (!OidIsValid(operoid))
-			elog(ERROR, "missing operator %d(%u,%u) in opfamily %u",
-				 strategy, key->partopcintype[col], key->partopcintype[col],
-				 key->partopfamily[col]);
-		*need_relabel = true;
-	}
-	else
-		*need_relabel = false;
+	*need_relabel = (key->parttypid[col] != key->partopcintype[col] &&
+					 key->partopcintype[col] != RECORDOID &&
+					 !IsPolymorphicType(key->partopcintype[col]));
 
 	return operoid;
 }
diff --git a/src/test/regress/expected/create_table.out b/src/test/regress/expected/create_table.out
index 672719e5d5e..8fdbca13458 100644
--- a/src/test/regress/expected/create_table.out
+++ b/src/test/regress/expected/create_table.out
@@ -882,7 +882,7 @@ CREATE TABLE arrlp12 PARTITION OF arrlp FOR VALUES IN ('{1}', '{2}');
 --------+-----------+-----------+----------+---------+----------+--------------+-------------
  a      | integer[] |           |          |         | extended |              | 
 Partition of: arrlp FOR VALUES IN ('{1}', '{2}')
-Partition constraint: ((a IS NOT NULL) AND (((a)::anyarray OPERATOR(pg_catalog.=) '{1}'::integer[]) OR ((a)::anyarray OPERATOR(pg_catalog.=) '{2}'::integer[])))
+Partition constraint: ((a IS NOT NULL) AND ((a = '{1}'::integer[]) OR (a = '{2}'::integer[])))
 
 DROP TABLE arrlp;
 -- partition on boolean column
diff --git a/src/test/regress/expected/partition_prune.out b/src/test/regress/expected/partition_prune.out
index 9059147e171..d15f1d37f13 100644
--- a/src/test/regress/expected/partition_prune.out
+++ b/src/test/regress/expected/partition_prune.out
@@ -2817,6 +2817,24 @@ explain (costs off) select * from pp_arrpart where a in ('{4, 5}', '{1}');
          Filter: ((a = '{4,5}'::integer[]) OR (a = '{1}'::integer[]))
 (5 rows)
 
+explain (costs off) update pp_arrpart set a = a where a = '{1}';
+               QUERY PLAN               
+----------------------------------------
+ Update on pp_arrpart
+   Update on pp_arrpart1
+   ->  Seq Scan on pp_arrpart1
+         Filter: (a = '{1}'::integer[])
+(4 rows)
+
+explain (costs off) delete from pp_arrpart where a = '{1}';
+               QUERY PLAN               
+----------------------------------------
+ Delete on pp_arrpart
+   Delete on pp_arrpart1
+   ->  Seq Scan on pp_arrpart1
+         Filter: (a = '{1}'::integer[])
+(4 rows)
+
 drop table pp_arrpart;
 -- array type hash partition key
 create table pph_arrpart (a int[]) partition by hash (a);
diff --git a/src/test/regress/sql/partition_prune.sql b/src/test/regress/sql/partition_prune.sql
index 11b92bfada6..b8e823d5620 100644
--- a/src/test/regress/sql/partition_prune.sql
+++ b/src/test/regress/sql/partition_prune.sql
@@ -721,6 +721,8 @@ create table pp_arrpart2 partition of pp_arrpart for values in ('{2, 3}', '{4, 5
 explain (costs off) select * from pp_arrpart where a = '{1}';
 explain (costs off) select * from pp_arrpart where a = '{1, 2}';
 explain (costs off) select * from pp_arrpart where a in ('{4, 5}', '{1}');
+explain (costs off) update pp_arrpart set a = a where a = '{1}';
+explain (costs off) delete from pp_arrpart where a = '{1}';
 drop table pp_arrpart;
 
 -- array type hash partition key