diff --git a/src/backend/utils/adt/partitionfuncs.c b/src/backend/utils/adt/partitionfuncs.c
index ffd66b6439..36d9f69cbc 100644
--- a/src/backend/utils/adt/partitionfuncs.c
+++ b/src/backend/utils/adt/partitionfuncs.c
@@ -69,9 +69,6 @@ pg_partition_tree(PG_FUNCTION_ARGS)
 	FuncCallContext *funcctx;
 	ListCell  **next;
 
-	if (!check_rel_can_be_partition(rootrelid))
-		PG_RETURN_NULL();
-
 	/* stuff done only on the first call of the function */
 	if (SRF_IS_FIRSTCALL())
 	{
@@ -82,6 +79,9 @@ pg_partition_tree(PG_FUNCTION_ARGS)
 		/* create a function context for cross-call persistence */
 		funcctx = SRF_FIRSTCALL_INIT();
 
+		if (!check_rel_can_be_partition(rootrelid))
+			SRF_RETURN_DONE(funcctx);
+
 		/* switch to memory context appropriate for multiple function calls */
 		oldcxt = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
 
diff --git a/src/test/regress/expected/partition_info.out b/src/test/regress/expected/partition_info.out
index a884df976f..73269ffd09 100644
--- a/src/test/regress/expected/partition_info.out
+++ b/src/test/regress/expected/partition_info.out
@@ -9,8 +9,7 @@ SELECT * FROM pg_partition_tree(NULL);
 SELECT * FROM pg_partition_tree(0);
  relid | parentrelid | isleaf | level 
 -------+-------------+--------+-------
-       |             |        |      
-(1 row)
+(0 rows)
 
 SELECT pg_partition_root(NULL);
  pg_partition_root 
@@ -163,14 +162,12 @@ CREATE MATERIALIZED VIEW ptif_test_matview AS SELECT 1;
 SELECT * FROM pg_partition_tree('ptif_test_view');
  relid | parentrelid | isleaf | level 
 -------+-------------+--------+-------
-       |             |        |      
-(1 row)
+(0 rows)
 
 SELECT * FROM pg_partition_tree('ptif_test_matview');
  relid | parentrelid | isleaf | level 
 -------+-------------+--------+-------
-       |             |        |      
-(1 row)
+(0 rows)
 
 SELECT pg_partition_root('ptif_test_view');
  pg_partition_root