diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..6e42511 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,99 @@ +name: Release wheels + +on: + push: + tags: + - "v*" + workflow_dispatch: + +permissions: + contents: read + +jobs: + linux-wheels: + name: Linux x86-64 wheels + runs-on: ubuntu-latest + timeout-minutes: 90 + + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + + - name: Check tag matches package version + if: github.event_name == 'push' + env: + RELEASE_TAG: ${{ github.ref_name }} + run: | + python - <<'PY' + import os + import tomllib + + with open("pyproject.toml", "rb") as file: + version = tomllib.load(file)["project"]["version"] + + expected_tag = f"v{version}" + if os.environ["RELEASE_TAG"] != expected_tag: + raise SystemExit( + f"release tag {os.environ['RELEASE_TAG']!r} does not match " + f"package version {version!r}; expected {expected_tag!r}" + ) + PY + + - name: Build and test wheels + uses: pypa/cibuildwheel@v4.2.0 + with: + output-dir: wheelhouse + + - name: Store wheel artifacts + uses: actions/upload-artifact@v4 + with: + name: linux-x86_64-wheels + path: wheelhouse/*.whl + if-no-files-found: error + + pypi-publish: + name: Publish wheels to PyPI + if: github.event_name == 'push' + needs: linux-wheels + runs-on: ubuntu-latest + environment: + name: pypi + url: https://pypi.org/project/rabitqlib/ + permissions: + id-token: write + + steps: + - name: Download wheel artifacts + uses: actions/download-artifact@v4 + with: + name: linux-x86_64-wheels + path: wheelhouse + + - name: Publish wheels + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: wheelhouse + + github-release: + name: Create GitHub release + if: github.event_name == 'push' + needs: [linux-wheels, pypi-publish] + runs-on: ubuntu-latest + permissions: + contents: write + + steps: + - name: Download wheel artifacts + uses: actions/download-artifact@v4 + with: + name: linux-x86_64-wheels + path: wheelhouse + + - name: Create GitHub release + env: + GH_TOKEN: ${{ github.token }} + RELEASE_TAG: ${{ github.ref_name }} + run: gh release create "$RELEASE_TAG" wheelhouse/*.whl --generate-notes --verify-tag diff --git a/README.md b/README.md index a750565..ddf9538 100644 --- a/README.md +++ b/README.md @@ -1,55 +1,81 @@ -# RaBitQ Library +
-[![C++ tests](https://github.com/VectorDB-NTU/RaBitQ-Library/actions/workflows/test.yaml/badge.svg)](https://github.com/VectorDB-NTU/RaBitQ-Library/actions/workflows/test.yaml) -[![Python tests](https://github.com/VectorDB-NTU/RaBitQ-Library/actions/workflows/python.yml/badge.svg)](https://github.com/VectorDB-NTU/RaBitQ-Library/actions/workflows/python.yml) -[![License](https://img.shields.io/badge/license-Apache--2.0-blue.svg)](LICENSE) +

RaBitQ Library

-RaBitQ Library is a C++17 library with Python bindings for compact, accurate -vector quantization and approximate nearest-neighbor search. It provides: +

Compact vectors. Accurate distances. Fast ANN search.

-- the [1-bit](https://arxiv.org/abs/2405.12497) and - [multi-bit](https://arxiv.org/abs/2409.09913) RaBitQ quantizers; -- IVF, HNSW, and [SymphonyQG](https://dl.acm.org/doi/abs/10.1145/3709730) - indexes powered by RaBitQ; -- Euclidean distance and inner-product search (cosine search is - available by normalizing vectors before using inner product); and -- optimized AVX2 and AVX-512 kernels with runtime CPU dispatch. +

+ A research-backed C++17 library with Python bindings for 1-bit and multi-bit
+ vector quantization, IVF, HNSW, and SymphonyQG. +

-RaBitQ is developed by the -[VectorDB group](https://vectordb-ntu.github.io/) at Nanyang Technological -University, Singapore. A GPU implementation is also available in -[cuvs_rabitq](https://github.com/Stardust-SJF/cuvs_rabitq/tree/cuvs_ivf_rabitq). +

+ PyPI + Python versions + C++ tests + Python tests + Documentation + Paper DOI + License +

-## Quick start +

+ Documentation · + Python package · + Paper · + Releases +

-### Python +
-#### Requirements +## Install -- Python 3.9 or newer -- a C++17 compiler -- CMake 3.15 or newer -- OpenMP -- an x86-64 CPU supported by the selected kernels: most paths accept either - AVX2 with FMA or AVX-512F/BW/DQ with FMA +```bash +pip install rabitqlib +``` -Most SIMD entry points select AVX-512 kernels when AVX-512F, AVX-512BW, and -AVX-512DQ are detected; otherwise they use AVX2 when AVX2 and FMA are -available. AVX-512 VPOPCNTDQ enables additional popcount kernels. The HNSW -AVX-512 core path also checks for AVX2 and FMA, and otherwise uses its AVX2 -path when available. AVX-512 translation units are compiled with FMA enabled. +Prebuilt wheels support Linux x86-64 and CPython 3.9–3.14. AVX2 + FMA is the +portable CPU baseline; supported AVX-512 kernels are selected at runtime. -On Ubuntu or Debian, install the system build tools and then install RaBitQ -from the repository: +## Adopted across the vector-search ecosystem -```bash -sudo apt-get update -sudo apt-get install -y build-essential cmake libomp-dev +[Milvus](https://github.com/milvus-io/milvus) · +[Faiss](https://github.com/facebookresearch/faiss) · +[VSAG](https://github.com/antgroup/vsag) · +[VectorChord](https://github.com/tensorchord/VectorChord) · +[Volcengine OpenSearch](https://www.volcengine.com/docs/6465/1553583) · +[CockroachDB](https://github.com/cockroachdb/cockroach) · +[Elasticsearch](https://github.com/elastic/elasticsearch) · +[Lucene](https://github.com/apache/lucene) · +[turbopuffer](https://turbopuffer.com/blog/ann-v3#:~:text=ANN%20v3%20employs%20the%20RaBitQ) · +[Zvec](https://github.com/alibaba/zvec) -git clone https://github.com/VectorDB-NTU/RaBitQ-Library.git -cd RaBitQ-Library -python -m pip install . -``` +## Accuracy at a glance + +![RaBitQ estimation error benchmark across MSong, YouTube, OpenAI embeddings, Word2Vec, and GIST](docs/docs/assets/img/acc_bench.png) + +*Average and maximum relative estimation error across six datasets; lower is +better. Results from the +[SIGMOD camera-ready paper](https://doi.org/10.1145/3725413).* + +## Why RaBitQ? + +| | | +| --- | --- | +| **Compact by design** | Choose [1-bit](https://arxiv.org/abs/2405.12497) or [multi-bit](https://doi.org/10.1145/3725413) codes to match your memory and accuracy target. | +| **Accurate estimates** | An asymptotically optimal theoretical error bound supports reliable ordering and reranking. | +| **Fast on x86-64** | Dedicated AVX2 and AVX-512 kernels are selected through runtime CPU dispatch. | +| **Ready for ANN search** | Use the quantizer directly or build complete IVF, HNSW, and [SymphonyQG](https://dl.acm.org/doi/abs/10.1145/3709730) indexes. | + +The library supports Euclidean distance and inner product. Cosine search is +available by normalizing vectors before using inner product. + +RaBitQ is developed by the +[VectorDB group](https://vectordb-ntu.github.io/) at Nanyang Technological +University, Singapore. A GPU implementation is also available in +[cuvs_rabitq](https://github.com/Stardust-SJF/cuvs_rabitq/tree/cuvs_ivf_rabitq). + +## Python quick start The following complete example builds a small IVF index and searches it. It uses deterministic synthetic data, so no dataset download is required. @@ -86,20 +112,41 @@ Python bindings are also available for `HnswIndex` and `SymqgIndex`. See the [Python examples](sample/python/) for index construction, querying, and index persistence. -### C++ +
+Build the Python bindings from source + +Source builds require a C++17 compiler, CMake 3.15 or newer, and OpenMP. On +Ubuntu or Debian: + +```bash +sudo apt-get update +sudo apt-get install -y build-essential cmake libomp-dev +git clone https://github.com/VectorDB-NTU/RaBitQ-Library.git +cd RaBitQ-Library +python -m pip install . +``` + +
+ +## C++ quick start -#### Requirements +### Requirements - CMake 3.15 or newer - a C++17 compiler with OpenMP support - an x86-64 CPU supported by the selected kernels: most paths accept either AVX2 with FMA or AVX-512F/BW/DQ with FMA +
+CPU dispatch details + Most SIMD entry points select AVX-512 kernels when AVX-512F, AVX-512BW, and AVX-512DQ are detected; otherwise they use AVX2 when AVX2 and FMA are available. AVX-512 VPOPCNTDQ enables additional popcount kernels. The HNSW -AVX-512 core path also checks for AVX2 and FMA, and otherwise uses its AVX2 -path when available. AVX-512 translation units are compiled with FMA enabled. +AVX-512 core path also checks for AVX2 and FMA. AVX-512 translation units are +compiled with FMA enabled. + +
Clone and build the library and example programs: @@ -139,45 +186,22 @@ GoogleTest is downloaded during test configuration. For a full benchmark on the GIST dataset, see [`example.sh`](example.sh). More detailed API and algorithm guidance is available in the [documentation](docs/docs/index.md). -## Contributing +## Choose the right building block -Contributions are welcome. See the [contributing guide](CONTRIBUTING.md) for -the build, formatting, pre-commit, and static-analysis workflows. +| Component | Best fit | Storage and search profile | +| --- | --- | --- | +| **Quantizer** | Integrating RaBitQ into an existing system | Low-level 1-bit or multi-bit encoding and distance estimation. | +| **IVF** | Memory-efficient partitioned search | Stores quantized codes without retaining the raw dataset. | +| **HNSW** | Graph search with compact vectors | Adds graph links and searches directly from quantized codes. | +| **SymphonyQG** | Query speed when more memory is available | Retains raw vectors and stores per-neighborhood quantization data. | -## Why RaBitQ? - -- **High accuracy with tiny codes.** RaBitQ provides state-of-the-art similarity - estimation across different bit widths and remains effective with a - one-bit code per padded dimension plus per-vector factors. -- **Fast distance estimation.** IVF and SymphonyQG use - [FastScan](https://arxiv.org/abs/1704.07355) for batched estimates, while - HNSW uses single-code AVX2 or AVX-512 kernels. -- **Theoretical error bounds.** RaBitQ provides an asymptotically optimal error - bound that can support reliable ordering and reranking. -- **Multiple index trade-offs.** IVF stores quantized codes without the raw - dataset. HNSW adds graph links but also searches from quantized codes. - SymphonyQG retains raw vectors and stores per-neighborhood quantization data - to improve its access pattern. +IVF and SymphonyQG use [FastScan](https://arxiv.org/abs/1704.07355) for batched +estimates, while HNSW uses single-code AVX2 or AVX-512 kernels. In typical workloads, 4-bit, 5-bit, and 7-bit quantization can achieve roughly 90%, 95%, and 99% recall, respectively, without reranking. Actual results depend on the dataset, index configuration, and search parameters. -## RaBitQ in industry - -RaBitQ has been adopted by vector databases, search engines, and libraries: - -- [Milvus](https://github.com/milvus-io/milvus) — IVF + RaBitQ (C++) -- [Faiss](https://github.com/facebookresearch/faiss) — IVF + RaBitQ (C++) -- [VSAG](https://github.com/antgroup/vsag) — HGraph + RaBitQ (C++) -- [VectorChord](https://github.com/tensorchord/VectorChord) — IVF + RaBitQ (Rust) -- [Volcengine OpenSearch](https://www.volcengine.com/docs/6465/1553583) — DiskANN + RaBitQ -- [CockroachDB](https://github.com/cockroachdb/cockroach) — CSPANN + RaBitQ (Go) -- [Elasticsearch](https://github.com/elastic/elasticsearch) — HNSW + BBQ, a modified RaBitQ implementation (Java) -- [Lucene](https://github.com/apache/lucene) — HNSW + BBQ, a modified RaBitQ implementation (Java) -- [turbopuffer](https://turbopuffer.com/blog/ann-v3#:~:text=ANN%20v3%20employs%20the%20RaBitQ) — SPFresh + RaBitQ -- [Zvec](https://github.com/alibaba/zvec) — HNSW/IVF + RaBitQ (C++) - ## Citation If RaBitQ helps your research or system, please cite: @@ -185,7 +209,13 @@ If RaBitQ helps your research or system, please cite: > Jianyang Gao, Yutong Gou, Yuexuan Xu, Yongyi Yang, Cheng Long, and Raymond > Chi-Wing Wong. “Practical and Asymptotically Optimal Quantization of > High-Dimensional Vectors in Euclidean Space for Approximate Nearest Neighbor -> Search.” SIGMOD 2025. [arXiv:2409.09913](https://arxiv.org/abs/2409.09913). +> Search.” *Proceedings of the ACM on Management of Data* 3, 3, Article 202 +> (June 2025), 26 pages. [https://doi.org/10.1145/3725413](https://doi.org/10.1145/3725413). + +## Contributing + +Contributions are welcome. See the [contributing guide](CONTRIBUTING.md) for +the build, formatting, pre-commit, and static-analysis workflows. ## Acknowledgements diff --git a/docs/docs/assets/img/acc_bench.png b/docs/docs/assets/img/acc_bench.png new file mode 100644 index 0000000..11a9c35 Binary files /dev/null and b/docs/docs/assets/img/acc_bench.png differ diff --git a/docs/docs/index.md b/docs/docs/index.md index ff8b4ab..0a54086 100644 --- a/docs/docs/index.md +++ b/docs/docs/index.md @@ -33,6 +33,14 @@ indexes backed by optimized AVX2 and AVX-512 kernels. +## Accuracy at a glance + +![RaBitQ estimation error benchmark across MSong, YouTube, OpenAI embeddings, Word2Vec, and GIST](assets/img/acc_bench.png) + +*Average and maximum relative estimation error across six datasets; lower is +better. Results from the +[SIGMOD camera-ready paper](https://doi.org/10.1145/3725413).* + ## Start with Python Install the latest release from PyPI: @@ -103,7 +111,7 @@ degree, and search parameters. The library supports Euclidean distance and inner product. Cosine similarity can be implemented by normalizing vectors and using inner product. It implements the [1-bit RaBitQ](https://arxiv.org/abs/2405.12497) and -[multi-bit RaBitQ](https://arxiv.org/abs/2409.09913) research from the +[multi-bit RaBitQ](https://doi.org/10.1145/3725413) research from the [VectorDB Group](https://vectordb-ntu.github.io/) at Nanyang Technological University. @@ -114,6 +122,7 @@ RaBitQ has been adopted by projects including [Faiss](https://github.com/facebookresearch/faiss), [VSAG](https://github.com/antgroup/vsag), [VectorChord](https://github.com/tensorchord/VectorChord), +[Volcengine OpenSearch](https://www.volcengine.com/docs/6465/1553583), [CockroachDB](https://github.com/cockroachdb/cockroach), [Elasticsearch](https://github.com/elastic/elasticsearch), [Lucene](https://github.com/apache/lucene), @@ -127,4 +136,5 @@ If RaBitQ helps your research or system, please cite: > Jianyang Gao, Yutong Gou, Yuexuan Xu, Yongyi Yang, Cheng Long, and Raymond > Chi-Wing Wong. “Practical and Asymptotically Optimal Quantization of > High-Dimensional Vectors in Euclidean Space for Approximate Nearest Neighbor -> Search.” SIGMOD 2025. [arXiv:2409.09913](https://arxiv.org/abs/2409.09913). +> Search.” *Proceedings of the ACM on Management of Data* 3, 3, Article 202 +> (June 2025), 26 pages. [https://doi.org/10.1145/3725413](https://doi.org/10.1145/3725413). diff --git a/docs/docs/rabitq/quantizer.md b/docs/docs/rabitq/quantizer.md index 0efc8a6..4d747d6 100644 --- a/docs/docs/rabitq/quantizer.md +++ b/docs/docs/rabitq/quantizer.md @@ -39,7 +39,7 @@ The binary code is easily computed by the function `one_bit_code` in `rabitq_imp The computation of ex-codes includes two versions of implementation. -In the first implementation, we compute the ex-codes of RaBitQ based on the algorithm described in the RaBitQ [paper](https://arxiv.org/abs/2409.09913) (Section 3.2.2). For a vector, to minimize the quantization error, the algorithm tries many different rescaling factors. For each rescaling factor, it rescales the vector and performs rounding (i.e., scalar quantization) to generate a quantization code. Then it finds out the factor and codes which minimizes the quantization error. Note that in the library, the range of enumeration is approriately shrinked, which brings better efficiency without affecting the accuracy. +In the first implementation, we compute the ex-codes of RaBitQ based on the algorithm described in the RaBitQ [SIGMOD paper](https://doi.org/10.1145/3725413) (Section 3.2.2). For a vector, to minimize the quantization error, the algorithm tries many different rescaling factors. For each rescaling factor, it rescales the vector and performs rounding (i.e., scalar quantization) to generate a quantization code. Then it finds out the factor and codes which minimizes the quantization error. Note that in the library, the range of enumeration is approriately shrinked, which brings better efficiency without affecting the accuracy. In the second implementation, instead of enumerating different rescaling factors, it directly rounds every vector based on the **expected optimal factor**. Specifically, recall that all data vectors are randomly rotated before quantization. The expected optimal factor is computed as follows. We sample several random vectors which follow uniform distribution on the unit sphere and use the first implementation to quantize them. We record the optimal factor for each and take the average of the optimal factors as the expected optimal factor. This implementation introduces some accuracy decrease while significantly speeds up the quantization. diff --git a/docs/docs/rabitq/rabitq.md b/docs/docs/rabitq/rabitq.md index 589941d..06206ec 100644 --- a/docs/docs/rabitq/rabitq.md +++ b/docs/docs/rabitq/rabitq.md @@ -1,6 +1,6 @@ # RaBitQ -The RaBitQ algorithm is a drop-in replacement of binary quantization and (uniform) scalar quantization, with its [1-bit version](https://arxiv.org/abs/2405.12497) (released in May 2024) and [multi-bit version](https://arxiv.org/abs/2409.09913) (released in Sep 2024), respectively. +The RaBitQ algorithm is a drop-in replacement of binary quantization and (uniform) scalar quantization, with its [1-bit version](https://arxiv.org/abs/2405.12497) (released in May 2024) and [multi-bit version](https://doi.org/10.1145/3725413) (published at SIGMOD 2025), respectively. diff --git a/pyproject.toml b/pyproject.toml index b94342c..e5fb4f0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "rabitqlib" -version = "0.1.2" +version = "0.2.0" description = "RaBitQ Python bindings for HNSW, IVF, and SymQG" readme = "README.md" requires-python = ">=3.9" @@ -23,6 +23,13 @@ cmake.args = [ "-DCMAKE_BUILD_TYPE=Release" ] +[tool.cibuildwheel] +build = "cp39-manylinux_x86_64 cp310-manylinux_x86_64 cp311-manylinux_x86_64 cp312-manylinux_x86_64 cp313-manylinux_x86_64 cp314-manylinux_x86_64" +build-verbosity = 1 +environment = { CMAKE_ARGS = "-DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF" } +test-requires = "pytest" +test-command = "pytest {project}/tests/python --tb=short -q" + [tool.pytest.ini_options] testpaths = ["tests/python"]