Faster exact row counts in YugabyteDB using range sharding and tablet-level parallelism

				
					set -euo pipefail

# yb_range_row_count.sh
# Exact row count for a RANGE-sharded YSQL table by counting each tablet key-range in parallel.
#
# Requires yb_tablet_ranges(regclass) from:
# https://yugabytedb.tips/decode-yb_local_tablets-split-ranges-into-human-readable-bounds/

usage() {
  cat <<'EOF'
Usage:
  ./yb_range_row_count.sh -t schema.table -h host -d db -U user [options]

Required:
  -t  schema.table
  -h  host
  -d  database
  -U  user

Options:
  -p  port            (default: 5433)
  -P  parallelism     (default: 8)
  --ysqlsh PATH       (default: ysqlsh)
  --where SQL         (extra predicate appended with AND)
  --quiet             print ONLY total row count (machine-friendly)
  --time-count        print count-phase elapsed time (ms) to stderr
  --no-validate       skip PK/type/HASH checks (faster setup, less safety)

Examples:
  PGPASSWORD=yugabyte ./yb_range_row_count.sh -t public.test -h 127.0.0.1 -d yugabyte -U yugabyte
  PGPASSWORD=yugabyte ./yb_range_row_count.sh -t public.test -h 127.0.0.1 -d yugabyte -U yugabyte --time-count -P 10
  PGPASSWORD=yugabyte ./yb_range_row_count.sh -t public.test -h 127.0.0.1 -d yugabyte -U yugabyte --quiet --time-count -P 10
EOF
}

# Defaults
PORT="5433"
PARALLELISM="8"
YSQLSH_BIN="ysqlsh"
TABLE=""
HOST=""
DB=""
DB_USER=""
EXTRA_WHERE=""
QUIET="0"
TIME_COUNT="0"
NO_VALIDATE="0"

while [[ $# -gt 0 ]]; do
  case "$1" in
    -t) TABLE="${2:-}"; shift 2;;
    -h) HOST="${2:-}"; shift 2;;
    -p) PORT="${2:-}"; shift 2;;
    -d) DB="${2:-}"; shift 2;;
    -U) DB_USER="${2:-}"; shift 2;;
    -P) PARALLELISM="${2:-}"; shift 2;;
    --ysqlsh) YSQLSH_BIN="${2:-}"; shift 2;;
    --where) EXTRA_WHERE="${2:-}"; shift 2;;
    --quiet) QUIET="1"; shift 1;;
    --time-count) TIME_COUNT="1"; shift 1;;
    --no-validate) NO_VALIDATE="1"; shift 1;;
    -?|--help) usage; exit 0;;
    *) echo "Unknown arg: $1" >&2; usage; exit 2;;
  esac
done

if [[ -z "$TABLE" || -z "$HOST" || -z "$DB" || -z "$DB_USER" ]]; then
  usage
  exit 2
fi

if [[ "$TABLE" != *.* ]]; then
  echo "ERROR: -t must be schema.table (got: $TABLE)" >&2
  exit 2
fi

SCHEMA="${TABLE%%.*}"
TNAME="${TABLE#*.}"

YSQLSH_FLAGS=(-h "$HOST" -p "$PORT" -U "$DB_USER" -d "$DB" -X -q -A -t -v ON_ERROR_STOP=1)
y() { "$YSQLSH_BIN" "${YSQLSH_FLAGS[@]}" -c "$1"; }

TMPDIR="$(mktemp -d)"
cleanup() { rm -rf "$TMPDIR"; }
trap cleanup EXIT

# Minimal checks (fast)
tc="$(y "SELECT COUNT(*) FROM pg_catalog.pg_tables WHERE schemaname='${SCHEMA}' AND tablename='${TNAME}';" | tr -d '[:space:]')"
if [[ "$tc" != "1" ]]; then
  echo "ERROR: Table not found: ${TABLE}" >&2
  exit 1
fi

has_fn="$(y "SELECT COUNT(*) FROM pg_proc WHERE proname='yb_tablet_ranges';" | tr -d '[:space:]')"
if [[ "$has_fn" == "0" ]]; then
  echo "ERROR: Missing function yb_tablet_ranges(regclass). See decode tip." >&2
  exit 1
fi

# Find PK column name (single-col only). Keep this even in --no-validate because we need the json key.
pk_cols="$(
  y "
    SELECT a.attname
    FROM pg_index i
    JOIN pg_class c ON c.oid=i.indrelid
    JOIN pg_namespace n ON n.oid=c.relnamespace
    JOIN pg_attribute a ON a.attrelid=c.oid AND a.attnum = ANY(i.indkey)
    WHERE n.nspname='${SCHEMA}' AND c.relname='${TNAME}' AND i.indisprimary
    ORDER BY array_position(i.indkey, a.attnum)
    LIMIT 2;
  "
)"
pk_cols_count="$(printf "%s\n" "$pk_cols" | sed '/^\s*$/d' | wc -l | tr -d '[:space:]')"
if [[ "$pk_cols_count" != "1" ]]; then
  echo "ERROR: single-column PRIMARY KEY required (found ${pk_cols_count})." >&2
  exit 1
fi
PK_COL="$(printf "%s\n" "$pk_cols" | head -n1 | tr -d '[:space:]')"

pk_type=""
if [[ "$NO_VALIDATE" == "0" ]]; then
  pk_type="$(y "
    SELECT format_type(a.atttypid, a.atttypmod)
    FROM pg_attribute a
    JOIN pg_class c ON c.oid=a.attrelid
    JOIN pg_namespace n ON n.oid=c.relnamespace
    WHERE n.nspname='${SCHEMA}' AND c.relname='${TNAME}' AND a.attname='${PK_COL}';
  " | tr -d '[:space:]')"
  if [[ "$pk_type" != "bigint" && "$pk_type" != "int8" ]]; then
    echo "ERROR: PK column ${PK_COL} must be BIGINT/int8 for this script (found ${pk_type})." >&2
    exit 1
  fi

  pk_def="$(y "
    SELECT pg_get_indexdef(i.indexrelid)
    FROM pg_index i
    JOIN pg_class c ON c.oid=i.indrelid
    JOIN pg_namespace n ON n.oid=c.relnamespace
    WHERE n.nspname='${SCHEMA}' AND c.relname='${TNAME}' AND i.indisprimary;
  ")"
  if echo "$pk_def" | grep -qi "HASH"; then
    echo "ERROR: table appears HASH-sharded; this script is for RANGE PKs." >&2
    exit 1
  fi
fi

# Pull ranges (tablet_id|start|end|split_range)
ranges_file="${TMPDIR}/ranges.txt"
y "
  SELECT
    tablet_id || '|' ||
    COALESCE(range_start->>'${PK_COL}', '') || '|' ||
    COALESCE(range_end->>'${PK_COL}', '')   || '|' ||
    split_range
  FROM yb_tablet_ranges('${SCHEMA}.${TNAME}'::regclass)
  ORDER BY split_range;
" > "$ranges_file"

num_ranges="$(grep -c '|' "$ranges_file" || true)"
if [[ "$num_ranges" -le 0 ]]; then
  echo "ERROR: No ranges found via yb_tablet_ranges()." >&2
  exit 1
fi

# Verbose header
if [[ "$QUIET" == "0" ]]; then
  echo "Table: ${TABLE}"
  if [[ -n "$pk_type" ]]; then
    echo "PK:    ${PK_COL} (${pk_type})"
  else
    echo "PK:    ${PK_COL}"
  fi
  echo "Host:  ${HOST}:${PORT}  DB: ${DB}  User: ${DB_USER}"
  echo "Par:   ${PARALLELISM}"
  if [[ -n "$EXTRA_WHERE" ]]; then
    echo "Where: ${EXTRA_WHERE}"
  fi
  echo
  echo "Ranges: ${num_ranges}"
  if [[ "$num_ranges" -eq 1 ]]; then
    echo "NOTE: Only 1 tablet range found (-inf..+inf). Nothing to parallelize."
  fi
  echo
fi

# Worker: outputs either "count" (quiet) or "tablet_id|count|split_range" (verbose)
export YSQLSH_BIN HOST PORT DB DB_USER TABLE PK_COL EXTRA_WHERE QUIET
count_one() {
  local line="$1"
  local tid start end split
  IFS="|" read -r tid start end split <<< "$line"

  local w="TRUE"
  [[ -n "$start" ]] && w="${PK_COL} >= ${start}"
  [[ -n "$end"   ]] && w="${w} AND ${PK_COL} < ${end}"
  [[ -n "$EXTRA_WHERE" ]] && w="${w} AND (${EXTRA_WHERE})"

  if [[ "$QUIET" == "1" ]]; then
    # Return just the count (one number)
    "$YSQLSH_BIN" -h "$HOST" -p "$PORT" -U "$DB_USER" -d "$DB" -X -q -A -t -v ON_ERROR_STOP=1 \
      -c "SELECT COUNT(*) FROM ${TABLE} WHERE ${w};"
  else
    # Return tablet_id|count|split_range for breakdown
    "$YSQLSH_BIN" -h "$HOST" -p "$PORT" -U "$DB_USER" -d "$DB" -X -q -A -t -v ON_ERROR_STOP=1 \
      -c "SELECT '${tid}' AS tablet_id, COUNT(*) AS c, '${split}' AS split_range FROM ${TABLE} WHERE ${w};"
  fi
}
export -f count_one

# Time only the count phase if requested
start_ns=0
if [[ "$TIME_COUNT" == "1" ]]; then
  start_ns="$(date +%s%N)"
fi

if [[ "$QUIET" == "1" ]]; then
  # Run counts in parallel, sum raw numbers
  total="$(
    cat "$ranges_file" \
      | xargs -I {} -P "$PARALLELISM" bash -c 'count_one "$1"' _ {} \
      | awk '{s += $1} END {printf "%.0f\n", s}'
  )"
else
  # Run counts in parallel, capture breakdown, sum counts
  counts_file="${TMPDIR}/counts.txt"
  cat "$ranges_file" \
    | xargs -I {} -P "$PARALLELISM" bash -c 'count_one "$1"' _ {} \
    | tee "$counts_file" >/dev/null

  total="$(awk -F'|' '{s += $2} END {printf "%.0f\n", s}' "$counts_file")"
fi

if [[ "$TIME_COUNT" == "1" ]]; then
  end_ns="$(date +%s%N)"
  ms="$(( (end_ns - start_ns) / 1000000 ))"
  echo "count_phase_ms=${ms}" >&2
fi

# Output
if [[ "$QUIET" == "1" ]]; then
  echo "$total"
else
  echo
  echo "Total row count: ${total}"
  echo
  if [[ -n "${counts_file:-}" && -s "${counts_file:-}" ]]; then
    echo "Per-range breakdown (tablet_id | count | split_range):"
    column -t -s'|' "$counts_file" | sed 's/^/  /'
  fi
fi