Show distributed table sizes in \d+

replace-pg_table_size.sql

-- Replace the pg_table_size function with one that supports distributed tables
BEGIN;
ALTER FUNCTION pg_catalog.pg_table_size(regclass) RENAME TO orig_table_size;

CREATE OR REPLACE FUNCTION pg_catalog.pg_table_size(table_name regclass)
RETURNS bigint
LANGUAGE plpgsql
AS $function$
BEGIN
        IF EXISTS (SELECT 1 FROM pg_dist_partition WHERE logicalrelid = table_name) THEN
                RETURN pg_catalog.citus_table_size(table_name);
        ELSE
                RETURN pg_catalog.orig_table_size(table_name);
        END IF;
END;
$function$;
COMMIT;

restore-pg_table_size.sql

-- Restore original pg_table_size
BEGIN;
DROP FUNCTION pg_catalog.pg_table_size(regclass);
ALTER FUNCTION pg_catalog.orig_table_size(regclass) RENAME TO pg_table_size;
COMMIT;

Example:

# CREATE TABLE test (x int, y serial);
CREATE TABLE
postgres=# SELECT create_distributed_table('test','x');
 create_distributed_table
--------------------------

(1 row)

postgres=# INSERT INTO test (x) SELECT s FROM generate_series(1,1000000) s;
INSERT 0 1000000
postgres=# \d+
                         List of relations
 Schema |    Name    |   Type   | Owner |    Size    | Description 
--------+------------+----------+-------+------------+-------------
 public | test       | table    | marco | 0 bytes    | 
 public | test_y_seq | sequence | marco | 8192 bytes | 
(2 rows)

postgres=# \i replace-pg_table_size.sql
BEGIN
ALTER FUNCTION
CREATE FUNCTION
COMMIT
postgres=# \d+
                         List of relations
 Schema |    Name    |   Type   | Owner |    Size    | Description
--------+------------+----------+-------+------------+-------------
 public | test       | table    | marco | 35 MB      |
 public | test_y_seq | sequence | marco | 8192 bytes |
(2 rows)