From ec60e592dc653651f410d5cc7927d110a90d3b92 Mon Sep 17 00:00:00 2001
From: prajnamort <prajnamort@gmail.com>
Date: Wed, 25 Mar 2020 15:50:01 +0800
Subject: [PATCH] Check operator when creating UNIQUE index on PARTITION table

Currently, we only check if every column in the partition key
is included in the unique index. But this is not enough when
custom opclass exists. Since there can be cases, that some
values belong to different partitions should be considered
equal according to unique index's opclass.

So, we should check the compatiblity of equality operator
between partition key's opclass and index column's opclass.
Only if they are the same operator should we consider them as
compatible.
---
 src/backend/commands/indexcmds.c                | 34 +++++++++++-
 src/test/regress/expected/partition_opclass.out | 69 ++++++++++++++++++++++++
 src/test/regress/parallel_schedule              |  2 +-
 src/test/regress/serial_schedule                |  1 +
 src/test/regress/sql/partition_opclass.sql      | 70 +++++++++++++++++++++++++
 5 files changed, 173 insertions(+), 3 deletions(-)
 create mode 100644 src/test/regress/expected/partition_opclass.out
 create mode 100644 src/test/regress/sql/partition_opclass.sql

diff --git a/src/backend/commands/indexcmds.c b/src/backend/commands/indexcmds.c
index 4e8263af4b..e7c3960011 100644
--- a/src/backend/commands/indexcmds.c
+++ b/src/backend/commands/indexcmds.c
@@ -877,8 +877,38 @@ DefineIndex(Oid relationId,
 			{
 				if (key->partattrs[i] == indexInfo->ii_IndexAttrNumbers[j])
 				{
-					found = true;
-					break;
+					/*
+					 * Check the equality operator in both the index's operator
+					 * class and the partition key's operator class. They must
+					 * be the same operator for us to consider them compatible.
+					 */
+					int ptkey_eq_strategy;
+					Oid ptkey_eqop;
+					Oid idxcol_opclass;
+					Oid idxcol_opfamily;
+					Oid idxcol_opclass_intype;
+					Oid idxcol_eqop;
+
+					if (key->strategy == PARTITION_STRATEGY_HASH)
+						ptkey_eq_strategy = HTEqualStrategyNumber;
+					else
+						ptkey_eq_strategy = BTEqualStrategyNumber;
+					ptkey_eqop = get_opfamily_member(key->partopfamily[i],
+													 key->partopcintype[i],
+													 key->partopcintype[i],
+													 ptkey_eq_strategy);
+					idxcol_opclass = classObjectId[j];
+					idxcol_opfamily = get_opclass_family(idxcol_opclass);
+					idxcol_opclass_intype = get_opclass_input_type(idxcol_opclass);
+					idxcol_eqop = get_opfamily_member(idxcol_opfamily,
+													  idxcol_opclass_intype,
+													  idxcol_opclass_intype,
+													  BTEqualStrategyNumber);
+					if (ptkey_eqop == idxcol_eqop)
+					{
+						found = true;
+						break;
+					}
 				}
 			}
 			if (!found)
diff --git a/src/test/regress/expected/partition_opclass.out b/src/test/regress/expected/partition_opclass.out
new file mode 100644
index 0000000000..27e0664de3
--- /dev/null
+++ b/src/test/regress/expected/partition_opclass.out
@@ -0,0 +1,69 @@
+--
+-- Test for checking operator class compatibility between partition key and
+-- unique index. If they are imcompatible, the unique index should not be
+-- allowed to created.
+--
+CREATE FUNCTION abscmp(int, int) RETURNS INT AS $$
+    begin return btint4cmp(abs($1),abs($2)); end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE FUNCTION abslt(int, int) RETURNS BOOL AS $$
+    begin return abscmp($1,$2) < 0; end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE FUNCTION abslte(int, int) RETURNS BOOL AS $$
+    begin return abscmp($1,$2) <= 0; end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE FUNCTION abseq(int, int) RETURNS BOOL AS $$
+    begin return abscmp($1,$2) = 0; end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE FUNCTION absgte(int, int) RETURNS BOOL AS $$
+    begin return abscmp($1,$2) >= 0; end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE FUNCTION absgt(int, int) RETURNS BOOL AS $$
+    begin return abscmp($1,$2) > 0; end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE OPERATOR |<| (PROCEDURE = abslt, LEFTARG = int, RIGHTARG = int, COMMUTATOR = |>|, HASHES);
+CREATE OPERATOR |<=| (PROCEDURE = abslt, LEFTARG = int, RIGHTARG = int, COMMUTATOR = |>=|, HASHES);
+CREATE OPERATOR |=| (PROCEDURE = abseq, LEFTARG = int, RIGHTARG = int, COMMUTATOR = |=|, HASHES, MERGES);
+CREATE OPERATOR |>=| (PROCEDURE = abslt, LEFTARG = int, RIGHTARG = int, COMMUTATOR = |<=|, HASHES);
+CREATE OPERATOR |>| (PROCEDURE = abslt, LEFTARG = int, RIGHTARG = int, COMMUTATOR = |<|, HASHES);
+CREATE OPERATOR CLASS abs_int_btree_ops FOR TYPE int4
+    USING btree AS
+    OPERATOR 1 |<|,
+    OPERATOR 2 |<=|,
+    OPERATOR 3 |=|,
+    OPERATOR 4 |>=|,
+    OPERATOR 5 |>|,
+    FUNCTION 1 abscmp(int, int);
+CREATE TABLE ptop_test (a int, b int, c int) PARTITION BY LIST (a);
+CREATE TABLE ptop_test_0 PARTITION OF ptop_test FOR VALUES IN ('0');
+CREATE TABLE ptop_test_p1 PARTITION OF ptop_test FOR VALUES IN ('1');
+CREATE TABLE ptop_test_m1 PARTITION OF ptop_test FOR VALUES IN ('-1');
+CREATE UNIQUE INDEX ptop_test_unq_abs_a ON ptop_test (a abs_int_btree_ops);  -- should fail
+ERROR:  insufficient columns in UNIQUE constraint definition
+DETAIL:  UNIQUE constraint on table "ptop_test" lacks column "a" which is part of the partition key.
+CREATE UNIQUE INDEX ptop_test_unq_a ON ptop_test (a);
+DROP INDEX ptop_test_unq_a;
+DROP TABLE ptop_test;
+CREATE TABLE ptop_test (a int, b int, c int) PARTITION BY LIST (a abs_int_btree_ops);
+CREATE TABLE ptop_test_0 PARTITION OF ptop_test FOR VALUES IN ('0');
+CREATE TABLE ptop_test_p1 PARTITION OF ptop_test FOR VALUES IN ('1', '-1');
+CREATE UNIQUE INDEX ptop_test_unq_a ON ptop_test (a);  -- should fail
+ERROR:  insufficient columns in UNIQUE constraint definition
+DETAIL:  UNIQUE constraint on table "ptop_test" lacks column "a" which is part of the partition key.
+CREATE UNIQUE INDEX ptop_test_unq_abs_a ON ptop_test (a abs_int_btree_ops);
+DROP INDEX ptop_test_unq_abs_a;
+DROP TABLE ptop_test;
+-- Clean up
+DROP OPERATOR CLASS abs_int_btree_ops USING btree;
+DROP OPERATOR FAMILY abs_int_btree_ops USING btree;
+DROP OPERATOR |<| (int, int);
+DROP OPERATOR |<=| (int, int);
+DROP OPERATOR |=| (int, int);
+DROP OPERATOR |>=| (int, int);
+DROP OPERATOR |>| (int, int);
+DROP FUNCTION abscmp(int, int);
+DROP FUNCTION abslt(int, int);
+DROP FUNCTION abslte(int, int);
+DROP FUNCTION abseq(int, int);
+DROP FUNCTION absgte(int, int);
+DROP FUNCTION absgt(int, int);
diff --git a/src/test/regress/parallel_schedule b/src/test/regress/parallel_schedule
index d2b17dd3ea..ac6919a8bc 100644
--- a/src/test/regress/parallel_schedule
+++ b/src/test/regress/parallel_schedule
@@ -112,7 +112,7 @@ test: plancache limit plpgsql copy2 temp domain rangefuncs prepare conversion tr
 # ----------
 # Another group of parallel tests
 # ----------
-test: partition_join partition_prune reloptions hash_part indexing partition_aggregate partition_info tuplesort explain
+test: partition_join partition_prune reloptions hash_part indexing partition_aggregate partition_info tuplesort explain partition_opclass
 
 # event triggers cannot run concurrently with any test that runs DDL
 test: event_trigger
diff --git a/src/test/regress/serial_schedule b/src/test/regress/serial_schedule
index acba391332..b63bff2a05 100644
--- a/src/test/regress/serial_schedule
+++ b/src/test/regress/serial_schedule
@@ -197,3 +197,4 @@ test: explain
 test: event_trigger
 test: fast_default
 test: stats
+test: partition_opclass
\ No newline at end of file
diff --git a/src/test/regress/sql/partition_opclass.sql b/src/test/regress/sql/partition_opclass.sql
new file mode 100644
index 0000000000..1ad279f531
--- /dev/null
+++ b/src/test/regress/sql/partition_opclass.sql
@@ -0,0 +1,70 @@
+--
+-- Test for checking operator class compatibility between partition key and
+-- unique index. If they are imcompatible, the unique index should not be
+-- allowed to created.
+--
+
+CREATE FUNCTION abscmp(int, int) RETURNS INT AS $$
+    begin return btint4cmp(abs($1),abs($2)); end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE FUNCTION abslt(int, int) RETURNS BOOL AS $$
+    begin return abscmp($1,$2) < 0; end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE FUNCTION abslte(int, int) RETURNS BOOL AS $$
+    begin return abscmp($1,$2) <= 0; end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE FUNCTION abseq(int, int) RETURNS BOOL AS $$
+    begin return abscmp($1,$2) = 0; end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE FUNCTION absgte(int, int) RETURNS BOOL AS $$
+    begin return abscmp($1,$2) >= 0; end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE FUNCTION absgt(int, int) RETURNS BOOL AS $$
+    begin return abscmp($1,$2) > 0; end;
+$$ LANGUAGE plpgsql STRICT IMMUTABLE;
+CREATE OPERATOR |<| (PROCEDURE = abslt, LEFTARG = int, RIGHTARG = int, COMMUTATOR = |>|, HASHES);
+CREATE OPERATOR |<=| (PROCEDURE = abslt, LEFTARG = int, RIGHTARG = int, COMMUTATOR = |>=|, HASHES);
+CREATE OPERATOR |=| (PROCEDURE = abseq, LEFTARG = int, RIGHTARG = int, COMMUTATOR = |=|, HASHES, MERGES);
+CREATE OPERATOR |>=| (PROCEDURE = abslt, LEFTARG = int, RIGHTARG = int, COMMUTATOR = |<=|, HASHES);
+CREATE OPERATOR |>| (PROCEDURE = abslt, LEFTARG = int, RIGHTARG = int, COMMUTATOR = |<|, HASHES);
+CREATE OPERATOR CLASS abs_int_btree_ops FOR TYPE int4
+    USING btree AS
+    OPERATOR 1 |<|,
+    OPERATOR 2 |<=|,
+    OPERATOR 3 |=|,
+    OPERATOR 4 |>=|,
+    OPERATOR 5 |>|,
+    FUNCTION 1 abscmp(int, int);
+
+CREATE TABLE ptop_test (a int, b int, c int) PARTITION BY LIST (a);
+CREATE TABLE ptop_test_0 PARTITION OF ptop_test FOR VALUES IN ('0');
+CREATE TABLE ptop_test_p1 PARTITION OF ptop_test FOR VALUES IN ('1');
+CREATE TABLE ptop_test_m1 PARTITION OF ptop_test FOR VALUES IN ('-1');
+CREATE UNIQUE INDEX ptop_test_unq_abs_a ON ptop_test (a abs_int_btree_ops);  -- should fail
+CREATE UNIQUE INDEX ptop_test_unq_a ON ptop_test (a);
+DROP INDEX ptop_test_unq_a;
+DROP TABLE ptop_test;
+
+CREATE TABLE ptop_test (a int, b int, c int) PARTITION BY LIST (a abs_int_btree_ops);
+CREATE TABLE ptop_test_0 PARTITION OF ptop_test FOR VALUES IN ('0');
+CREATE TABLE ptop_test_p1 PARTITION OF ptop_test FOR VALUES IN ('1', '-1');
+CREATE UNIQUE INDEX ptop_test_unq_a ON ptop_test (a);  -- should fail
+CREATE UNIQUE INDEX ptop_test_unq_abs_a ON ptop_test (a abs_int_btree_ops);
+DROP INDEX ptop_test_unq_abs_a;
+DROP TABLE ptop_test;
+
+
+-- Clean up
+DROP OPERATOR CLASS abs_int_btree_ops USING btree;
+DROP OPERATOR FAMILY abs_int_btree_ops USING btree;
+DROP OPERATOR |<| (int, int);
+DROP OPERATOR |<=| (int, int);
+DROP OPERATOR |=| (int, int);
+DROP OPERATOR |>=| (int, int);
+DROP OPERATOR |>| (int, int);
+DROP FUNCTION abscmp(int, int);
+DROP FUNCTION abslt(int, int);
+DROP FUNCTION abslte(int, int);
+DROP FUNCTION abseq(int, int);
+DROP FUNCTION absgte(int, int);
+DROP FUNCTION absgt(int, int);
-- 
2.11.0

