Support RabitQ in pgvector
Summary
I noticed pgvector maintainer mentioned supporting RaBitQ in #989.
We have implemented RaBitQ quantization for both the IVFFlat and HNSW access method in pgvector. The implementation supports cosine distance and does not change the existing query flow or SQL interface.
It has seen some internal use in IVF RaBitQ, supporting cosine-distance search on single heap table containing hundreds of millions of vectors.
We would like to ask whether the maintainers are interested in accepting this work into pgvector.
Benchmark
Our ann-benchmark results show that RaBitQ reduces index storage and improves query performance while maintaining competitive recall.
- After each index was built,
pg_prewarmwas used to load it into the PostgreSQL buffer cache before running search benchmarks. shared_bufferswas configured to be large enough to hold the indexes. This minimizes storage I/O during search and makes the benchmark focus on the distance-computation speedup provided by RaBitQ.
Environment
- Hardware:
- CPU: Intel Xeon Platinum 8269CY @ 2.50GHz
- OS kernel version: Linux 5.10.134
- Software:
- PostgreSQL
- version: v17.9
- shared_buffers = 10GB
- max_parallel_workers = 4
- pgvector base version: v0.8.0
- PostgreSQL
- Benchmark configuration
- TopK = 10
- dataset: dbpedia-openai-1000k-angular (1536D1M)
- IVFFlat configuration:
- lists = 1000
- probes = 1, 2, 5, 10, 20, 50, 100
- quantization_bit = 1
- HNSW configuration:
- m = 16
- ef_construction = 64
- ef_search = 40, 60, 80, 100, 120, 140, 160, 180
- quantization_bit = 7
IVFFlat
Build
| Implementation | Build time | Index size |
|---|---|---|
| IVFFLAT baseline | 144.9s | 7820MB |
| IVF RaBitQ | 120s | 281MB |
Search performance
probes |
Baseline QPS | Baseline Recall@10 | RaBitQ QPS | RaBitQ Recall@10 | Speed up |
|---|---|---|---|---|---|
| 1 | 224.957 | 68.9% | 690.172 | 69.2% | 3.07x |
| 2 | 141.526 | 81.1% | 499.285 | 81.0% | 3.53x |
| 5 | 76.083 | 90.6% | 297.533 | 90.1% | 3.91x |
| 10 | 58.398 | 94.3% | 198.78 | 93.9% | 3.40x |
| 20 | 22.794 | 96.4% | 139.59 | 96.3% | 6.12x |
| 50 | 9.066 | 98.3% | 78.50 | 98.1% | 8.66x |
| 100 | 5.681 | 99.0% | 45.83 | 98.9% | 8.07x |
HNSW
Build
| Implementation | Build time | Index size |
|---|---|---|
| HNSW baseline | 359.76s | 7813MB |
| HNSW RaBitQ | 714.22s | 2604MB |
Search performance
ef_search |
Baseline QPS | Baseline Recall@10 | RaBitQ QPS | RaBitQ Recall@10 | Speed up |
|---|---|---|---|---|---|
| 40 | 395.243 | 99.0% | 447.133 | 98.7% | 1.13x |
| 60 | 295.626 | 99.2% | 353.140 | 99.2% | 1.19x |
| 80 | 277.763 | 99.4% | 300.351 | 99.3% | 1.08x |
| 100 | 225.656 | 99.4% | 245.344 | 99.4% | 1.09x |
| 120 | 187.349 | 99.4% | 213.504 | 99.4% | 1.14x |
| 140 | 185.545 | 99.4% | 208.332 | 99.3% | 1.12x |
| 160 | 167.264 | 99.5% | 193.393 | 99.3% | 1.16x |
| 180 | 151.726 | 99.5% | 175.822 | 99.4% | 1.16x |
Next steps
If the maintainers are interested, we are willing to prepare the code for an upstream contribution and do the additional work needed, including:
- Providing more detailed and reproducible benchmarks
- Discussing design details, including the GUCs, index reloptions, operator classes, and extension upgrade objects requires by the design
- Organizing the implemented features into focused, independent commits to make review easier
- Improving the implementation based on maintainer feedback
- Adding support for other distance type, such as L2 distance support
- Adding compatibility and fallback support for different CPU architectures
- Adding documentation
We would appreciate guidance on whether this feature would be of interest and how the maintainers would prefer us to proceed.
Source: pgvector/pgvector