2 files changed, 171 insertions, 0 deletions
diff --git a/src/test/regress/expected/stats_import.out b/src/test/regress/expected/stats_import.out
index fb50da1cd83..0e8491131e3 100644
--- a/src/test/regress/expected/stats_import.out
+++ b/src/test/regress/expected/stats_import.out
@@ -85,6 +85,44 @@ WHERE oid = 'stats_import.test'::regclass;
        17 |       400 |             4
 (1 row)
 
+CREATE INDEX test_i ON stats_import.test(id);
+BEGIN;
+-- regular indexes have special case locking rules
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.test_i'::regclass,
+        relpages => 18::integer);
+ pg_set_relation_stats 
+-----------------------
+ 
+(1 row)
+
+SELECT mode FROM pg_locks
+WHERE relation = 'stats_import.test'::regclass AND
+      pid = pg_backend_pid() AND granted;
+           mode           
+--------------------------
+ ShareUpdateExclusiveLock
+(1 row)
+
+SELECT mode FROM pg_locks
+WHERE relation = 'stats_import.test_i'::regclass AND
+      pid = pg_backend_pid() AND granted;
+      mode       
+-----------------
+ AccessShareLock
+(1 row)
+
+COMMIT;
+SELECT
+    pg_catalog.pg_restore_relation_stats(
+        'relation', 'stats_import.test_i'::regclass,
+        'relpages', 19::integer );
+ pg_restore_relation_stats 
+---------------------------
+ t
+(1 row)
+
 -- positional arguments
 SELECT
     pg_catalog.pg_set_relation_stats(
@@ -182,6 +220,7 @@ CREATE TABLE stats_import.part_child_1
   PARTITION OF stats_import.part_parent
   FOR VALUES FROM (0) TO (10)
   WITH (autovacuum_enabled = false);
+CREATE INDEX part_parent_i ON stats_import.part_parent(i);
 ANALYZE stats_import.part_parent;
 SELECT relpages
 FROM pg_class
@@ -195,6 +234,15 @@ WHERE oid = 'stats_import.part_parent'::regclass;
 -- positive value is still allowed
 SELECT
     pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.part_parent_i'::regclass,
+        relpages => 2::integer);
+ pg_set_relation_stats 
+-----------------------
+ 
+(1 row)
+
+SELECT
+    pg_catalog.pg_set_relation_stats(
         relation => 'stats_import.part_parent'::regclass,
         relpages => 2::integer);
  pg_set_relation_stats 
@@ -202,6 +250,48 @@ SELECT
  
 (1 row)
 
+--
+-- Partitioned indexes aren't analyzed but it is possible to set
+-- stats. The locking rules are different from normal indexes due to
+-- the rules for in-place updates: both the partitioned table and the
+-- partitioned index are locked in ShareUpdateExclusive mode.
+--
+BEGIN;
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.part_parent_i'::regclass,
+        relpages => 2::integer);
+ pg_set_relation_stats 
+-----------------------
+ 
+(1 row)
+
+SELECT mode FROM pg_locks
+WHERE relation = 'stats_import.part_parent'::regclass AND
+      pid = pg_backend_pid() AND granted;
+           mode           
+--------------------------
+ ShareUpdateExclusiveLock
+(1 row)
+
+SELECT mode FROM pg_locks
+WHERE relation = 'stats_import.part_parent_i'::regclass AND
+      pid = pg_backend_pid() AND granted;
+           mode           
+--------------------------
+ ShareUpdateExclusiveLock
+(1 row)
+
+COMMIT;
+SELECT
+    pg_catalog.pg_restore_relation_stats(
+        'relation', 'stats_import.part_parent_i'::regclass,
+        'relpages', 2::integer);
+ pg_restore_relation_stats 
+---------------------------
+ t
+(1 row)
+
 -- nothing stops us from setting it to -1
 SELECT
     pg_catalog.pg_set_relation_stats(
@@ -1414,6 +1504,19 @@ SELECT 3, 'tre', (3, 3.3, 'TRE', '2003-03-03', NULL)::stats_import.complex_type,
 UNION ALL
 SELECT 4, 'four', NULL, int4range(0,100), NULL;
 CREATE INDEX is_odd ON stats_import.test(((comp).a % 2 = 1));
+-- restoring stats on index
+SELECT * FROM pg_catalog.pg_restore_relation_stats(
+    'relation', 'stats_import.is_odd'::regclass,
+    'version', '180000'::integer,
+    'relpages', '11'::integer,
+    'reltuples', '10000'::real,
+    'relallvisible', '0'::integer
+);
+ pg_restore_relation_stats 
+---------------------------
+ t
+(1 row)
+
 -- Generate statistics on table with data
 ANALYZE stats_import.test;
 CREATE TABLE stats_import.test_clone ( LIKE stats_import.test )
diff --git a/src/test/regress/sql/stats_import.sql b/src/test/regress/sql/stats_import.sql
index d3058bf8f6b..5e24c779d80 100644
--- a/src/test/regress/sql/stats_import.sql
+++ b/src/test/regress/sql/stats_import.sql
@@ -64,6 +64,30 @@ SELECT relpages, reltuples, relallvisible
 FROM pg_class
 WHERE oid = 'stats_import.test'::regclass;
 
+CREATE INDEX test_i ON stats_import.test(id);
+
+BEGIN;
+-- regular indexes have special case locking rules
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.test_i'::regclass,
+        relpages => 18::integer);
+
+SELECT mode FROM pg_locks
+WHERE relation = 'stats_import.test'::regclass AND
+      pid = pg_backend_pid() AND granted;
+
+SELECT mode FROM pg_locks
+WHERE relation = 'stats_import.test_i'::regclass AND
+      pid = pg_backend_pid() AND granted;
+
+COMMIT;
+
+SELECT
+    pg_catalog.pg_restore_relation_stats(
+        'relation', 'stats_import.test_i'::regclass,
+        'relpages', 19::integer );
+
 -- positional arguments
 SELECT
     pg_catalog.pg_set_relation_stats(
@@ -127,6 +151,8 @@ CREATE TABLE stats_import.part_child_1
   FOR VALUES FROM (0) TO (10)
   WITH (autovacuum_enabled = false);
 
+CREATE INDEX part_parent_i ON stats_import.part_parent(i);
+
 ANALYZE stats_import.part_parent;
 
 SELECT relpages
@@ -137,9 +163,42 @@ WHERE oid = 'stats_import.part_parent'::regclass;
 -- positive value is still allowed
 SELECT
     pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.part_parent_i'::regclass,
+        relpages => 2::integer);
+
+SELECT
+    pg_catalog.pg_set_relation_stats(
         relation => 'stats_import.part_parent'::regclass,
         relpages => 2::integer);
 
+--
+-- Partitioned indexes aren't analyzed but it is possible to set
+-- stats. The locking rules are different from normal indexes due to
+-- the rules for in-place updates: both the partitioned table and the
+-- partitioned index are locked in ShareUpdateExclusive mode.
+--
+BEGIN;
+
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.part_parent_i'::regclass,
+        relpages => 2::integer);
+
+SELECT mode FROM pg_locks
+WHERE relation = 'stats_import.part_parent'::regclass AND
+      pid = pg_backend_pid() AND granted;
+
+SELECT mode FROM pg_locks
+WHERE relation = 'stats_import.part_parent_i'::regclass AND
+      pid = pg_backend_pid() AND granted;
+
+COMMIT;
+
+SELECT
+    pg_catalog.pg_restore_relation_stats(
+        'relation', 'stats_import.part_parent_i'::regclass,
+        'relpages', 2::integer);
+
 -- nothing stops us from setting it to -1
 SELECT
     pg_catalog.pg_set_relation_stats(
@@ -1062,6 +1121,15 @@ SELECT 4, 'four', NULL, int4range(0,100), NULL;
 
 CREATE INDEX is_odd ON stats_import.test(((comp).a % 2 = 1));
 
+-- restoring stats on index
+SELECT * FROM pg_catalog.pg_restore_relation_stats(
+    'relation', 'stats_import.is_odd'::regclass,
+    'version', '180000'::integer,
+    'relpages', '11'::integer,
+    'reltuples', '10000'::real,
+    'relallvisible', '0'::integer
+);
+
 -- Generate statistics on table with data
 ANALYZE stats_import.test;