Implementing Sharding with Citus

Introduction

In the era of big data, when storage capacity reaches the Terabyte threshold or the number of requests exceeds the physical limit of a single server (Vertical Scaling / Scale-up), Sharding is the optimal solution for Horizontal Scaling / Scale-out
It is achieved by dividing a massive data table (billions of rows) into multiple small, independent and self-managed parts called Shards
Each Shard is a separate physical database located on a different physical Server.
The key point is that the data in the Shards does not overlap, but when combined, it forms a complete dataset.

Horizontal Partitioning vs Sharding

Horizontal Partitioning: Splitting a large table into smaller tables (such as by month) but all of these child tables still reside on the same physical server.
Sharding: Distributing those child tables after partitioning across multiple different physical servers. Therefore, Sharding is the architecture of horizontal data partitioning in a distributed environment.

Sharding Key

To know which Shard a row will be stored in, the system relies on a specific data column called Sharding Key (or Distribution Key), which is used as a routing landmark.

Common routing algorithms include:

Range-based Sharding: Partitioning based on a range of key values
- Example: IDs from 1 - 10,000 in Shard 1, from 10,001 - 20,000 in Shard 2
- Disadvantage: Easily leads to load imbalance if new data is concentrated within a certain range.
Directory-based Sharding: Using a separate configuration storage service to map which Key value goes to which Shard.
- Hash-based Sharding: This is the most commonly used method
- The system applies a Hash Function to the Sharding Key according to a formula to determine the location: Shard ID = Hash(Sharding Key) mod N (where N is the number of Shards)
- Advantage: Data is distributed extremely evenly among servers. Hash-based Sharding

Query Routing

When there is a SQL statement, it connects to an intermediate node (Router or Coordinator). This Node analyzes the statement, relies on the Sharding Key to calculate the Shard location and then sends the statement directly to the correct Shard storing the data without needing to scan others
If the statement needs to aggregate from multiple Shards, the Coordinator will run the subqueries in parallel and merge the results before returning them to the Client.

Components

Citus

Citus is a powerful open-source extension that transforms Postgres from a monolithic Database into a distributed Database. The architecture consists of 2 main components:

Node Worker
- It is the physical Server, which is actually a regular Postgres database containing multiple small physical tables (the logical Shards).
- For example, in a system with 2 Workers and 32 Shards: Worker 1 will contain 16 actual Postgres tables representing Shards 1-16 and Worker 2 will contain 16 tables representing Shards 17-32.
- The Workers directly perform calculations, filtering and sorting of data according to commands from the Coordinator
- They have independent processing performance, allowing linear query speedups when adding Workers.
Node Coordinator
- This is the sole communication gateway with the application, it does not contain actual data but only stores the metadata of the Node Workers including IP, host, port and shard information.
- When you query data, the Coordinator plans the execution of the distributed SQL statement and coordinates by opening direct Socket Connections to the Workers to request data processing, then returns the result to it

Prerequisites

If your goal is to improve query performance on large datasets, you should first consider Query Tuning, using appropriate Index types like Partial Index or Table Partitioning as primary choices
If none of the above solutions help because the system data size is too large, only then should you think about Sharding

Detail

First, create a docker-compose.yml file with the following content

services:
  citus_master:
    image: citusdata/citus:alpine
    container_name: citus_coordinator
    ports:
      - "5432:5432"
    environment:
      &citus-env
      POSTGRES_USER: ${POSTGRES_USER}
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
      POSTGRES_DB: ${POSTGRES_DB}

  citus_worker_1:
    image: citusdata/citus:alpine
    container_name: citus_worker_1
    environment: *citus-env
    depends_on:
      - citus_master

  citus_worker_2:
    image: citusdata/citus:alpine
    container_name: citus_worker_2
    environment: *citus-env
    depends_on:
      - citus_master

  pgadmin:
    image: dpage/pgadmin4
    container_name: citus_pgadmin
    ports:
      - "8080:80"
    environment:
      PGADMIN_DEFAULT_EMAIL: ${PGADMIN_DEFAULT_EMAIL}
      PGADMIN_DEFAULT_PASSWORD: ${PGADMIN_DEFAULT_PASSWORD}
    depends_on:
      - citus_master

You can see that I have created citus_master and 2 workers citus_worker_1/citus_worker_2, along with pgadmin
You can change the Postgres value information accordingly in the .env file

Then connect with pgAdmin to citus_master as follows

Next, add the workers to the Coordinator with 5432 as the port defined in docker-compose

SELECT citus_add_node('citus_worker_1', 5432);
SELECT citus_add_node('citus_worker_2', 5432);

You can check the successfully added worker information and the metadata of the Coordinator as follows

SELECT * FROM citus_get_active_worker_nodes();
SELECT * FROM pg_dist_node;

After that, create the tables and enable the Sharding feature as follows

CREATE TABLE products (
    product_id SERIAL PRIMARY KEY,
    product_name VARCHAR(150) NOT NULL,
    price NUMERIC(12, 2) NOT NULL CHECK (price >= 0),
    stock_quantity INT NOT NULL DEFAULT 0,
    status VARCHAR(20) DEFAULT 'active'
);

CREATE TABLE IF NOT EXISTS customers (
    id SERIAL PRIMARY KEY,
    fullname VARCHAR(100) NOT NULL,
    email VARCHAR(100) NOT NULL,
    phone VARCHAR(15),
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

SELECT create_distributed_table('products', 'product_id');
SELECT create_distributed_table('customers', 'id', shard_count := 16);

SHOW citus.shard_count;
SET citus.shard_count = 8;

You can see that I use the create_distributed_table function to create the distributed table products with the Sharding key as product_id
By default, it will be divided into 32 Shards, which you can check with SHOW citus.shard_count
If you want to adjust it, use SET citus.shard_count with the appropriate value or pass shard_count directly into the create_distributed_table function
The number of Shards will be divided equally among Workers. For example, if we have 2 Workers, each will have 16 Shards. If you adjust shard_count := 16, each Worker will have 8 Shards

After seeding data, you can check directly on the Coordinator and it will retrieve data from Workers to aggregate and return

SELECT * FROM products;
SELECT * FROM customers;

To view specific Shard information

-- Query 1
SELECT 
    table_name,
    shardid,
    nodename,
    nodeport,
    shard_size
FROM citus_shards
ORDER BY table_name, shardid;

-- Query 2
SELECT 
    shardid AS shard_name, 
    result::bigint AS row_count
FROM run_command_on_shards('products', 'SELECT count(*) FROM %s');

-- Query 3
SELECT 
    shardid,
    result AS row_data
FROM run_command_on_shards(
    'products', 
    'SELECT json_agg(t) FROM (SELECT * FROM %s LIMIT 5) t'
)
WHERE shardid = 102008;

Query 1: Shard information within each Worker
Query 2: Execute query SELECT count(*) FROM %s on each Shard to count data
Query 3: Query data on a specific Shard, namely 102008