#!/usr/bin/env bash
# run.sh — hstore vs jsonb in PostgreSQL, measured.
#
#   docker run -d --name hstore-bench -e POSTGRES_PASSWORD=x postgres:17
#   ./run.sh
#   docker rm -f hstore-bench
#
# Environment:
#   CONTAINER  container name             (default hstore-bench)
#   ROWS       rows per table             (default 1000000)
#   STEPS      which measurements to run  (default "1 2 3 4 5 7 8")
#
# Output: raw/<pgNN>/ (every psql session verbatim), results-<pgNN>.csv, medians-<pgNN>.txt
#
# Method
#   - Default server settings, nothing tuned. Data is generated once (setup.sql, setseed)
#     into a typed source table; the three document tables are loaded from it.
#   - Every timed statement runs in its own psql session with \timing (client-side
#     wall clock). Three runs per variant, interleaved: run 1 of every variant, then
#     run 2, ... so drift hits all variants alike. The median is reported.
#   - Queries get one untimed warm-up per variant first (warm cache).
#   - Step order: load (2) -> VACUUM ANALYZE -> storage (1) -> seq-scan reads (3, 7a)
#     -> GIN indexes (4) -> expression indexes (7b) -> drop secondary indexes ->
#     updates (5) -> type behaviour (8). Updates run last, without secondary indexes,
#     because a rolled-back update leaves dead tuples behind.
set -euo pipefail
export LC_ALL=C
cd "$(dirname "$0")"

C=${CONTAINER:-hstore-bench}
ROWS=${ROWS:-1000000}
STEPS=${STEPS:-"1 2 3 4 5 7 8"}
PSQL=(docker exec -i "$C" psql -U postgres -X -q -v ON_ERROR_STOP=1)
has() { [[ " $STEPS " == *" $1 "* ]]; }

until docker exec "$C" pg_isready -U postgres -q; do sleep 1; done
sleep 2   # the image restarts the server once after initdb

TAG=pg$("${PSQL[@]}" -Atc "select current_setting('server_version_num')::int / 10000")
RAW=raw/$TAG
CSV=results-$TAG.csv
mkdir -p "$RAW"
echo "version,step,variant,run,ms" > "$CSV"

"${PSQL[@]}" -Atc "select version()" | tee "$RAW/version.txt"
"${PSQL[@]}" -c "select name, setting, unit from pg_settings
  where name in ('shared_buffers','work_mem','maintenance_work_mem','max_parallel_workers_per_gather',
                 'max_parallel_maintenance_workers','jit','effective_cache_size','random_page_cost')
  order by name" > "$RAW/settings.txt"
echo "setup ($ROWS rows)..."
"${PSQL[@]}" -v rows="$ROWS" < setup.sql > "$RAW/setup.out" 2>&1

# timed <step> <variant> <run> <sql before> <timed sql: one statement> [<sql after>]
timed() {
  local out="$RAW/${1}_${2}_run${3}.txt"
  "${PSQL[@]}" > "$out" 2>&1 <<SQL
$4
\\timing on
$5
\\timing off
${6:-}
SQL
  local ms; ms=$(grep -m1 '^Time:' "$out" | awk '{print $2}')
  echo "$TAG,$1,$2,$3,$ms" >> "$CSV"
  printf '%-14s %-16s run%s %10s ms\n' "$1" "$2" "$3" "$ms"
}
# sql <file> <sql>: untimed session, output kept
sql() { "${PSQL[@]}" -v ON_ERROR_STOP=0 > "$RAW/$1.txt" 2>&1 <<< "$2" || true; }

# ---------------------------------------------------------------- 2. load
declare -A LOAD
LOAD[h]="insert into docs_h
select id, hstore(array['color','size','sku','note','brand','qty','stock','rank','price','active'],
                  array[color, size, sku, note, brand, qty::text, stock::text, rank::text, price::text, active::text])
           || case when rare is null then ''::hstore else hstore('rare_key', rare) end
  from src;"
LOAD[j]="insert into docs_j
select id, jsonb_build_object('color', color, 'size', size, 'sku', sku, 'note', note, 'brand', brand,
                              'qty', qty, 'stock', stock, 'rank', rank, 'price', price, 'active', active)
           || case when rare is null then '{}'::jsonb else jsonb_build_object('rare_key', rare) end
  from src;"
LOAD[jt]="insert into docs_j_text
select id, jsonb_build_object('color', color, 'size', size, 'sku', sku, 'note', note, 'brand', brand,
                              'qty', qty::text, 'stock', stock::text, 'rank', rank::text,
                              'price', price::text, 'active', active::text)
           || case when rare is null then '{}'::jsonb else jsonb_build_object('rare_key', rare) end
  from src;"
declare -A TBL=([h]=docs_h [j]=docs_j [jt]=docs_j_text)

if has 2; then
  for run in 1 2 3; do
    for v in h j jt; do
      timed 2_load "$v" "$run" "truncate ${TBL[$v]}; checkpoint;" "${LOAD[$v]}"
    done
  done
else
  for v in h j jt; do "${PSQL[@]}" -c "truncate ${TBL[$v]}" -c "${LOAD[$v]}"; done
fi
"${PSQL[@]}" -c "vacuum analyze docs_h, docs_j, docs_j_text" -c "checkpoint"
sql sample_rows "select id, attrs from docs_h where id in (1, 2);
select id, attrs from docs_j where id in (1, 2);
select id, attrs from docs_j_text where id in (1, 2);"

# ---------------------------------------------------------------- 1. storage
if has 1; then
  sql 1_storage "select t as table_name,
       pg_total_relation_size(t) as total_bytes, pg_size_pretty(pg_total_relation_size(t)) as total,
       pg_relation_size(t) as heap_bytes, pg_size_pretty(pg_relation_size(t)) as heap,
       coalesce(pg_total_relation_size(nullif(c.reltoastrelid, 0)), 0) as toast_bytes,
       pg_indexes_size(t) as index_bytes, c.relpages as heap_pages
  from unnest(array['docs_h','docs_j','docs_j_text']::regclass[]) t join pg_class c on c.oid = t;
select 'docs_h' as table_name, round(avg(pg_column_size(attrs)), 1) as avg_attrs_bytes,
       min(pg_column_size(attrs)), max(pg_column_size(attrs)) from docs_h
union all select 'docs_j', round(avg(pg_column_size(attrs)), 1), min(pg_column_size(attrs)), max(pg_column_size(attrs)) from docs_j
union all select 'docs_j_text', round(avg(pg_column_size(attrs)), 1), min(pg_column_size(attrs)), max(pg_column_size(attrs)) from docs_j_text;"
  cat "$RAW/1_storage.txt"
fi

# query runner: warm-up, 3 interleaved runs, one EXPLAIN (ANALYZE, BUFFERS) per variant
# queries <step> <variant>=<sql> ...
queries() {
  local step=$1; shift
  local -a names=() sqls=()
  for kv in "$@"; do names+=("${kv%%=*}"); sqls+=("${kv#*=}"); done
  for i in "${!names[@]}"; do
    "${PSQL[@]}" > "$RAW/${step}_${names[$i]}_warmup.txt" 2>&1 <<< "${sqls[$i]}" || true
  done
  for run in 1 2 3; do
    for i in "${!names[@]}"; do timed "$step" "${names[$i]}" "$run" "" "${sqls[$i]}"; done
  done
  for i in "${!names[@]}"; do
    sql "${step}_${names[$i]}_explain" "explain (analyze, buffers) ${sqls[$i]}"
  done
}

# ---------------------------------------------------------------- 3. key read, no index
if has 3; then
  queries 3_key_seqscan \
    "h=select count(*) from docs_h where attrs -> 'color' = 'red';" \
    "j=select count(*) from docs_j where attrs ->> 'color' = 'red';" \
    "jt=select count(*) from docs_j_text where attrs ->> 'color' = 'red';"
  # containment without an index, as the baseline for step 4
  queries 3b_contain_seqscan \
    "h=select count(*) from docs_h where attrs @> 'color=>red, size=>L';" \
    "j=select count(*) from docs_j where attrs @> '{\"color\":\"red\",\"size\":\"L\"}';" \
    "jt=select count(*) from docs_j_text where attrs @> '{\"color\":\"red\",\"size\":\"L\"}';"
fi

# ---------------------------------------------------------------- 7a. numeric filter, no index
if has 7; then
  queries 7a_num_seqscan \
    "h_int=select count(*) from docs_h where (attrs -> 'qty')::int > 899;" \
    "j_text_int=select count(*) from docs_j where (attrs ->> 'qty')::int > 899;" \
    "j_int=select count(*) from docs_j where (attrs -> 'qty')::int > 899;" \
    "j_numeric=select count(*) from docs_j where (attrs -> 'qty')::numeric > 899;" \
    "jt_text_int=select count(*) from docs_j_text where (attrs ->> 'qty')::int > 899;"
fi

# ---------------------------------------------------------------- 4. GIN
# build_index <step> <variant> <drop sql> <create sql>: 3 builds per variant, interleaved
build_indexes() {
  local step=$1; shift
  local -a names=() drops=() creates=()
  while (($#)); do names+=("$1"); drops+=("$2"); creates+=("$3"); shift 3; done
  for run in 1 2 3; do
    for i in "${!names[@]}"; do
      timed "$step" "${names[$i]}" "$run" "${drops[$i]} checkpoint;" "${creates[$i]}"
    done
  done
}
index_sizes() {
  sql "$1" "select indexrelid::regclass as index_name, indrelid::regclass as table_name,
       pg_relation_size(indexrelid) as bytes, pg_size_pretty(pg_relation_size(indexrelid)) as size
  from pg_index where indrelid in ('docs_h'::regclass, 'docs_j'::regclass, 'docs_j_text'::regclass)
 order by 2, 1;"
  cat "$RAW/$1.txt"
}
CONTAIN_H="attrs @> 'color=>red, size=>L'"
CONTAIN_J="attrs @> '{\"color\":\"red\",\"size\":\"L\"}'"

if has 4; then
  # phase A: hstore gin_hstore_ops, jsonb jsonb_ops (the defaults)
  build_indexes 4_build \
    h  "drop index if exists docs_h_gin;"      "create index docs_h_gin on docs_h using gin (attrs);" \
    j  "drop index if exists docs_j_gin;"      "create index docs_j_gin on docs_j using gin (attrs);" \
    jt "drop index if exists docs_j_text_gin;" "create index docs_j_text_gin on docs_j_text using gin (attrs);"
  "${PSQL[@]}" -c "analyze docs_h, docs_j, docs_j_text"
  index_sizes 4_index_sizes_a
  queries 4_contain_gin \
    "h=select count(*) from docs_h where $CONTAIN_H;" \
    "j=select count(*) from docs_j where $CONTAIN_J;" \
    "jt=select count(*) from docs_j_text where $CONTAIN_J;"
  queries 4_exists_gin \
    "h=select count(*) from docs_h where attrs ? 'rare_key';" \
    "j=select count(*) from docs_j where attrs ? 'rare_key';" \
    "jt=select count(*) from docs_j_text where attrs ? 'rare_key';"

  # phase B: jsonb_path_ops instead of jsonb_ops (hstore keeps its index, as reference)
  "${PSQL[@]}" -c "drop index docs_j_gin" -c "drop index docs_j_text_gin"
  build_indexes 4_build \
    j_po  "drop index if exists docs_j_gin_po;"      "create index docs_j_gin_po on docs_j using gin (attrs jsonb_path_ops);" \
    jt_po "drop index if exists docs_j_text_gin_po;" "create index docs_j_text_gin_po on docs_j_text using gin (attrs jsonb_path_ops);"
  "${PSQL[@]}" -c "analyze docs_j, docs_j_text"
  index_sizes 4_index_sizes_b
  queries 4_contain_gin \
    "j_po=select count(*) from docs_j where $CONTAIN_J;" \
    "jt_po=select count(*) from docs_j_text where $CONTAIN_J;"
  # jsonb_path_ops has no operator for ?: the planner cannot use the index
  queries 4_exists_gin \
    "j_po=select count(*) from docs_j where attrs ? 'rare_key';" \
    "jt_po=select count(*) from docs_j_text where attrs ? 'rare_key';"
  "${PSQL[@]}" -c "drop index docs_h_gin" -c "drop index docs_j_gin_po" -c "drop index docs_j_text_gin_po"
fi

# ---------------------------------------------------------------- 7b. expression index on a number
if has 7; then
  # hstore has no ->> operator; this shows what happens if you try
  sql 7b_hstore_arrow_text_index "create index docs_h_qty_txt on docs_h (((attrs ->> 'qty')::int));"
  build_indexes 7b_build \
    h_int      "drop index if exists docs_h_qty;"      "create index docs_h_qty on docs_h (((attrs -> 'qty')::int));" \
    j_text_int "drop index if exists docs_j_qty;"      "create index docs_j_qty on docs_j (((attrs ->> 'qty')::int));" \
    jt_text_int "drop index if exists docs_j_text_qty;" "create index docs_j_text_qty on docs_j_text (((attrs ->> 'qty')::int));"
  "${PSQL[@]}" -c "analyze docs_h, docs_j, docs_j_text"
  index_sizes 7b_index_sizes
  queries 7b_num_index \
    "h_int=select count(*) from docs_h where (attrs -> 'qty')::int between 500 and 509;" \
    "j_text_int=select count(*) from docs_j where (attrs ->> 'qty')::int between 500 and 509;" \
    "jt_text_int=select count(*) from docs_j_text where (attrs ->> 'qty')::int between 500 and 509;"
  # a query written with a different expression than the index
  sql 7b_j_other_expr_explain "explain (analyze, buffers) select count(*) from docs_j where (attrs -> 'qty')::int between 500 and 509;"
  "${PSQL[@]}" -c "drop index docs_h_qty" -c "drop index docs_j_qty" -c "drop index docs_j_text_qty"
fi

# ---------------------------------------------------------------- 5. update one key in 100k rows
if has 5; then
  for run in 1 2 3; do
    for v in h j_concat j_set jt_concat jt_set; do
      case $v in
        h)         t=docs_h;      expr="attrs || hstore('note', 'x')";;
        j_concat)  t=docs_j;      expr="attrs || '{\"note\":\"x\"}'";;
        j_set)     t=docs_j;      expr="jsonb_set(attrs, '{note}', '\"x\"')";;
        jt_concat) t=docs_j_text; expr="attrs || '{\"note\":\"x\"}'";;
        jt_set)    t=docs_j_text; expr="jsonb_set(attrs, '{note}', '\"x\"')";;
      esac
      timed 5_update "$v" "$run" "vacuum $t; checkpoint; begin;" \
        "update $t set attrs = $expr where id <= 100000;" "rollback;"
    done
  done
fi

# ---------------------------------------------------------------- 8. type behaviour
if has 8; then
  sql 8_types "\\x on
select attrs as hstore_value, hstore_to_jsonb(attrs), hstore_to_jsonb_loose(attrs) from docs_h where id = 1;
select attrs as jsonb_value from docs_j where id = 1;
\\x off
select hstore_to_jsonb_loose('active=>true, a=>t, b=>f, qty=>042, price=>123.40, n=>1e3, s=>NULL') as loose_examples;
select max(attrs -> 'stock') as hstore_max_text, max((attrs -> 'stock')::int) as hstore_max_int from docs_h;
select max(attrs ->> 'stock') as jsonb_max_text, max((attrs -> 'stock')::int) as jsonb_max_int from docs_j;
select max(attrs ->> 'stock') as jsonb_text_max_text, max((attrs ->> 'stock')::int) as jsonb_text_max_int from docs_j_text;
select jsonb_typeof(attrs -> 'stock') as docs_j, (select jsonb_typeof(attrs -> 'stock') from docs_j_text where id = 1) as docs_j_text from docs_j where id = 1;
select max(attrs -> 'stock') as jsonb_max_jsonb from docs_j;
select (attrs -> 'qty')::numeric from docs_j_text where id = 1;"
  cat "$RAW/8_types.txt"
fi

# ---------------------------------------------------------------- medians
sort -t, -k2,2 -k3,3 -k5,5g <(tail -n +2 "$CSV") | awk -F, '
  function flush() { if (n) printf "%-20s %-12s median %10.3f ms   runs %s\n", k2, k3, (n%2 ? a[(n+1)/2] : (a[n/2]+a[n/2+1])/2), runs }
  { if ($2 != k2 || $3 != k3) { flush(); n = 0; runs = "" ; k2 = $2; k3 = $3 }
    a[++n] = $5 }
  END { flush() }' > /tmp/medians.$$
# show runs in run order, not sorted order
while read -r step var _ med _ _ _; do
  runs=$(awk -F, -v s="$step" -v v="$var" '$2==s && $3==v {printf "%s%s", sep, $5; sep=" / "}' "$CSV")
  printf "%-20s %-12s median %10s ms   runs %s\n" "$step" "$var" "$med" "$runs"
done < /tmp/medians.$$ > "medians-$TAG.txt"
rm -f /tmp/medians.$$
cat "medians-$TAG.txt"