barebox/benchmarks/run-pgvector.sh at main · Cooperation-org/barebox · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
#!/bin/bash
# pgvector HNSW benchmark runner.
# Tests index build time and search throughput at varying concurrency.
#
# Configure via environment:
#   POSTGRES_VM_IP  (default: 10.0.0.100)
#   PG_PORT         (default: 5432)
#   PG_USER         (default: postgres)
#   PG_DB           (default: pgvector_bench)
#   NUM_VECTORS     (default: 100000)
#   DIMENSIONS      (default: 1536)
#   DURATION        (default: 30 seconds per run)
#   CLIENT_COUNTS   (default: "1 2 4 8 16 32 64")
#
# Note: Inserting 100K vectors at 1536 dimensions via SQL is slow (~10-20 min).
# For faster loading, consider generating vectors in Python and using COPY.
#
# Usage: ./run-pgvector.sh
set -euo pipefail

PG_HOST="${POSTGRES_VM_IP:-10.0.0.100}"
PG_PORT="${PG_PORT:-5432}"
PG_USER="${PG_USER:-postgres}"
PG_DB="${PG_DB:-pgvector_bench}"
NUM_VECTORS=${NUM_VECTORS:-100000}
DIMENSIONS=${DIMENSIONS:-1536}
DURATION=${DURATION:-30}
CLIENT_COUNTS="${CLIENT_COUNTS:-1 2 4 8 16 32 64}"

echo "=== pgvector HNSW Benchmark ==="
echo "Host: $PG_HOST:$PG_PORT"
echo "Vectors: $NUM_VECTORS at $DIMENSIONS dimensions"
echo "Duration: ${DURATION}s per search run"
echo ""

# Setup database and extension
echo "--- Setting up database ---"
createdb -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" "$PG_DB" 2>/dev/null || true
psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" -d "$PG_DB" <<SQL
CREATE EXTENSION IF NOT EXISTS vector;
DROP TABLE IF EXISTS bench_vectors;
CREATE TABLE bench_vectors (
  id serial PRIMARY KEY,
  embedding vector($DIMENSIONS)
);
SQL

# Insert random vectors in batches
echo ""
echo "--- Inserting $NUM_VECTORS random vectors (this may take a while) ---"
BATCH_SIZE=1000
INSERTED=0
while [ $INSERTED -lt $NUM_VECTORS ]; do
  REMAINING=$(( NUM_VECTORS - INSERTED ))
  CURRENT_BATCH=$(( REMAINING < BATCH_SIZE ? REMAINING : BATCH_SIZE ))
  psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" -d "$PG_DB" -q <<SQL
INSERT INTO bench_vectors (embedding)
SELECT
  (SELECT array_agg(random())::vector($DIMENSIONS)
   FROM generate_series(1, $DIMENSIONS))
FROM generate_series(1, $CURRENT_BATCH);
SQL
  INSERTED=$(( INSERTED + CURRENT_BATCH ))
  echo "  Inserted $INSERTED / $NUM_VECTORS"
done

# Build HNSW index (timed)
echo ""
echo "========================================="
echo "=== HNSW Index Build (m=16, ef_construction=64) ==="
echo "========================================="
time psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" -d "$PG_DB" <<SQL
SET maintenance_work_mem = '2GB';
CREATE INDEX ON bench_vectors
  USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);
SQL

# Generate a random query vector
echo ""
echo "--- Generating query vector ---"
QUERY_VEC=$(psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" -d "$PG_DB" -t -A -c \
  "SELECT (SELECT array_agg(random())::vector($DIMENSIONS) FROM generate_series(1, $DIMENSIONS));")

# Create pgbench custom script for vector search
BENCH_SCRIPT=$(mktemp /tmp/pgvector-bench.XXXXXX)
cat > "$BENCH_SCRIPT" <<SQL
SET hnsw.ef_search = 40;
SELECT id FROM bench_vectors
  ORDER BY embedding <=> '$QUERY_VEC'
  LIMIT 10;
SQL

# Run search benchmarks
echo ""
echo "========================================="
echo "=== Search Throughput                 ==="
echo "========================================="
for c in $CLIENT_COUNTS; do
  echo ""
  echo "--- $c clients ---"
  pgbench -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" \
    -c "$c" -j "$c" -T "$DURATION" -f "$BENCH_SCRIPT" "$PG_DB"
done

rm -f "$BENCH_SCRIPT"
echo ""
echo "=== Done ==="