#!/bin/bash
# Benchmarks TOAST compression methods, master against v4 of "ZSTD TOAST
# compression, and an extensible compression method encoding".
#   bench.sh BUILD METHOD [SIZES]
#     BUILD   base | v4           (see build.sh)
#     METHOD  pglz | lz4 | zstd   (zstd only exists in v4)
#     SIZES   value sizes in bytes (default: 4000 32000 1000000)
# One fresh cluster per run.  Three corpora, the same bytes for every build
# and method:
#   docs    the PostgreSQL SGML documentation, cut into values of SIZE bytes
#   json    generated JSON documents with realistic, repetitive keys
#   random  base64 of random bytes (setseed), nearly incompressible
# Storage is the default (EXTENDED), so small values are compressed inline
# and large ones compressed and moved out of line.  Prints one "key=value"
# line per measurement.
set -eu
BUILD=$1 METHOD=$2
SIZES=${3:-4000 32000 1000000}
REPS=${REPS:-5}
W=$HOME/pgzstd; B=$W/i-$BUILD/bin
DOCS=$W/src-base/doc/src/sgml
D=$(mktemp -d /tmp/claude-1000/zs.XXXX); P=$((56300 + RANDOM % 100))
# SQL_ASCII and C: substr() then works on bytes and fetches only the chunks
# it needs from the 10 MB source; compression sees the same bytes either way.
"$B/initdb" -D $D -A trust --no-sync -U postgres -E SQL_ASCII --locale=C >/dev/null
cat >> $D/postgresql.conf <<EOF
port = $P
unix_socket_directories = '/tmp'
shared_buffers = 2GB
max_wal_size = 32GB
checkpoint_timeout = 30min
autovacuum = off
jit = off
default_toast_compression = $METHOD
EOF
"$B/pg_ctl" -D $D -l $D/log -w start >/dev/null
q()  { "$B/psql" -X -qAt -h /tmp -p $P -U postgres -d postgres -v ON_ERROR_STOP=1 "$@"; }
ms() { python3 -c "import time; print(int(time.time()*1000))"; }
q -c "CREATE EXTENSION pg_prewarm"

# The documentation as one text, about 10 MB (file order is fixed).
cat $(ls $DOCS/*.sgml $DOCS/ref/*.sgml | sort) > $D/docs.txt
# Stored uncompressed, so substr() fetches only the chunks it needs instead
# of decompressing all 10 MB for every value.
q -c "CREATE TABLE docs_src (t text)"
q -c "ALTER TABLE docs_src ALTER COLUMN t SET STORAGE EXTERNAL"
q -c "INSERT INTO docs_src SELECT pg_read_file('$D/docs.txt')"
DOCS_LEN=$(q -c "SELECT length(t) FROM docs_src")

out() { echo "build=$BUILD method=$METHOD corpus=$CORPUS size=$SIZE $*"; }

for SIZE in $SIZES; do
  # About 200 MB of raw data per corpus and size.
  N=$(( 200000000 / SIZE )); [ $N -gt 50000 ] && N=50000
  for CORPUS in docs json random; do
    case $CORPUS in
      docs)   src="SELECT g, substr((SELECT t FROM docs_src), 1 + (g * 7919) % $((DOCS_LEN - SIZE)), $SIZE)
                     FROM generate_series(1, $N) g" ;;
      json)   src="SELECT g, left(string_agg(format(
                       '{\"id\": %s, \"customer\": {\"name\": \"customer %s\", \"tier\": \"%s\"}, \"status\": \"%s\", \"items\": [{\"sku\": \"SKU-%s\", \"qty\": %s, \"price\": %s}], \"tags\": [\"%s\", \"%s\"]}',
                       g * 1000 + i, (g * 31 + i) % 5000, (ARRAY['gold','silver','bronze'])[1 + i % 3],
                       (ARRAY['paid','shipped','pending','refunded'])[1 + (g + i) % 4],
                       (g * 17 + i) % 100000, 1 + i % 9, round(((g * 13 + i) % 100000) / 100.0, 2),
                       (ARRAY['web','store','app'])[1 + g % 3], (ARRAY['promo','none'])[1 + i % 2]), ','), $SIZE)
                     FROM generate_series(1, $N) g,
                          generate_series(1, $SIZE / 150 + 1) i
                    GROUP BY g" ;;
      random) src="SELECT g, left(string_agg(md5(random()::text) || encode(sha256(random()::text::bytea), 'base64'), ''), $SIZE)
                     FROM generate_series(1, $N) g, generate_series(1, $SIZE / 76 + 1) i
                    GROUP BY g" ;;
    esac
    q -c "DROP TABLE IF EXISTS t, raw"
    q -c "SELECT setseed(0.42)"
    q -c "CREATE UNLOGGED TABLE raw (id int, v text)"
    q -c "ALTER TABLE raw ALTER COLUMN v SET STORAGE EXTERNAL"
    q -c "INSERT INTO raw $src"
    q -c "CREATE TABLE t (id int PRIMARY KEY, v text)"
    q -c "SELECT pg_prewarm('raw'), pg_prewarm((SELECT reltoastrelid FROM pg_class WHERE relname = 'raw'))" >/dev/null
    q -c "CHECKPOINT"

    # 1. Load: compression happens here.
    lsn0=$(q -c "SELECT pg_current_wal_insert_lsn()")
    t0=$(ms); q -c "INSERT INTO t SELECT id, v FROM raw"; t1=$(ms)
    wal=$(q -c "SELECT pg_wal_lsn_diff(pg_current_wal_insert_lsn(), '$lsn0')")
    sizes=$(q -c "SELECT format('rows=%s raw_mb=%s heap_kb=%s toast_kb=%s compressed_pct=%s',
                   count(*), round(sum(octet_length(v)) / 1e6, 1),
                   pg_relation_size('t') / 1024,
                   pg_relation_size((SELECT reltoastrelid FROM pg_class WHERE relname = 't')) / 1024,
                   round(100.0 * count(*) FILTER (WHERE pg_column_compression(v) IS NOT NULL) / count(*), 1))
                  FROM t")
    out "op=load ms=$((t1 - t0)) wal_bytes=$wal $sizes"
    q -c "VACUUM ANALYZE t"; q -c "CHECKPOINT"
    q -c "SELECT pg_prewarm('t'), pg_prewarm((SELECT reltoastrelid FROM pg_class WHERE relname = 't'))" >/dev/null

    # 2. Full read, decompressing every value (hot).
    for rep in $(seq $REPS); do
      t0=$(ms); q -c "SELECT sum(length(v || 'x')) FROM t" >/dev/null; t1=$(ms)
      out "op=read_all rep=$rep ms=$((t1 - t0))"
    done
    # 3. A 100-byte slice from the middle of every value (hot).
    for rep in $(seq $REPS); do
      t0=$(ms); q -c "SELECT sum(length(substr(v, $SIZE / 2, 100))) FROM t" >/dev/null; t1=$(ms)
      out "op=slice rep=$rep ms=$((t1 - t0))"
    done
    # 4. A 100-byte prefix of every value (hot): pglz can stop early here.
    for rep in $(seq $REPS); do
      t0=$(ms); q -c "SELECT sum(length(substr(v, 1, 100))) FROM t" >/dev/null; t1=$(ms)
      out "op=prefix rep=$rep ms=$((t1 - t0))"
    done
  done
done
"$B/pg_ctl" -D $D -m fast -w stop >/dev/null
rm -rf $D