test_partition_effects.sh

#!/bin/bash

set -eu -o pipefail

export PGHOST=localhost
export PGPORT=5432
export PGDATABASE=postgres
export PGUSER=postgres
export PGPASSWORD=postgres
export PATH=/usr/lib/postgresql/15/bin:$PATH


PGBENCH_SCALES="5000" # ~5x RAM
PGBENCH_INIT_FLAGS="--unlogged"
PGBENCH_PARTITIONS="0 16 64 256 1024 4096"
PGBENCH_DURATION=3600
PGBENCH_CACHE_WARMUP_DURATION=300
PROTOCOLS="simple"
PGBENCH_CLIENTS=2
PGBENCH_PART_METHODS="range hash"
HOSTNAME=`hostname`

SQL_PGSS_SETUP="CREATE EXTENSION IF NOT EXISTS pg_stat_statements SCHEMA public;"
SQL_PGSS_RESULTS_SETUP="CREATE TABLE IF NOT EXISTS public.pgss_results_planning_test AS SELECT ''::text AS hostname, now() AS test_start_time, now() AS created_on, ''::text as test_name, 0 AS scale, 0 AS duration, 0 AS clients, ''::text AS protocol, ''::text as part_method, 0 AS partitions, mean_plan_time, stddev_plan_time, mean_exec_time, stddev_exec_time, calls, rows, shared_blks_hit, shared_blks_read, query FROM public.pg_stat_statements WHERE false;"
SQL_PGSS_RESET="SELECT public.pg_stat_statements_reset();"

SQL_ZIPF=$(cat <<-"EOF"
select count(*) * 1e5 as aid_max from pgbench_branches \gset
\set aid random_zipfian(1,:aid_max,1.001)
SELECT * FROM pgbench_accounts WHERE aid = :aid;
EOF
)

SQL_RAND=$(cat <<-"EOF"
select count(*) * 1e5 as aid_max from pgbench_branches \gset
\set aid random(1,:aid_max)
SELECT * FROM pgbench_accounts WHERE aid = :aid;
EOF
)

declare -a QUERY_MODES
declare -a QUERIES
QUERIES+=("$SQL_ZIPF")
QUERY_MODES+=("zipfian_access")
QUERIES+=("$SQL_RAND")
QUERY_MODES+=("random_access")

START_TIME=`date +%s`
START_TIME_PG=`psql -qAXtc "select now();"`
date

psql -Xc "$SQL_PGSS_SETUP"
psql -Xc "$SQL_PGSS_RESULTS_SETUP"
psql -Xc "select * from pg_stat_statements limit 1" &>/dev/null # test if pg_stat_statements in shared_preload_libraries


for SCALE in $PGBENCH_SCALES ; do

for PARTS in $PGBENCH_PARTITIONS ; do

for PART_METHOD in $PGBENCH_PART_METHODS ; do

echo -e "\n*** SCALE $SCALE ***\n"

echo "Creating test data using pgbench ..."
echo "pgbench -i -q $PGBENCH_INIT_FLAGS -s $SCALE --partitions $PARTS --partition-method=$PART_METHOD &>/dev/null"
if [ "$PARTS" -gt 0 ]; then
    pgbench -i -q $PGBENCH_INIT_FLAGS -s $SCALE --partitions $PARTS --partition-method=$PART_METHOD &>/dev/null
else
    pgbench -i -q $PGBENCH_INIT_FLAGS -s $SCALE --partitions $PARTS &>/dev/null
fi
echo "Create an extra dummy indexes on aid to give the planner a tiny bit more work ..."
echo "CREATE INDEX ON pgbench_accounts (aid) INCLUDE (abalance) WHERE abalance > 0"
psql -Xc "CREATE INDEX ON pgbench_accounts (aid) INCLUDE (abalance) WHERE abalance > 0"

DBSIZE=`psql -XAtqc "select pg_size_pretty(pg_database_size(current_database()))"`
echo "DB size = $DBSIZE"

for PROTOCOL in $PROTOCOLS ; do

i=0
for QUERY_MODE in "${QUERY_MODES[@]}" ; do
QUERY=${QUERIES[i]}
i=$((i+1))

echo "Doing cache warmup for $PGBENCH_CACHE_WARMUP_DURATION seconds..."
echo "pgbench -S -c $PGBENCH_CLIENTS -T $PGBENCH_CACHE_WARMUP_DURATION &>/dev/null"
pgbench -S -c $PGBENCH_CLIENTS -T $PGBENCH_CACHE_WARMUP_DURATION &>/dev/null

echo "Starting the test loop ..."

echo "Reseting pg_stat_statements..."
psql -Xc "$SQL_PGSS_RESET"

echo "Running the timed query test"
echo -e "echo $QUERY | pgbench -n -f- --random-seed 666 -M $PROTOCOL -c $PGBENCH_CLIENTS -T $PGBENCH_DURATION\n"
echo "$QUERY" | pgbench -n -f- --random-seed 666 -M $PROTOCOL -c $PGBENCH_CLIENTS -T $PGBENCH_DURATION

echo -e "\nStoring results from current run to public.pgss_results_planning_test ..."
psql -Xc "INSERT INTO public.pgss_results_planning_test SELECT '${HOSTNAME}', '${START_TIME_PG}', now(), '${QUERY_MODE}', ${SCALE}, ${PGBENCH_DURATION}, ${PGBENCH_CLIENTS}, '${PROTOCOL}', '${PART_METHOD}',  $PARTS, mean_plan_time, stddev_plan_time, mean_exec_time, stddev_exec_time, calls, rows, shared_blks_hit, shared_blks_read, query from public.pg_stat_statements where query ~* '(INSERT|UPDATE|SELECT).*pgbench_accounts' order by calls desc limit 1"


echo "Done with QUERY_MODE $QUERY_MODE"
done

echo "Done with PROTOCOL $PROTOCOL"
done

echo "Done with PART_METHOD $PART_METHOD"
done

echo "Done with PARTS $PARTS"
done

echo "Done with SCALE $SCALE"
done


date
END_TIME=`date +%s`
echo -e "\nDONE in $((END_TIME-START_TIME)) s"