#!/usr/bin/env bash
# The first statement of a new backend over a column index of many segments: the blocks its planning and its execution read, and the milliseconds of planning. A segment
# is added by every promotion of the write buffer; the leaf statistics a plan reads (the sections of every segment) are folded once for the server in a record of shared
# memory, under a stamp of the segments, so what a new backend pays is the segments the record does not cover. COUNTS are the measure (EXPLAIN BUFFERS: Shared Read and Hit
# Blocks of Planning and of the statement; a tmpfs cluster has no device latency); the milliseconds are secondary.
#
#   PGHOST=.. PGPORT=.. bench/steps/first_statement_planning.sh [SEGMENTS [KEYS [ROWS_PER_KEY]]] > records.tsv
#   LOAD=0 reuses the table of an earlier run
#
# Needs a scratch cluster with a small shared_buffers (16-64 MB) and pg_sift. One table of events loaded in SEGMENTS batches (sift.col_auto_merge off, one promotion each),
# a table of padding rows larger than shared_buffers whose scan evicts the pool, a statement over 200 keys. Two arms (sift: the planner's own choice; planner: every sift
# planner feature off, so the plan is the one PostgreSQL makes, which still costs the column index and so reads its statistics) in two states of the record:
#   warm      the statistics were folded by another backend and the pool was evicted since
#   promoted  a small batch was promoted after that (the stamp moved by one segment), the pool was evicted, a new backend runs the statement
set -euo pipefail
SEGS=${1:-60} ENT=${2:-300} PER=${3:-1000}
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_sift;
drop table if exists events, 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, measures jsonb not null);
create index events_sift on events using sift_col (id, dim_a_id, kind_id, event_date, amount, label, 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;
create table src as select e as dim_a_id, 1 + (e * 7 + g) % 19 as kind_id, date '2024-01-01' + (random() * 900)::int as event_date, round((random() * 100)::numeric, 2)::float8 as amount,
  'label_' || (random() * 500)::int as label, jsonb_build_object('m1', round((random() * 10)::numeric, 1), 'm2', (random() * 5)::int) as measures, (random() * 1000000)::int as r
  from generate_series(1, $ENT) e, generate_series(1, $PER) g;
set sift.col_auto_merge = off;
do \$\$ declare i int; begin
  for i in 0..$SEGS - 1 loop
    insert into events (dim_a_id, kind_id, event_date, amount, label, measures) select dim_a_id, kind_id, event_date, amount, label, measures from src where r % $SEGS = i;
    perform sift.col_promote('events_sift');
    commit;
  end loop;
end \$\$;
drop table src;
vacuum (freeze, analyze, index_cleanup off) events, bulk;
SQL

step=$((ENT / 200)); [ "$step" -ge 1 ] || step=1
keys=$(seq 1 "$step" "$ENT" | head -n 200 | paste -sd, -)
Q="select kind_id, count(*), sum(amount) from events where dim_a_id in ($keys) and event_date between '2024-03-01' and '2025-03-01' group by kind_id order by 1"
evict() { $P -c "set sift.enable_agg = off; set enable_seqscan = off; set enable_indexscan = off; select count(*), sum(length(pad)) from bulk where k < 900000000" >/dev/null; }
segs() { $P -c "select count(*) from sift.col_segments('events_sift') where kind = 'immutable'"; }
measure() {  # ARM STATE: a new backend runs the statement as its first
  local arm=$1 state=$2 SEGN=$3 set="set sift.enable_agg = on;"
  [ "$arm" = planner ] && set="set sift.enable_agg = off; set sift.enable_col_scan = off; set sift.enable_scan = off; set sift.enable_topn = off; set sift.enable_count = off;"
  { echo "set max_parallel_workers_per_gather = 0; $set"; echo '\o /dev/stdout'; echo "explain (analyze, buffers, costs off, format json) $Q;"; } | $P | python3 -I -c "
import json, sys
j = json.load(sys.stdin)[0]
p, plan = j['Planning'], j['Plan']
print('\t'.join(str(x) for x in ['$arm', '$state', '$SEGN', p['Shared Read Blocks'], p['Shared Hit Blocks'], round(j['Planning Time'], 2), plan['Shared Read Blocks'], plan['Shared Hit Blocks'], round(j['Execution Time'], 2)]))"
}
printf 'arm\tstate\tsegments\tplan_read_blocks\tplan_hit_blocks\tplan_ms\texec_read_blocks\texec_hit_blocks\texec_ms\n'
for arm in sift planner; do
  # another backend folds the statistics of the segments (the record), then the pool is evicted
  $P -c "set sift.enable_agg = on; $Q" >/dev/null; n=$(segs); evict
  measure $arm warm "$n"
  # a small batch is promoted: one segment more, the stamp of the record is not that of the directory
  $P -c "set sift.col_auto_merge = off; insert into events (dim_a_id, kind_id, event_date, amount, label, measures) select dim_a_id, kind_id, event_date, amount, label, measures from events where id % 400 = 7" -c "select sift.col_promote('events_sift')" >/dev/null
  n=$(segs); evict
  measure $arm promoted "$n"
done