3 files changed, 71 insertions, 1 deletions
diff --git a/src/backend/utils/adt/partitionfuncs.c b/src/backend/utils/adt/partitionfuncs.c
index 98c8ef77f17..714a0242a18 100644
--- a/src/backend/utils/adt/partitionfuncs.c
+++ b/src/backend/utils/adt/partitionfuncs.c
@@ -189,8 +189,16 @@ pg_partition_root(PG_FUNCTION_ARGS)
 	if (!check_rel_can_be_partition(relid))
 		PG_RETURN_NULL();
 
-	/* Fetch the top-most parent */
+	/* fetch the list of ancestors */
 	ancestors = get_partition_ancestors(relid);
+
+	/*
+	 * If the input relation is already the top-most parent, just return
+	 * itself.
+	 */
+	if (ancestors == NIL)
+		PG_RETURN_OID(relid);
+
 	rootrelid = llast_oid(ancestors);
 	list_free(ancestors);
 
diff --git a/src/test/regress/expected/partition_info.out b/src/test/regress/expected/partition_info.out
index 92aa6abe3f9..42b6bc77cad 100644
--- a/src/test/regress/expected/partition_info.out
+++ b/src/test/regress/expected/partition_info.out
@@ -46,6 +46,31 @@ CREATE TABLE ptif_test2 PARTITION OF ptif_test
 -- This partitioned table should remain with no partitions.
 CREATE TABLE ptif_test3 PARTITION OF ptif_test
   FOR VALUES FROM (200) TO (maxvalue) PARTITION BY list (b);
+-- Test pg_partition_root for tables
+SELECT pg_partition_root('ptif_test');
+ pg_partition_root 
+-------------------
+ ptif_test
+(1 row)
+
+SELECT pg_partition_root('ptif_test0');
+ pg_partition_root 
+-------------------
+ ptif_test
+(1 row)
+
+SELECT pg_partition_root('ptif_test01');
+ pg_partition_root 
+-------------------
+ ptif_test
+(1 row)
+
+SELECT pg_partition_root('ptif_test3');
+ pg_partition_root 
+-------------------
+ ptif_test
+(1 row)
+
 -- Test index partition tree
 CREATE INDEX ptif_test_index ON ONLY ptif_test (a);
 CREATE INDEX ptif_test0_index ON ONLY ptif_test0 (a);
@@ -60,6 +85,31 @@ CREATE INDEX ptif_test2_index ON ptif_test2 (a);
 ALTER INDEX ptif_test_index ATTACH PARTITION ptif_test2_index;
 CREATE INDEX ptif_test3_index ON ptif_test3 (a);
 ALTER INDEX ptif_test_index ATTACH PARTITION ptif_test3_index;
+-- Test pg_partition_root for indexes
+SELECT pg_partition_root('ptif_test_index');
+ pg_partition_root 
+-------------------
+ ptif_test_index
+(1 row)
+
+SELECT pg_partition_root('ptif_test0_index');
+ pg_partition_root 
+-------------------
+ ptif_test_index
+(1 row)
+
+SELECT pg_partition_root('ptif_test01_index');
+ pg_partition_root 
+-------------------
+ ptif_test_index
+(1 row)
+
+SELECT pg_partition_root('ptif_test3_index');
+ pg_partition_root 
+-------------------
+ ptif_test_index
+(1 row)
+
 -- List all tables members of the tree
 SELECT relid, parentrelid, level, isleaf
   FROM pg_partition_tree('ptif_test');
diff --git a/src/test/regress/sql/partition_info.sql b/src/test/regress/sql/partition_info.sql
index adaca232324..b5060bec7f0 100644
--- a/src/test/regress/sql/partition_info.sql
+++ b/src/test/regress/sql/partition_info.sql
@@ -22,6 +22,12 @@ CREATE TABLE ptif_test2 PARTITION OF ptif_test
 CREATE TABLE ptif_test3 PARTITION OF ptif_test
   FOR VALUES FROM (200) TO (maxvalue) PARTITION BY list (b);
 
+-- Test pg_partition_root for tables
+SELECT pg_partition_root('ptif_test');
+SELECT pg_partition_root('ptif_test0');
+SELECT pg_partition_root('ptif_test01');
+SELECT pg_partition_root('ptif_test3');
+
 -- Test index partition tree
 CREATE INDEX ptif_test_index ON ONLY ptif_test (a);
 CREATE INDEX ptif_test0_index ON ONLY ptif_test0 (a);
@@ -37,6 +43,12 @@ ALTER INDEX ptif_test_index ATTACH PARTITION ptif_test2_index;
 CREATE INDEX ptif_test3_index ON ptif_test3 (a);
 ALTER INDEX ptif_test_index ATTACH PARTITION ptif_test3_index;
 
+-- Test pg_partition_root for indexes
+SELECT pg_partition_root('ptif_test_index');
+SELECT pg_partition_root('ptif_test0_index');
+SELECT pg_partition_root('ptif_test01_index');
+SELECT pg_partition_root('ptif_test3_index');
+
 -- List all tables members of the tree
 SELECT relid, parentrelid, level, isleaf
   FROM pg_partition_tree('ptif_test');