[BACKPORT 2.20][#20302,#20914] YSQL: fix colocation option issues rel…

…ated to table rewrite and table partition

Summary:
This diff fixes two issues of the colocation option in table creation.
(1) Table rewrite
YB table rewrite operations like ALTER TABLE ADD PRIMARY KEY and REFRESH MATERIALIZED VIEW create a new YB table under the hood.
`reloptions` field in `pg_class` catalog table is read to maintain the property of tables.
Option value in `reloptions` syntax like `colocation=0` is parsed as integer type by the PG parser,
but when untransformed from `reloptions` field, its type is string.
To solve this issue, extend function `defGetBoolean` to process string `0` and `1` as false and true, respectively.

(2) Table partition
We have a check to prevent tablegroup from using with colocation option mainly serving as a grammar guard,
and this check is placed in an inappropriate place causing table partition creation to fail when the colocation option is specified.
Solve this issue by moving the check to an earlier place in the CREATE TABLE execution path.
Jira: DB-9268, DB-9896

Original commit: 2cbb49d / D33001

Test Plan: ./yb_build.sh --java-test 'org.yb.pgsql.TestPgRegressColocation'

Reviewers: tverona, fizaa

Reviewed By: fizaa

Subscribers: yql

Tags: #jenkins-ready

Differential Revision: https://phorge.dev.yugabyte.com/D33346

src/postgres/src/backend/commands/define.c

@@ -42,6 +42,9 @@
 #include "parser/scansup.h"
 #include "utils/builtins.h"
 
+/* YB Includes */
+#include "pg_yb_utils.h"
+
 /*
  * Extract a string value (otherwise uninterpreted) from a DefElem.
  */
@@ -149,6 +152,30 @@ defGetBoolean(DefElem *def)
 					return true;
 				if (pg_strcasecmp(sval, "off") == 0)
 					return false;
+				/*
+				 * YB: PG function untransformRelOptions untransform text-array
+				 * format of reloptions stored in pg_class to a list of DefElem
+				 * whose arg has T_String type.
+				 * For example, reloption colocation=0 is parsed to a DefElem
+				 * whose arg 0 has T_Integer by parser.
+				 * However, text format colocation=0 stored in pg_class is
+				 * converted to a DefElem whose arg 0 has T_String type by
+				 * untransformRelOptions.
+				 * We use untransformRelOptions to support table rewrite
+				 * operations like ALTER TABLE, REFRESH MATERIALIZED VIEW, so
+				 * for YB's table boolean reloptions, we allow T_String type
+				 * "0" and "1" to be extracted as boolean false and true,
+				 * respectively.
+				 */
+				if (IsYugaByteEnabled() &&
+					(strcmp(def->defname, "colocated") == 0 ||
+					 strcmp(def->defname, "colocation") == 0))
+				{
+					if (pg_strcasecmp(sval, "1") == 0)
+						return true;
+					if (pg_strcasecmp(sval, "0") == 0)
+						return false;
+				}
 			}
 			break;
 	}

src/postgres/src/backend/commands/tablecmds.c

@@ -591,6 +591,7 @@ DefineRelation(CreateStmt *stmt, char relkind, Oid ownerId,
 	Oid			rowTypeId = InvalidOid;
 	bool		relisshared = false;
 	bool		use_initdb_acl = false;
+	ListCell   *opt_cell;
 
 	/*
 	 * Truncate relname to appropriate length (probably a waste of time, as
@@ -799,6 +800,27 @@ DefineRelation(CreateStmt *stmt, char relkind, Oid ownerId,
 		? get_tablegroup_oid(stmt->tablegroupname, false)
 		: InvalidOid;
 
+	if (IsYugaByteEnabled())
+	{
+		foreach(opt_cell, stmt->options)
+		{
+			DefElem *def = (DefElem *) lfirst(opt_cell);
+
+			/*
+			 * A check in parse_utilcmd.c makes sure only one of these two options
+			 * can be specified.
+			 */
+			if (strcmp(def->defname, "colocated") == 0 ||
+				strcmp(def->defname, "colocation") == 0)
+			{
+				if (OidIsValid(tablegroupId))
+					ereport(ERROR,
+							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							 errmsg("cannot use \'colocation=true/false\' with tablegroup")));
+			}
+		}
+	}
+
 	/*
 	 * Check permissions for tablegroup. To create a table within a tablegroup, a user must
 	 * either be a superuser, the owner of the tablegroup, or have create perms on it.
@@ -11943,7 +11965,7 @@ ATPrepSetTableSpace(AlteredTableInfo *tab, Relation rel, const char *tablespacen
 				 errmsg("cannot set tablespaces for temporary tables")));
 	}
 
-	if (IsYugaByteEnabled() && tablespacename && 
+	if (IsYugaByteEnabled() && tablespacename &&
 		rel->rd_index &&
 		rel->rd_index->indisprimary) {
 		/*
@@ -12438,7 +12460,7 @@ ATExecSetTableSpaceNoStorage(Relation rel, Oid newTableSpace)
 
 	/* Notify the user that this command is async */
 	ereport(NOTICE,
-			(errmsg("Data movement for table %s is successfully initiated.", 
+			(errmsg("Data movement for table %s is successfully initiated.",
 					RelationGetRelationName(rel)),
 			 errdetail("Data movement is a long running asynchronous process "
 					   "and can be monitored by checking the tablet placement "

src/postgres/src/backend/commands/ybccmds.c

@@ -657,11 +657,6 @@ YBCCreateTable(CreateStmt *stmt, char relkind, TupleDesc desc,
 		if (strcmp(def->defname, "colocated") == 0 ||
 			strcmp(def->defname, "colocation") == 0)
 		{
-			if (OidIsValid(tablegroupId))
-				ereport(ERROR,
-						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-						 errmsg("cannot use \'colocation=true/false\' with tablegroup")));
-
 			bool colocated_relopt = defGetBoolean(def);
 			if (MyDatabaseColocated)
 				is_colocated_via_database = colocated_relopt;

src/postgres/src/test/regress/expected/yb_feature_colocation.out

@@ -948,3 +948,110 @@ select is_colocated from yb_table_properties('m5'::regclass);
 --------------
  f
 (1 row)
+
+-- Test Table Rewrite operations with table colocation option
+CREATE TABLE non_col_tbl_without_pk (k INT) WITH (colocation=0);
+INSERT INTO non_col_tbl_without_pk SELECT generate_series(1, 100);
+ALTER TABLE non_col_tbl_without_pk ADD PRIMARY KEY(k);
+NOTICE:  table rewrite may lead to inconsistencies
+DETAIL:  Concurrent DMLs may not be reflected in the new table.
+HINT:  See https://github.com/yugabyte/yugabyte-db/issues/19860. Set 'ysql_suppress_unsafe_alter_notice' yb-tserver gflag to true to suppress this notice.
+\d non_col_tbl_without_pk
+       Table "public.non_col_tbl_without_pk"
+ Column |  Type   | Collation | Nullable | Default
+--------+---------+-----------+----------+---------
+ k      | integer |           | not null |
+Indexes:
+    "non_col_tbl_without_pk_pkey" PRIMARY KEY, lsm (k HASH)
+
+SELECT min(k), max(k) FROM non_col_tbl_without_pk;
+ min | max
+-----+-----
+   1 | 100
+(1 row)
+
+CREATE TABLE col_tbl_without_pk (k INT) WITH (colocation=1);
+INSERT INTO col_tbl_without_pk SELECT generate_series(1, 10);
+ALTER TABLE col_tbl_without_pk ADD PRIMARY KEY(k);
+NOTICE:  table rewrite may lead to inconsistencies
+DETAIL:  Concurrent DMLs may not be reflected in the new table.
+HINT:  See https://github.com/yugabyte/yugabyte-db/issues/19860. Set 'ysql_suppress_unsafe_alter_notice' yb-tserver gflag to true to suppress this notice.
+\d col_tbl_without_pk
+         Table "public.col_tbl_without_pk"
+ Column |  Type   | Collation | Nullable | Default
+--------+---------+-----------+----------+---------
+ k      | integer |           | not null |
+Indexes:
+    "col_tbl_without_pk_pkey" PRIMARY KEY, lsm (k ASC), colocation: true
+Colocation: true
+
+SELECT sum(k) FROM col_tbl_without_pk;
+ sum
+-----
+  55
+(1 row)
+
+CREATE MATERIALIZED VIEW non_col_mv WITH (colocation=0) AS SELECT * FROM col_tbl_without_pk WHERE k % 2 = 0;
+CREATE MATERIALIZED VIEW col_mv WITH (colocation=1) AS SELECT * FROM col_tbl_without_pk WHERE k % 2 = 1;
+SELECT sum(k) FROM non_col_mv;
+ sum
+-----
+  30
+(1 row)
+
+SELECT sum(k) FROM col_mv;
+ sum
+-----
+  25
+(1 row)
+
+INSERT INTO col_tbl_without_pk VALUES (11);
+REFRESH MATERIALIZED VIEW non_col_mv;
+REFRESH MATERIALIZED VIEW col_mv;
+SELECT sum(k) FROM non_col_mv;
+ sum
+-----
+  30
+(1 row)
+
+SELECT sum(k) FROM col_mv;
+ sum
+-----
+  36
+(1 row)
+
+-- Test table partitions with colocation option
+CREATE TABLE partitioned_table (k INT) PARTITION BY LIST (k);
+CREATE TABLE col_table_partition PARTITION OF partitioned_table FOR VALUES IN (2) WITH (COLOCATION=TRUE);
+\d col_table_partition
+        Table "public.col_table_partition"
+ Column |  Type   | Collation | Nullable | Default
+--------+---------+-----------+----------+---------
+ k      | integer |           |          |
+Partition of: partitioned_table FOR VALUES IN (2)
+Colocation: true
+
+CREATE TABLE non_col_table_partition PARTITION OF partitioned_table FOR VALUES IN (4) WITH (COLOCATION=FALSE);
+\d non_col_table_partition
+      Table "public.non_col_table_partition"
+ Column |  Type   | Collation | Nullable | Default
+--------+---------+-----------+----------+---------
+ k      | integer |           |          |
+Partition of: partitioned_table FOR VALUES IN (4)
+Colocation: true
+
+INSERT INTO partitioned_table VALUES (2), (4);
+SELECT * FROM col_table_partition;
+ k
+---
+ 2
+(1 row)
+
+SELECT * FROM non_col_table_partition;
+ k
+---
+ 4
+(1 row)
+
+\c yugabyte
+DROP DATABASE colocation_test;

src/postgres/src/test/regress/sql/yb_feature_colocation.sql

@@ -351,7 +351,7 @@ CREATE MATERIALIZED VIEW m3 with (colocation = false) AS SELECT * FROM t1;
 \c yugabyte
 DROP DATABASE colocation_test;
 
--- Test Colocated Materialized View 
+-- Test Colocated Materialized View
 CREATE DATABASE colocation_test WITH colocation = true;
 \c colocation_test
 CREATE TABLE t1 (a INT PRIMARY KEY) WITH (colocation = true);
@@ -368,3 +368,40 @@ select is_colocated from yb_table_properties('m2'::regclass);
 select is_colocated from yb_table_properties('m3'::regclass);
 select is_colocated from yb_table_properties('m4'::regclass);
 select is_colocated from yb_table_properties('m5'::regclass);
+
+-- Test Table Rewrite operations with table colocation option
+CREATE TABLE non_col_tbl_without_pk (k INT) WITH (colocation=0);
+INSERT INTO non_col_tbl_without_pk SELECT generate_series(1, 100);
+ALTER TABLE non_col_tbl_without_pk ADD PRIMARY KEY(k);
+\d non_col_tbl_without_pk
+SELECT min(k), max(k) FROM non_col_tbl_without_pk;
+
+CREATE TABLE col_tbl_without_pk (k INT) WITH (colocation=1);
+INSERT INTO col_tbl_without_pk SELECT generate_series(1, 10);
+ALTER TABLE col_tbl_without_pk ADD PRIMARY KEY(k);
+\d col_tbl_without_pk
+SELECT sum(k) FROM col_tbl_without_pk;
+
+CREATE MATERIALIZED VIEW non_col_mv WITH (colocation=0) AS SELECT * FROM col_tbl_without_pk WHERE k % 2 = 0;
+CREATE MATERIALIZED VIEW col_mv WITH (colocation=1) AS SELECT * FROM col_tbl_without_pk WHERE k % 2 = 1;
+SELECT sum(k) FROM non_col_mv;
+SELECT sum(k) FROM col_mv;
+INSERT INTO col_tbl_without_pk VALUES (11);
+REFRESH MATERIALIZED VIEW non_col_mv;
+REFRESH MATERIALIZED VIEW col_mv;
+SELECT sum(k) FROM non_col_mv;
+SELECT sum(k) FROM col_mv;
+
+-- Test table partitions with colocation option
+CREATE TABLE partitioned_table (k INT) PARTITION BY LIST (k);
+CREATE TABLE col_table_partition PARTITION OF partitioned_table FOR VALUES IN (2) WITH (COLOCATION=TRUE);
+\d col_table_partition
+CREATE TABLE non_col_table_partition PARTITION OF partitioned_table FOR VALUES IN (4) WITH (COLOCATION=FALSE);
+\d non_col_table_partition
+
+INSERT INTO partitioned_table VALUES (2), (4);
+SELECT * FROM col_table_partition;
+SELECT * FROM non_col_table_partition;
+
+\c yugabyte
+DROP DATABASE colocation_test;