From 9f08325983cbc37f9a68ecf4cdab201f1b75dca0 Mon Sep 17 00:00:00 2001
From: "tender.wang" <tender.wang@openpie.com>
Date: Wed, 6 Mar 2024 12:20:00 +0800
Subject: [PATCH v4] Fix CREATE INDEX failed due to parallel unsafe function.

Currently, the expression passed to is_parallel_safe() is created from
RelationGetIndexExpressions(). In RelationGetIndexExpressions(), it will
flatten the raw index expression if possible. It will ok in most cases, because
index expression usually only include column reference.

If index expressions include function, which is defined with PARALLEL UNSAFE label.
And the UDF is so simple that RelationGetIndexExpressions() will flatten it. It will
report error if building index in parallel mode. For example, the UDF contain EXCEPTION
statement.

The solution is to pass the raw index expression to is_parallel_safe(). The raw
expression should get from syscache not relcache.
---
 src/backend/optimizer/plan/planner.c      |  9 ++-
 src/backend/utils/cache/lsyscache.c       | 76 +++++++++++++++++++++++
 src/include/utils/lsyscache.h             |  2 +
 src/test/regress/expected/btree_index.out | 17 +++++
 src/test/regress/sql/btree_index.sql      | 21 +++++++
 5 files changed, 123 insertions(+), 2 deletions(-)

diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index ac97575453..0876179e83 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -6667,10 +6667,15 @@ plan_create_index_workers(Oid tableOid, Oid indexOid)
 	 * Currently, parallel workers can't access the leader's temporary tables.
 	 * Furthermore, any index predicate or index expressions must be parallel
 	 * safe.
+	 *
+	 * Note that we must get index expressions or predicate from syscache not
+	 * relcache. The expressions stored in relcache may have been flattened in
+	 * RelationGetIndexExpression(). The flattened expressions would lose PARALLEL
+	 * UNSAFE label, which it is not what we want.
 	 */
 	if (heap->rd_rel->relpersistence == RELPERSISTENCE_TEMP ||
-		!is_parallel_safe(root, (Node *) RelationGetIndexExpressions(index)) ||
-		!is_parallel_safe(root, (Node *) RelationGetIndexPredicate(index)))
+		!is_parallel_safe(root, (Node *) get_index_expression(indexOid)) ||
+		!is_parallel_safe(root, (Node *) get_index_predicate(indexOid)))
 	{
 		parallel_workers = 0;
 		goto done;
diff --git a/src/backend/utils/cache/lsyscache.c b/src/backend/utils/cache/lsyscache.c
index 6418d1c6eb..378764b0a8 100644
--- a/src/backend/utils/cache/lsyscache.c
+++ b/src/backend/utils/cache/lsyscache.c
@@ -3502,6 +3502,82 @@ get_index_column_opclass(Oid index_oid, int attno)
 	return opclass;
 }
 
+/*
+ * get_index_expression
+ *
+ *		Given the index OID, return the index expression.
+ */
+List *
+get_index_expression(Oid index_oid)
+{
+	List	  *result;
+	HeapTuple tuple;
+	Datum	  exprDatum;
+	bool	  isnull;
+	char	  *exprString;
+
+	/* Fetch the pg_index tuple by the Oid of the index */
+	tuple = SearchSysCache1(INDEXRELID, ObjectIdGetDatum(index_oid));
+	if (!HeapTupleIsValid(tuple))
+		elog(ERROR, "cache lookup failed for index %u", index_oid);
+
+	/* Extract raw index expression from index tuple */
+	exprDatum = SysCacheGetAttr(INDEXRELID, tuple,
+								Anum_pg_index_indexprs, &isnull);
+	if (isnull)
+	{
+		ReleaseSysCache(tuple);
+		return NIL;
+	}
+
+	exprString = TextDatumGetCString(exprDatum);
+	result = (List *) stringToNode(exprString);
+	pfree(exprString);
+
+	/* Clean up */
+	ReleaseSysCache(tuple);
+
+	return result;
+}
+
+/*
+ * get_index_predicate
+ *
+ *		Given the index OID, return index predicate.
+ */
+List *
+get_index_predicate(Oid index_oid)
+{
+	List	  *result;
+	HeapTuple tuple;
+	Datum	  predDatum;
+	bool	  isnull;
+	char	  *predString;
+
+	/* Fetch the pg_index tuple by the Oid of the index */
+	tuple = SearchSysCache1(INDEXRELID, ObjectIdGetDatum(index_oid));
+	if (!HeapTupleIsValid(tuple))
+		elog(ERROR, "cache lookup failed for index %u", index_oid);
+
+	/* Extract raw index predicate from index tuple */
+	predDatum = SysCacheGetAttr(INDEXRELID, tuple,
+								Anum_pg_index_indpred, &isnull);
+	if (isnull)
+	{
+		ReleaseSysCache(tuple);
+		return NIL;
+	}
+
+	predString = TextDatumGetCString(predDatum);
+	result = (List *) stringToNode(predString);
+	pfree(predString);
+
+	/* Clean up */
+	ReleaseSysCache(tuple);
+
+	return result;
+}
+
 /*
  * get_index_isreplident
  *
diff --git a/src/include/utils/lsyscache.h b/src/include/utils/lsyscache.h
index e4a200b00e..cde6bb76ac 100644
--- a/src/include/utils/lsyscache.h
+++ b/src/include/utils/lsyscache.h
@@ -195,6 +195,8 @@ extern Oid	get_range_collation(Oid rangeOid);
 extern Oid	get_range_multirange(Oid rangeOid);
 extern Oid	get_multirange_range(Oid multirangeOid);
 extern Oid	get_index_column_opclass(Oid index_oid, int attno);
+extern List *get_index_expression(Oid index_oid);
+extern List *get_index_predicate(Oid index_oid);
 extern bool get_index_isreplident(Oid index_oid);
 extern bool get_index_isvalid(Oid index_oid);
 extern bool get_index_isclustered(Oid index_oid);
diff --git a/src/test/regress/expected/btree_index.out b/src/test/regress/expected/btree_index.out
index 8311a03c3d..29a2903d80 100644
--- a/src/test/regress/expected/btree_index.out
+++ b/src/test/regress/expected/btree_index.out
@@ -434,3 +434,20 @@ ALTER INDEX btree_part_idx ALTER COLUMN id SET (n_distinct=100);
 ERROR:  ALTER action ALTER COLUMN ... SET cannot be performed on relation "btree_part_idx"
 DETAIL:  This operation is not supported for partitioned indexes.
 DROP TABLE btree_part;
+-- Test index expression includes parallel unsafe function
+CREATE FUNCTION para_unsafe_f() RETURNS int IMMUTABLE PARALLEL UNSAFE
+AS $$
+BEGIN
+    RETURN 0;
+EXCEPTION WHEN OTHERS THEN
+    RETURN 1;
+END$$ LANGUAGE plpgsql;
+CREATE TABLE btree_para_bld(i int);
+ALTER TABLE btree_para_bld SET (parallel_workers = 4);
+SET max_parallel_maintenance_workers TO 4;
+CREATE INDEX ON btree_para_bld((i + para_unsafe_f()));
+-- Test index predicate includes parallel unsafe function
+CREATE INDEX ON btree_para_bld(i) WHERE i > para_unsafe_f();
+RESET max_parallel_maintenance_workers;
+DROP TABLE btree_para_bld;
+DROP FUNCTION para_unsafe_f;
diff --git a/src/test/regress/sql/btree_index.sql b/src/test/regress/sql/btree_index.sql
index ef84354234..a054d106fb 100644
--- a/src/test/regress/sql/btree_index.sql
+++ b/src/test/regress/sql/btree_index.sql
@@ -267,3 +267,24 @@ CREATE TABLE btree_part (id int4) PARTITION BY RANGE (id);
 CREATE INDEX btree_part_idx ON btree_part(id);
 ALTER INDEX btree_part_idx ALTER COLUMN id SET (n_distinct=100);
 DROP TABLE btree_part;
+
+-- Test index expression includes parallel unsafe function
+CREATE FUNCTION para_unsafe_f() RETURNS int IMMUTABLE PARALLEL UNSAFE
+AS $$
+BEGIN
+    RETURN 0;
+EXCEPTION WHEN OTHERS THEN
+    RETURN 1;
+END$$ LANGUAGE plpgsql;
+
+CREATE TABLE btree_para_bld(i int);
+ALTER TABLE btree_para_bld SET (parallel_workers = 4);
+SET max_parallel_maintenance_workers TO 4;
+CREATE INDEX ON btree_para_bld((i + para_unsafe_f()));
+
+-- Test index predicate includes parallel unsafe function
+CREATE INDEX ON btree_para_bld(i) WHERE i > para_unsafe_f();
+
+RESET max_parallel_maintenance_workers;
+DROP TABLE btree_para_bld;
+DROP FUNCTION para_unsafe_f;
-- 
2.25.1

