#!/usr/bin/env bash
# The first statement of a backend after its blocks left shared_buffers: how many blocks it reads from storage, in how many read requests, of what size.
# These are COUNTS (pg_stat_io reads and read_bytes, the blocks a segment read waited for against those read-ahead brought in); a tmpfs cluster gives
# counts and sequentiality, never device latency, and the milliseconds printed are secondary, not a proof of an IO gain.
#
#   PGHOST=.. PGPORT=.. bench/steps/first_statement_reads.sh [KEYS [ROWS_PER_KEY]] > records.tsv
#   LOAD=0 reuses the tables of an earlier run (the same rows for two builds of the library)
#
# Needs a scratch cluster with a small shared_buffers (16-64 MB), pg_buffercache and a pg_sift built with SIFT_TESTING (sift._test_segment_reads).
# One table of events in random insertion order with a column index cut by the key column, a copy with a b-tree on (dim_a_id, event_date), and a
# table of padding rows larger than shared_buffers whose bitmap scan evicts them. Three shapes (below) over key sets of 25, 200 and 2,000 keys in three regimes:
#   evict   the evicting scan ran first (everything is out, the catalogs too: the requests include their re-reads by the planner)
#   cold    the relation's own buffers are evicted (pg_buffercache_evict_relation)
#   steady  the statement ran a moment before
set -euo pipefail
ENT=${1:-20000} PER=${2:-50}
P="psql -X -q -At -v ON_ERROR_STOP=1"

[ "${LOAD:-1}" = 1 ] && $P <<SQL >/dev/null
set client_min_messages = warning;
create extension if not exists pg_buffercache; create extension if not exists pg_sift; create extension if not exists pg_sift_test;
drop table if exists events, events_n, bulk cascade;
create table events (id bigint generated by default as identity primary key, dim_a_id int not null, kind_id int not null, event_date date not null,
  amount float8 not null, label text not null, v1 float8, v2 float8, v3 float8, v4 float8, v5 float8, v6 float8, v7 float8, v8 float8, v9 float8,
  v10 float8, v11 float8, v12 float8, measures jsonb not null);
insert into events (dim_a_id, kind_id, event_date, amount, label, v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11, v12, measures)
select e, 1 + (e * 7 + g) % 19, date '2024-01-01' + (random() * 900)::int, round((random() * 100)::numeric, 2), 'label_' || (random() * 500)::int,
  random(), random(), random(), random(), random(), random(), random(), random(), random(), random(), random(), random(),
  jsonb_build_object('m1', round((random() * 10)::numeric, 1), 'm2', round((random() * 10)::numeric, 1), 'm3', (random() * 5)::int)
from generate_series(1, $ENT) e, generate_series(1, $PER) g order by random();
create table events_n as select * from events;
create index events_n_key on events_n (dim_a_id, event_date);
create index events_sift on events using sift_col (id, dim_a_id, kind_id, event_date, amount, label, v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11, v12, measures)
  with (key_cut = dim_a_id, paths = 'measures.* float8');
create table bulk (k int primary key, pad text);
insert into bulk select g, repeat('x', 900) from generate_series(1, 24 * (select setting::int from pg_settings where name = 'shared_buffers')) g;
vacuum (analyze) events, events_n, bulk;
SQL

# query SHAPE TABLE N -> the statement of a shape for N keys spread over the key column:
#   group      a key list, a date window, a group by (the units are read for the same columns)
#   selective  the same keys, a filter that keeps a few rows of a unit, then another column for them (not every unit reads it)
#   limit      the newest 20 rows of one key (a scan that stops early)
query() {
  local shape=$1 t=$2 n=$3 step=$((ENT / $3))
  local keys; keys=$(seq 1 "$step" "$ENT" | head -n "$n" | paste -sd, -)
  case $shape in
    group) echo "select kind_id, count(*), sum(amount) from $t where dim_a_id in ($keys) and event_date between '2024-03-01' and '2025-03-01' group by kind_id order by 1" ;;
    selective) echo "select sum(v1), count(*) from $t where dim_a_id in ($keys) and amount > 99.5" ;;
    limit) echo "select * from $t where dim_a_id = $((ENT / 2)) order by event_date desc, id limit 20" ;;
  esac
}

IO="select coalesce(sum(reads), 0), coalesce(sum(read_bytes), 0) from pg_stat_io where object = 'relation' and context in ('normal', 'bulkread', 'vacuum')"
W=$(mktemp -d)
trap 'rm -rf "$W"' EXIT
printf 'shape\tkeys\tarm\tregime\tblocks_read\tread_requests\tavg_request_kb\tdemand_blocks\tahead_blocks\texec_ms\n'
for case in group:25 group:200 group:2000 selective:200 limit:1; do
  shape=${case%%:*} n=${case##*:}
  for arm in sift native; do
    t=events; idx=events_sift; on=on
    [ $arm = native ] && { t=events_n; idx="events_n events_n_key"; on=off; }
    for regime in evict cold steady; do
      {
        echo "set client_min_messages = warning; set sift.enable_agg = $on; set max_parallel_workers_per_gather = 0;"
        echo '\o /dev/null'
        case $regime in
          evict) echo "set sift.enable_agg = off; set enable_seqscan = off; set enable_indexscan = off; select count(*), sum(length(pad)) from bulk where k < 900000000;"
                 echo "set sift.enable_agg = $on; reset enable_seqscan; reset enable_indexscan;" ;;
          cold) for r in $idx; do echo "select count(pg_buffercache_evict_relation('$r'::regclass));"; done ;;
          steady) echo "$(query $shape $t "$n");" ;;
        esac
        echo "select pg_stat_force_next_flush();"
        echo "\\o $W/a"; echo "$IO;"
        echo "select sift._test_segment_reads(true);"
        echo "\\o $W/p"; echo "explain (analyze, buffers, costs off, summary off, format json) $(query $shape $t "$n");"
        echo "\\o $W/c"; echo "select demand_blocks, ahead_blocks from sift._test_segment_reads();"
        echo "\\o /dev/null"; echo "select pg_stat_force_next_flush();"
        echo "\\o $W/b"; echo "$IO;"
      } | $P -F' ' >/dev/null
      python3 -I - "$W" "$shape" "$n" "$arm" "$regime" <<'PY'
import json, os, sys
w, shape, n, arm, regime = sys.argv[1:6]
a = open(w + '/a').read().split(); b = open(w + '/b').read().split()
demand, ahead = open(w + '/c').read().split() if arm == 'sift' else ('0', '0')
plan = json.load(open(w + '/p'))[0]['Plan']
reads, kb = float(b[0]) - float(a[0]), (float(b[1]) - float(a[1])) / 1024
print(f"{shape}\t{n}\t{arm}\t{regime}\t{plan['Shared Read Blocks']}\t{reads:.0f}\t{kb / max(reads, 1):.1f}\t{demand}\t{ahead}\t{plan['Actual Total Time']:.2f}")
PY
      rm -f "$W"/?
    done
  done
done