Exploring Tablet Distribution and Leader Placement in YugabyteDB

				
					CREATE OR REPLACE FUNCTION public.yb_tablet_placement_details(
  schema_name_p text,
  table_name_p  text,
  leaders_only_p boolean DEFAULT false
)
RETURNS TABLE(
  schema_name text,
  table_name text,
  ysql_oid oid,
  tablet_id text,
  tablet_hash_start integer,
  tablet_hash_end integer,
  replica_role text,
  node_host text,
  node_port integer,
  node_public_ip text,
  node_uuid text,
  node_type text,
  cloud text,
  region text,
  zone text,
  tablet_row_count bigint,
  tablet_pct_of_table numeric
)
LANGUAGE plpgsql
AS $function$
DECLARE
  v_relid             oid;
  v_relkind           "char";
  v_base_relid        oid;
  v_base_schema       text;
  v_base_relname      text;

  v_pk_index_oid      oid;
  v_pk_def            text;

  v_idx_def           text;
  v_hash_cols_raw     text;
  v_hash_cols_arr     text[];
  v_hash_expr         text;

  v_has_bounds        boolean;
  v_role_filter       text := '';
  v_sql               text;
BEGIN
  -- Resolve relation OID (table or index)
  SELECT to_regclass(format('%I.%I', schema_name_p, table_name_p))::oid
    INTO v_relid;

  IF v_relid IS NULL THEN
    RAISE EXCEPTION 'Relation %.% not found', schema_name_p, table_name_p;
  END IF;

  SELECT c.relkind
    INTO v_relkind
  FROM pg_class c
  WHERE c.oid = v_relid;

  IF leaders_only_p THEN
    v_role_filter := 'WHERE replica_role = ''LEADER''';
  END IF;

  -- Do we have hash boundaries for THIS relation (table or index)?
  SELECT EXISTS (
    SELECT 1
    FROM yb_tablet_metadata ytm
    WHERE ytm.db_name = current_database()
      AND ytm.oid = v_relid
      AND (ytm.start_hash_code IS NOT NULL OR ytm.end_hash_code IS NOT NULL)
  )
  INTO v_has_bounds;

  /*
    Determine hash expression:
    - If TABLE: use PK hash cols
    - If INDEX: use index HASH cols, but computed from the BASE TABLE alias "t"
  */
  IF v_has_bounds THEN
    IF v_relkind = 'i' THEN
      -- Base table for this index
      SELECT i.indrelid
        INTO v_base_relid
      FROM pg_index i
      WHERE i.indexrelid = v_relid;

      IF v_base_relid IS NULL THEN
        RAISE EXCEPTION 'Could not resolve base table for index %.%', schema_name_p, table_name_p;
      END IF;

      SELECT n.nspname, c.relname
        INTO v_base_schema, v_base_relname
      FROM pg_class c
      JOIN pg_namespace n ON n.oid = c.relnamespace
      WHERE c.oid = v_base_relid;

      -- Extract HASH columns from the index definition (lsm (... HASH ...))
      v_idx_def := pg_get_indexdef(v_relid);

      SELECT string_agg(quote_ident(regexp_replace(trim(part), '\s+HASH.*$', '', 'i')), ',' )
        INTO v_hash_cols_raw
      FROM regexp_split_to_table(
             regexp_replace(v_idx_def, '^.*\((.*)\).*$','\1'),
             ','
           ) AS part
      WHERE part ~* '\mHASH\M';

      IF v_hash_cols_raw IS NULL OR length(btrim(v_hash_cols_raw)) = 0 THEN
        RAISE EXCEPTION 'Index %.% does not appear to have HASH columns (indexdef=%)',
          schema_name_p, table_name_p, v_idx_def;
      END IF;

      SELECT array_agg(btrim(x))
        INTO v_hash_cols_arr
      FROM unnest(string_to_array(v_hash_cols_raw, ',')) AS x;

      -- Build: yb_hash_code(t.col1, t.col2, ...)
      SELECT
        'yb_hash_code(' ||
        string_agg('t.' || col, ', ') ||
        ')'
      INTO v_hash_expr
      FROM unnest(v_hash_cols_arr) AS col;

    ELSE
      -- TABLE path: PK-based hash expression
      SELECT i.indexrelid
        INTO v_pk_index_oid
      FROM pg_index i
      WHERE i.indrelid = v_relid
        AND i.indisprimary
      LIMIT 1;

      IF v_pk_index_oid IS NULL THEN
        RAISE EXCEPTION 'Table %.% has no primary key (required to map rows to tablets)', schema_name_p, table_name_p;
      END IF;

      v_pk_def := pg_get_indexdef(v_pk_index_oid);

      -- Try to extract explicit HASH columns from PK definition: "((c1, c2) HASH ...)"
      SELECT (regexp_match(v_pk_def, '\(\((.+?)\)\s+HASH', 'i'))[1]
        INTO v_hash_cols_raw;

      IF v_hash_cols_raw IS NOT NULL THEN
        SELECT array_agg(btrim(x))
          INTO v_hash_cols_arr
        FROM unnest(string_to_array(v_hash_cols_raw, ',')) AS x;
      ELSE
        -- Fallback: all PK columns in index order
        SELECT array_agg(quote_ident(a.attname) ORDER BY s.ord)
          INTO v_hash_cols_arr
        FROM pg_index i
        JOIN LATERAL unnest(i.indkey) WITH ORDINALITY AS s(attnum, ord) ON true
        JOIN pg_attribute a
          ON a.attrelid = i.indrelid
         AND a.attnum   = s.attnum
        WHERE i.indexrelid = v_pk_index_oid;
      END IF;

      IF v_hash_cols_arr IS NULL OR array_length(v_hash_cols_arr, 1) = 0 THEN
        RAISE EXCEPTION 'Could not determine hash key columns for %.% (pk_def=%)', schema_name_p, table_name_p, v_pk_def;
      END IF;

      SELECT
        'yb_hash_code(' ||
        string_agg('t.' || quote_ident(replace(col,'"','')), ', ') ||
        ')'
      INTO v_hash_expr
      FROM unnest(v_hash_cols_arr) AS col;
    END IF;
  END IF;

  /*
    Build the query
  */
  IF v_has_bounds THEN
    v_sql := format($q$
      WITH tablet_map AS (
        SELECT
          ytm.tablet_id,
          ytm.oid,
          COALESCE(ytm.start_hash_code, 0)::int     AS start_hash_code_eff,
          COALESCE(ytm.end_hash_code, 65536)::int   AS end_hash_code_eff,
          ytm.leader,
          ytm.replicas
        FROM yb_tablet_metadata ytm
        WHERE ytm.db_name = current_database()
          AND ytm.oid = %s
      ),
      -- ✅ production nicety: count starting from tablet_map so empty tablets appear with 0
      counts AS (
        SELECT
          tm.tablet_id,
          COUNT(t.*)::bigint AS tablet_row_count
        FROM tablet_map tm
        LEFT JOIN %I.%I t
          ON %s >= tm.start_hash_code_eff
         AND %s <  tm.end_hash_code_eff
        GROUP BY tm.tablet_id
      ),
      totals AS (
        SELECT COALESCE(SUM(tablet_row_count)::numeric, 0::numeric) AS total_rows
        FROM counts
      ),
      replica_rows AS (
        SELECT
          tm.tablet_id,
          tm.oid,
          tm.start_hash_code_eff,
          tm.end_hash_code_eff,
          rep AS replica_hostport,
          tm.leader,
          split_part(rep, ':', 1) AS rep_host,
          split_part(rep, ':', 2)::int AS rep_port
        FROM tablet_map tm
        CROSS JOIN LATERAL unnest(tm.replicas) AS rep
      ),
      base AS (
        SELECT
          %L::text  AS schema_name,
          %L::text  AS table_name,
          rr.oid    AS ysql_oid,
          rr.tablet_id,
          rr.start_hash_code_eff AS tablet_hash_start,
          rr.end_hash_code_eff   AS tablet_hash_end,
          CASE
            WHEN rr.replica_hostport = rr.leader THEN 'LEADER'
            ELSE 'FOLLOWER'
          END AS replica_role,
          ys.host        AS node_host,
          ys.port::int   AS node_port,
          ys.public_ip   AS node_public_ip,
          ys.uuid::text  AS node_uuid,
          ys.node_type   AS node_type,
          ys.cloud,
          ys.region,
          ys.zone,
          COALESCE(c.tablet_row_count, 0)::bigint AS tablet_row_count,
          CASE
            WHEN (SELECT total_rows FROM totals) = 0
              THEN NULL::numeric(6,2)
            ELSE ROUND((COALESCE(c.tablet_row_count, 0)::numeric / (SELECT total_rows FROM totals)) * 100, 2)
          END AS tablet_pct_of_table
        FROM replica_rows rr
        JOIN yb_servers() ys
          ON rr.rep_port = ys.port::int
         AND (rr.rep_host = ys.host OR rr.rep_host = ys.public_ip)
        LEFT JOIN counts c
          ON c.tablet_id = rr.tablet_id
      )
      SELECT *
      FROM base
      %s
      ORDER BY tablet_hash_start, tablet_id, replica_role DESC, node_host, node_port;
    $q$,
      v_relid,
      -- scan base table if index, else scan requested table
      (CASE WHEN v_relkind = 'i' THEN v_base_schema ELSE schema_name_p END),
      (CASE WHEN v_relkind = 'i' THEN v_base_relname ELSE table_name_p END),
      v_hash_expr, v_hash_expr,
      schema_name_p, table_name_p,
      v_role_filter
    );
  ELSE
    v_sql := format($q$
      WITH tablet_map AS (
        SELECT
          ytm.tablet_id,
          ytm.oid,
          ytm.leader,
          ytm.replicas
        FROM yb_tablet_metadata ytm
        WHERE ytm.db_name = current_database()
          AND ytm.oid = %s
      ),
      replica_rows AS (
        SELECT
          tm.tablet_id,
          tm.oid,
          rep AS replica_hostport,
          tm.leader,
          split_part(rep, ':', 1) AS rep_host,
          split_part(rep, ':', 2)::int AS rep_port
        FROM tablet_map tm
        CROSS JOIN LATERAL unnest(tm.replicas) AS rep
      ),
      base AS (
        SELECT
          %L::text  AS schema_name,
          %L::text  AS table_name,
          rr.oid    AS ysql_oid,
          rr.tablet_id,
          NULL::int AS tablet_hash_start,
          NULL::int AS tablet_hash_end,
          CASE
            WHEN rr.replica_hostport = rr.leader THEN 'LEADER'
            ELSE 'FOLLOWER'
          END AS replica_role,
          ys.host        AS node_host,
          ys.port::int   AS node_port,
          ys.public_ip   AS node_public_ip,
          ys.uuid::text  AS node_uuid,
          ys.node_type   AS node_type,
          ys.cloud,
          ys.region,
          ys.zone,
          NULL::bigint   AS tablet_row_count,
          NULL::numeric(6,2) AS tablet_pct_of_table
        FROM replica_rows rr
        JOIN yb_servers() ys
          ON rr.rep_port = ys.port::int
         AND (rr.rep_host = ys.host OR rr.rep_host = ys.public_ip)
      )
      SELECT *
      FROM base
      %s
      ORDER BY tablet_id, replica_role DESC, node_host, node_port;
    $q$,
      v_relid,
      schema_name_p, table_name_p,
      v_role_filter
    );
  END IF;

  RETURN QUERY EXECUTE v_sql;
END;
$function$;