Vector similarity search with pgvector

types_session.py

"""pgvector 0.8.7 on PostgreSQL 18.6: the vector type, three distance operators that disagree, and the size of a value."""

from pg_extensions import extension_settings
from pg_session import Server

# pgvector may sit next to PostgreSQL or in a prefix of its own: pg_extensions.py finds such a prefix and returns
# the two settings PostgreSQL 18 needs for it (none when pgvector is installed with PostgreSQL).
SETTINGS = {"autovacuum": "off", "max_parallel_workers_per_gather": 0, "jit": "off", "timezone": "'UTC'"}
SETTINGS.update(extension_settings())

with Server(settings=SETTINGS) as pg:
    pg.up()
    pg.psql("types.sql", errors=1)

$ psql -X -a -f types.sql
-- pgvector on PostgreSQL 18.6: a vector column, three distances that rank five vectors differently,
-- what each type takes on disk, and a value of the wrong length.
CREATE EXTENSION vector;
CREATE EXTENSION
SELECT extversion FROM pg_extension WHERE extname = 'vector';
 extversion 
------------
 0.8.7
(1 row)

-- Two numbers per vector so the distances can be checked by hand; an embedding has hundreds.
CREATE TABLE items (id int PRIMARY KEY, label text, emb vector(2));
CREATE TABLE
INSERT INTO items VALUES
  (1, 'close',          '[0.9, 1.1]'),
  (2, 'same direction', '[2, 2]'),
  (3, 'long east',      '[5, 0]'),
  (4, 'north',          '[0.2, 1]'),
  (5, 'opposite',       '[-1, -1]');
INSERT 0 5
-- Each distance to the query vector [1, 1]; the smallest value is the nearest.
SELECT label,
       round((emb <-> '[1,1]')::numeric, 3) AS l2,
       round((emb <=> '[1,1]')::numeric, 3) AS cosine,
       round((emb <#> '[1,1]')::numeric, 3) AS neg_inner
FROM items ORDER BY emb <-> '[1,1]';
     label      |  l2   | cosine | neg_inner 
----------------+-------+--------+-----------
 close          | 0.141 |  0.005 |    -2.000
 north          | 0.800 |  0.168 |    -1.200
 same direction | 1.414 |  0.000 |    -4.000
 opposite       | 2.828 |  2.000 |     2.000
 long east      | 4.123 |  0.293 |    -5.000
(5 rows)

-- The nearest vector by each operator.
SELECT '<-> l2' AS operator, (SELECT label FROM items ORDER BY emb <-> '[1,1]' LIMIT 1) AS nearest
UNION ALL
SELECT '<=> cosine', (SELECT label FROM items ORDER BY emb <=> '[1,1]' LIMIT 1)
UNION ALL
SELECT '<#> inner product', (SELECT label FROM items ORDER BY emb <#> '[1,1]' LIMIT 1);
     operator      |    nearest     
-------------------+----------------
 <-> l2            | close
 <=> cosine        | same direction
 <#> inner product | long east
(3 rows)

-- Bytes per value: 4 per dimension plus 8, 2 plus 8, and 8 per non-zero element plus 16.
SELECT pg_column_size(array_fill(0.5, ARRAY[384])::vector(384))  AS vector_384,
       pg_column_size(array_fill(0.5, ARRAY[384])::halfvec(384)) AS halfvec_384,
       pg_column_size('{7:0.5, 9000:1.5}/30000'::sparsevec)       AS sparsevec_2_of_30000;
 vector_384 | halfvec_384 | sparsevec_2_of_30000 
------------+-------------+----------------------
       1544 |         776 |                   32
(1 row)

-- A value with three numbers does not fit vector(2).
INSERT INTO items VALUES (6, 'too long', '[1, 2, 3]');
psql:types.sql:35: ERROR:  expected 2 dimensions, not 3

types.sql

-- pgvector on PostgreSQL 18.6: a vector column, three distances that rank five vectors differently,
-- what each type takes on disk, and a value of the wrong length.
CREATE EXTENSION vector;
SELECT extversion FROM pg_extension WHERE extname = 'vector';

-- Two numbers per vector so the distances can be checked by hand; an embedding has hundreds.
CREATE TABLE items (id int PRIMARY KEY, label text, emb vector(2));
INSERT INTO items VALUES
  (1, 'close',          '[0.9, 1.1]'),
  (2, 'same direction', '[2, 2]'),
  (3, 'long east',      '[5, 0]'),
  (4, 'north',          '[0.2, 1]'),
  (5, 'opposite',       '[-1, -1]');

-- Each distance to the query vector [1, 1]; the smallest value is the nearest.
SELECT label,
       round((emb <-> '[1,1]')::numeric, 3) AS l2,
       round((emb <=> '[1,1]')::numeric, 3) AS cosine,
       round((emb <#> '[1,1]')::numeric, 3) AS neg_inner
FROM items ORDER BY emb <-> '[1,1]';

-- The nearest vector by each operator.
SELECT '<-> l2' AS operator, (SELECT label FROM items ORDER BY emb <-> '[1,1]' LIMIT 1) AS nearest
UNION ALL
SELECT '<=> cosine', (SELECT label FROM items ORDER BY emb <=> '[1,1]' LIMIT 1)
UNION ALL
SELECT '<#> inner product', (SELECT label FROM items ORDER BY emb <#> '[1,1]' LIMIT 1);

-- Bytes per value: 4 per dimension plus 8, 2 plus 8, and 8 per non-zero element plus 16.
SELECT pg_column_size(array_fill(0.5, ARRAY[384])::vector(384))  AS vector_384,
       pg_column_size(array_fill(0.5, ARRAY[384])::halfvec(384)) AS halfvec_384,
       pg_column_size('{7:0.5, 9000:1.5}/30000'::sparsevec)       AS sparsevec_2_of_30000;

-- A value with three numbers does not fit vector(2).
INSERT INTO items VALUES (6, 'too long', '[1, 2, 3]');

pg_session.py

"""Throwaway PostgreSQL 18.6 servers for the examples of this lesson.

Every session creates a cluster of its own with initdb in a new folder inside the examples folder, starts it, runs
psql scripts and server programs on it, prints what they printed, and stops and deletes it again, also when something
fails. The server is the pinned PostgreSQL 18.6 from the mise installations next to the Python that runs this file.
It listens only on 127.0.0.1, an address no other computer can reach, on a port the operating system reports as
unused, with no Unix socket, and it asks for a password (scram-sha-256) made up for the run and never printed.
Besides PostgreSQL's own programs, the sessions read only this folder (and an extension's own folder where a lesson
needs one), and they connect to nothing but 127.0.0.1.

The output is exactly what the programs printed, with four substitutions so that a record never shows the computer
it ran on and is the same on every run: the run's folder is shown as ".", the folder of PostgreSQL's programs is left
out of their paths, the operating-system user's name is shown as "learner", and the dots pg_ctl prints while it waits
(one about every tenth of a second) as three. A session may name more patterns (mask=) for numbers that change on
every run, such as the rates and times ANALYZE VERBOSE prints or the random key of pg_dump's \\restrict line; the
session's own file says which, and its record shows a fixed word in their place. The sessions never print the port.
A session stops with status 1 when a command fails that should have worked, or reports a different number of errors
than the session expects, so a run that went wrong is never recorded as if it had worked.
"""

import os
import pwd
import re
import secrets
import shlex
import shutil
import socket
import subprocess
import sys
import tempfile
import textwrap
import time
from pathlib import Path

PG_VERSION = "18.6"
# .../installs/python/3.14.8/bin/python3 -> .../installs
INSTALLS = Path(sys.executable).resolve().parents[3]
BIN = INSTALLS / "postgres" / PG_VERSION / "bin"
OS_USER = pwd.getpwuid(os.geteuid()).pw_name


def free_port():
    """A TCP port on 127.0.0.1 that nothing listens on now."""
    with socket.socket() as probe:
        probe.bind(("127.0.0.1", 0))
        return probe.getsockname()[1]


class Server:
    """One throwaway cluster: initdb, settings, start; stopped and deleted when the with block ends."""

    def __init__(self, settings=None, initdb=(), name="data"):
        self.settings = dict(settings or {})
        self.initdb_args = list(initdb)
        self.name = name
        self.running = []

    def __enter__(self):
        if not (BIN / "postgres").exists():
            sys.exit(f"PostgreSQL {PG_VERSION} was not found next to this Python")
        self.base = Path(tempfile.mkdtemp(prefix="pg", dir=".")).resolve()
        self.port = free_port()
        self.password = secrets.token_urlsafe(24)
        (self.base / "pwfile").write_text(self.password + "\n")
        self.env = {**os.environ, "PGHOST": "127.0.0.1", "PGPORT": str(self.port), "PGUSER": "postgres",
                    "PGPASSWORD": self.password, "PGDATABASE": "postgres", "PGCONNECT_TIMEOUT": "60"}
        for key in ("PGSERVICE", "PGSERVICEFILE", "PGOPTIONS", "PSQLRC"):
            self.env.pop(key, None)
        for script in Path(".").glob("*.sql"):
            shutil.copy(script, self.base / script.name)
        for data in Path(".").glob("*.csv"):
            shutil.copy(data, self.base / data.name)
        return self

    def __exit__(self, *exc):
        for folder, _port in reversed(self.running):
            self._call(["pg_ctl", "-D", folder, "-m", "fast", "-w", "-t", "60", "stop"])
        self.running = []
        shutil.rmtree(self.base, ignore_errors=True)

    # ---- the cluster ---------------------------------------------------------------------------------------------
    def initdb(self, show=False):
        args = ["initdb", "-D", self.name, "-U", "postgres", "--auth=scram-sha-256", "--pwfile=pwfile",
                *self.initdb_args]
        if not show:
            args.append("--no-instructions")
        self._expect(args, show=show)
        with open(self.base / self.name / "postgresql.conf", "a") as conf:
            conf.write("\n# Added by the session: 127.0.0.1 only, no Unix socket.\n")
            conf.write(f"listen_addresses = '127.0.0.1'\nport = {self.port}\nunix_socket_directories = ''\n")
            for key, value in self.settings.items():
                conf.write(f"{key} = {value}\n")

    def start(self, show=False, folder=None, port=None):
        folder = folder or self.name
        self._expect(["pg_ctl", "-D", folder, "-l", f"{folder}.log", "-w", "-t", "60", "start"], show=show)
        self.running.append((folder, port or self.port))

    def up(self):
        """initdb and start, quietly."""
        self.initdb()
        self.start()
        return self

    def standby(self, folder="standby"):
        """A streaming standby of this server, made with pg_basebackup -R, started on a port of its own."""
        self._expect(["pg_basebackup", "-D", folder, "-R", "-X", "stream", "-c", "fast"])
        port = free_port()
        with open(self.base / folder / "postgresql.conf", "a") as conf:
            conf.write(f"port = {port}\n")
        self.start(folder=folder, port=port)
        for _ in range(300):
            if self.query("SELECT count(*) FROM pg_stat_replication WHERE state = 'streaming'") == "1":
                return port
            time.sleep(0.1)
        sys.exit("the standby did not start streaming")

    # ---- running things ------------------------------------------------------------------------------------------
    def _call(self, args, env=None):
        args = [str(a) for a in args]
        done = subprocess.run([str(BIN / args[0]), *args[1:]], stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                              text=True, cwd=self.base, env=env or self.env)
        return done.returncode, self.clean(done.stdout)

    def _expect(self, args, show=False, status=0, mask=()):
        code, out = self._call(args)
        out = _masked(out, mask)
        if show:
            _show(args, out)
        if code != status:
            if not show:
                print(out)
            sys.exit(f"{args[0]} ended with status {code}, not {status}")
        return out

    def clean(self, text):
        """The run's folder as ".", the user as learner, pg_ctl's waiting dots as three (see the docstring)."""
        text = text.replace(str(BIN) + "/", "").replace(str(self.base) + "/", "").replace(str(self.base), ".")
        text = re.sub(r"(waiting for server to (?:start|shut down))\.+", r"\1...", text)
        return text.replace(f'"{OS_USER}"', '"learner"')

    def wait_for(self, sql, value):
        """Waits (unprinted) until a query returns the value, for example after a reload."""
        for _ in range(300):
            if self.query(sql) == value:
                return
            time.sleep(0.1)
        sys.exit(f"{sql} never returned {value}")

    def show_file(self, path):
        """Prints a text file of the cluster, as `cat` would."""
        print(f"$ cat {path}")
        print((self.base / path).read_text().rstrip("\n"))
        print()

    def query(self, sql, port=None):
        """One value from one query, unprinted (the session uses it to wait or to decide)."""
        env = {**self.env, "PGPORT": str(port or self.port)}
        code, out = self._call(["psql", "-X", "-At", "-c", sql], env=env)
        if code != 0:
            sys.exit(f"the query failed: {out}")
        return out.strip()

    def silent(self, sql):
        """Runs a statement without printing it (used only to give a role the run's password)."""
        self.query(sql)

    def psql(self, script, errors=0, status=None, user=None, db=None, port=None, flags=("-a",), note=None, mask=()):
        """Runs a psql script with -X (no ~/.psqlrc) and the given flags, and prints the command and its output.
        `errors` is how many ERROR lines the script must print; `status` the exit status it must end with."""
        env = {**self.env, "PGPORT": str(port or self.port)}
        if user:
            env["PGUSER"] = user
        if db:
            env["PGDATABASE"] = db
        args = ["psql", "-X", *flags, "-f", script]
        code, out = self._call(args, env=env)
        out = _masked(out, mask)
        shown = ["psql", "-X", *flags, "-f", script]
        if user:
            shown[1:1] = ["-U", user]
        if db:
            shown[1:1] = ["-d", db]
        if note:
            print(f"-- {note}")
        _show(shown, out)
        found = sum(1 for line in out.splitlines() if "ERROR:" in line or ": error:" in line)
        expected_status = 0 if status is None else status
        if found != errors or code != expected_status:
            sys.exit(f"{script}: {found} errors and status {code}; the session expects {errors} and {expected_status}")

    def sessions(self, count):
        """Opens `count` idle psql connections (each one a client backend) until the with block ends."""
        return _Idle(self, count)

    def shell(self, *args, status=0, mask=()):
        """Runs one of PostgreSQL's programs and prints the command and its output."""
        return self._expect(list(args), show=True, status=status, mask=mask)

    def listing(self, folder, prefix=""):
        """Prints the names in a folder of the cluster (sorted, folders with a slash), as `ls -p` would list them."""
        path = self.base / folder
        names = sorted(n + ("/" if (path / n).is_dir() else "") for n in os.listdir(path) if n.startswith(prefix))
        print(f"-- in {folder}" + (f", the names that start with {prefix}" if prefix else "") + ":")
        print(textwrap.fill("  ".join(names), width=78, initial_indent="   ", subsequent_indent="   ",
                            break_on_hyphens=False))
        print()
        return names


def _masked(text, mask):
    """Replaces what differs on every run (a rate, a time, a random key) with fixed text (see the docstring)."""
    for pattern, replacement in mask:
        text = re.sub(pattern, replacement, text, flags=re.M)
    return text


def _arg(a):
    """A command-line argument as a reader would type it: double quotes around text with single quotes in it."""
    a = str(a)
    if re.fullmatch(r"[\w@%+=:,./-]+", a):
        return a
    if not re.search(r"[\"$`\\!]", a):
        return f'"{a}"'
    return shlex.quote(a)


def _show(args, out):
    print("$ " + " ".join(_arg(a) for a in args))
    if out.strip():
        print(out.rstrip("\n"))
    print()


class _Idle:
    """Idle psql connections, closed again when the with block ends."""

    def __init__(self, server, count):
        self.server, self.count, self.procs = server, count, []

    def __enter__(self):
        before = int(self.server.query("SELECT count(*) FROM pg_stat_activity WHERE backend_type = 'client backend'"))
        for _ in range(self.count):
            self.procs.append(subprocess.Popen([str(BIN / "psql"), "-X", "-q"], stdin=subprocess.PIPE,
                                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                                               env=self.server.env, cwd=self.server.base))
        sql = "SELECT count(*) FROM pg_stat_activity WHERE backend_type = 'client backend'"
        for _ in range(300):
            # the query's own connection is one more client backend
            if int(self.server.query(sql)) == before + self.count:
                return self
            time.sleep(0.1)
        sys.exit("the idle connections did not open")

    def __exit__(self, *exc):
        for p in self.procs:
            p.stdin.close()
            p.wait(timeout=30)

pg_extensions.py

"""Where the sessions of this lesson find an extension that is not installed next to PostgreSQL.

An extension installed together with PostgreSQL (by a package manager, for example) needs nothing from this file:
extension_settings() then returns no settings. An extension built into a prefix of its own, a folder that holds
share/extension/ and lib/, is found by PostgreSQL 18 through two settings, extension_control_path and
dynamic_library_path. This file builds them from the prefixes it finds, in this order:

1. PG_EXTENSION_DIR, a colon-separated list of such prefixes, when it is set;
2. otherwise the entries of a folder named pg-extensions in the session's folder or in a folder above it: each entry
   is a prefix, or a link to one.

The settings go into the throwaway cluster's configuration only. The sessions never print them, because they hold
paths of the computer that ran them.
"""

import os
from pathlib import Path


def extension_prefixes():
    """The prefixes to add, or [] when the extension sits next to PostgreSQL."""
    listed = [p for p in os.environ.get("PG_EXTENSION_DIR", "").split(":") if p]
    if listed:
        return listed
    here = Path.cwd()
    for folder in (here, *here.parents):
        holder = folder / "pg-extensions"
        if holder.is_dir():
            return sorted(str(p.resolve()) for p in holder.iterdir() if (p / "share" / "extension").is_dir())
    return []


def extension_settings():
    """Settings for Server(settings=...): empty when no prefix is needed."""
    prefixes = extension_prefixes()
    if not prefixes:
        return {}
    return {
        "extension_control_path": "'" + ":".join(f"{p}/share" for p in prefixes) + ":$system'",
        "dynamic_library_path": "'" + ":".join(f"{p}/lib" for p in prefixes) + ":$libdir'",
    }

plans_session.py

"""pgvector 0.8.7 on PostgreSQL 18.6: exact search reads every row, an HNSW index serves only its own operator."""

from pg_extensions import extension_settings
from pg_session import Server

# pg_extensions.py finds pgvector when it sits in a prefix of its own (see its docstring). Index builds run in the
# session's own process, with no parallel workers. They are random all the same (pgvector draws from PostgreSQL's
# per-process generator, which setseed() does not reach), so the script prints only what every build agrees on.
SETTINGS = {"autovacuum": "off", "max_parallel_workers_per_gather": 0, "jit": "off", "timezone": "'UTC'",
            "max_parallel_maintenance_workers": 0}
SETTINGS.update(extension_settings())

with Server(settings=SETTINGS) as pg:
    pg.up()
    pg.psql("docs.sql")
    pg.psql("plans.sql")

$ psql -X -a -f docs.sql
-- 5,000 documents and 200 queries, each a vector of 64 numbers between -1 and 1 from random(); the seed makes every
-- run identical. A real system gets such vectors from an embedding model outside the database.
CREATE EXTENSION vector;
CREATE EXTENSION
SELECT setseed(0.42) AS seeded \gset
CREATE TABLE docs AS
  SELECT i AS id, array_agg(random() * 2 - 1 ORDER BY d)::vector(64) AS emb
  FROM generate_series(1, 5000) AS i, generate_series(1, 64) AS d
  GROUP BY i;
SELECT 5000
CREATE TABLE queries AS
  SELECT q AS qid, array_agg(random() * 2 - 1 ORDER BY d)::vector(64) AS emb
  FROM generate_series(1, 200) AS q, generate_series(1, 64) AS d
  GROUP BY q;
SELECT 200
ALTER TABLE docs ADD PRIMARY KEY (id);
ALTER TABLE
ALTER TABLE queries ADD PRIMARY KEY (qid);
ALTER TABLE
ANALYZE docs, queries;
ANALYZE

$ psql -X -a -f plans.sql
-- The 10 documents nearest to query 1 by cosine distance, before any vector index exists.
EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF, BUFFERS OFF)
SELECT id FROM docs ORDER BY emb <=> (SELECT emb FROM queries WHERE qid = 1) LIMIT 10;
                                 QUERY PLAN                                  
-----------------------------------------------------------------------------
 Limit (actual rows=10.00 loops=1)
   InitPlan 1
     ->  Index Scan using queries_pkey on queries (actual rows=1.00 loops=1)
           Index Cond: (qid = 1)
           Index Searches: 1
   ->  Sort (actual rows=10.00 loops=1)
         Sort Key: ((docs.emb <=> (InitPlan 1).col1))
         Sort Method: top-N heapsort  Memory: 25kB
         ->  Seq Scan on docs (actual rows=5000.00 loops=1)
(9 rows)

-- An HNSW index for cosine distance, with the default m = 16 and ef_construction = 64.
CREATE INDEX docs_emb_idx ON docs USING hnsw (emb vector_cosine_ops);
CREATE INDEX
EXPLAIN (COSTS OFF)
SELECT id FROM docs ORDER BY emb <=> (SELECT emb FROM queries WHERE qid = 1) LIMIT 10;
                    QUERY PLAN                    
--------------------------------------------------
 Limit
   InitPlan 1
     ->  Index Scan using queries_pkey on queries
           Index Cond: (qid = 1)
   ->  Index Scan using docs_emb_idx on docs
         Order By: (emb <=> (InitPlan 1).col1)
(6 rows)

-- The same question by L2 distance: the index was built for cosine distance.
EXPLAIN (COSTS OFF)
SELECT id FROM docs ORDER BY emb <-> (SELECT emb FROM queries WHERE qid = 1) LIMIT 10;
                      QUERY PLAN                      
------------------------------------------------------
 Limit
   InitPlan 1
     ->  Index Scan using queries_pkey on queries
           Index Cond: (qid = 1)
   ->  Sort
         Sort Key: ((docs.emb <-> (InitPlan 1).col1))
         ->  Seq Scan on docs
(7 rows)

-- Cosine similarity, highest first: an expression around the operator, in descending order.
EXPLAIN (COSTS OFF)
SELECT id FROM docs ORDER BY 1 - (emb <=> (SELECT emb FROM queries WHERE qid = 1)) DESC LIMIT 10;
                                     QUERY PLAN                                      
-------------------------------------------------------------------------------------
 Limit
   InitPlan 1
     ->  Index Scan using queries_pkey on queries
           Index Cond: (qid = 1)
   ->  Sort
         Sort Key: (('1'::double precision - (docs.emb <=> (InitPlan 1).col1))) DESC
         ->  Seq Scan on docs
(7 rows)

docs.sql

-- 5,000 documents and 200 queries, each a vector of 64 numbers between -1 and 1 from random(); the seed makes every
-- run identical. A real system gets such vectors from an embedding model outside the database.
CREATE EXTENSION vector;
SELECT setseed(0.42) AS seeded \gset
CREATE TABLE docs AS
  SELECT i AS id, array_agg(random() * 2 - 1 ORDER BY d)::vector(64) AS emb
  FROM generate_series(1, 5000) AS i, generate_series(1, 64) AS d
  GROUP BY i;
CREATE TABLE queries AS
  SELECT q AS qid, array_agg(random() * 2 - 1 ORDER BY d)::vector(64) AS emb
  FROM generate_series(1, 200) AS q, generate_series(1, 64) AS d
  GROUP BY q;
ALTER TABLE docs ADD PRIMARY KEY (id);
ALTER TABLE queries ADD PRIMARY KEY (qid);
ANALYZE docs, queries;

plans.sql

-- The 10 documents nearest to query 1 by cosine distance, before any vector index exists.
EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF, BUFFERS OFF)
SELECT id FROM docs ORDER BY emb <=> (SELECT emb FROM queries WHERE qid = 1) LIMIT 10;

-- An HNSW index for cosine distance, with the default m = 16 and ef_construction = 64.
CREATE INDEX docs_emb_idx ON docs USING hnsw (emb vector_cosine_ops);

EXPLAIN (COSTS OFF)
SELECT id FROM docs ORDER BY emb <=> (SELECT emb FROM queries WHERE qid = 1) LIMIT 10;

-- The same question by L2 distance: the index was built for cosine distance.
EXPLAIN (COSTS OFF)
SELECT id FROM docs ORDER BY emb <-> (SELECT emb FROM queries WHERE qid = 1) LIMIT 10;

-- Cosine similarity, highest first: an expression around the operator, in descending order.
EXPLAIN (COSTS OFF)
SELECT id FROM docs ORDER BY 1 - (emb <=> (SELECT emb FROM queries WHERE qid = 1)) DESC LIMIT 10;