Skip to content

Instantly share code, notes, and snippets.

@onderkalaci
Created August 19, 2016 08:01
Show Gist options
  • Save onderkalaci/b7737864beb0fd865fc4a84a547a411c to your computer and use it in GitHub Desktop.
Save onderkalaci/b7737864beb0fd865fc4a84a547a411c to your computer and use it in GitHub Desktop.
```
CREATE TABLE pg_dist_colocation
(shardCount INT, workerNodeList [TEXT], replicationFactor INT, colocationId SERIAL,
UNIQUE (shardCount, workerNodeList))
```
ALTER TABLE pg_dist_partition ADD COLUMN colocationId INT;
def create_hash_partitioned_table(tableName, partitionColumn):
shardCount = citus.shardCount
workerNodeList = WorkerNodeList()
colocationId = GetOrInsertColocationId(shardCount, workerNodeList, replicationFactor
)
CreateDistributedTable('tableName'::regclass, DISTRIBUTED_BY_HASH, partitionColumn, colocationId);
CreateWorkerShards('tableName'::regclass, shardCount, replicationFactor);
def GetOrInsertColocationId(shardCountIn, workerNodeListIn, replicationFactorIn):
```
INSERT INTO
pg_dist_colocation (shardCount, workerNodeList, replicationFactor) VALUES
(shardCountIn, workerNodeListIn, replicationFactorIn)
ON CONFLICT (shardCount, workerNodeList) DO NOTHING
RETURNING colocationId;
```
def rebalance_colocated_tables(tableName):
colocationId = ColocationId(tableName)
AcquireExlusiveLock(colocationId)
colocatedTableList = ColocatedTableList(colocationId)
for colocatedTable in colocatedTableList:
sortedShardPlacementList = SortShardPlacementList()
def SortShardPlacementList(relationId):
shardPlacementList = ShardPlcementList(relationId)
# sort the placements wrt shardId, nodeHost, and nodePort
sortedShardPlacementList = SortList (shardPlacementList, [shardId, nodeHost, nodePort])
return sortedShardPlacementList
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment