Skip to content

Speed up IVFFlat scans with direct kernel calls and candidate array - #1028

Open
kirito632 wants to merge 1 commit into
pgvector:masterfrom
kirito632:ivfflat-fastscan
Open

kirito632 wants to merge 1 commit into
pgvector:masterfrom
kirito632:ivfflat-fastscan

Conversation

@kirito632

Copy link
Copy Markdown

Motivation

Profiling an IVFFlat index scan at moderate probes shows that the distance
computation itself is only a small fraction of scan CPU. Most time goes to
per-tuple machinery in GetScanItems:

  1. fmgr dispatch per candidate row — the opclass distance function is
    called through FunctionCallInvoke for every candidate, paying full
    argument marshalling (each Vector * re-validated, dim re-checked).
  2. Tuple slot + tuplesort per batch — every candidate is stored in a
    heap tuple and run through tuplesort, even though only the top
    k (bounded by work_mem) can survive.

Both costs scale with the number of scanned candidates (probes x cluster
size), not with result size, so they dominate exactly when recall settings
are aggressive.

What this PR changes

  • ivfflatbeginscan resolves the opclass distance proc (FUNCTION 1) to
    one of three exported C kernels — vector_l2_squared_distance,
    vector_negative_inner_product (shared by ip_ops and cosine_ops) —
    and stores the function pointer on the scan. Kernels take
    (int dim, const float *ax, const float *bx).
  • IvfflatVectorParts() resolves both varlena header layouts inline:
    index tuples for dim <= 30 use a 1-byte short header, and the previous
    direct (Vector *) cast was only valid for 4-byte headers. (This was the
    source of out-of-bounds reads on small-dim data — see Testing.)
  • GetScanItems fast path accumulates {distance, tid} pairs into a
    work_mem-bounded candidate array, sorted once per batch. Ties are
    broken by TID.
  • Fallbacks are preserved: capacity overflow, dimension mismatch, or an
    unrecognized opclass proc all revert to the original slot + tuplesort
    path. NULL order-by values keep using the original path.

Benchmark

Frozen configuration: deterministic clustered-Gaussian synthetic data, N=100k, dim=128, 1,000 queries,
P2 matrix (3 operators x 14 lists/probes cells), R=5 repetitions per arm,
same-machine A/B (OpenTenBase / PostgreSQL 19beta3, gcc 13.3, x86_64).
Recall = mean Recall@10 per cell; QPS = 1,000 queries / freeze-robust mean wall seconds (concurrency=1); wall time = median of repetitions with
freeze outliers (>5x median) excluded.

  • Recall: unchanged on every acceptance cell — worst |delta| = 0.0011
    against a 0.02 tolerance.
  • Wall time: mean 1.43x / median 1.44x faster across the 42-cell grid;
    largest single-cell gain 2.93x (trend cells at low lists).
  • pgvector regression suite: 15/15 (14 existing + 1 new).

Testing

  • New test/sql/ivfflat_fastscan.sql covers both varlena header forms
    (dim 8 short-header, dim 64 4-byte-header) for L2/IP/cosine, the
    work_mem overflow fallback, and lateral rescans. The dim<=30 short
    header case is a regression guard: the naive (Vector *) cast reads a
    garbage dim there and crashes.
  • Adversarial smoke tests outside the suite additionally verified
    full-probes equality with exact scans (diff = 0) on both header forms.
  • One behavioral note for reviewers: the fast path breaks distance ties by
    TID, whereas the tuplesort path's tie order is unspecified. The new test
    pins the fallback query with a secondary sort key where ties occur.

Notes

  • No user-visible defaults or GUCs change; the fast path is transparent.
  • The kernels are also the natural seam for future SIMD work (the current
    build already auto-vectorizes the inner loops with AVX2).

Bypass fmgr dispatch in GetScanItems by resolving the opclass distance
proc (FUNCTION 1: vector_l2_squared_distance / vector_negative_inner_product)
to exported C kernels, and replace the per-tuple slot + tuplesort
machinery with a work_mem-bounded {distance, tid} candidate array that is
sorted once per batch. Ties are broken by TID; capacity overflow or any
unrecognized kernel falls back to the original path.

Index tuples for dim <= 30 use a 1-byte short varlena header; kernels take
(dim, ax, bx) and IvfflatVectorParts resolves both header layouts inline.

A/B on the frozen benchmark (3 operators x 14 lists/probes cells, R=5,
N=100k, dim=128): recall unchanged (worst delta 0.0011 on acceptance
cells), wall time mean 1.43x / median 1.44x faster.

Adds test/sql/ivfflat_fastscan.sql covering short-header (dim 8) and
4-byte-header (dim 64) scans for L2/IP/cosine, the work_mem overflow
fallback, and lateral rescans. Full suite: 15/15.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Development

Successfully merging this pull request may close these issues.

1 participant