1 2SET citus.next_shard_id TO 1300000; 3ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 4; 4 5-- =================================================================== 6-- create test utility function 7-- =================================================================== 8 9CREATE SEQUENCE colocation_test_seq 10 MINVALUE 1000 11 NO CYCLE; 12 13/* a very simple UDF that only sets the colocation ids the same 14 * DO NOT USE THIS FUNCTION IN PRODUCTION. It manually sets colocationid column of 15 * pg_dist_partition and it does not check anything about pyshical state about shards. 16 */ 17CREATE OR REPLACE FUNCTION colocation_test_colocate_tables(source_table regclass, target_table regclass) 18 RETURNS BOOL 19 LANGUAGE plpgsql 20 AS $colocate_tables$ 21DECLARE nextid INTEGER; 22BEGIN 23 SELECT nextval('colocation_test_seq') INTO nextid; 24 25 UPDATE pg_dist_partition SET colocationId = nextid 26 WHERE logicalrelid IN 27 ( 28 (SELECT p1.logicalrelid 29 FROM pg_dist_partition p1, pg_dist_partition p2 30 WHERE 31 p2.logicalrelid = source_table AND 32 (p1.logicalrelid = source_table OR 33 (p1.colocationId = p2.colocationId AND p1.colocationId != 0))) 34 UNION 35 (SELECT target_table) 36 ); 37 RETURN TRUE; 38END; 39$colocate_tables$; 40 41-- =================================================================== 42-- create test functions 43-- =================================================================== 44 45CREATE FUNCTION get_table_colocation_id(regclass) 46 RETURNS INTEGER 47 AS 'citus' 48 LANGUAGE C STRICT; 49 50CREATE FUNCTION tables_colocated(regclass, regclass) 51 RETURNS bool 52 AS 'citus' 53 LANGUAGE C; 54 55CREATE FUNCTION shards_colocated(bigint, bigint) 56 RETURNS bool 57 AS 'citus' 58 LANGUAGE C STRICT; 59 60CREATE FUNCTION find_shard_interval_index(bigint) 61 RETURNS int 62 AS 'citus' 63 LANGUAGE C STRICT; 64 65-- =================================================================== 66-- test co-location util functions 67-- =================================================================== 68 69-- create distributed table observe shard pruning 70CREATE TABLE table1_group1 ( id int ); 71SELECT master_create_distributed_table('table1_group1', 'id', 'hash'); 72SELECT master_create_worker_shards('table1_group1', 4, 2); 73 74CREATE TABLE table2_group1 ( id int ); 75SELECT master_create_distributed_table('table2_group1', 'id', 'hash'); 76SELECT master_create_worker_shards('table2_group1', 4, 2); 77 78CREATE TABLE table3_group2 ( id int ); 79SELECT master_create_distributed_table('table3_group2', 'id', 'hash'); 80SELECT master_create_worker_shards('table3_group2', 4, 2); 81 82CREATE TABLE table4_group2 ( id int ); 83SELECT master_create_distributed_table('table4_group2', 'id', 'hash'); 84SELECT master_create_worker_shards('table4_group2', 4, 2); 85 86CREATE TABLE table5_groupX ( id int ); 87SELECT master_create_distributed_table('table5_groupX', 'id', 'hash'); 88SELECT master_create_worker_shards('table5_groupX', 4, 2); 89 90CREATE TABLE table6_append ( id int ); 91SELECT master_create_distributed_table('table6_append', 'id', 'append'); 92SELECT master_create_empty_shard('table6_append'); 93SELECT master_create_empty_shard('table6_append'); 94 95 96-- make table1_group1 and table2_group1 co-located manually 97SELECT colocation_test_colocate_tables('table1_group1', 'table2_group1'); 98 99-- check co-location id 100SELECT get_table_colocation_id('table1_group1'); 101SELECT get_table_colocation_id('table5_groupX'); 102SELECT get_table_colocation_id('table6_append'); 103 104-- check self table co-location 105SELECT tables_colocated('table1_group1', 'table1_group1'); 106SELECT tables_colocated('table5_groupX', 'table5_groupX'); 107SELECT tables_colocated('table6_append', 'table6_append'); 108 109-- check table co-location with same co-location group 110SELECT tables_colocated('table1_group1', 'table2_group1'); 111 112-- check table co-location with different co-location group 113SELECT tables_colocated('table1_group1', 'table3_group2'); 114 115-- check table co-location with invalid co-location group 116SELECT tables_colocated('table1_group1', 'table5_groupX'); 117SELECT tables_colocated('table1_group1', 'table6_append'); 118 119-- check self shard co-location 120SELECT shards_colocated(1300000, 1300000); 121SELECT shards_colocated(1300016, 1300016); 122SELECT shards_colocated(1300020, 1300020); 123 124-- check shard co-location with same co-location group 125SELECT shards_colocated(1300000, 1300004); 126 127-- check shard co-location with same table different co-location group 128SELECT shards_colocated(1300000, 1300001); 129 130-- check shard co-location with different co-location group 131SELECT shards_colocated(1300000, 1300005); 132 133-- check shard co-location with invalid co-location group 134SELECT shards_colocated(1300000, 1300016); 135SELECT shards_colocated(1300000, 1300020); 136 137-- check co-located table list 138SELECT UNNEST(get_colocated_table_array('table1_group1'))::regclass ORDER BY 1; 139SELECT UNNEST(get_colocated_table_array('table5_groupX'))::regclass ORDER BY 1; 140SELECT UNNEST(get_colocated_table_array('table6_append'))::regclass ORDER BY 1; 141 142-- check co-located shard list 143SELECT UNNEST(get_colocated_shard_array(1300000))::regclass ORDER BY 1; 144SELECT UNNEST(get_colocated_shard_array(1300016))::regclass ORDER BY 1; 145SELECT UNNEST(get_colocated_shard_array(1300020))::regclass ORDER BY 1; 146 147-- check FindShardIntervalIndex function 148SELECT find_shard_interval_index(1300000); 149SELECT find_shard_interval_index(1300001); 150SELECT find_shard_interval_index(1300002); 151SELECT find_shard_interval_index(1300003); 152SELECT find_shard_interval_index(1300016); 153 154-- check external colocation API 155 156SELECT count(*) FROM pg_dist_partition WHERE colocationid = 4; 157DELETE FROM pg_dist_colocation WHERE colocationid = 4; 158 159SET citus.shard_count = 2; 160 161CREATE TABLE table1_groupA ( id int ); 162SELECT create_distributed_table('table1_groupA', 'id'); 163 164CREATE TABLE table2_groupA ( id int ); 165SELECT create_distributed_table('table2_groupA', 'id'); 166 167-- change shard replication factor 168SET citus.shard_replication_factor = 1; 169 170CREATE TABLE table1_groupB ( id int ); 171SELECT create_distributed_table('table1_groupB', 'id'); 172 173CREATE TABLE table2_groupB ( id int ); 174SELECT create_distributed_table('table2_groupB', 'id'); 175 176UPDATE pg_dist_partition SET repmodel='c' WHERE logicalrelid='table1_groupB'::regclass; 177UPDATE pg_dist_partition SET repmodel='c' WHERE logicalrelid='table2_groupB'::regclass; 178 179-- revert back to default shard replication factor 180SET citus.shard_replication_factor to DEFAULT; 181 182-- change partition column type 183CREATE TABLE table1_groupC ( id text ); 184SELECT create_distributed_table('table1_groupC', 'id'); 185 186CREATE TABLE table2_groupC ( id text ); 187SELECT create_distributed_table('table2_groupC', 'id'); 188 189-- change shard count 190SET citus.shard_count = 8; 191 192CREATE TABLE table1_groupD ( id int ); 193SELECT create_distributed_table('table1_groupD', 'id'); 194 195CREATE TABLE table2_groupD ( id int ); 196SELECT create_distributed_table('table2_groupD', 'id'); 197 198-- try other distribution methods 199CREATE TABLE table_append ( id int ); 200SELECT create_distributed_table('table_append', 'id', 'append'); 201 202CREATE TABLE table_range ( id int ); 203SELECT create_distributed_table('table_range', 'id', 'range'); 204 205-- test foreign table creation 206CREATE FOREIGN TABLE table3_groupD ( id int ) SERVER fake_fdw_server; 207SELECT create_distributed_table('table3_groupD', 'id'); 208 209-- check metadata 210SELECT * FROM pg_dist_colocation 211 WHERE colocationid >= 1 AND colocationid < 1000 212 ORDER BY colocationid; 213 214SELECT logicalrelid, colocationid FROM pg_dist_partition 215 WHERE colocationid >= 1 AND colocationid < 1000 216 ORDER BY logicalrelid; 217 218-- check effects of dropping tables 219DROP TABLE table1_groupA; 220SELECT * FROM pg_dist_colocation WHERE colocationid = 4; 221 222-- dropping all tables in a colocation group also deletes the colocation group 223DROP TABLE table2_groupA; 224SELECT * FROM pg_dist_colocation WHERE colocationid = 4; 225 226-- create dropped colocation group again 227SET citus.shard_count = 2; 228 229CREATE TABLE table1_groupE ( id int ); 230SELECT create_distributed_table('table1_groupE', 'id'); 231 232CREATE TABLE table2_groupE ( id int ); 233SELECT create_distributed_table('table2_groupE', 'id'); 234 235-- test different table DDL 236CREATE TABLE table3_groupE ( dummy_column text, id int ); 237SELECT create_distributed_table('table3_groupE', 'id'); 238 239-- test different schema 240CREATE SCHEMA schema_colocation; 241 242CREATE TABLE schema_colocation.table4_groupE ( id int ); 243SELECT create_distributed_table('schema_colocation.table4_groupE', 'id'); 244 245-- test colocate_with option 246CREATE TABLE table1_group_none_1 ( id int ); 247SELECT create_distributed_table('table1_group_none_1', 'id', colocate_with => 'none'); 248 249CREATE TABLE table2_group_none_1 ( id int ); 250SELECT create_distributed_table('table2_group_none_1', 'id', colocate_with => 'table1_group_none_1'); 251 252CREATE TABLE table1_group_none_2 ( id int ); 253SELECT create_distributed_table('table1_group_none_2', 'id', colocate_with => 'none'); 254 255CREATE TABLE table4_groupE ( id int ); 256SELECT create_distributed_table('table4_groupE', 'id', colocate_with => 'default'); 257 258SET citus.shard_count = 3; 259 260-- check that this new configuration does not have a default group 261CREATE TABLE table1_group_none_3 ( id int ); 262SELECT create_distributed_table('table1_group_none_3', 'id', colocate_with => 'NONE'); 263 264-- a new table does not use a non-default group 265CREATE TABLE table1_group_default ( id int ); 266SELECT create_distributed_table('table1_group_default', 'id', colocate_with => 'DEFAULT'); 267 268-- check metadata 269SELECT * FROM pg_dist_colocation 270 WHERE colocationid >= 1 AND colocationid < 1000 271 ORDER BY colocationid; 272 273SELECT logicalrelid, colocationid FROM pg_dist_partition 274 WHERE colocationid >= 1 AND colocationid < 1000 275 ORDER BY colocationid, logicalrelid; 276 277-- check failing colocate_with options 278CREATE TABLE table_postgresql( id int ); 279CREATE TABLE table_failing ( id int ); 280 281SELECT create_distributed_table('table_failing', 'id', colocate_with => 'table_append'); 282SELECT create_distributed_table('table_failing', 'id', 'append', 'table1_groupE'); 283SELECT create_distributed_table('table_failing', 'id', colocate_with => 'table_postgresql'); 284SELECT create_distributed_table('table_failing', 'id', colocate_with => 'no_table'); 285SELECT create_distributed_table('table_failing', 'id', colocate_with => ''); 286SELECT create_distributed_table('table_failing', 'id', colocate_with => NULL); 287 288-- check with different distribution column types 289CREATE TABLE table_bigint ( id bigint ); 290SELECT create_distributed_table('table_bigint', 'id', colocate_with => 'table1_groupE'); 291-- check worker table schemas 292\c - - - :worker_1_port 293SELECT "Column", "Type", "Modifiers" FROM table_desc WHERE relid='public.table3_groupE_1300062'::regclass; 294SELECT "Column", "Type", "Modifiers" FROM table_desc WHERE relid='schema_colocation.table4_groupE_1300064'::regclass; 295 296\c - - - :master_port 297SET citus.next_shard_id TO 1300080; 298 299CREATE TABLE table1_groupF ( id int ); 300SELECT create_reference_table('table1_groupF'); 301 302CREATE TABLE table2_groupF ( id int ); 303SELECT create_reference_table('table2_groupF'); 304 305-- check metadata 306SELECT * FROM pg_dist_colocation 307 WHERE colocationid >= 1 AND colocationid < 1000 308 ORDER BY colocationid; 309 310-- cross check with internal colocation API 311SELECT 312 p1.logicalrelid::regclass AS table1, 313 p2.logicalrelid::regclass AS table2, 314 tables_colocated(p1.logicalrelid , p2.logicalrelid) AS colocated 315FROM 316 pg_dist_partition p1, 317 pg_dist_partition p2 318WHERE 319 p1.logicalrelid < p2.logicalrelid AND 320 p1.colocationid != 0 AND 321 p2.colocationid != 0 AND 322 tables_colocated(p1.logicalrelid , p2.logicalrelid) is TRUE 323ORDER BY 324 table1, 325 table2; 326 327-- check created shards 328SELECT 329 logicalrelid, 330 pg_dist_shard.shardid AS shardid, 331 shardstorage, 332 nodeport, 333 shardminvalue, 334 shardmaxvalue 335FROM 336 pg_dist_shard, 337 pg_dist_shard_placement 338WHERE 339 pg_dist_shard.shardid = pg_dist_shard_placement.shardid AND 340 pg_dist_shard.shardid >= 1300026 341ORDER BY 342 logicalrelid, 343 shardmaxvalue::integer, 344 shardid, 345 nodeport; 346 347-- reset colocation ids to test update_distributed_table_colocation 348ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 1; 349DELETE FROM pg_dist_colocation 350 WHERE colocationid >= 1 AND colocationid < 1000; 351UPDATE pg_dist_partition SET colocationid = 0 352 WHERE colocationid >= 1 AND colocationid < 1000; 353 354-- check metadata 355SELECT * FROM pg_dist_colocation 356 WHERE colocationid >= 1 AND colocationid < 1000 357 ORDER BY colocationid; 358 359SELECT logicalrelid, colocationid FROM pg_dist_partition 360 WHERE colocationid >= 1 AND colocationid < 1000 361 ORDER BY colocationid, logicalrelid; 362 363-- first check failing cases 364SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupC'); 365SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupD'); 366SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupE'); 367SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupF'); 368SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupD'); 369 370-- check metadata to see failing calls didn't have any side effects 371SELECT * FROM pg_dist_colocation 372 WHERE colocationid >= 1 AND colocationid < 1000 373 ORDER BY colocationid; 374 375SELECT logicalrelid, colocationid FROM pg_dist_partition 376 WHERE colocationid >= 1 AND colocationid < 1000 377 ORDER BY colocationid, logicalrelid; 378 379-- check successfully cololated tables 380SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table2_groupB'); 381SELECT update_distributed_table_colocation('table1_groupC', colocate_with => 'table2_groupC'); 382SELECT update_distributed_table_colocation('table1_groupD', colocate_with => 'table2_groupD'); 383SELECT update_distributed_table_colocation('table1_groupE', colocate_with => 'table2_groupE'); 384SELECT update_distributed_table_colocation('table1_groupE', colocate_with => 'table3_groupE'); 385 386-- check to colocate with itself 387SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupB'); 388 389SET citus.shard_count = 2; 390 391CREATE TABLE table1_group_none ( id int ); 392SELECT create_distributed_table('table1_group_none', 'id', colocate_with => 'NONE'); 393 394CREATE TABLE table2_group_none ( id int ); 395SELECT create_distributed_table('table2_group_none', 'id', colocate_with => 'NONE'); 396 397-- check metadata to see colocation groups are created successfully 398SELECT * FROM pg_dist_colocation 399 WHERE colocationid >= 1 AND colocationid < 1000 400 ORDER BY colocationid; 401 402SELECT logicalrelid, colocationid FROM pg_dist_partition 403 WHERE colocationid >= 1 AND colocationid < 1000 404 ORDER BY colocationid, logicalrelid; 405 406-- move the all tables in colocation group 5 to colocation group 7 407SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table1_groupE'); 408SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table2_groupE'); 409SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table3_groupE'); 410 411-- move a table with a colocation id which is already not in pg_dist_colocation 412SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table2_group_none'); 413 414-- check metadata to see that unused colocation group is deleted 415SELECT * FROM pg_dist_colocation 416 WHERE colocationid >= 1 AND colocationid < 1000 417 ORDER BY colocationid; 418 419SELECT logicalrelid, colocationid FROM pg_dist_partition 420 WHERE colocationid >= 1 AND colocationid < 1000 421 ORDER BY colocationid, logicalrelid; 422 423-- try to colocate different replication models 424CREATE TABLE table1_groupG ( id int ); 425SELECT create_distributed_table('table1_groupG', 'id'); 426 427-- update replication model 428UPDATE pg_dist_partition SET repmodel = 's' WHERE logicalrelid = 'table1_groupG'::regclass; 429 430CREATE TABLE table2_groupG ( id int ); 431SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'table1_groupG'); 432DROP TABLE table2_groupG; 433 434CREATE TABLE table2_groupG ( id int ); 435SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'NONE'); 436 437SELECT update_distributed_table_colocation('table1_groupG', colocate_with => 'table2_groupG'); 438 439CREATE TABLE d1(a int, b int); 440CREATE TABLE d2(a int, b int); 441CREATE TABLE d3(a int, b int); 442CREATE TABLE d4(a int, b int); 443CREATE TABLE different_d1(ch char); 444CREATE TABLE append_table(a int, b int); 445CREATE TABLE range_table(a int, b int); 446-- special keyword none 447CREATE TABLE none(a int, b int); 448CREATE TABLE ref(a int); 449CREATE TABLE local_table(a int); 450 451SELECT create_distributed_table('d1', 'a'); 452SELECT create_distributed_table('d2', 'a', colocate_with => 'd1'); 453SELECT create_distributed_table('d3', 'a', colocate_with => 'd2'); 454SELECT create_distributed_table('d4', 'a', colocate_with => 'd3'); 455SELECT create_distributed_table('none', 'a', colocate_with => 'd4'); 456SELECT create_distributed_table('different_d1', 'ch'); 457SELECT create_distributed_table('append_table', 'a', 'append'); 458SELECT create_distributed_table('range_table', 'a', 'range'); 459 460 461SELECT create_reference_table('ref'); 462 463SELECT tables_colocated('d1', 'd2'); 464SELECT tables_colocated('d2', 'd3'); 465SELECT tables_colocated('d2', 'd4'); 466SELECT tables_colocated('d3', 'd4'); 467SELECT tables_colocated('d1', 'd3'); 468SELECT tables_colocated('d1', 'd4'); 469 470-- break colocation of d2 471SELECT update_distributed_table_colocation('d2', colocate_with => 'none'); 472 473SELECT tables_colocated('d1', 'd2'); 474SELECT tables_colocated('d2', 'd3'); 475SELECT tables_colocated('d1', 'd3'); 476SELECT tables_colocated('d1', 'd4'); 477 478-- break colocation of d2 479-- update colocation should not error if d2 doesn't have any colocated table. 480SELECT update_distributed_table_colocation('d2', colocate_with => 'none'); 481 482SELECT tables_colocated('d1', 'd2'); 483SELECT tables_colocated('d2', 'd3'); 484SELECT tables_colocated('d1', 'd3'); 485SELECT tables_colocated('d1', 'd4'); 486 487SELECT update_distributed_table_colocation('d3', colocate_with => 'd2'); 488 489SELECT tables_colocated('d1', 'd2'); 490SELECT tables_colocated('d2', 'd3'); 491SELECT tables_colocated('d1', 'd3'); 492SELECT tables_colocated('d1', 'd4'); 493 494-- special case, colocate with a table named "none". 495SELECT update_distributed_table_colocation('d3', colocate_with => '"none"'); 496 497SELECT tables_colocated('d1', 'd2'); 498SELECT tables_colocated('d2', 'd3'); 499SELECT tables_colocated('d1', 'd3'); 500SELECT tables_colocated('d1', 'd4'); 501SELECT tables_colocated('d1', 'none'); 502SELECT tables_colocated('d4', 'none'); 503SELECT tables_colocated('d3', 'none'); 504SELECT tables_colocated('d2', 'none'); 505 506-- make sure reference and local tables return an error. 507SELECT update_distributed_table_colocation('ref', colocate_with => 'none'); 508SELECT update_distributed_table_colocation('local_table', colocate_with => 'none'); 509 510-- make sure that different types cannot be colocated 511SELECT update_distributed_table_colocation('different_d1', colocate_with => 'd1'); 512SELECT update_distributed_table_colocation('d1', colocate_with => 'different_d1'); 513 514-- make sure that append distributed tables cannot be colocated 515SELECT update_distributed_table_colocation('append_table', colocate_with => 'd1'); 516SELECT update_distributed_table_colocation('d1', colocate_with => 'append_table'); 517SELECT update_distributed_table_colocation('range_table', colocate_with => 'd1'); 518SELECT update_distributed_table_colocation('d1', colocate_with => 'range_table'); 519 520 521 522 523-- drop tables to clean test space 524DROP TABLE table1_groupb; 525DROP TABLE table2_groupb; 526DROP TABLE table1_groupc; 527DROP TABLE table2_groupc; 528DROP TABLE table1_groupd; 529DROP TABLE table2_groupd; 530DROP TABLE table1_groupf; 531DROP TABLE table2_groupf; 532DROP TABLE table1_groupg; 533DROP TABLE table2_groupg; 534DROP TABLE table1_groupe; 535DROP TABLE table2_groupe; 536DROP TABLE table3_groupe; 537DROP TABLE table4_groupe; 538DROP TABLE schema_colocation.table4_groupe; 539DROP TABLE table1_group_none_1; 540DROP TABLE table2_group_none_1; 541DROP TABLE table1_group_none_2; 542DROP TABLE table1_group_none_3; 543DROP TABLE table1_group_none; 544DROP TABLE table2_group_none; 545DROP TABLE table1_group_default; 546DROP TABLE d1; 547DROP TABLE d2; 548DROP TABLE d3; 549DROP TABLE d4; 550DROP TABLE different_d1; 551DROP TABLE append_table; 552DROP TABLE range_table; 553DROP TABLE none; 554DROP TABLE ref; 555DROP TABLE local_table; 556 557 558