Exact Per-Tablet Row Counts in YugabyteDB – Hash and Range Sharding in One SQL Function

				
					CREATE OR REPLACE FUNCTION public.yb_tablet_key_counts(
    p_table regclass
)
RETURNS TABLE (
    tablet_ordinal int,
    tablet_id text,
    partitioning_scheme text,
    start_hash_code int,
    end_hash_code int,
    start_range_value jsonb,
    end_range_value jsonb,
    split_range text,
    leader text,
    replicas text[],
    row_count bigint,
    distinct_key_count bigint
)
LANGUAGE plpgsql
AS $$
DECLARE
    v_schema_name text;
    v_table_name  text;
    v_indexdef    text;
    v_is_hash     boolean := false;

    v_hash_cols_text text;
    v_hash_expr      text;
    v_distinct_expr  text;

    v_pk_cols text[];
    v_key_tuple text;
    v_range_key_tuple text;
    v_range_key_col_count int;

    v_union_sql text := '';
    v_branch_sql text;

    r record;

    v_tablet_ordinal int := 0;

    v_split_clause text;
    v_values_text text;
    v_inner text;
    v_split_values text[] := ARRAY[]::text[];
    v_normalized_split_values text[] := ARRAY[]::text[];

    v_current text := '';
    v_ch text;
    v_depth int := 0;
    v_in_quote boolean := false;
    v_pos int := 1;
    v_len int := 0;
    v_next_ch text;

    v_start_value text;
    v_end_value text;
    v_tablet_count int;
    v_range_idx int := 0;

    v_predicate text;

    /*
      Helpers for normalizing split tuples like:

        ('uuid-value', MINVALUE)

      into:

        ('uuid-value')
    */
    v_split_value text;
    v_tuple_inner text;
    v_elem_values text[];
    v_elem_current text;
    v_elem text;
    v_norm_elems text[];
    v_norm_tuple text;
    v_elem_count int;
    i int;
BEGIN
    /*
      Resolve schema/table name.
    */
    SELECT n.nspname, c.relname
    INTO v_schema_name, v_table_name
    FROM pg_class c
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE c.oid = p_table;

    IF v_schema_name IS NULL THEN
        RAISE EXCEPTION 'Table % not found', p_table;
    END IF;

    /*
      Get primary key definition.
    */
    SELECT pg_get_indexdef(idx.indexrelid)
    INTO v_indexdef
    FROM pg_index idx
    WHERE idx.indrelid = p_table
      AND idx.indisprimary;

    IF v_indexdef IS NULL THEN
        RAISE EXCEPTION 'Table % has no primary key', p_table;
    END IF;

    v_is_hash := v_indexdef ~* '\sHASH';

    /*
      Get primary key columns in PK order.
    */
    SELECT array_agg(a.attname ORDER BY k.ord)
    INTO v_pk_cols
    FROM pg_index idx
    CROSS JOIN LATERAL generate_subscripts(idx.indkey, 1) AS k(ord)
    JOIN pg_attribute a
      ON a.attrelid = idx.indrelid
     AND a.attnum = idx.indkey[k.ord]
    WHERE idx.indrelid = p_table
      AND idx.indisprimary
    GROUP BY idx.indexrelid;

    IF v_pk_cols IS NULL THEN
        RAISE EXCEPTION 'Could not determine primary key columns for %', p_table;
    END IF;

    SELECT '(' || string_agg(format('%I', x.col), ', ') || ')'
    INTO v_key_tuple
    FROM unnest(v_pk_cols) AS x(col);

    /*
      ----------------------------------------------------------------------
      HASH-SHARDED TABLE
      ----------------------------------------------------------------------
    */
    IF v_is_hash THEN

        /*
          Extract the HASH key expression.

          Handles:
            lsm (c1 HASH)
            lsm ((c1, c2) HASH, c3 ASC)
        */
        v_hash_cols_text := substring(v_indexdef FROM '\(\((.*?)\)\s+HASH');

        IF v_hash_cols_text IS NULL THEN
            v_hash_cols_text := substring(v_indexdef FROM '\(([[:alnum:]_".]+)\s+HASH');
        END IF;

        IF v_hash_cols_text IS NULL THEN
            RAISE EXCEPTION 'Could not extract HASH columns from primary key definition: %', v_indexdef;
        END IF;

        v_hash_expr := format('yb_hash_code(%s)', v_hash_cols_text);
        v_distinct_expr := format('(%s)', v_hash_cols_text);

        v_tablet_ordinal := 0;

        FOR r IN
            SELECT
                m.tablet_id,
                m.start_hash_code,
                m.end_hash_code,
                m.leader,
                m.replicas
            FROM yb_tablet_metadata m
            WHERE m.oid = p_table::oid
            ORDER BY m.start_hash_code
        LOOP
            v_tablet_ordinal := v_tablet_ordinal + 1;

            v_branch_sql := format(
                $fmt$
                SELECT
                    %s::int AS tablet_ordinal,
                    %L::text AS tablet_id,
                    'hash'::text AS partitioning_scheme,
                    %s::int AS start_hash_code,
                    %s::int AS end_hash_code,
                    NULL::jsonb AS start_range_value,
                    NULL::jsonb AS end_range_value,
                    NULL::text AS split_range,
                    %L::text AS leader,
                    %L::text[] AS replicas,
                    COUNT(*)::bigint AS row_count,
                    COUNT(DISTINCT %s)::bigint AS distinct_key_count
                FROM %s
                WHERE %s >= %s
                  AND %s <  %s
                $fmt$,
                v_tablet_ordinal,
                r.tablet_id,
                r.start_hash_code,
                r.end_hash_code,
                r.leader,
                r.replicas::text,
                v_distinct_expr,
                p_table,
                v_hash_expr,
                r.start_hash_code,
                v_hash_expr,
                r.end_hash_code
            );

            v_union_sql :=
                v_union_sql ||
                CASE WHEN v_union_sql = '' THEN '' ELSE E'\nUNION ALL\n' END ||
                v_branch_sql;
        END LOOP;

    /*
      ----------------------------------------------------------------------
      RANGE-SHARDED TABLE
      ----------------------------------------------------------------------

      Uses YugabyteDB's reconstructed SPLIT AT VALUES clause instead of
      decoding yb_local_tablets().partition_key_start/end bytea values.
    */
    ELSE
        v_split_clause := NULL;

        BEGIN
            SELECT yb_get_range_split_clause(p_table)
            INTO v_split_clause;
        EXCEPTION
            WHEN OTHERS THEN
                v_split_clause := NULL;
        END;

        IF v_split_clause IS NULL OR btrim(v_split_clause) = '' THEN
            RAISE EXCEPTION
                'Could not get SPLIT AT VALUES metadata for range-sharded table %. This function currently requires declared range splits.',
                p_table;
        END IF;

        v_split_clause := btrim(v_split_clause);
        v_split_clause := regexp_replace(v_split_clause, ';\s*$', '');

        /*
          Strip the leading "SPLIT AT VALUES".

          Example:

            SPLIT AT VALUES (('4000...'), ('8000...'))

          Becomes:

            (('4000...'), ('8000...'))
        */
        v_values_text :=
            btrim(
                regexp_replace(
                    v_split_clause,
                    '^\s*SPLIT\s+AT\s+VALUES\s*',
                    '',
                    'i'
                )
            );

        /*
          Remove the outer wrapper parentheses.

          Example:

            (('4000...'), ('8000...'))

          Becomes:

            ('4000...'), ('8000...')
        */
        IF left(v_values_text, 1) = '('
           AND right(v_values_text, 1) = ')' THEN
            v_inner := substring(v_values_text FROM 2 FOR length(v_values_text) - 2);
        ELSE
            RAISE EXCEPTION 'Could not parse SPLIT AT VALUES clause: %', v_split_clause;
        END IF;

        /*
          Parse top-level split tuples:

            (250000), (500000), (750000)
            ('uuid', MINVALUE), ('uuid2', MINVALUE)
        */
        v_split_values := ARRAY[]::text[];
        v_current := '';
        v_depth := 0;
        v_in_quote := false;
        v_pos := 1;
        v_len := length(v_inner);

        WHILE v_pos <= v_len LOOP
            v_ch := substr(v_inner, v_pos, 1);
            v_next_ch :=
                CASE
                    WHEN v_pos < v_len THEN substr(v_inner, v_pos + 1, 1)
                    ELSE NULL
                END;

            IF v_ch = '''' THEN
                v_current := v_current || v_ch;

                IF v_in_quote AND v_next_ch = '''' THEN
                    /*
                      Escaped quote inside SQL literal.
                    */
                    v_current := v_current || v_next_ch;
                    v_pos := v_pos + 2;
                    CONTINUE;
                ELSE
                    v_in_quote := NOT v_in_quote;
                    v_pos := v_pos + 1;
                    CONTINUE;
                END IF;
            END IF;

            IF NOT v_in_quote THEN
                IF v_ch = '(' THEN
                    v_depth := v_depth + 1;
                    v_current := v_current || v_ch;
                    v_pos := v_pos + 1;
                    CONTINUE;
                ELSIF v_ch = ')' THEN
                    v_depth := v_depth - 1;
                    v_current := v_current || v_ch;
                    v_pos := v_pos + 1;
                    CONTINUE;
                ELSIF v_ch = ',' AND v_depth = 0 THEN
                    IF btrim(v_current) <> '' THEN
                        v_split_values := array_append(v_split_values, btrim(v_current));
                    END IF;

                    v_current := '';
                    v_pos := v_pos + 1;
                    CONTINUE;
                END IF;
            END IF;

            v_current := v_current || v_ch;
            v_pos := v_pos + 1;
        END LOOP;

        IF btrim(v_current) <> '' THEN
            v_split_values := array_append(v_split_values, btrim(v_current));
        END IF;

        IF array_length(v_split_values, 1) IS NULL THEN
            RAISE EXCEPTION 'No split values parsed from clause: %', v_split_clause;
        END IF;

        /*
          Normalize each split tuple by stripping trailing MINVALUE/MAXVALUE.

          Example:

            ('4000...', MINVALUE)

          becomes:

            ('4000...')

          This allows predicates to use only the leading range key columns.
        */
        v_normalized_split_values := ARRAY[]::text[];
        v_range_key_col_count := NULL;

        FOREACH v_split_value IN ARRAY v_split_values LOOP
            v_split_value := btrim(v_split_value);

            IF left(v_split_value, 1) = '('
               AND right(v_split_value, 1) = ')' THEN
                v_tuple_inner := substring(v_split_value FROM 2 FOR length(v_split_value) - 2);
            ELSE
                RAISE EXCEPTION 'Could not parse split tuple: %', v_split_value;
            END IF;

            /*
              Parse elements inside one split tuple.

              Example:

                '4000...', MINVALUE

              becomes array:

                ['4000...', 'MINVALUE']
            */
            v_elem_values := ARRAY[]::text[];
            v_elem_current := '';
            v_depth := 0;
            v_in_quote := false;
            v_pos := 1;
            v_len := length(v_tuple_inner);

            WHILE v_pos <= v_len LOOP
                v_ch := substr(v_tuple_inner, v_pos, 1);
                v_next_ch :=
                    CASE
                        WHEN v_pos < v_len THEN substr(v_tuple_inner, v_pos + 1, 1)
                        ELSE NULL
                    END;

                IF v_ch = '''' THEN
                    v_elem_current := v_elem_current || v_ch;

                    IF v_in_quote AND v_next_ch = '''' THEN
                        v_elem_current := v_elem_current || v_next_ch;
                        v_pos := v_pos + 2;
                        CONTINUE;
                    ELSE
                        v_in_quote := NOT v_in_quote;
                        v_pos := v_pos + 1;
                        CONTINUE;
                    END IF;
                END IF;

                IF NOT v_in_quote THEN
                    IF v_ch = '(' THEN
                        v_depth := v_depth + 1;
                        v_elem_current := v_elem_current || v_ch;
                        v_pos := v_pos + 1;
                        CONTINUE;
                    ELSIF v_ch = ')' THEN
                        v_depth := v_depth - 1;
                        v_elem_current := v_elem_current || v_ch;
                        v_pos := v_pos + 1;
                        CONTINUE;
                    ELSIF v_ch = ',' AND v_depth = 0 THEN
                        v_elem_values := array_append(v_elem_values, btrim(v_elem_current));
                        v_elem_current := '';
                        v_pos := v_pos + 1;
                        CONTINUE;
                    END IF;
                END IF;

                v_elem_current := v_elem_current || v_ch;
                v_pos := v_pos + 1;
            END LOOP;

            IF btrim(v_elem_current) <> '' THEN
                v_elem_values := array_append(v_elem_values, btrim(v_elem_current));
            END IF;

            /*
              Keep values until MINVALUE/MAXVALUE.

              Example:

                ('4000...', MINVALUE)

              becomes:

                ('4000...')
            */
            v_norm_elems := ARRAY[]::text[];

            FOREACH v_elem IN ARRAY v_elem_values LOOP
                IF upper(btrim(v_elem)) IN ('MINVALUE', 'MAXVALUE') THEN
                    EXIT;
                END IF;

                v_norm_elems := array_append(v_norm_elems, btrim(v_elem));
            END LOOP;

            v_elem_count := COALESCE(array_length(v_norm_elems, 1), 0);

            IF v_elem_count = 0 THEN
                RAISE EXCEPTION 'Split tuple % did not contain any usable values before MINVALUE/MAXVALUE', v_split_value;
            END IF;

            IF v_range_key_col_count IS NULL THEN
                v_range_key_col_count := v_elem_count;
            ELSIF v_range_key_col_count <> v_elem_count THEN
                RAISE EXCEPTION
                    'Inconsistent split tuple widths after MINVALUE/MAXVALUE normalization. Expected %, found % in %',
                    v_range_key_col_count,
                    v_elem_count,
                    v_split_value;
            END IF;

            IF v_range_key_col_count > array_length(v_pk_cols, 1) THEN
                RAISE EXCEPTION
                    'Split tuple has more columns (%) than primary key columns (%)',
                    v_range_key_col_count,
                    array_length(v_pk_cols, 1);
            END IF;

            v_norm_tuple := '(' || array_to_string(v_norm_elems, ', ') || ')';
            v_normalized_split_values := array_append(v_normalized_split_values, v_norm_tuple);
        END LOOP;

        /*
          Build the tuple of leading PK columns used by the split values.

          For:

            PRIMARY KEY (customer_id ASC, order_id ASC)
            SPLIT AT VALUES (('uuid', MINVALUE))

          this becomes:

            (customer_id)
        */
        SELECT '(' || string_agg(format('%I', v_pk_cols[g.col_num]), ', ' ORDER BY g.col_num) || ')'
        INTO v_range_key_tuple
        FROM generate_series(1, v_range_key_col_count) AS g(col_num);

        SELECT COUNT(*)
        INTO v_tablet_count
        FROM yb_local_tablets()
        WHERE table_name = v_table_name;

        IF v_tablet_count <> array_length(v_normalized_split_values, 1) + 1 THEN
            RAISE NOTICE
                'Range tablet count (%) does not equal parsed split count + 1 (%). Counts may not reflect the current physical tablet layout.',
                v_tablet_count,
                array_length(v_normalized_split_values, 1) + 1;
        END IF;

        v_range_idx := 0;

        /*
          The ordinal is assigned in physical range order based on
          partition_key_start. The caller should ORDER BY tablet_ordinal.
        */
        FOR r IN
            SELECT
                l.tablet_id,
                l.partition_key_start,
                m.leader,
                m.replicas
            FROM yb_local_tablets() l
            LEFT JOIN yb_tablet_metadata m
              ON m.tablet_id = l.tablet_id
             AND m.oid = p_table::oid
            WHERE l.table_name = v_table_name
            ORDER BY l.partition_key_start NULLS FIRST
        LOOP
            v_range_idx := v_range_idx + 1;

            v_start_value :=
                CASE
                    WHEN v_range_idx = 1 THEN NULL
                    ELSE v_normalized_split_values[v_range_idx - 1]
                END;

            v_end_value :=
                CASE
                    WHEN v_range_idx <= array_length(v_normalized_split_values, 1)
                    THEN v_normalized_split_values[v_range_idx]
                    ELSE NULL
                END;

            v_predicate := 'TRUE';

            IF v_start_value IS NOT NULL THEN
                v_predicate :=
                    v_predicate || format(' AND %s >= %s', v_range_key_tuple, v_start_value);
            END IF;

            IF v_end_value IS NOT NULL THEN
                v_predicate :=
                    v_predicate || format(' AND %s < %s', v_range_key_tuple, v_end_value);
            END IF;

            v_branch_sql := format(
                $fmt$
                SELECT
                    %s::int AS tablet_ordinal,
                    %L::text AS tablet_id,
                    'range'::text AS partitioning_scheme,
                    NULL::int AS start_hash_code,
                    NULL::int AS end_hash_code,
                    %s::jsonb AS start_range_value,
                    %s::jsonb AS end_range_value,
                    %L::text AS split_range,
                    %L::text AS leader,
                    %L::text[] AS replicas,
                    COUNT(*)::bigint AS row_count,
                    COUNT(DISTINCT %s)::bigint AS distinct_key_count
                FROM %s
                WHERE %s
                $fmt$,
                v_range_idx,
                r.tablet_id,
                CASE
                    WHEN v_start_value IS NULL THEN 'NULL'
                    ELSE quote_literal(jsonb_build_object('split_at_value', v_start_value)::text)
                END,
                CASE
                    WHEN v_end_value IS NULL THEN 'NULL'
                    ELSE quote_literal(jsonb_build_object('split_at_value', v_end_value)::text)
                END,
                format(
                    '%s .. %s',
                    COALESCE(v_start_value, '-inf'),
                    COALESCE(v_end_value, '+inf')
                ),
                r.leader,
                r.replicas::text,
                v_range_key_tuple,
                p_table,
                v_predicate
            );

            v_union_sql :=
                v_union_sql ||
                CASE WHEN v_union_sql = '' THEN '' ELSE E'\nUNION ALL\n' END ||
                v_branch_sql;
        END LOOP;
    END IF;

    IF v_union_sql = '' THEN
        RAISE EXCEPTION 'No tablet branches generated for table %', p_table;
    END IF;

    RETURN QUERY EXECUTE v_union_sql;
END;
$$;