diff --git a/src/backend/distributed/sql/downgrades/citus--10.1-1--10.0-3.sql b/src/backend/distributed/sql/downgrades/citus--10.1-1--10.0-3.sql
index b12b27e5e..2cd802ef9 100644
--- a/src/backend/distributed/sql/downgrades/citus--10.1-1--10.0-3.sql
+++ b/src/backend/distributed/sql/downgrades/citus--10.1-1--10.0-3.sql
@@ -16,8 +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);
+DROP FUNCTION pg_catalog.worker_partitioned_relation_total_size(regclass);
+DROP FUNCTION pg_catalog.worker_partitioned_relation_size(regclass);
+DROP FUNCTION pg_catalog.worker_partitioned_table_size(regclass);
 
 #include "../udfs/citus_finish_pg_upgrade/10.0-1.sql"
diff --git a/src/backend/distributed/sql/udfs/worker_partitioned_relation_size/10.1-1.sql b/src/backend/distributed/sql/udfs/worker_partitioned_relation_size/10.1-1.sql
index 751c93f38..8eee5ec0c 100644
--- a/src/backend/distributed/sql/udfs/worker_partitioned_relation_size/10.1-1.sql
+++ b/src/backend/distributed/sql/udfs/worker_partitioned_relation_size/10.1-1.sql
@@ -1,7 +1,7 @@
-CREATE OR REPLACE FUNCTION worker_partitioned_relation_size(relation text)
+CREATE OR REPLACE FUNCTION worker_partitioned_relation_size(relation regclass)
 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)
+COMMENT ON FUNCTION pg_catalog.worker_partitioned_relation_size(regclass)
     IS 'Calculates and returns the size of a partitioned relation';
diff --git a/src/backend/distributed/sql/udfs/worker_partitioned_relation_size/latest.sql b/src/backend/distributed/sql/udfs/worker_partitioned_relation_size/latest.sql
index 751c93f38..8eee5ec0c 100644
--- a/src/backend/distributed/sql/udfs/worker_partitioned_relation_size/latest.sql
+++ b/src/backend/distributed/sql/udfs/worker_partitioned_relation_size/latest.sql
@@ -1,7 +1,7 @@
-CREATE OR REPLACE FUNCTION worker_partitioned_relation_size(relation text)
+CREATE OR REPLACE FUNCTION worker_partitioned_relation_size(relation regclass)
 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)
+COMMENT ON FUNCTION pg_catalog.worker_partitioned_relation_size(regclass)
     IS 'Calculates and returns the size of a partitioned relation';
diff --git a/src/backend/distributed/sql/udfs/worker_partitioned_relation_total_size/10.1-1.sql b/src/backend/distributed/sql/udfs/worker_partitioned_relation_total_size/10.1-1.sql
index 4bde3f9e6..e0859c216 100644
--- a/src/backend/distributed/sql/udfs/worker_partitioned_relation_total_size/10.1-1.sql
+++ b/src/backend/distributed/sql/udfs/worker_partitioned_relation_total_size/10.1-1.sql
@@ -1,7 +1,7 @@
-CREATE OR REPLACE FUNCTION worker_partitioned_relation_total_size(relation text)
+CREATE OR REPLACE FUNCTION worker_partitioned_relation_total_size(relation regclass)
 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)
+COMMENT ON FUNCTION worker_partitioned_relation_total_size(regclass)
 	IS 'Calculates and returns the total size of a partitioned relation';
diff --git a/src/backend/distributed/sql/udfs/worker_partitioned_relation_total_size/latest.sql b/src/backend/distributed/sql/udfs/worker_partitioned_relation_total_size/latest.sql
index 4bde3f9e6..e0859c216 100644
--- a/src/backend/distributed/sql/udfs/worker_partitioned_relation_total_size/latest.sql
+++ b/src/backend/distributed/sql/udfs/worker_partitioned_relation_total_size/latest.sql
@@ -1,7 +1,7 @@
-CREATE OR REPLACE FUNCTION worker_partitioned_relation_total_size(relation text)
+CREATE OR REPLACE FUNCTION worker_partitioned_relation_total_size(relation regclass)
 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)
+COMMENT ON FUNCTION worker_partitioned_relation_total_size(regclass)
 	IS 'Calculates and returns the total size of a partitioned relation';
diff --git a/src/backend/distributed/sql/udfs/worker_partitioned_table_size/10.1-1.sql b/src/backend/distributed/sql/udfs/worker_partitioned_table_size/10.1-1.sql
index d2be892a2..0907fffad 100644
--- a/src/backend/distributed/sql/udfs/worker_partitioned_table_size/10.1-1.sql
+++ b/src/backend/distributed/sql/udfs/worker_partitioned_table_size/10.1-1.sql
@@ -1,7 +1,7 @@
-CREATE OR REPLACE FUNCTION worker_partitioned_table_size(relation text)
+CREATE OR REPLACE FUNCTION worker_partitioned_table_size(relation regclass)
 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)
+COMMENT ON FUNCTION pg_catalog.worker_partitioned_table_size(regclass)
     IS 'Calculates and returns the size of a partitioned table';
diff --git a/src/backend/distributed/sql/udfs/worker_partitioned_table_size/latest.sql b/src/backend/distributed/sql/udfs/worker_partitioned_table_size/latest.sql
index d2be892a2..0907fffad 100644
--- a/src/backend/distributed/sql/udfs/worker_partitioned_table_size/latest.sql
+++ b/src/backend/distributed/sql/udfs/worker_partitioned_table_size/latest.sql
@@ -1,7 +1,7 @@
-CREATE OR REPLACE FUNCTION worker_partitioned_table_size(relation text)
+CREATE OR REPLACE FUNCTION worker_partitioned_table_size(relation regclass)
 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)
+COMMENT ON FUNCTION pg_catalog.worker_partitioned_table_size(regclass)
     IS 'Calculates and returns the size of a partitioned table';
diff --git a/src/test/regress/expected/multi_extension.out b/src/test/regress/expected/multi_extension.out
index 5cbdf0113..29ce306f2 100644
--- a/src/test/regress/expected/multi_extension.out
+++ b/src/test/regress/expected/multi_extension.out
@@ -565,9 +565,9 @@ 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)
-                                                                               | function worker_partitioned_relation_size(text)
-                                                                               | function worker_partitioned_relation_total_size(text)
-                                                                               | function worker_partitioned_table_size(text)
+                                                                               | function worker_partitioned_relation_size(regclass)
+                                                                               | function worker_partitioned_relation_total_size(regclass)
+                                                                               | function worker_partitioned_table_size(regclass)
 (6 rows)
 
 DROP TABLE prev_objects, extension_diff;
diff --git a/src/test/regress/expected/multi_partitioning.out b/src/test/regress/expected/multi_partitioning.out
index 7d99c754b..ce7386a6a 100644
--- a/src/test/regress/expected/multi_partitioning.out
+++ b/src/test/regress/expected/multi_partitioning.out
@@ -2088,7 +2088,7 @@ SELECT create_distributed_table('"events.Energy Added"', 'user_id', colocate_wit
 
 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');
 -- test shard cost by disk size function
-SELECT citus_shard_cost_by_disk_size(1660207);
+SELECT citus_shard_cost_by_disk_size(shardid) FROM pg_dist_shard WHERE logicalrelid = '"events.Energy Added"'::regclass ORDER BY shardid LIMIT 1;
  citus_shard_cost_by_disk_size
 ---------------------------------------------------------------------
                          16384
@@ -2097,21 +2097,21 @@ SELECT citus_shard_cost_by_disk_size(1660207);
 CREATE INDEX idx_btree_hobbies ON "events.Energy Added" USING BTREE ((data->>'location'));
  \c - - - :worker_1_port
 -- should not be zero because of TOAST, vm, fms
-SELECT worker_partitioned_table_size('"events.Energy Added_1660207"');
+SELECT worker_partitioned_table_size(oid) FROM pg_class WHERE relname LIKE '%events.Energy Added%' ORDER BY relname LIMIT 1;
  worker_partitioned_table_size
 ---------------------------------------------------------------------
                           8192
 (1 row)
 
 -- should be zero since no data
-SELECT worker_partitioned_relation_size('"events.Energy Added_1660207"');
+SELECT worker_partitioned_relation_size(oid) FROM pg_class WHERE relname LIKE '%events.Energy Added%' ORDER BY relname LIMIT 1;
  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"');
+SELECT worker_partitioned_relation_total_size(oid) FROM pg_class WHERE relname LIKE '%events.Energy Added%' ORDER BY relname LIMIT 1;
  worker_partitioned_relation_total_size
 ---------------------------------------------------------------------
                                   24576
diff --git a/src/test/regress/expected/upgrade_list_citus_objects.out b/src/test/regress/expected/upgrade_list_citus_objects.out
index 7bfb5a90b..786ffad09 100644
--- a/src/test/regress/expected/upgrade_list_citus_objects.out
+++ b/src/test/regress/expected/upgrade_list_citus_objects.out
@@ -199,9 +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_partitioned_relation_size(regclass)
+ function worker_partitioned_relation_total_size(regclass)
+ function worker_partitioned_table_size(regclass)
  function worker_range_partition_table(bigint,integer,text,text,oid,anyarray)
  function worker_record_sequence_dependency(regclass,regclass,name)
  function worker_repartition_cleanup(bigint)
diff --git a/src/test/regress/sql/multi_partitioning.sql b/src/test/regress/sql/multi_partitioning.sql
index 98aa5b962..f086a1393 100644
--- a/src/test/regress/sql/multi_partitioning.sql
+++ b/src/test/regress/sql/multi_partitioning.sql
@@ -1230,15 +1230,15 @@ SELECT create_distributed_table('"events.Energy Added"', 'user_id', colocate_wit
 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');
 
 -- test shard cost by disk size function
-SELECT citus_shard_cost_by_disk_size(1660207);
+SELECT citus_shard_cost_by_disk_size(shardid) FROM pg_dist_shard WHERE logicalrelid = '"events.Energy Added"'::regclass ORDER BY shardid LIMIT 1;
 CREATE INDEX idx_btree_hobbies ON "events.Energy Added" USING BTREE ((data->>'location'));
  \c - - - :worker_1_port
 -- should not be zero because of TOAST, vm, fms
-SELECT worker_partitioned_table_size('"events.Energy Added_1660207"');
+SELECT worker_partitioned_table_size(oid) FROM pg_class WHERE relname LIKE '%events.Energy Added%' ORDER BY relname LIMIT 1;
 -- should be zero since no data
-SELECT worker_partitioned_relation_size('"events.Energy Added_1660207"');
+SELECT worker_partitioned_relation_size(oid) FROM pg_class WHERE relname LIKE '%events.Energy Added%' ORDER BY relname LIMIT 1;
 -- should not be zero because of indexes + pg_table_size()
-SELECT worker_partitioned_relation_total_size('"events.Energy Added_1660207"');
+SELECT worker_partitioned_relation_total_size(oid) FROM pg_class WHERE relname LIKE '%events.Energy Added%' ORDER BY relname LIMIT 1;
  \c - - - :master_port
 DROP TABLE "events.Energy Added";