diff --git a/src/backend/distributed/sql/citus--9.0-2--9.1-1.sql b/src/backend/distributed/sql/citus--9.0-2--9.1-1.sql
index 7ff5d972c..6c55b16fe 100644
--- a/src/backend/distributed/sql/citus--9.0-2--9.1-1.sql
+++ b/src/backend/distributed/sql/citus--9.0-2--9.1-1.sql
@@ -4,6 +4,8 @@ COMMENT ON COLUMN pg_catalog.pg_dist_node.shouldhaveshards IS
 
 #include "udfs/master_set_node_property/9.1-1.sql"
 #include "udfs/master_drain_node/9.1-1.sql"
+#include "udfs/rebalance_table_shards/9.1-1.sql"
+#include "udfs/get_rebalance_table_shards_plan/9.1-1.sql"
 
 -- we don't maintain replication factor of reference tables anymore and just
 -- use -1 instead.
diff --git a/src/backend/distributed/sql/udfs/get_rebalance_table_shards_plan/9.1-1.sql b/src/backend/distributed/sql/udfs/get_rebalance_table_shards_plan/9.1-1.sql
new file mode 100644
index 000000000..68fb854e5
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/get_rebalance_table_shards_plan/9.1-1.sql
@@ -0,0 +1,23 @@
+-- get_rebalance_table_shards_plan shows the actual events that will be performed
+-- if a rebalance operation will be performed with the same arguments, which allows users
+-- to understand the impact of the change overall availability of the application and
+-- network trafic.
+--
+DROP FUNCTION pg_catalog.get_rebalance_table_shards_plan;
+CREATE OR REPLACE FUNCTION pg_catalog.get_rebalance_table_shards_plan(
+        relation regclass default NULL,
+        threshold float4 default 0,
+        max_shard_moves int default 1000000,
+        excluded_shard_list bigint[] default '{}',
+        drain_only boolean default false)
+    RETURNS TABLE (table_name regclass,
+                   shardid bigint,
+                   shard_size bigint,
+                   sourcename text,
+                   sourceport int,
+                   targetname text,
+                   targetport int)
+    AS 'MODULE_PATHNAME'
+    LANGUAGE C VOLATILE;
+COMMENT ON FUNCTION pg_catalog.get_rebalance_table_shards_plan(regclass, float4, int, bigint[], boolean)
+    IS 'returns the list of shard placement moves to be done on a rebalance operation';
diff --git a/src/backend/distributed/sql/udfs/get_rebalance_table_shards_plan/latest.sql b/src/backend/distributed/sql/udfs/get_rebalance_table_shards_plan/latest.sql
index 1079be2b4..68fb854e5 100644
--- a/src/backend/distributed/sql/udfs/get_rebalance_table_shards_plan/latest.sql
+++ b/src/backend/distributed/sql/udfs/get_rebalance_table_shards_plan/latest.sql
@@ -3,11 +3,13 @@
 -- to understand the impact of the change overall availability of the application and
 -- network trafic.
 --
+DROP FUNCTION pg_catalog.get_rebalance_table_shards_plan;
 CREATE OR REPLACE FUNCTION pg_catalog.get_rebalance_table_shards_plan(
-        relation regclass,
+        relation regclass default NULL,
         threshold float4 default 0,
         max_shard_moves int default 1000000,
-        excluded_shard_list bigint[] default '{}')
+        excluded_shard_list bigint[] default '{}',
+        drain_only boolean default false)
     RETURNS TABLE (table_name regclass,
                    shardid bigint,
                    shard_size bigint,
@@ -16,6 +18,6 @@ CREATE OR REPLACE FUNCTION pg_catalog.get_rebalance_table_shards_plan(
                    targetname text,
                    targetport int)
     AS 'MODULE_PATHNAME'
-    LANGUAGE C STRICT VOLATILE;
-COMMENT ON FUNCTION pg_catalog.get_rebalance_table_shards_plan(regclass, float4, int, bigint[])
+    LANGUAGE C VOLATILE;
+COMMENT ON FUNCTION pg_catalog.get_rebalance_table_shards_plan(regclass, float4, int, bigint[], boolean)
     IS 'returns the list of shard placement moves to be done on a rebalance operation';
diff --git a/src/backend/distributed/sql/udfs/master_drain_node/9.1-1.sql b/src/backend/distributed/sql/udfs/master_drain_node/9.1-1.sql
index 44ab5b03d..02b536f9a 100644
--- a/src/backend/distributed/sql/udfs/master_drain_node/9.1-1.sql
+++ b/src/backend/distributed/sql/udfs/master_drain_node/9.1-1.sql
@@ -1,14 +1,11 @@
 CREATE FUNCTION pg_catalog.master_drain_node(
     nodename text,
     nodeport integer,
-    threshold float4 default 0,
-    max_shard_moves int default 1000000,
-    excluded_shard_list bigint[] default '{}',
     shard_transfer_mode citus.shard_transfer_mode default 'auto')
   RETURNS VOID
   LANGUAGE C STRICT
   AS 'MODULE_PATHNAME', $$master_drain_node$$;
-COMMENT ON FUNCTION pg_catalog.master_drain_node(text,int,float4,int,bigint[],citus.shard_transfer_mode)
+COMMENT ON FUNCTION pg_catalog.master_drain_node(text,int,citus.shard_transfer_mode)
   IS 'mark a node to be drained of data and actually drain it as well';
 
-REVOKE ALL ON FUNCTION pg_catalog.master_drain_node(text,int,float4,int,bigint[],citus.shard_transfer_mode) FROM PUBLIC;
+REVOKE ALL ON FUNCTION pg_catalog.master_drain_node(text,int,citus.shard_transfer_mode) FROM PUBLIC;
diff --git a/src/backend/distributed/sql/udfs/master_drain_node/latest.sql b/src/backend/distributed/sql/udfs/master_drain_node/latest.sql
index 44ab5b03d..02b536f9a 100644
--- a/src/backend/distributed/sql/udfs/master_drain_node/latest.sql
+++ b/src/backend/distributed/sql/udfs/master_drain_node/latest.sql
@@ -1,14 +1,11 @@
 CREATE FUNCTION pg_catalog.master_drain_node(
     nodename text,
     nodeport integer,
-    threshold float4 default 0,
-    max_shard_moves int default 1000000,
-    excluded_shard_list bigint[] default '{}',
     shard_transfer_mode citus.shard_transfer_mode default 'auto')
   RETURNS VOID
   LANGUAGE C STRICT
   AS 'MODULE_PATHNAME', $$master_drain_node$$;
-COMMENT ON FUNCTION pg_catalog.master_drain_node(text,int,float4,int,bigint[],citus.shard_transfer_mode)
+COMMENT ON FUNCTION pg_catalog.master_drain_node(text,int,citus.shard_transfer_mode)
   IS 'mark a node to be drained of data and actually drain it as well';
 
-REVOKE ALL ON FUNCTION pg_catalog.master_drain_node(text,int,float4,int,bigint[],citus.shard_transfer_mode) FROM PUBLIC;
+REVOKE ALL ON FUNCTION pg_catalog.master_drain_node(text,int,citus.shard_transfer_mode) FROM PUBLIC;
diff --git a/src/backend/distributed/sql/udfs/rebalance_table_shards/9.1-1.sql b/src/backend/distributed/sql/udfs/rebalance_table_shards/9.1-1.sql
new file mode 100644
index 000000000..434d36ad1
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/rebalance_table_shards/9.1-1.sql
@@ -0,0 +1,16 @@
+-- rebalance_table_shards uses the shard rebalancer's C UDF functions to rebalance
+-- shards of the given relation.
+--
+DROP FUNCTION pg_catalog.rebalance_table_shards;
+CREATE OR REPLACE FUNCTION pg_catalog.rebalance_table_shards(
+        relation regclass default NULL,
+        threshold float4 default 0,
+        max_shard_moves int default 1000000,
+        excluded_shard_list bigint[] default '{}',
+        shard_transfer_mode citus.shard_transfer_mode default 'auto',
+        drain_only boolean default false)
+    RETURNS VOID
+    AS 'MODULE_PATHNAME'
+    LANGUAGE C VOLATILE;
+COMMENT ON FUNCTION pg_catalog.rebalance_table_shards(regclass, float4, int, bigint[], citus.shard_transfer_mode, boolean)
+    IS 'rebalance the shards of the given table across the worker nodes (including colocated shards of other tables)';
diff --git a/src/backend/distributed/sql/udfs/rebalance_table_shards/latest.sql b/src/backend/distributed/sql/udfs/rebalance_table_shards/latest.sql
index c714235ed..434d36ad1 100644
--- a/src/backend/distributed/sql/udfs/rebalance_table_shards/latest.sql
+++ b/src/backend/distributed/sql/udfs/rebalance_table_shards/latest.sql
@@ -1,14 +1,16 @@
 -- rebalance_table_shards uses the shard rebalancer's C UDF functions to rebalance
 -- shards of the given relation.
 --
+DROP FUNCTION pg_catalog.rebalance_table_shards;
 CREATE OR REPLACE FUNCTION pg_catalog.rebalance_table_shards(
-        relation regclass,
+        relation regclass default NULL,
         threshold float4 default 0,
         max_shard_moves int default 1000000,
         excluded_shard_list bigint[] default '{}',
-        shard_transfer_mode citus.shard_transfer_mode default 'auto')
+        shard_transfer_mode citus.shard_transfer_mode default 'auto',
+        drain_only boolean default false)
     RETURNS VOID
     AS 'MODULE_PATHNAME'
-    LANGUAGE C STRICT VOLATILE;
-COMMENT ON FUNCTION pg_catalog.rebalance_table_shards(regclass, float4, int, bigint[], citus.shard_transfer_mode)
+    LANGUAGE C VOLATILE;
+COMMENT ON FUNCTION pg_catalog.rebalance_table_shards(regclass, float4, int, bigint[], citus.shard_transfer_mode, boolean)
     IS 'rebalance the shards of the given table across the worker nodes (including colocated shards of other tables)';