Introduce 3 partitioned size udfs (#4899)

* Introduce 3 partitioned size udfs

* Add tests for new partition size udfs

* Fix type incompatibilities

* Convert UDFs into pure sql functions

* Fix function comment

src/backend/distributed/metadata/metadata_utility.c

@@ -596,7 +596,8 @@ ShardIntervalsOnWorkerGroup(WorkerNode *workerNode, Oid relationId)
  * size of multiple tables. Note that, different size functions supported by PG
  * are also supported by this function changing the size query type given as the
  * last parameter to function. Depending on the sizeQueryType enum parameter, the
- * generated query will be pg_relation_size or pg_total_relation_size.
+ * generated query will call one of the functions: pg_relation_size,
+ * pg_total_relation_size, pg_table_size and cstore_table_size.
  */
 StringInfo
 GenerateSizeQueryOnMultiplePlacements(List *shardIntervalList,

src/backend/distributed/sql/citus--10.0-3--10.1-1.sql

@@ -2,4 +2,7 @@
 
 #include "../../columnar/sql/columnar--10.0-3--10.1-1.sql"
 #include "udfs/create_distributed_table/10.1-1.sql";
+#include "udfs/worker_partitioned_relation_total_size/10.1-1.sql"
+#include "udfs/worker_partitioned_relation_size/10.1-1.sql"
+#include "udfs/worker_partitioned_table_size/10.1-1.sql"
 #include "udfs/citus_finish_pg_upgrade/10.1-1.sql"

src/backend/distributed/sql/downgrades/citus--10.1-1--10.0-3.sql

@@ -1,4 +1,4 @@
--- citus--10.1-1--10.0-2
+-- citus--10.1-1--10.0-3
 
 #include "../../../columnar/sql/downgrades/columnar--10.1-1--10.0-3.sql"
 
@@ -16,4 +16,8 @@ COMMENT ON FUNCTION create_distributed_table(table_name regclass,
 											 colocate_with text)
     IS 'creates a distributed table';
 
+DROP FUNCTION pg_catalog.worker_partitioned_relation_total_size(text);
+DROP FUNCTION pg_catalog.worker_partitioned_relation_size(text);
+DROP FUNCTION pg_catalog.worker_partitioned_table_size(text);
+
 #include "../udfs/citus_finish_pg_upgrade/10.0-1.sql"

src/backend/distributed/sql/udfs/worker_partitioned_relation_size/10.1-1.sql

@@ -0,0 +1,7 @@
+CREATE OR REPLACE FUNCTION worker_partitioned_relation_size(relation text)
+RETURNS bigint AS $$
+    SELECT sum(pg_relation_size(relid))::bigint
+					  FROM (SELECT relid from pg_partition_tree(relation)) partition_tree;
+$$ LANGUAGE SQL;
+COMMENT ON FUNCTION pg_catalog.worker_partitioned_relation_size(text)
+    IS 'Calculates and returns the size of a partitioned relation';

src/backend/distributed/sql/udfs/worker_partitioned_relation_size/latest.sql

@@ -0,0 +1,7 @@
+CREATE OR REPLACE FUNCTION worker_partitioned_relation_size(relation text)
+RETURNS bigint AS $$
+    SELECT sum(pg_relation_size(relid))::bigint
+					  FROM (SELECT relid from pg_partition_tree(relation)) partition_tree;
+$$ LANGUAGE SQL;
+COMMENT ON FUNCTION pg_catalog.worker_partitioned_relation_size(text)
+    IS 'Calculates and returns the size of a partitioned relation';

src/backend/distributed/sql/udfs/worker_partitioned_relation_total_size/10.1-1.sql

@@ -0,0 +1,7 @@
+CREATE OR REPLACE FUNCTION worker_partitioned_relation_total_size(relation text)
+RETURNS bigint AS $$
+    SELECT sum(pg_total_relation_size(relid))::bigint
+					  FROM (SELECT relid from pg_partition_tree(relation)) partition_tree;
+$$ LANGUAGE SQL;
+COMMENT ON FUNCTION worker_partitioned_relation_total_size(text)
+	IS 'Calculates and returns the total size of a partitioned relation';

src/backend/distributed/sql/udfs/worker_partitioned_relation_total_size/latest.sql

@@ -0,0 +1,7 @@
+CREATE OR REPLACE FUNCTION worker_partitioned_relation_total_size(relation text)
+RETURNS bigint AS $$
+    SELECT sum(pg_total_relation_size(relid))::bigint
+					  FROM (SELECT relid from pg_partition_tree(relation)) partition_tree;
+$$ LANGUAGE SQL;
+COMMENT ON FUNCTION worker_partitioned_relation_total_size(text)
+	IS 'Calculates and returns the total size of a partitioned relation';

src/backend/distributed/sql/udfs/worker_partitioned_table_size/10.1-1.sql

@@ -0,0 +1,7 @@
+CREATE OR REPLACE FUNCTION worker_partitioned_table_size(relation text)
+RETURNS bigint AS $$
+    SELECT sum(pg_table_size(relid))::bigint
+					  FROM (SELECT relid from pg_partition_tree(relation)) partition_tree;
+$$ LANGUAGE SQL;
+COMMENT ON FUNCTION pg_catalog.worker_partitioned_table_size(text)
+    IS 'Calculates and returns the size of a partitioned table';

src/backend/distributed/sql/udfs/worker_partitioned_table_size/latest.sql

@@ -0,0 +1,7 @@
+CREATE OR REPLACE FUNCTION worker_partitioned_table_size(relation text)
+RETURNS bigint AS $$
+    SELECT sum(pg_table_size(relid))::bigint
+					  FROM (SELECT relid from pg_partition_tree(relation)) partition_tree;
+$$ LANGUAGE SQL;
+COMMENT ON FUNCTION pg_catalog.worker_partitioned_table_size(text)
+    IS 'Calculates and returns the size of a partitioned table';

src/test/regress/expected/multi_extension.out

@@ -565,7 +565,10 @@ SELECT * FROM print_extension_changes();
  function citus_internal.columnar_ensure_objects_exist()                       |
  function create_distributed_table(regclass,text,citus.distribution_type,text) |
                                                                                | function create_distributed_table(regclass,text,citus.distribution_type,text,integer)
-(3 rows)
+                                                                               | function worker_partitioned_relation_size(text)
+                                                                               | function worker_partitioned_relation_total_size(text)
+                                                                               | function worker_partitioned_table_size(text)
+(6 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version

src/test/regress/expected/multi_partitioning.out

@@ -2078,6 +2078,40 @@ SELECT create_distributed_table('test_inheritance','a');
 CREATE TABLE local_inheritance (k int) INHERITS (test_inheritance);
 ERROR:  non-distributed tables cannot inherit distributed tables
 DROP TABLE test_inheritance;
+-- test worker partitioned table size functions
+CREATE TABLE "events.Energy Added" (user_id int, time timestamp with time zone, data jsonb, PRIMARY KEY (user_id, time )) PARTITION BY RANGE ("time");
+ CREATE INDEX idx_btree_hobbies ON "events.Energy Added" USING BTREE ((data->>'location'));
+ SELECT create_distributed_table('"events.Energy Added"', 'user_id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE "Energy Added_17634"  PARTITION OF "events.Energy Added" FOR VALUES  FROM ('2018-04-13 00:00:00+00') TO ('2018-04-14 00:00:00+00');
+ \c - - - :worker_1_port
+-- should not be zero because of TOAST, vm, fms
+SELECT worker_partitioned_table_size('"events.Energy Added_1660207"');
+ worker_partitioned_table_size
+---------------------------------------------------------------------
+                          8192
+(1 row)
+
+-- should be zero since no data
+SELECT worker_partitioned_relation_size('"events.Energy Added_1660207"');
+ worker_partitioned_relation_size
+---------------------------------------------------------------------
+                                0
+(1 row)
+
+-- should not be zero because of indexes + pg_table_size()
+SELECT worker_partitioned_relation_total_size('"events.Energy Added_1660207"');
+ worker_partitioned_relation_total_size
+---------------------------------------------------------------------
+                                  24576
+(1 row)
+
+ \c - - - :master_port
+DROP TABLE "events.Energy Added";
 DROP SCHEMA partitioning_schema CASCADE;
 NOTICE:  drop cascades to table partitioning_schema."schema-test"
 DROP TABLE IF EXISTS

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -199,6 +199,9 @@ ORDER BY 1;
  function worker_partial_agg_ffunc(internal)
  function worker_partial_agg_sfunc(internal,oid,anyelement)
  function worker_partition_query_result(text,text,integer,citus.distribution_type,text[],text[],boolean)
+ function worker_partitioned_relation_size(text)
+ function worker_partitioned_relation_total_size(text)
+ function worker_partitioned_table_size(text)
  function worker_range_partition_table(bigint,integer,text,text,oid,anyarray)
  function worker_record_sequence_dependency(regclass,regclass,name)
  function worker_repartition_cleanup(bigint)
@@ -242,5 +245,5 @@ ORDER BY 1;
  view citus_worker_stat_activity
  view pg_dist_shard_placement
  view time_partitions
-(226 rows)
+(229 rows)
 

src/test/regress/sql/multi_partitioning.sql

@@ -1225,6 +1225,21 @@ SELECT create_distributed_table('test_inheritance','a');
 CREATE TABLE local_inheritance (k int) INHERITS (test_inheritance);
 DROP TABLE test_inheritance;
 
+-- test worker partitioned table size functions
+CREATE TABLE "events.Energy Added" (user_id int, time timestamp with time zone, data jsonb, PRIMARY KEY (user_id, time )) PARTITION BY RANGE ("time");
+ CREATE INDEX idx_btree_hobbies ON "events.Energy Added" USING BTREE ((data->>'location'));
+ SELECT create_distributed_table('"events.Energy Added"', 'user_id');
+CREATE TABLE "Energy Added_17634"  PARTITION OF "events.Energy Added" FOR VALUES  FROM ('2018-04-13 00:00:00+00') TO ('2018-04-14 00:00:00+00');
+ \c - - - :worker_1_port
+-- should not be zero because of TOAST, vm, fms
+SELECT worker_partitioned_table_size('"events.Energy Added_1660207"');
+-- should be zero since no data
+SELECT worker_partitioned_relation_size('"events.Energy Added_1660207"');
+-- should not be zero because of indexes + pg_table_size()
+SELECT worker_partitioned_relation_total_size('"events.Energy Added_1660207"');
+ \c - - - :master_port
+DROP TABLE "events.Energy Added";
+
 DROP SCHEMA partitioning_schema CASCADE;
 DROP TABLE IF EXISTS
 	partitioning_hash_test,