scripts/benchmark_postgres_fts.py
Rama main · réplica de lectura del repositorio original.
"""Reproduce large-project FTS costs against a disposable local PostgreSQL database.
Example (dedicated server; do not point this at a development database):
docker run -d --name bm-fts-bench --cpus 1 --memory 2g \
-e POSTGRES_PASSWORD=benchmark -e POSTGRES_DB=benchmark \
-p 127.0.0.1:55491:5432 pgvector/pgvector:pg17 \
-c shared_buffers=128MB -c work_mem=4MB \
-c max_parallel_workers_per_gather=0 -c jit=off
uv run python scripts/benchmark_postgres_fts.py --seed
uv run python scripts/benchmark_postgres_fts.py --output /tmp/fts-before.json
The default corpus has 32k documents and 320k TOAST-heavy chunks, using the actual
search DDL and query compiler. Five percent of documents match the selective query,
with its words split across chunks. No embedding model or customer data is needed.
Use --restart-container bm-fts-bench to restart the dedicated server before each
statement, then measure it twice. This clears PostgreSQL shared buffers, NOT the
Docker VM's OS page cache or Neon's file cache. Treat buffer counts as the primary
comparison; these timings do not establish cold Neon latency. Never run benchmarks
concurrently. The JSON keeps complete EXPLAIN (ANALYZE, BUFFERS) plans.
"""
from __future__ import annotations
import argparse
import asyncio
import json
import subprocess
import time
from pathlib import Path
from typing import Any
from sqlalchemy import text
from sqlalchemy.engine import make_url
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from basic_memory.models import search as ddl
from basic_memory.repository.postgres_search_query import compile_fts_filter
from basic_memory.repository.search_query import PreparedSearchQuery
from basic_memory.repository.search_scope import ProjectScope
from basic_memory.schemas.search import SearchItemType
DEFAULT_URL = "postgresql+asyncpg://postgres:benchmark@127.0.0.1:55491/benchmark"
QUERIES = ("meeting budget review", "meeting", "archive", "missingword budget")
async def seed(engine: AsyncEngine, notes: int) -> None:
"""Create a fresh fixture; fail if the target already contains these tables."""
async with engine.begin() as connection:
await connection.execute(text("CREATE TABLE project (id integer PRIMARY KEY)"))
await connection.execute(text("INSERT INTO project VALUES (1), (2), (3), (4)"))
for statement in (
ddl.CREATE_POSTGRES_SEARCH_INDEX_TABLE,
ddl.CREATE_POSTGRES_SEARCH_INDEX_FTS_CHUNKS_TABLE,
):
await connection.execute(statement)
await connection.execute(
text("""
INSERT INTO search_index (
id, project_id, type, title, content_stems, content_snippet,
entity_id, file_path, permalink, metadata, created_at, updated_at
)
SELECT i, 1 + (i % 4), 'entity', 'Document ' || i, 'archive', 'archive',
i, 'notes/' || i || '.md', 'notes/' || i, '{}', now(), now()
FROM generate_series(1, :notes) i
"""),
{"notes": notes},
)
# Distinct lexemes defeat compression and force both text and vectors into
# TOAST. Reusing 1000 vocabularies keeps seed generation manageable.
await connection.execute(
text("""
CREATE TABLE filler AS
SELECT i AS filler_id, string_agg(md5(i::text || '-' || j::text), ' ') AS body
FROM generate_series(1,1000) i CROSS JOIN generate_series(1,128) j
GROUP BY i
""")
)
for start in range(0, notes, 1000):
async with engine.begin() as connection:
await connection.execute(
text("""
INSERT INTO search_index_fts_chunks (
project_id, search_index_id, search_index_type, chunk_index, chunk_text
)
SELECT 1 + (i % 4), i, 'entity', j,
CASE WHEN i % 20 = 0 THEN
CASE j % 3 WHEN 0 THEN 'meeting '
WHEN 1 THEN 'budget ' ELSE 'review ' END
ELSE 'archive ' END || f.body
FROM generate_series(:start + 1, :stop) i
CROSS JOIN generate_series(0,9) j
JOIN filler f ON f.filler_id = 1 + (i % 1000)
"""),
{"start": start, "stop": min(start + 1000, notes)},
)
print(f"Seeded {min(start + 1000, notes) * 10} chunks", flush=True)
async with engine.begin() as connection:
for statement in (
ddl.CREATE_POSTGRES_SEARCH_INDEX_FTS,
ddl.CREATE_POSTGRES_SEARCH_INDEX_FTS_CHUNKS_INDEX,
ddl.CREATE_POSTGRES_SEARCH_INDEX_SCRIPT_NGRAMS_FTS,
ddl.CREATE_POSTGRES_SEARCH_INDEX_FTS_CHUNKS_SCRIPT_NGRAMS_INDEX,
):
await connection.execute(statement)
async with engine.connect() as connection:
await connection.execution_options(isolation_level="AUTOCOMMIT")
await connection.execute(text("VACUUM ANALYZE"))
def restart(container: str) -> None:
"""Clear shared buffers on the explicitly named disposable benchmark server."""
subprocess.run(["docker", "restart", container], check=True, capture_output=True)
for _ in range(60):
ready = subprocess.run(
["docker", "exec", container, "pg_isready", "-U", "postgres", "-d", "benchmark"],
capture_output=True,
)
if ready.returncode == 0:
return
time.sleep(0.5)
raise RuntimeError("Benchmark PostgreSQL did not restart within 30 seconds")
async def measure(engine: AsyncEngine, container: str | None) -> dict[str, Any]:
measurements: dict[str, Any] = {}
for query in QUERIES:
compiled = compile_fts_filter(
ProjectScope.of([1, 2, 3, 4]),
PreparedSearchQuery(search_text=query, search_item_types=[SearchItemType.ENTITY]),
allow_relaxed=True,
)
statements = {
"page": (
f"SELECT search_index.id, {compiled.score_expression} AS score "
f"FROM {compiled.from_clause} WHERE {compiled.where_clause} "
"ORDER BY score DESC, search_index.id ASC LIMIT 130"
),
"count": (f"SELECT count(*) FROM {compiled.from_clause} WHERE {compiled.where_clause}"),
}
for kind, sql in statements.items():
if container:
await engine.dispose()
await asyncio.to_thread(restart, container)
for attempt in range(2):
async with engine.begin() as connection:
# Match execute_fts_statement's planner setting for text queries.
await connection.execute(text("SET LOCAL enable_seqscan = off"))
await connection.execute(text("SET LOCAL statement_timeout = '120s'"))
result = await connection.execute(
text("EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) " + sql), compiled.params
)
explained = result.scalar_one()[0]
label = f"{query}/{kind}/{attempt + 1}"
measurements[label] = explained
plan = explained["Plan"]
print(
f"{label}: {explained['Execution Time']:.1f} ms; "
f"{plan.get('Shared Hit Blocks', 0)} hits, "
f"{plan.get('Shared Read Blocks', 0)} reads",
flush=True,
)
return measurements
async def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--url", default=DEFAULT_URL)
parser.add_argument("--seed", action="store_true")
parser.add_argument("--notes", type=int, default=32000)
parser.add_argument("--restart-container")
parser.add_argument("--output", type=Path, default=Path("/tmp/postgres-fts-benchmark.json"))
args = parser.parse_args()
url = make_url(args.url)
if url.host not in {"localhost", "127.0.0.1", "::1"}:
parser.error("Use a dedicated local PostgreSQL server")
if args.notes <= 0:
parser.error("--notes must be positive")
engine = create_async_engine(url)
try:
if args.seed:
await seed(engine, args.notes)
else:
results = await measure(engine, args.restart_container)
args.output.write_text(json.dumps(results, indent=2) + "\n")
finally:
await engine.dispose()
if __name__ == "__main__":
asyncio.run(main())