diff --git a/src/backend/distributed/utils/colocation_utils.c b/src/backend/distributed/utils/colocation_utils.c
index b1038fd5f..d1307c551 100644
--- a/src/backend/distributed/utils/colocation_utils.c
+++ b/src/backend/distributed/utils/colocation_utils.c
@@ -113,17 +113,17 @@ update_distributed_table_colocation(PG_FUNCTION_ARGS)
 	CheckCitusVersion(ERROR);
 	EnsureCoordinator();
 	EnsureTableOwner(targetRelationId);
-	EnsureHashDistributedTable(targetRelationId);
 
 	char *colocateWithTableName = text_to_cstring(colocateWithTableNameText);
 	if (IsColocateWithNone(colocateWithTableName))
 	{
+		EnsureHashDistributedTable(targetRelationId);
 		BreakColocation(targetRelationId);
 	}
 	else
 	{
 		Oid colocateWithTableId = ResolveRelationId(colocateWithTableNameText, false);
-		EnsureHashDistributedTable(colocateWithTableId);
+		EnsureTableOwner(colocateWithTableId);
 		MarkTablesColocated(colocateWithTableId, targetRelationId);
 	}
 	PG_RETURN_VOID();
@@ -248,6 +248,8 @@ MarkTablesColocated(Oid sourceRelationId, Oid targetRelationId)
 							   "other tables")));
 	}
 
+	EnsureHashDistributedTable(sourceRelationId);
+	EnsureHashDistributedTable(targetRelationId);
 	CheckReplicationModel(sourceRelationId, targetRelationId);
 	CheckDistributionColumnType(sourceRelationId, targetRelationId);
 
diff --git a/src/test/regress/expected/citus_local_tables.out b/src/test/regress/expected/citus_local_tables.out
index 65dc6fe1d..45beceeef 100644
--- a/src/test/regress/expected/citus_local_tables.out
+++ b/src/test/regress/expected/citus_local_tables.out
@@ -295,17 +295,17 @@ SELECT create_reference_table('reference_table');
 
 -- show that colociation of citus local tables are not supported for now
 -- between citus local tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['citus_local_table_2']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'citus_local_table_2');
 ERROR:  citus local tables cannot be colocated with other tables
 -- between citus local tables and reference tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['reference_table']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'reference_table');
 ERROR:  citus local tables cannot be colocated with other tables
-SELECT mark_tables_colocated('reference_table', ARRAY['citus_local_table_1']);
+SELECT update_distributed_table_colocation('reference_table', colocate_with => 'citus_local_table_1');
 ERROR:  citus local tables cannot be colocated with other tables
 -- between citus local tables and distributed tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['distributed_table']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'distributed_table');
 ERROR:  citus local tables cannot be colocated with other tables
-SELECT mark_tables_colocated('distributed_table', ARRAY['citus_local_table_1']);
+SELECT update_distributed_table_colocation('distributed_table', colocate_with => 'citus_local_table_1');
 ERROR:  citus local tables cannot be colocated with other tables
 -- master_create_empty_shard is not supported
 SELECT master_create_empty_shard('citus_local_table_1');
diff --git a/src/test/regress/expected/multi_colocation_utils.out b/src/test/regress/expected/multi_colocation_utils.out
index cea7003e9..dfc26c7a2 100644
--- a/src/test/regress/expected/multi_colocation_utils.out
+++ b/src/test/regress/expected/multi_colocation_utils.out
@@ -833,7 +833,7 @@ ORDER BY
  table2_groupf                   | 1300081 | t            |    57638 |               |
 (108 rows)
 
--- reset colocation ids to test mark_tables_colocated
+-- reset colocation ids to test update_distributed_table_colocation
 ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 1;
 DELETE FROM pg_dist_colocation
     WHERE colocationid >= 1 AND colocationid < 1000;
@@ -855,21 +855,20 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
 (0 rows)
 
 -- first check failing cases
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupC']);
-ERROR:  cannot colocate tables table1_groupb and table1_groupc
-DETAIL:  Distribution column types don't match for table1_groupb and table1_groupc.
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupD']);
-ERROR:  cannot colocate tables table1_groupb and table1_groupd
-DETAIL:  Shard counts don't match for table1_groupb and table1_groupd.
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupE']);
-ERROR:  cannot colocate tables table1_groupb and table1_groupe
-DETAIL:  Shard 1300026 of table1_groupb and shard xxxxx of table1_groupe have different number of shard placements.
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupF']);
-ERROR:  cannot colocate tables table1_groupb and table1_groupf
-DETAIL:  Replication models don't match for table1_groupb and table1_groupf.
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table2_groupB', 'table1_groupD']);
-ERROR:  cannot colocate tables table1_groupb and table1_groupd
-DETAIL:  Shard counts don't match for table1_groupb and table1_groupd.
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupC');
+ERROR:  cannot colocate tables table1_groupc and table1_groupb
+DETAIL:  Distribution column types don't match for table1_groupc and table1_groupb.
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupD');
+ERROR:  cannot colocate tables table1_groupd and table1_groupb
+DETAIL:  Shard counts don't match for table1_groupd and table1_groupb.
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupE');
+ERROR:  cannot colocate tables table1_groupe and table1_groupb
+DETAIL:  Shard 1300058 of table1_groupe and shard xxxxx of table1_groupb have different number of shard placements.
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupF');
+ERROR:  relation table1_groupf should be a hash distributed table
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupD');
+ERROR:  cannot colocate tables table1_groupd and table1_groupb
+DETAIL:  Shard counts don't match for table1_groupd and table1_groupb.
 -- check metadata to see failing calls didn't have any side effects
 SELECT * FROM pg_dist_colocation
     WHERE colocationid >= 1 AND colocationid < 1000
@@ -886,39 +885,39 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
 (0 rows)
 
 -- check successfully cololated tables
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table2_groupB']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table2_groupB');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
-SELECT mark_tables_colocated('table1_groupC', ARRAY['table2_groupC']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupC', colocate_with => 'table2_groupC');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
-SELECT mark_tables_colocated('table1_groupD', ARRAY['table2_groupD']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupD', colocate_with => 'table2_groupD');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
-SELECT mark_tables_colocated('table1_groupE', ARRAY['table2_groupE', 'table3_groupE']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupE', colocate_with => 'table2_groupE');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
-SELECT mark_tables_colocated('table1_groupF', ARRAY['table2_groupF']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupE', colocate_with => 'table3_groupE');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
 -- check to colocate with itself
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupB']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupB');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
@@ -944,40 +943,53 @@ SELECT * FROM pg_dist_colocation
     ORDER BY colocationid;
  colocationid | shardcount | replicationfactor | distributioncolumntype | distributioncolumncollation
 ---------------------------------------------------------------------
-            2 |          2 |                 1 |                     23 |                           0
-            3 |          2 |                 2 |                     25 |                         100
-            4 |          8 |                 2 |                     23 |                           0
+            1 |          2 |                 1 |                     23 |                           0
+            2 |          2 |                 2 |                     25 |                         100
+            3 |          8 |                 2 |                     23 |                           0
+            4 |          2 |                 2 |                     23 |                           0
             5 |          2 |                 2 |                     23 |                           0
-(4 rows)
+(5 rows)
 
 SELECT logicalrelid, colocationid FROM pg_dist_partition
     WHERE colocationid >= 1 AND colocationid < 1000
     ORDER BY colocationid, logicalrelid;
    logicalrelid    | colocationid
 ---------------------------------------------------------------------
- table1_groupb     |            2
- table2_groupb     |            2
- table1_groupc     |            3
- table2_groupc     |            3
- table1_groupd     |            4
- table2_groupd     |            4
+ table1_groupb     |            1
+ table2_groupb     |            1
+ table1_groupc     |            2
+ table2_groupc     |            2
+ table1_groupd     |            3
+ table2_groupd     |            3
+ table2_groupe     |            4
  table1_groupe     |            5
- table2_groupe     |            5
  table3_groupe     |            5
  table1_group_none |            6
  table2_group_none |            7
 (11 rows)
 
 -- move the all tables in colocation group 5 to colocation group 7
-SELECT mark_tables_colocated('table1_group_none', ARRAY['table1_groupE', 'table2_groupE', 'table3_groupE']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table1_groupE');
+ update_distributed_table_colocation
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table2_groupE');
+ update_distributed_table_colocation
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table3_groupE');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
 -- move a table with a colocation id which is already not in pg_dist_colocation
-SELECT mark_tables_colocated('table1_group_none', ARRAY['table2_group_none']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table2_group_none');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
@@ -988,27 +1000,29 @@ SELECT * FROM pg_dist_colocation
     ORDER BY colocationid;
  colocationid | shardcount | replicationfactor | distributioncolumntype | distributioncolumncollation
 ---------------------------------------------------------------------
-            2 |          2 |                 1 |                     23 |                           0
-            3 |          2 |                 2 |                     25 |                         100
-            4 |          8 |                 2 |                     23 |                           0
-(3 rows)
+            1 |          2 |                 1 |                     23 |                           0
+            2 |          2 |                 2 |                     25 |                         100
+            3 |          8 |                 2 |                     23 |                           0
+            4 |          2 |                 2 |                     23 |                           0
+            5 |          2 |                 2 |                     23 |                           0
+(5 rows)
 
 SELECT logicalrelid, colocationid FROM pg_dist_partition
     WHERE colocationid >= 1 AND colocationid < 1000
     ORDER BY colocationid, logicalrelid;
    logicalrelid    | colocationid
 ---------------------------------------------------------------------
- table1_groupb     |            2
- table2_groupb     |            2
- table1_groupc     |            3
- table2_groupc     |            3
- table1_groupd     |            4
- table2_groupd     |            4
- table1_groupe     |            6
- table2_groupe     |            6
- table3_groupe     |            6
- table1_group_none |            6
- table2_group_none |            6
+ table1_groupb     |            1
+ table2_groupb     |            1
+ table1_groupc     |            2
+ table2_groupc     |            2
+ table1_groupd     |            3
+ table2_groupd     |            3
+ table2_groupe     |            4
+ table1_groupe     |            5
+ table3_groupe     |            5
+ table1_group_none |            7
+ table2_group_none |            7
 (11 rows)
 
 -- try to colocate different replication models
@@ -1033,9 +1047,9 @@ SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'NONE');
 
 (1 row)
 
-SELECT mark_tables_colocated('table1_groupG', ARRAY['table2_groupG']);
-ERROR:  cannot colocate tables table1_groupg and table2_groupg
-DETAIL:  Replication models don't match for table1_groupg and table2_groupg.
+SELECT update_distributed_table_colocation('table1_groupG', colocate_with => 'table2_groupG');
+ERROR:  cannot colocate tables table2_groupg and table1_groupg
+DETAIL:  Replication models don't match for table2_groupg and table1_groupg.
 CREATE TABLE d1(a int, b int);
 CREATE TABLE d2(a int, b int);
 CREATE TABLE d3(a int, b int);
@@ -1101,31 +1115,6 @@ SELECT create_reference_table('ref');
 
 (1 row)
 
--- check d1, d2, d3 and d4 has the same colocation id => they are colocated.
-SELECT get_table_colocation_id('d1');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d2');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d3');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d4');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
 SELECT tables_colocated('d1', 'd2');
  tables_colocated
 ---------------------------------------------------------------------
@@ -1169,31 +1158,6 @@ SELECT update_distributed_table_colocation('d2', colocate_with => 'none');
 
 (1 row)
 
--- d1 and d3 and d4 should be colocated, d2 should have a new colocation id.
-SELECT get_table_colocation_id('d1');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d2');
- get_table_colocation_id
----------------------------------------------------------------------
-                      11
-(1 row)
-
-SELECT get_table_colocation_id('d3');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d4');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
 SELECT tables_colocated('d1', 'd2');
  tables_colocated
 ---------------------------------------------------------------------
@@ -1226,31 +1190,6 @@ SELECT update_distributed_table_colocation('d2', colocate_with => 'none');
 
 (1 row)
 
--- d1 and d3 and d4 should be colocated, d2 should have a new colocation id.
-SELECT get_table_colocation_id('d1');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d2');
- get_table_colocation_id
----------------------------------------------------------------------
-                      12
-(1 row)
-
-SELECT get_table_colocation_id('d3');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d4');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
 SELECT tables_colocated('d1', 'd2');
  tables_colocated
 ---------------------------------------------------------------------
@@ -1281,31 +1220,6 @@ SELECT update_distributed_table_colocation('d3', colocate_with => 'd2');
 
 (1 row)
 
--- d1 and d4 should be colocated, d2 and d3 should be colocated.
-SELECT get_table_colocation_id('d1');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d2');
- get_table_colocation_id
----------------------------------------------------------------------
-                      12
-(1 row)
-
-SELECT get_table_colocation_id('d3');
- get_table_colocation_id
----------------------------------------------------------------------
-                      12
-(1 row)
-
-SELECT get_table_colocation_id('d4');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
 SELECT tables_colocated('d1', 'd2');
  tables_colocated
 ---------------------------------------------------------------------
@@ -1337,37 +1251,6 @@ SELECT update_distributed_table_colocation('d3', colocate_with => '"none"');
 
 (1 row)
 
--- d1, d4, d3 and "none" should be colocated;
-SELECT get_table_colocation_id('d1');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d2');
- get_table_colocation_id
----------------------------------------------------------------------
-                      12
-(1 row)
-
-SELECT get_table_colocation_id('d3');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d4');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('none');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
 SELECT tables_colocated('d1', 'd2');
  tables_colocated
 ---------------------------------------------------------------------
diff --git a/src/test/regress/expected/multi_drop_extension.out b/src/test/regress/expected/multi_drop_extension.out
index cd3098ee8..45ae5e9eb 100644
--- a/src/test/regress/expected/multi_drop_extension.out
+++ b/src/test/regress/expected/multi_drop_extension.out
@@ -16,7 +16,6 @@ SELECT create_distributed_table('testtableddl', 'distributecol', 'append');
 -- around.
 SET client_min_messages TO 'WARNING';
 DROP FUNCTION pg_catalog.master_create_worker_shards(text, integer, integer);
-DROP FUNCTION pg_catalog.mark_tables_colocated(regclass, regclass[]);
 DROP EXTENSION citus CASCADE;
 RESET client_min_messages;
 BEGIN;
diff --git a/src/test/regress/expected/multi_metadata_sync.out b/src/test/regress/expected/multi_metadata_sync.out
index c457b5a6b..8d1452be9 100644
--- a/src/test/regress/expected/multi_metadata_sync.out
+++ b/src/test/regress/expected/multi_metadata_sync.out
@@ -808,13 +808,8 @@ SELECT "Constraint", "Definition" FROM table_fkeys WHERE relid='mx_test_schema_1
  mx_fk_constraint_2 | FOREIGN KEY (col1) REFERENCES mx_test_schema_2.mx_table_2(col1)
 (1 row)
 
--- Check that mark_tables_colocated call propagates the changes to the workers
+-- Check that update_distributed_table_colocation call propagates the changes to the workers
 \c - - - :master_port
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
 SELECT nextval('pg_catalog.pg_dist_colocationid_seq') AS last_colocation_id \gset
 ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 10000;
 SET citus.shard_count TO 7;
@@ -865,9 +860,9 @@ SET
 WHERE
 	logicalrelid = 'mx_colocation_test_1'::regclass
 	OR logicalrelid = 'mx_colocation_test_2'::regclass;
--- Mark tables colocated and see the changes on the master and the worker
-SELECT mark_tables_colocated('mx_colocation_test_1', ARRAY['mx_colocation_test_2']);
- mark_tables_colocated
+-- Update colocation and see the changes on the master and the worker
+SELECT update_distributed_table_colocation('mx_colocation_test_1', colocate_with => 'mx_colocation_test_2');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
@@ -881,8 +876,8 @@ WHERE
 	OR logicalrelid = 'mx_colocation_test_2'::regclass;
      logicalrelid     | colocationid
 ---------------------------------------------------------------------
- mx_colocation_test_1 |        10001
  mx_colocation_test_2 |        10001
+ mx_colocation_test_1 |        10001
 (2 rows)
 
 \c - - - :worker_1_port
@@ -895,8 +890,8 @@ WHERE
 	OR logicalrelid = 'mx_colocation_test_2'::regclass;
      logicalrelid     | colocationid
 ---------------------------------------------------------------------
- mx_colocation_test_1 |        10001
  mx_colocation_test_2 |        10001
+ mx_colocation_test_1 |        10001
 (2 rows)
 
 \c - - - :master_port
diff --git a/src/test/regress/expected/multi_multiuser.out b/src/test/regress/expected/multi_multiuser.out
index daf86bad7..a0480d8a0 100644
--- a/src/test/regress/expected/multi_multiuser.out
+++ b/src/test/regress/expected/multi_multiuser.out
@@ -313,7 +313,7 @@ ABORT;
 SELECT * FROM citus_stat_statements_reset();
 ERROR:  permission denied for function citus_stat_statements_reset
 -- should not be allowed to co-located tables
-SELECT mark_tables_colocated('test', ARRAY['test_coloc'::regclass]);
+SELECT update_distributed_table_colocation('test', colocate_with => 'test_coloc');
 ERROR:  must be owner of table test
 -- should not be allowed to take any locks
 BEGIN;
diff --git a/src/test/regress/expected/multi_reference_table.out b/src/test/regress/expected/multi_reference_table.out
index 83e7e8c84..22ab25eca 100644
--- a/src/test/regress/expected/multi_reference_table.out
+++ b/src/test/regress/expected/multi_reference_table.out
@@ -1196,23 +1196,12 @@ DEBUG:  cannot perform distributed INSERT INTO ... SELECT because the partition
 DETAIL:  The target table's partition column should correspond to a partition column in the subquery.
 DEBUG:  performing repartitioned INSERT ... SELECT
 RESET client_min_messages;
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
--- some tests for mark_tables_colocated
--- should error out
-SELECT mark_tables_colocated('colocated_table_test_2', ARRAY['reference_table_test']);
-ERROR:  cannot colocate tables colocated_table_test_2 and reference_table_test
-DETAIL:  Replication models don't match for colocated_table_test_2 and reference_table_test.
--- should work sliently
-SELECT mark_tables_colocated('reference_table_test', ARRAY['reference_table_test_fifth']);
- mark_tables_colocated
----------------------------------------------------------------------
-
-(1 row)
-
+-- some tests for mark_tables_colocated	
+-- should error out	
+SELECT update_distributed_table_colocation('colocated_table_test_2', colocate_with => 'reference_table_test');
+ERROR:  relation reference_table_test should be a hash distributed table
+SELECT update_distributed_table_colocation('reference_table_test', colocate_with => 'reference_table_test_fifth');
+ERROR:  relation reference_table_test_fifth should be a hash distributed table
 -- ensure that reference tables on
 -- different queries works as expected
 CREATE SCHEMA reference_schema;
diff --git a/src/test/regress/expected/multi_unsupported_worker_operations.out b/src/test/regress/expected/multi_unsupported_worker_operations.out
index 7ec8728bd..ceb34e252 100644
--- a/src/test/regress/expected/multi_unsupported_worker_operations.out
+++ b/src/test/regress/expected/multi_unsupported_worker_operations.out
@@ -220,14 +220,8 @@ SELECT master_remove_node('localhost', 5432);
 (1 row)
 
 \c - - - :worker_1_port
--- mark_tables_colocated
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
 UPDATE pg_dist_partition SET colocationid = 0 WHERE logicalrelid='mx_table_2'::regclass;
-SELECT mark_tables_colocated('mx_table', ARRAY['mx_table_2']);
+SELECT update_distributed_table_colocation('mx_table', colocate_with => 'mx_table_2');
 ERROR:  operation is not allowed on this node
 HINT:  Connect to the coordinator and run it again.
 SELECT colocationid FROM pg_dist_partition WHERE logicalrelid='mx_table_2'::regclass;
diff --git a/src/test/regress/sql/citus_local_tables.sql b/src/test/regress/sql/citus_local_tables.sql
index b4e7a7c26..e01534f34 100644
--- a/src/test/regress/sql/citus_local_tables.sql
+++ b/src/test/regress/sql/citus_local_tables.sql
@@ -220,15 +220,15 @@ SELECT create_reference_table('reference_table');
 
 -- show that colociation of citus local tables are not supported for now
 -- between citus local tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['citus_local_table_2']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'citus_local_table_2');
 
 -- between citus local tables and reference tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['reference_table']);
-SELECT mark_tables_colocated('reference_table', ARRAY['citus_local_table_1']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'reference_table');
+SELECT update_distributed_table_colocation('reference_table', colocate_with => 'citus_local_table_1');
 
 -- between citus local tables and distributed tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['distributed_table']);
-SELECT mark_tables_colocated('distributed_table', ARRAY['citus_local_table_1']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'distributed_table');
+SELECT update_distributed_table_colocation('distributed_table', colocate_with => 'citus_local_table_1');
 
 -- master_create_empty_shard is not supported
 SELECT master_create_empty_shard('citus_local_table_1');
diff --git a/src/test/regress/sql/multi_colocation_utils.sql b/src/test/regress/sql/multi_colocation_utils.sql
index 55f563234..4959087f8 100644
--- a/src/test/regress/sql/multi_colocation_utils.sql
+++ b/src/test/regress/sql/multi_colocation_utils.sql
@@ -344,7 +344,7 @@ ORDER BY
     shardid,
     nodeport;
 
--- reset colocation ids to test mark_tables_colocated
+-- reset colocation ids to test update_distributed_table_colocation
 ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 1;
 DELETE FROM pg_dist_colocation
     WHERE colocationid >= 1 AND colocationid < 1000;
@@ -361,11 +361,11 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
     ORDER BY colocationid, logicalrelid;
 
 -- first check failing cases
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupC']);
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupD']);
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupE']);
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupF']);
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table2_groupB', 'table1_groupD']);
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupC');
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupD');
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupE');
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupF');
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupD');
 
 -- check metadata to see failing calls didn't have any side effects
 SELECT * FROM pg_dist_colocation
@@ -377,14 +377,14 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
     ORDER BY colocationid, logicalrelid;
 
 -- check successfully cololated tables
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table2_groupB']);
-SELECT mark_tables_colocated('table1_groupC', ARRAY['table2_groupC']);
-SELECT mark_tables_colocated('table1_groupD', ARRAY['table2_groupD']);
-SELECT mark_tables_colocated('table1_groupE', ARRAY['table2_groupE', 'table3_groupE']);
-SELECT mark_tables_colocated('table1_groupF', ARRAY['table2_groupF']);
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table2_groupB');
+SELECT update_distributed_table_colocation('table1_groupC', colocate_with => 'table2_groupC');
+SELECT update_distributed_table_colocation('table1_groupD', colocate_with => 'table2_groupD');
+SELECT update_distributed_table_colocation('table1_groupE', colocate_with => 'table2_groupE');
+SELECT update_distributed_table_colocation('table1_groupE', colocate_with => 'table3_groupE');
 
 -- check to colocate with itself
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupB']);
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupB');
 
 SET citus.shard_count = 2;
 
@@ -404,10 +404,12 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
     ORDER BY colocationid, logicalrelid;
 
 -- move the all tables in colocation group 5 to colocation group 7
-SELECT mark_tables_colocated('table1_group_none', ARRAY['table1_groupE', 'table2_groupE', 'table3_groupE']);
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table1_groupE');
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table2_groupE');
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table3_groupE');
 
 -- move a table with a colocation id which is already not in pg_dist_colocation
-SELECT mark_tables_colocated('table1_group_none', ARRAY['table2_group_none']);
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table2_group_none');
 
 -- check metadata to see that unused colocation group is deleted
 SELECT * FROM pg_dist_colocation
@@ -431,7 +433,7 @@ SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'table1_
 CREATE TABLE table2_groupG ( id int );
 SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'NONE');
 
-SELECT mark_tables_colocated('table1_groupG', ARRAY['table2_groupG']);
+SELECT update_distributed_table_colocation('table1_groupG', colocate_with => 'table2_groupG');
 
 CREATE TABLE d1(a int, b int);
 CREATE TABLE d2(a int, b int);
@@ -457,14 +459,6 @@ SELECT create_distributed_table('range_table', 'a', 'range');
 
 SELECT create_reference_table('ref');
 
-
--- check d1, d2, d3 and d4 has the same colocation id => they are colocated.
-SELECT get_table_colocation_id('d1');
-SELECT get_table_colocation_id('d2');
-SELECT get_table_colocation_id('d3');
-SELECT get_table_colocation_id('d4');
-
-
 SELECT tables_colocated('d1', 'd2');
 SELECT tables_colocated('d2', 'd3');
 SELECT tables_colocated('d2', 'd4');
@@ -475,12 +469,6 @@ SELECT tables_colocated('d1', 'd4');
 -- break colocation of d2
 SELECT update_distributed_table_colocation('d2', colocate_with => 'none');
 
--- d1 and d3 and d4 should be colocated, d2 should have a new colocation id.
-SELECT get_table_colocation_id('d1');
-SELECT get_table_colocation_id('d2');
-SELECT get_table_colocation_id('d3');
-SELECT get_table_colocation_id('d4');
-
 SELECT tables_colocated('d1', 'd2');
 SELECT tables_colocated('d2', 'd3');
 SELECT tables_colocated('d1', 'd3');
@@ -490,12 +478,6 @@ SELECT tables_colocated('d1', 'd4');
 -- update colocation should not error if d2 doesn't have any colocated table.
 SELECT update_distributed_table_colocation('d2', colocate_with => 'none');
 
--- d1 and d3 and d4 should be colocated, d2 should have a new colocation id.
-SELECT get_table_colocation_id('d1');
-SELECT get_table_colocation_id('d2');
-SELECT get_table_colocation_id('d3');
-SELECT get_table_colocation_id('d4');
-
 SELECT tables_colocated('d1', 'd2');
 SELECT tables_colocated('d2', 'd3');
 SELECT tables_colocated('d1', 'd3');
@@ -503,12 +485,6 @@ SELECT tables_colocated('d1', 'd4');
 
 SELECT update_distributed_table_colocation('d3', colocate_with => 'd2');
 
--- d1 and d4 should be colocated, d2 and d3 should be colocated.
-SELECT get_table_colocation_id('d1');
-SELECT get_table_colocation_id('d2');
-SELECT get_table_colocation_id('d3');
-SELECT get_table_colocation_id('d4');
-
 SELECT tables_colocated('d1', 'd2');
 SELECT tables_colocated('d2', 'd3');
 SELECT tables_colocated('d1', 'd3');
@@ -516,12 +492,6 @@ SELECT tables_colocated('d1', 'd4');
 
 -- special case, colocate with a table named "none".
 SELECT update_distributed_table_colocation('d3', colocate_with => '"none"');
--- d1, d4, d3 and "none" should be colocated;
-SELECT get_table_colocation_id('d1');
-SELECT get_table_colocation_id('d2');
-SELECT get_table_colocation_id('d3');
-SELECT get_table_colocation_id('d4');
-SELECT get_table_colocation_id('none');
 
 SELECT tables_colocated('d1', 'd2');
 SELECT tables_colocated('d2', 'd3');
diff --git a/src/test/regress/sql/multi_drop_extension.sql b/src/test/regress/sql/multi_drop_extension.sql
index ec1bee234..0598d31ed 100644
--- a/src/test/regress/sql/multi_drop_extension.sql
+++ b/src/test/regress/sql/multi_drop_extension.sql
@@ -16,7 +16,6 @@ SELECT create_distributed_table('testtableddl', 'distributecol', 'append');
 -- around.
 SET client_min_messages TO 'WARNING';
 DROP FUNCTION pg_catalog.master_create_worker_shards(text, integer, integer);
-DROP FUNCTION pg_catalog.mark_tables_colocated(regclass, regclass[]);
 DROP EXTENSION citus CASCADE;
 RESET client_min_messages;
 
diff --git a/src/test/regress/sql/multi_metadata_sync.sql b/src/test/regress/sql/multi_metadata_sync.sql
index a0ff53a7c..8a9b380af 100644
--- a/src/test/regress/sql/multi_metadata_sync.sql
+++ b/src/test/regress/sql/multi_metadata_sync.sql
@@ -354,13 +354,8 @@ NOT VALID;
 \c - - - :worker_1_port
 SELECT "Constraint", "Definition" FROM table_fkeys WHERE relid='mx_test_schema_1.mx_table_1'::regclass;
 
--- Check that mark_tables_colocated call propagates the changes to the workers
+-- Check that update_distributed_table_colocation call propagates the changes to the workers
 \c - - - :master_port
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
 SELECT nextval('pg_catalog.pg_dist_colocationid_seq') AS last_colocation_id \gset
 ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 10000;
 SET citus.shard_count TO 7;
@@ -400,8 +395,8 @@ WHERE
 	logicalrelid = 'mx_colocation_test_1'::regclass
 	OR logicalrelid = 'mx_colocation_test_2'::regclass;
 
--- Mark tables colocated and see the changes on the master and the worker
-SELECT mark_tables_colocated('mx_colocation_test_1', ARRAY['mx_colocation_test_2']);
+-- Update colocation and see the changes on the master and the worker
+SELECT update_distributed_table_colocation('mx_colocation_test_1', colocate_with => 'mx_colocation_test_2');
 SELECT
 	logicalrelid, colocationid
 FROM
diff --git a/src/test/regress/sql/multi_multiuser.sql b/src/test/regress/sql/multi_multiuser.sql
index 90cd28790..b88fadb03 100644
--- a/src/test/regress/sql/multi_multiuser.sql
+++ b/src/test/regress/sql/multi_multiuser.sql
@@ -193,7 +193,7 @@ ABORT;
 SELECT * FROM citus_stat_statements_reset();
 
 -- should not be allowed to co-located tables
-SELECT mark_tables_colocated('test', ARRAY['test_coloc'::regclass]);
+SELECT update_distributed_table_colocation('test', colocate_with => 'test_coloc');
 
 -- should not be allowed to take any locks
 BEGIN;
diff --git a/src/test/regress/sql/multi_reference_table.sql b/src/test/regress/sql/multi_reference_table.sql
index 34b9900db..a432730d5 100644
--- a/src/test/regress/sql/multi_reference_table.sql
+++ b/src/test/regress/sql/multi_reference_table.sql
@@ -764,18 +764,11 @@ WHERE
 
 RESET client_min_messages;
 
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
-
 -- some tests for mark_tables_colocated
 -- should error out
-SELECT mark_tables_colocated('colocated_table_test_2', ARRAY['reference_table_test']);
+SELECT update_distributed_table_colocation('colocated_table_test_2', colocate_with => 'reference_table_test');
 
--- should work sliently
-SELECT mark_tables_colocated('reference_table_test', ARRAY['reference_table_test_fifth']);
+SELECT update_distributed_table_colocation('reference_table_test', colocate_with => 'reference_table_test_fifth');
 
 -- ensure that reference tables on
 -- different queries works as expected
diff --git a/src/test/regress/sql/multi_unsupported_worker_operations.sql b/src/test/regress/sql/multi_unsupported_worker_operations.sql
index 69677d3c1..5b07f2cac 100644
--- a/src/test/regress/sql/multi_unsupported_worker_operations.sql
+++ b/src/test/regress/sql/multi_unsupported_worker_operations.sql
@@ -142,15 +142,9 @@ SELECT master_remove_node('localhost', 5432);
 
 \c - - - :worker_1_port
 
--- mark_tables_colocated
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
 UPDATE pg_dist_partition SET colocationid = 0 WHERE logicalrelid='mx_table_2'::regclass;
 
-SELECT mark_tables_colocated('mx_table', ARRAY['mx_table_2']);
+SELECT update_distributed_table_colocation('mx_table', colocate_with => 'mx_table_2');
 SELECT colocationid FROM pg_dist_partition WHERE logicalrelid='mx_table_2'::regclass;
 
 SELECT colocationid AS old_colocation_id