Fix bug: calculate AO table size in CTAS (#105)

The extension file with (segno=0, column=1) is not traversed by
ao_foreach_extent_file(), we need to handle the size of it additionally.

Co-authored-by: hzhang2 <hzhang2@vmware.com>
Co-authored-by: Xing Guo <higuoxing+github@gmail.com>

Author: 2 people

gpcontrib/diskquota/diskquota_utility.c

@@ -1252,6 +1252,12 @@ calculate_relation_size_all_forks(RelFileNodeBackend *rnode, char relstorage)
 		RelationFileStatCtx ctx = {0};
 		ctx.relation_path = relpathbackend(rnode->node, rnode->backend, MAIN_FORKNUM);
 		ctx.size = 0;
+		/*
+		 * Since the extension file with (segno=0, column=1) is not traversed by
+		 * ao_foreach_extent_file(), we need to handle the size of it additionally.
+		 * See comments in ao_foreach_extent_file() for details.
+		 */
+		relation_file_stat(0, &ctx);
 		ao_foreach_extent_file(relation_file_stat, &ctx);
 		return ctx.size;
 	}

gpcontrib/diskquota/tests/regress/expected/test_uncommitted_table_size.out

@@ -123,7 +123,7 @@ SELECT pg_table_size('ao');
 (1 row)
 
 commit;
--- AOCS table index
+-- AO table index
 begin;
 CREATE INDEX ao_idx on ao(i);
 SELECT pg_sleep(5);
@@ -144,6 +144,30 @@ SELECT pg_table_size('ao_idx');
        2490368
 (1 row)
 
+commit;
+DROP TABLE ao;
+-- AO table CTAS
+begin;
+CREATE TABLE ao WITH(appendonly=true) AS SELECT generate_series(1, 10000);
+NOTICE:  Table doesn't have 'DISTRIBUTED BY' clause. Creating a NULL policy entry.
+SELECT pg_sleep(5);
+ pg_sleep 
+----------
+ 
+(1 row)
+
+SELECT tableid::regclass, size, segid FROM diskquota.table_size WHERE tableid = 'ao'::regclass and segid = -1;
+ tableid |  size  | segid 
+---------+--------+-------
+ ao      | 329576 |    -1
+(1 row)
+
+SELECT pg_table_size('ao');
+ pg_table_size 
+---------------
+        329576
+(1 row)
+
 commit;
 DROP TABLE ao;
 -- AOCS table
@@ -194,3 +218,28 @@ SELECT pg_table_size('aocs_idx');
 
 commit;
 DROP TABLE aocs;
+-- AOCS table CTAS
+begin;
+CREATE TABLE aocs WITH(appendonly=true, orientation=column) AS SELECT i, array(select * from generate_series(1,1000)) FROM generate_series(1, 100) AS i;
+NOTICE:  Table doesn't have 'DISTRIBUTED BY' clause -- Using column(s) named 'i' as the Greenplum Database data distribution key for this table.
+HINT:  The 'DISTRIBUTED BY' clause determines the distribution of data. Make sure column(s) chosen are the optimal data distribution key to minimize skew.
+SELECT pg_sleep(5);
+ pg_sleep 
+----------
+ 
+(1 row)
+
+SELECT tableid::regclass, size, segid FROM diskquota.table_size WHERE tableid = 'aocs'::regclass and segid = -1;
+ tableid |  size  | segid 
+---------+--------+-------
+ aocs    | 763936 |    -1
+(1 row)
+
+SELECT pg_table_size('aocs');
+ pg_table_size 
+---------------
+        763936
+(1 row)
+
+commit;
+DROP TABLE aocs;

gpcontrib/diskquota/tests/regress/sql/test_uncommitted_table_size.sql

@@ -48,7 +48,7 @@ SELECT tableid::regclass, size, segid FROM diskquota.table_size WHERE tableid =
 SELECT pg_table_size('ao');
 commit;
 
--- AOCS table index
+-- AO table index
 begin;
 CREATE INDEX ao_idx on ao(i);
 SELECT pg_sleep(5);
@@ -58,6 +58,15 @@ commit;
 
 DROP TABLE ao;
 
+-- AO table CTAS
+begin;
+CREATE TABLE ao WITH(appendonly=true) AS SELECT generate_series(1, 10000);
+SELECT pg_sleep(5);
+SELECT tableid::regclass, size, segid FROM diskquota.table_size WHERE tableid = 'ao'::regclass and segid = -1;
+SELECT pg_table_size('ao');
+commit;
+DROP TABLE ao;
+
 -- AOCS table
 begin;
 CREATE TABLE aocs (i int, t text) WITH (appendonly=true, orientation=column);
@@ -76,3 +85,12 @@ SELECT pg_table_size('aocs_idx');
 commit;
 
 DROP TABLE aocs;
+
+-- AOCS table CTAS
+begin;
+CREATE TABLE aocs WITH(appendonly=true, orientation=column) AS SELECT i, array(select * from generate_series(1,1000)) FROM generate_series(1, 100) AS i;
+SELECT pg_sleep(5);
+SELECT tableid::regclass, size, segid FROM diskquota.table_size WHERE tableid = 'aocs'::regclass and segid = -1;
+SELECT pg_table_size('aocs');
+commit;
+DROP TABLE aocs;