Fix partitionned index + test

adding more tests suggested by Onur Tirtir, found and fixed a related bug.
citusdata · Oct 29, 2023 · 526cdaa · 526cdaa
1 parent 23235fb
commit 526cdaa
Show file tree

Hide file tree

Showing 3 changed files with 169 additions and 51 deletions.
diff --git a/src/backend/distributed/metadata/metadata_utility.c b/src/backend/distributed/metadata/metadata_utility.c
@@ -602,7 +602,8 @@ DistributedRelationSizeOnWorker(WorkerNode *workerNode, Oid relationId,
 
 	/* if the relation is an index, update relationId and define indexId */
 	Oid indexId = InvalidOid;
-	if (get_rel_relkind(relationId) == RELKIND_INDEX)
+	Oid relKind = get_rel_relkind(relationId);
+	if (relKind == RELKIND_INDEX || relKind == RELKIND_PARTITIONED_INDEX)
 	{
 		indexId = relationId;
 
@@ -1046,7 +1047,8 @@ ErrorIfNotSuitableToGetSize(Oid relationId)
 {
 	if (!IsCitusTable(relationId))
 	{
-		if (get_rel_relkind(relationId) != RELKIND_INDEX)
+		Oid relKind = get_rel_relkind(relationId);
+		if (relKind != RELKIND_INDEX && relKind != RELKIND_PARTITIONED_INDEX)
 		{
 			char *relationName = get_rel_name(relationId);
 			char *escapedQueryString = quote_literal_cstr(relationName);

diff --git a/src/test/regress/expected/multi_size_queries.out b/src/test/regress/expected/multi_size_queries.out
@@ -14,12 +14,14 @@ SELECT citus_total_relation_size(1);
 ERROR:  could not compute relation size: relation does not exist
 -- Tests with non-distributed table
 CREATE TABLE non_distributed_table (x int primary key);
+-- table
 SELECT citus_table_size('non_distributed_table');
 ERROR:  cannot calculate the size because relation 'non_distributed_table' is not distributed
 SELECT citus_relation_size('non_distributed_table');
 ERROR:  cannot calculate the size because relation 'non_distributed_table' is not distributed
 SELECT citus_total_relation_size('non_distributed_table');
 ERROR:  cannot calculate the size because relation 'non_distributed_table' is not distributed
+-- index
 SELECT citus_table_size('non_distributed_table_pkey');
 ERROR:  cannot calculate the size because table 'non_distributed_table' for index 'non_distributed_table_pkey' is not distributed
 SELECT citus_relation_size('non_distributed_table_pkey');
@@ -35,8 +37,10 @@ SELECT replicate_table_shards('lineitem_hash_part', shard_replication_factor:=2,
 
 (1 row)
 
+CREATE INDEX lineitem_hash_part_idx ON lineitem_hash_part(l_orderkey);
 -- Tests on distributed table with replication factor > 1
 VACUUM (FULL) lineitem_hash_part;
+-- table
 SELECT citus_table_size('lineitem_hash_part');
  citus_table_size
 ---------------------------------------------------------------------
@@ -52,42 +56,33 @@ SELECT citus_relation_size('lineitem_hash_part');
 SELECT citus_total_relation_size('lineitem_hash_part');
  citus_total_relation_size
 ---------------------------------------------------------------------
-                   3801088
+                   4259840
 (1 row)
 
-VACUUM (FULL) customer_copy_hash;
--- Tests on distributed tables with streaming replication.
-SELECT citus_table_size('customer_copy_hash');
+-- index
+SELECT citus_table_size('lineitem_hash_part_idx');
  citus_table_size
 ---------------------------------------------------------------------
-           548864
+           458752
 (1 row)
 
-SELECT citus_relation_size('customer_copy_hash');
+SELECT citus_relation_size('lineitem_hash_part_idx');
  citus_relation_size
 ---------------------------------------------------------------------
-              548864
+              458752
 (1 row)
 
-SELECT citus_total_relation_size('customer_copy_hash');
+SELECT citus_total_relation_size('lineitem_hash_part_idx');
  citus_total_relation_size
 ---------------------------------------------------------------------
-                   1597440
-(1 row)
-
-VACUUM (FULL) supplier;
--- Make sure we can get multiple sizes in a single query
-SELECT citus_table_size('customer_copy_hash'),
-       citus_table_size('customer_copy_hash'),
-       citus_table_size('supplier');
- citus_table_size | citus_table_size | citus_table_size
----------------------------------------------------------------------
-           548864 |           548864 |           565248
+                    458752
 (1 row)
 
-CREATE INDEX index_1 on customer_copy_hash(c_custkey);
+DROP INDEX lineitem_hash_part_idx;
+CREATE INDEX customer_copy_hash_idx on customer_copy_hash(c_custkey);
 VACUUM (FULL) customer_copy_hash;
--- Tests on distributed table with index.
+-- Tests on distributed tables with streaming replication.
+-- table
 SELECT citus_table_size('customer_copy_hash');
  citus_table_size
 ---------------------------------------------------------------------
@@ -106,27 +101,37 @@ SELECT citus_total_relation_size('customer_copy_hash');
                    2646016
 (1 row)
 
--- Tests on index on distributed table.
-SELECT citus_table_size('customer_copy_hash_pkey');
+-- index
+SELECT citus_table_size('customer_copy_hash_idx');
  citus_table_size
 ---------------------------------------------------------------------
           1048576
 (1 row)
 
-SELECT citus_relation_size('customer_copy_hash_pkey');
+SELECT citus_relation_size('customer_copy_hash_idx');
  citus_relation_size
 ---------------------------------------------------------------------
              1048576
 (1 row)
 
-SELECT citus_total_relation_size('customer_copy_hash_pkey');
+SELECT citus_total_relation_size('customer_copy_hash_idx');
  citus_total_relation_size
 ---------------------------------------------------------------------
                    1048576
 (1 row)
 
--- Tests on reference table
 VACUUM (FULL) supplier;
+-- Make sure we can get multiple sizes in a single query
+SELECT citus_table_size('customer_copy_hash'),
+       citus_table_size('customer_copy_hash'),
+       citus_table_size('supplier');
+ citus_table_size | citus_table_size | citus_table_size
+---------------------------------------------------------------------
+           548864 |           548864 |           565248
+(1 row)
+
+-- Tests on reference table
+-- table
 SELECT citus_table_size('supplier');
  citus_table_size
 ---------------------------------------------------------------------
@@ -145,8 +150,8 @@ SELECT citus_total_relation_size('supplier');
                     565248
 (1 row)
 
-CREATE INDEX index_2 on supplier(s_suppkey);
-VACUUM (FULL) supplier;
+CREATE INDEX supplier_idx on supplier(s_suppkey);
+-- table
 SELECT citus_table_size('supplier');
  citus_table_size
 ---------------------------------------------------------------------
@@ -165,6 +170,82 @@ SELECT citus_total_relation_size('supplier');
                     688128
 (1 row)
 
+-- index
+SELECT citus_table_size('supplier_idx');
+ citus_table_size
+---------------------------------------------------------------------
+           122880
+(1 row)
+
+SELECT citus_relation_size('supplier_idx');
+ citus_relation_size
+---------------------------------------------------------------------
+              122880
+(1 row)
+
+SELECT citus_total_relation_size('supplier_idx');
+ citus_total_relation_size
+---------------------------------------------------------------------
+                    122880
+(1 row)
+
+-- Test on partitioned table
+CREATE TABLE split_me (dist_col int, partition_col timestamp) PARTITION BY RANGE (partition_col);
+CREATE INDEX ON split_me(dist_col);
+-- create 2 partitions
+CREATE TABLE m PARTITION OF split_me FOR VALUES FROM ('2018-01-01') TO ('2019-01-01');
+CREATE TABLE e PARTITION OF split_me FOR VALUES FROM ('2019-01-01') TO ('2020-01-01');
+-- before citus
+SELECT citus_relation_size('split_me');
+ERROR:  cannot calculate the size because relation 'split_me' is not distributed
+SELECT citus_relation_size('split_me_dist_col_idx');
+ERROR:  cannot calculate the size because table 'split_me' for index 'split_me_dist_col_idx' is not distributed
+SELECT citus_relation_size('m');
+ERROR:  cannot calculate the size because relation 'm' is not distributed
+SELECT citus_relation_size('m_dist_col_idx');
+ERROR:  cannot calculate the size because table 'm' for index 'm_dist_col_idx' is not distributed
+-- distribute the table(s)
+SELECT create_distributed_table('split_me', 'dist_col');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- after citus
+SELECT citus_relation_size('split_me');
+ citus_relation_size
+---------------------------------------------------------------------
+                   0
+(1 row)
+
+SELECT citus_relation_size('split_me_dist_col_idx');
+ citus_relation_size
+---------------------------------------------------------------------
+                   0
+(1 row)
+
+SELECT citus_relation_size('m');
+ citus_relation_size
+---------------------------------------------------------------------
+                   0
+(1 row)
+
+SELECT citus_relation_size('m_dist_col_idx');
+ citus_relation_size
+---------------------------------------------------------------------
+               65536
+(1 row)
+
+-- And we should make sure that following always returns true:
+SELECT citus_relation_size('split_me_dist_col_idx')
+    < citus_relation_size('split_me_dist_col_idx')
+    + citus_relation_size('m_dist_col_idx')
+    + citus_relation_size('e_dist_col_idx');
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
 -- Test inside the transaction
 BEGIN;
 ALTER TABLE supplier ALTER COLUMN s_suppkey SET NOT NULL;
@@ -212,5 +293,5 @@ SELECT pg_reload_conf();
  t
 (1 row)
 
-DROP INDEX index_1;
-DROP INDEX index_2;
+DROP INDEX customer_copy_hash_idx;
+DROP INDEX supplier_idx;
diff --git a/src/test/regress/sql/multi_size_queries.sql b/src/test/regress/sql/multi_size_queries.sql
@@ -14,9 +14,11 @@ SELECT citus_total_relation_size(1);
 
 -- Tests with non-distributed table
 CREATE TABLE non_distributed_table (x int primary key);
+-- table
 SELECT citus_table_size('non_distributed_table');
 SELECT citus_relation_size('non_distributed_table');
 SELECT citus_total_relation_size('non_distributed_table');
+-- index
 SELECT citus_table_size('non_distributed_table_pkey');
 SELECT citus_relation_size('non_distributed_table_pkey');
 SELECT citus_total_relation_size('non_distributed_table_pkey');
@@ -26,19 +28,33 @@ DROP TABLE non_distributed_table;
 SET client_min_messages TO ERROR;
 SELECT replicate_table_shards('lineitem_hash_part', shard_replication_factor:=2, shard_transfer_mode:='block_writes');
 
+CREATE INDEX lineitem_hash_part_idx ON lineitem_hash_part(l_orderkey);
 -- Tests on distributed table with replication factor > 1
 VACUUM (FULL) lineitem_hash_part;
 
+-- table
 SELECT citus_table_size('lineitem_hash_part');
 SELECT citus_relation_size('lineitem_hash_part');
 SELECT citus_total_relation_size('lineitem_hash_part');
+-- index
+SELECT citus_table_size('lineitem_hash_part_idx');
+SELECT citus_relation_size('lineitem_hash_part_idx');
+SELECT citus_total_relation_size('lineitem_hash_part_idx');
 
+DROP INDEX lineitem_hash_part_idx;
+
+CREATE INDEX customer_copy_hash_idx on customer_copy_hash(c_custkey);
 VACUUM (FULL) customer_copy_hash;
 
 -- Tests on distributed tables with streaming replication.
+-- table
 SELECT citus_table_size('customer_copy_hash');
 SELECT citus_relation_size('customer_copy_hash');
 SELECT citus_total_relation_size('customer_copy_hash');
+-- index
+SELECT citus_table_size('customer_copy_hash_idx');
+SELECT citus_relation_size('customer_copy_hash_idx');
+SELECT citus_total_relation_size('customer_copy_hash_idx');
 
 VACUUM (FULL) supplier;
 
@@ -47,32 +63,51 @@ SELECT citus_table_size('customer_copy_hash'),
        citus_table_size('customer_copy_hash'),
        citus_table_size('supplier');
 
-CREATE INDEX index_1 on customer_copy_hash(c_custkey);
-VACUUM (FULL) customer_copy_hash;
-
--- Tests on distributed table with index.
-SELECT citus_table_size('customer_copy_hash');
-SELECT citus_relation_size('customer_copy_hash');
-SELECT citus_total_relation_size('customer_copy_hash');
-
--- Tests on index on distributed table.
-SELECT citus_table_size('customer_copy_hash_pkey');
-SELECT citus_relation_size('customer_copy_hash_pkey');
-SELECT citus_total_relation_size('customer_copy_hash_pkey');
-
 -- Tests on reference table
-VACUUM (FULL) supplier;
-
+-- table
 SELECT citus_table_size('supplier');
 SELECT citus_relation_size('supplier');
 SELECT citus_total_relation_size('supplier');
 
-CREATE INDEX index_2 on supplier(s_suppkey);
-VACUUM (FULL) supplier;
+CREATE INDEX supplier_idx on supplier(s_suppkey);
 
+-- table
 SELECT citus_table_size('supplier');
 SELECT citus_relation_size('supplier');
 SELECT citus_total_relation_size('supplier');
+-- index
+SELECT citus_table_size('supplier_idx');
+SELECT citus_relation_size('supplier_idx');
+SELECT citus_total_relation_size('supplier_idx');
+
+-- Test on partitioned table
+CREATE TABLE split_me (dist_col int, partition_col timestamp) PARTITION BY RANGE (partition_col);
+CREATE INDEX ON split_me(dist_col);
+
+-- create 2 partitions
+CREATE TABLE m PARTITION OF split_me FOR VALUES FROM ('2018-01-01') TO ('2019-01-01');
+CREATE TABLE e PARTITION OF split_me FOR VALUES FROM ('2019-01-01') TO ('2020-01-01');
+
+-- before citus
+SELECT citus_relation_size('split_me');
+SELECT citus_relation_size('split_me_dist_col_idx');
+SELECT citus_relation_size('m');
+SELECT citus_relation_size('m_dist_col_idx');
+
+-- distribute the table(s)
+SELECT create_distributed_table('split_me', 'dist_col');
+
+-- after citus
+SELECT citus_relation_size('split_me');
+SELECT citus_relation_size('split_me_dist_col_idx');
+SELECT citus_relation_size('m');
+SELECT citus_relation_size('m_dist_col_idx');
+
+-- And we should make sure that following always returns true:
+SELECT citus_relation_size('split_me_dist_col_idx')
+    < citus_relation_size('split_me_dist_col_idx')
+    + citus_relation_size('m_dist_col_idx')
+    + citus_relation_size('e_dist_col_idx');
 
 -- Test inside the transaction
 BEGIN;
@@ -94,5 +129,5 @@ SELECT citus_total_relation_size('customer_copy_hash');
 ALTER SYSTEM RESET citus.node_conninfo;
 SELECT pg_reload_conf();
 
-DROP INDEX index_1;
-DROP INDEX index_2;
+DROP INDEX customer_copy_hash_idx;
+DROP INDEX supplier_idx;