qann 0.1.0a1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (41) hide show
  1. qann-0.1.0a1/.github/workflows/ci.yml +68 -0
  2. qann-0.1.0a1/.github/workflows/wheels.yml +55 -0
  3. qann-0.1.0a1/.gitignore +42 -0
  4. qann-0.1.0a1/ARCHITECTURE.md +306 -0
  5. qann-0.1.0a1/CMakeLists.txt +174 -0
  6. qann-0.1.0a1/LICENSE +21 -0
  7. qann-0.1.0a1/PKG-INFO +73 -0
  8. qann-0.1.0a1/README.md +25 -0
  9. qann-0.1.0a1/benchmarks/baseline_siftsmall.json +98 -0
  10. qann-0.1.0a1/benchmarks/bench_distance.cpp +76 -0
  11. qann-0.1.0a1/benchmarks/bench_indexes.py +230 -0
  12. qann-0.1.0a1/benchmarks/compare.py +136 -0
  13. qann-0.1.0a1/benchmarks/profile_query.cpp +95 -0
  14. qann-0.1.0a1/bindings/python_bindings.cpp +264 -0
  15. qann-0.1.0a1/build.sh +19 -0
  16. qann-0.1.0a1/include/vecengine/aligned_allocator.hpp +82 -0
  17. qann-0.1.0a1/include/vecengine/distance.hpp +39 -0
  18. qann-0.1.0a1/include/vecengine/distance_inl.hpp +40 -0
  19. qann-0.1.0a1/include/vecengine/flat_index.hpp +31 -0
  20. qann-0.1.0a1/include/vecengine/index.hpp +38 -0
  21. qann-0.1.0a1/include/vecengine/index_factory.hpp +31 -0
  22. qann-0.1.0a1/include/vecengine/ivf_index.hpp +69 -0
  23. qann-0.1.0a1/include/vecengine/pq.hpp +102 -0
  24. qann-0.1.0a1/include/vecengine/refine_index.hpp +44 -0
  25. qann-0.1.0a1/include/vecengine/threads.hpp +20 -0
  26. qann-0.1.0a1/pyproject.toml +56 -0
  27. qann-0.1.0a1/src/distance.cpp +186 -0
  28. qann-0.1.0a1/src/flat_index.cpp +52 -0
  29. qann-0.1.0a1/src/index.cpp +28 -0
  30. qann-0.1.0a1/src/index_factory.cpp +22 -0
  31. qann-0.1.0a1/src/ivf_index.cpp +224 -0
  32. qann-0.1.0a1/src/kmeans.cpp +229 -0
  33. qann-0.1.0a1/src/kmeans.hpp +36 -0
  34. qann-0.1.0a1/src/parallel.hpp +63 -0
  35. qann-0.1.0a1/src/pq.cpp +127 -0
  36. qann-0.1.0a1/src/profile.hpp +28 -0
  37. qann-0.1.0a1/src/refine_index.cpp +74 -0
  38. qann-0.1.0a1/src/threads.cpp +22 -0
  39. qann-0.1.0a1/src/topk.hpp +68 -0
  40. qann-0.1.0a1/tests/test_distance.cpp +814 -0
  41. qann-0.1.0a1/tests/test_python.py +35 -0
@@ -0,0 +1,68 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+
8
+ jobs:
9
+ test:
10
+ name: build + test (${{ matrix.os }})
11
+ runs-on: ${{ matrix.os }}
12
+ strategy:
13
+ fail-fast: false
14
+ matrix:
15
+ # ubuntu exercises the find_package(BLAS) path with OpenBLAS;
16
+ # macos-14 (arm64) uses Accelerate like local development.
17
+ os: [ubuntu-latest, macos-14]
18
+ steps:
19
+ - uses: actions/checkout@v4
20
+
21
+ - uses: actions/setup-python@v5
22
+ with:
23
+ python-version: "3.12"
24
+
25
+ - name: Install OpenBLAS
26
+ if: runner.os == 'Linux'
27
+ run: sudo apt-get update && sudo apt-get install -y libopenblas-dev
28
+
29
+ - name: Install Python dependencies
30
+ run: python -m pip install numpy pytest
31
+
32
+ - name: Configure
33
+ run: cmake -B build -DCMAKE_BUILD_TYPE=Release -DPython_EXECUTABLE="$(which python)"
34
+
35
+ - name: Build
36
+ run: cmake --build build -j 4
37
+
38
+ - name: C++ tests
39
+ run: ctest --test-dir build --output-on-failure
40
+
41
+ - name: Python tests
42
+ env:
43
+ PYTHONPATH: build
44
+ run: python -m pytest tests
45
+
46
+ - name: Cache siftsmall
47
+ if: runner.os == 'macOS'
48
+ id: sift-cache
49
+ uses: actions/cache@v4
50
+ with:
51
+ path: benchmarks/data/siftsmall
52
+ key: siftsmall-v1
53
+
54
+ - name: Download siftsmall
55
+ if: runner.os == 'macOS' && steps.sift-cache.outputs.cache-hit != 'true'
56
+ run: |
57
+ mkdir -p benchmarks/data
58
+ curl --retry 5 --retry-delay 10 -fsSL ftp://ftp.irisa.fr/local/texmex/corpus/siftsmall.tar.gz \
59
+ | tar -xz -C benchmarks/data
60
+
61
+ # The baseline was recorded with Accelerate on Apple Silicon, so the
62
+ # gate runs on macOS only; OpenBLAS can round differently and shift
63
+ # k-means results.
64
+ - name: Recall regression gate (siftsmall)
65
+ if: runner.os == 'macOS'
66
+ env:
67
+ PYTHONPATH: build
68
+ run: python benchmarks/bench_indexes.py --check benchmarks/baseline_siftsmall.json
@@ -0,0 +1,55 @@
1
+ name: Wheels
2
+
3
+ # PRs touching packaging and manual runs build and test wheels only; pushing a
4
+ # v* tag also publishes to PyPI.
5
+ on:
6
+ workflow_dispatch:
7
+ pull_request:
8
+ paths: [pyproject.toml, CMakeLists.txt, .github/workflows/wheels.yml]
9
+ push:
10
+ tags: ["v*"]
11
+
12
+ jobs:
13
+ wheels:
14
+ name: wheels (${{ matrix.os }})
15
+ runs-on: ${{ matrix.os }}
16
+ strategy:
17
+ fail-fast: false
18
+ matrix:
19
+ os: [ubuntu-latest, ubuntu-24.04-arm, macos-14, macos-15-intel]
20
+ steps:
21
+ - uses: actions/checkout@v4
22
+
23
+ - uses: pypa/cibuildwheel@v4.3.0
24
+
25
+ - uses: actions/upload-artifact@v4
26
+ with:
27
+ name: wheels-${{ matrix.os }}
28
+ path: wheelhouse/*.whl
29
+
30
+ sdist:
31
+ runs-on: ubuntu-latest
32
+ steps:
33
+ - uses: actions/checkout@v4
34
+
35
+ - run: pipx run build --sdist
36
+
37
+ - uses: actions/upload-artifact@v4
38
+ with:
39
+ name: sdist
40
+ path: dist/*.tar.gz
41
+
42
+ publish:
43
+ if: startsWith(github.ref, 'refs/tags/v')
44
+ needs: [wheels, sdist]
45
+ runs-on: ubuntu-latest
46
+ environment: pypi
47
+ permissions:
48
+ id-token: write
49
+ steps:
50
+ - uses: actions/download-artifact@v4
51
+ with:
52
+ path: dist
53
+ merge-multiple: true
54
+
55
+ - uses: pypa/gh-action-pypi-publish@release/v1
@@ -0,0 +1,42 @@
1
+ # Build output
2
+ build/
3
+ build-prof/
4
+ cmake-build-*/
5
+
6
+ # CMake generated
7
+ CMakeCache.txt
8
+ CMakeFiles/
9
+ cmake_install.cmake
10
+ Makefile
11
+ CTestTestfile.cmake
12
+ compile_commands.json
13
+ _deps/
14
+
15
+ # Python artifacts
16
+ __pycache__/
17
+ *.py[cod]
18
+ *.egg-info/
19
+ dist/
20
+ *.egg
21
+ .eggs/
22
+
23
+ # nanobind / Python extension modules
24
+ *.so
25
+ *.dylib
26
+ *.pyd
27
+
28
+ # macOS
29
+ .DS_Store
30
+
31
+ # IDE
32
+ .vscode/
33
+ .idea/
34
+ *.swp
35
+ *.swo
36
+
37
+ # clangd index
38
+ .cache/
39
+
40
+ # benchmark datasets and outputs
41
+ benchmarks/data/
42
+ benchmarks/results/
@@ -0,0 +1,306 @@
1
+ # QANN: Architecture & File Reference
2
+
3
+ ## Overview
4
+
5
+ QANN (Quantized Approximate Nearest Neighbors) is a C++20 vector search library with a Python extension module, `qann`, built on nanobind. The C++ API lives in namespace `vecengine` under `include/vecengine/`. It has three layers:
6
+
7
+ 1. **Distance kernels**: squared L2 and cosine distance with scalar, AVX2+FMA and ARM NEON paths, selected at compile time.
8
+ 2. **Quantization and clustering**: a shared k-means helper (BLAS-accelerated when a CBLAS is available) and a Product Quantization (PQ) codebook.
9
+ 3. **Indexes**: `FlatIndex` (exact brute force), `IVFIndex` (inverted file, optionally PQ-compressed) and `RefineIndex` (exact re-ranking wrapper), behind a common `Index` interface.
10
+
11
+ All arithmetic lives in C++. Python gets zero-copy access to NumPy arrays via `nb::ndarray`. `benchmarks/bench_indexes.py` measures recall vs QPS on the SIFT datasets.
12
+
13
+ ---
14
+
15
+ ## Directory Structure
16
+
17
+ ```
18
+ vecengine/
19
+ ├── .github/workflows/ci.yml # build + test (Ubuntu, macOS), recall gate
20
+ ├── CMakeLists.txt # build orchestrator
21
+ ├── build.sh # configure + build + test in one command
22
+ ├── pyproject.toml # PEP 517 packaging (scikit-build-core)
23
+ ├── include/vecengine/ # public API
24
+ │ ├── aligned_allocator.hpp # std::allocator with 32/64-byte alignment
25
+ │ ├── distance.hpp # distance function declarations
26
+ │ ├── distance_inl.hpp # inline compile-time SIMD dispatch
27
+ │ ├── index.hpp # Index interface + Neighbor
28
+ │ ├── flat_index.hpp # exact brute-force index
29
+ │ ├── ivf_index.hpp # IVF index, optional PQ
30
+ │ ├── pq.hpp # PQCodebook + PQDistance
31
+ │ ├── refine_index.hpp # exact re-ranking wrapper
32
+ │ └── index_factory.hpp # make_index() + IndexOptions
33
+ ├── src/
34
+ │ ├── distance.cpp # scalar / AVX2 / NEON kernels
35
+ │ ├── kmeans.hpp / kmeans.cpp # internal k-means helper (not public API)
36
+ │ ├── parallel.hpp # internal parallel_for (not public API)
37
+ │ ├── index.cpp # default Index::add_batch / query_batch
38
+ │ ├── pq.cpp
39
+ │ ├── flat_index.cpp
40
+ │ ├── ivf_index.cpp
41
+ │ ├── refine_index.cpp
42
+ │ └── index_factory.cpp
43
+ ├── tests/test_distance.cpp # Catch2 v3 tests for every layer
44
+ ├── benchmarks/
45
+ │ ├── bench_distance.cpp # Google Benchmark distance kernels
46
+ │ ├── bench_indexes.py # recall vs QPS on SIFT
47
+ │ ├── compare.py # ref vs working tree: recall, QPS, build times
48
+ │ ├── baseline_siftsmall.json # recall baseline for the CI gate
49
+ │ ├── data/ # downloaded datasets (git-ignored)
50
+ │ └── results/ # CSV + PNG output (git-ignored)
51
+ └── bindings/python_bindings.cpp # nanobind module
52
+ ```
53
+
54
+ ---
55
+
56
+ ## Layer 1: Distance Kernels
57
+
58
+ ### `distance.hpp` / `distance_inl.hpp` / `distance.cpp`
59
+
60
+ ```cpp
61
+ float l2_distance(const float* a, const float* b, size_t n); // squared L2
62
+ float cosine_distance(const float* a, const float* b, size_t n); // 1 - cos(a, b)
63
+ ```
64
+
65
+ `l2_distance` returns **squared** L2. It ranks neighbors identically to true L2 and skips the `sqrt`. Every index and PQ distance in this repo is squared L2.
66
+
67
+ The unsuffixed functions dispatch at compile time in `distance_inl.hpp`:
68
+
69
+ | Build | Path |
70
+ |---|---|
71
+ | x86-64 Release (`-mavx2 -mfma`) | `_avx2`: 8 floats per iteration, `_mm256_fmadd_ps` |
72
+ | arm64 (Apple Silicon, Graviton) | `_neon`: 4 floats per iteration |
73
+ | anything else | `_scalar` |
74
+
75
+ The `_scalar`, `_avx2` and `_neon` symbols always exist (non-native paths forward to scalar) so tests and benchmarks can call any of them directly. There is no runtime CPU detection, function pointer or virtual call.
76
+
77
+ These kernels live in `distance.cpp` and are not inlined into other translation units, so each call has a fixed overhead. That is negligible on full-length vectors (128 floats) but dominates on 4 to 16 float PQ slices, which is why the PQ table builds below avoid calling them per centroid.
78
+
79
+ Cosine returns `1.0f` for zero-magnitude inputs rather than dividing by zero.
80
+
81
+ ### `aligned_allocator.hpp`
82
+
83
+ `AlignedAllocator<T, Align>` uses `posix_memalign` (or `_aligned_malloc` on MSVC). Aliases: `AVX2Vector<T>` (32-byte) and `AVX512Vector<T>` (64-byte). Used by the benchmarks; the indexes use plain `std::vector` and unaligned loads.
84
+
85
+ ---
86
+
87
+ ## Layer 2: Clustering and Quantization
88
+
89
+ ### `src/kmeans.hpp` / `kmeans.cpp` (`vecengine::detail`)
90
+
91
+ Shared by `IVFIndex` and `PQCodebook`, so Lloyd's algorithm exists exactly once.
92
+
93
+ ```cpp
94
+ KMeansResult kmeans(const float* data, size_t n, size_t stride, size_t d, size_t k, size_t max_iters);
95
+ void assign_nearest(const float* data, size_t n, size_t stride, const float* centroids, size_t k, size_t d, uint32_t* out);
96
+ size_t nearest_centroid(const float* point, const float* centroids, size_t k, size_t d);
97
+ ```
98
+
99
+ - Point `i` is the `d` floats at `data + i * stride`. `stride` and `d` are separate so PQ can cluster one subspace slice of every row in place: pass `data + s * sub_dim`, `stride = dim`, `d = sub_dim`. IVF passes `stride = d = dim`.
100
+ - Returned centroids are packed (`k * d`, centroid `c` at `c * d`), never strided.
101
+ - Init picks `k` distinct random points (partial Fisher-Yates). Empty clusters are handled by `reseed_empty_clusters`: one pass keeps the E furthest points (distance to their assigned centroid, skipping points within `1e-6`) in a size-E min-heap, and each of the E empty clusters gets a distinct one. If the heap runs out, the largest cluster is split FAISS-style: its centroid is copied and the two copies are scaled by `1 ± 1/1024`, and its size is halved so the next split picks another cluster.
102
+ - A final assignment pass runs after the last update, so `assignments` match the returned `centroids`. IVF relies on this to compute PQ training residuals.
103
+
104
+ **Assignment is ~99.9% of k-means time**, so it is done in bulk by `assign_nearest`:
105
+
106
+ - **With BLAS (`VECENGINE_USE_BLAS`)**: `argmin_c ‖x − c‖² = argmin_c (‖c‖² − 2 x·c)`, and the `x·c` terms for a block of 1024 points against all centroids are one `cblas_sgemm` call. On Apple that is Accelerate, which runs on the M-series matrix unit; elsewhere any BLAS that ships `cblas.h` (OpenBLAS, netlib CBLAS). `lda = stride` lets PQ subspaces use strided input without copying. The per-row argmin is branchless (conditional selects).
107
+ - **Without BLAS**: falls back to `nearest_centroid` per point.
108
+
109
+ Measured on SIFT's 100k learn set: IVF assignment (k = 1000) 1.54 s → 0.086 s (18×), PQ subspaces 2.9 to 3.7×, with identical assignments. The `‖c‖² − 2x·c` scores are only used for ranking, never as distances.
110
+
111
+ `nearest_centroid` is the single-point scan, still used by `IVFIndex::add`. `IVFIndex::add_batch` uses `assign_nearest`.
112
+
113
+ `kmeans` takes an optional `seed`; without one it seeds from `std::random_device`, so training is non-deterministic run to run. `IVFIndex::train` and `PQCodebook::train` take an optional seed too (Python: `train(x, seed=1)`): coarse k-means uses `seed`, PQ uses `seed + 1` onward, one per subspace. With the same seed and the same BLAS, results are bit-identical, which the recall gate relies on.
114
+
115
+ ### `pq.hpp` / `pq.cpp`: `PQCodebook`
116
+
117
+ Splits a `dim`-vector into `m = num_subspaces` slices of `sub_dim = dim / m` floats and quantizes each slice to one of `K = centroids_per_subspace` (max 256) centroids. A vector becomes `m` bytes: at `dim = 128`, `m = 16` that is 512 bytes down to 16.
118
+
119
+ | Member | Layout | Used by |
120
+ |---|---|---|
121
+ | `centroids_` | `m * K * sub_dim`, `[subspace][centroid][dim]` | reading whole centroids |
122
+ | `centroids_t_` | `m * sub_dim * K`, `[subspace][dim][centroid]` (transposed copy) | ADC and SDC table builds |
123
+ | `sdc_table_` | `m * K * K`: distance between centroids `c1`, `c2` of subspace `s` at `s*K*K + c1*K + c2` | `distance_sdc` |
124
+
125
+ All three are built in `train()`.
126
+
127
+ - `train(vectors, n)`: one `kmeans` call per subspace, run in parallel with `parallel_for`, then the transposed copy, then `sdc_table_`.
128
+ - `compute_adc_table(query, out)` + `distance_adc(table, code)`: **ADC**. The caller-provided `out` holds `m * K` distances from each query slice to every centroid (`table[s * K + c]`); the distance is `m` lookups. `distance_adc` is defined in the header so it inlines into the IVF scan loop.
129
+ - `encode(vec)`: builds the ADC table for `vec`, then takes the argmin of each subspace's row (a vectorized min pass, then a short search for its position).
130
+ - `distance_sdc(query_code, code)`: **SDC**. Both sides are codes; `m` lookups into `sdc_table_`, no float math.
131
+
132
+ **Why `centroids_t_`:** the table loops compute, for one dimension, `out[c] += (q_d − centroid_c[d])²` over all `K` centroids. In `[dim][centroid]` order those `K` values are contiguous, so the compiler vectorizes across centroids with no function calls or horizontal sums. The previous version made `K * m` separate `l2_distance` calls on 4 to 16 floats each (8,192 calls for m = 32). The change made ADC queries 2.0× (m = 8) to 3.6× (m = 32) faster on siftsmall with identical results.
133
+
134
+ Both distances work because squared L2 decomposes across disjoint slices: `‖a − b‖² = Σₛ ‖aₛ − bₛ‖²`. ADC keeps the query exact and is more accurate; SDC quantizes both sides.
135
+
136
+ `PQDistance { ADC, SDC }` lives here and selects the mode in `IVFIndex`.
137
+
138
+ ---
139
+
140
+ ## Layer 3: Indexes
141
+
142
+ ### `index.hpp`: `Index`
143
+
144
+ ```cpp
145
+ virtual void add(span<const float> vec) = 0;
146
+ virtual void add_batch(span<const float> vecs, size_t n);
147
+ virtual vector<Neighbor> query(span<const float> vec, size_t k) const = 0;
148
+ virtual vector<vector<Neighbor>> query_batch(span<const float> queries, size_t num_queries, size_t k) const;
149
+ virtual size_t size() const noexcept = 0;
150
+ virtual size_t dim() const noexcept = 0;
151
+ ```
152
+
153
+ `Neighbor { size_t index; float distance; }`. Results are sorted nearest first.
154
+
155
+ **Invariant:** `add()` assigns IDs in insertion order starting at 0, and `Neighbor::index` is that ID. `RefineIndex` depends on this. `add_batch` must keep it: row `i` of a batch gets the next ID in order.
156
+
157
+ Defaults in `src/index.cpp`: `add_batch` calls `add` per row; `query_batch` calls `query` for each row in parallel via `parallel_for`. This is safe because `query` is `const` and touches no shared mutable state, and each result goes to its own slot. Subclasses override `add_batch` only when a whole batch allows a faster path (`IVFIndex`, `RefineIndex`).
158
+
159
+ ### `src/parallel.hpp`: `parallel_for` (`vecengine::detail`)
160
+
161
+ `parallel_for(n, fn)` calls `fn(i)` once for each `i` in `[0, n)` on up to `num_threads()` `std::thread`s, the caller being one of them. With one thread (or `n <= 1`) it is a plain loop on the caller: no atomics, no threads, exceptions propagate directly. Workers take the next index from a shared atomic counter, so uneven per-item cost (IVF list sizes, M-series performance vs efficiency cores) doesn't leave threads idle. The first exception thrown by `fn` stops the other workers and is rethrown on the caller after all join. Threads are created per call; at batch and training granularity that cost is negligible, so there is no pool. Don't nest it: `fn` must not call `parallel_for` itself.
162
+
163
+ **Thread count** (`include/vecengine/threads.hpp`, `src/threads.cpp`): `set_num_threads(n)` sets one process-wide count in a relaxed `std::atomic`, safe to change while other threads query. `0` (the default) means one thread per core; `num_threads()` resolves `0` to `hardware_concurrency()` (at least 1) on every read without writing it back. It does not cap the BLAS library's own threads inside `sgemm`; use `VECLIB_MAXIMUM_THREADS` / `OPENBLAS_NUM_THREADS` for those.
164
+
165
+ ### `flat_index.hpp`: `FlatIndex`
166
+
167
+ Exact search. Vectors stored row-major in one `vector<float>`. `query` scans all of them with a size-`k` max-heap and drains it back to front for ascending order. This is the ground truth for recall tests.
168
+
169
+ ### `ivf_index.hpp`: `IVFIndex`
170
+
171
+ `IVFIndex(dim, nlist, nprobe = 10)`, with `nlist` in `[100, 65535]` (cluster IDs are `uint16_t`).
172
+
173
+ **Lifecycle:** `enable_pq(...)` (optional) → `train(sample)` → `add(v)` for each vector → `query`. `train` learns the model only; it stores no vectors. Training vectors are not searchable unless they are also added.
174
+
175
+ - **`train`**: one `kmeans` call for the `nlist` coarse centroids. With PQ on, computes residuals `v − centroid(assignment)` from `kmeans`'s assignments and trains the codebook on them.
176
+ - **`add`**: `nearest_centroid` picks the cluster; the vector's ID is appended to that cluster's `InvertedList`, plus either the raw floats (`vecs`) or the PQ code of its residual (`codes`).
177
+ - **`add_batch`**: three phases. (1) One `assign_nearest` call assigns every row. (2) With PQ on, `parallel_for` encodes each residual into its own slot of an `n * m` buffer. (3) A serial loop appends IDs and data to the lists in input order, because `unordered_map` and list growth are not thread-safe and the ID invariant needs input order. `assign_nearest` ranks by `‖c‖² − 2x·c`, so a near-tie can pick a different list than `add` would.
178
+ - **`query`**:
179
+ 1. Coarse: find the `nprobe` nearest coarse centroids.
180
+ 2. Fine: scan only those lists with a size-`k` max-heap. One ADC table buffer is allocated per query and reused for every list. Per list, all query-and-cluster-dependent work happens once, before the per-vector loop:
181
+ - PQ off: `l2_distance` per stored vector.
182
+ - ADC: residual `q − centroid(c)` → `compute_adc_table` → `distance_adc` per code.
183
+ - SDC: residual → `encode` → `distance_sdc` per code.
184
+
185
+ The residual is rebuilt per probed list because it depends on the cluster's centroid. Stored codes encode `x − c`, the query becomes `q − c`, and `(q − c) − (x − c) = q − x`, so residual distances equal true distances. Inside IVF, SDC costs about the same as ADC (both need a full pass over the centroids per list) and is less accurate. ADC is the default.
186
+
187
+ `nprobe` and `pq_distance` both have setters and can be changed at any time after training; they only affect queries.
188
+
189
+ ### `refine_index.hpp`: `RefineIndex`
190
+
191
+ `RefineIndex(base, k_factor = 10)` wraps an **empty, trained** approximate index by reference and keeps its own row-major copy of the raw vectors.
192
+
193
+ - `add(v)`: forwards to `base`, then appends the raw floats, so ID `i` is at `data_[i * dim]`.
194
+ - `add_batch(vecs, n)`: forwards to `base.add_batch` and appends all rows at once.
195
+ - `query(q, k)`: asks `base` for `min(k * k_factor, size())` candidates, replaces each approximate distance with the exact `l2_distance` against its raw vector, and returns the exact top `k`.
196
+
197
+ It only helps when `base` ranks with approximate distances (IVF + PQ). It restores recall but stores the raw vectors in RAM, so IVF + PQ + re-ranking uses more memory than plain IVF; in production the raw vectors would live on disk and only the candidates would be read. `k_factor` is a query-time setting; `k_factor = 1` returns the base's IDs with exact distances.
198
+
199
+ ### `index_factory.hpp`: `make_index`
200
+
201
+ `make_index(IndexType, dim, IndexOptions)` returns `unique_ptr<Index>`. `IndexOptions` holds `capacity` (Flat), `nlist`, `nprobe`, `pq_subspaces` (0 = PQ off) and `pq_centroids` (IVF). IVF-specific calls like `train()` require the concrete `IVFIndex`.
202
+
203
+ ---
204
+
205
+ ## Benchmarks (`benchmarks/bench_indexes.py`)
206
+
207
+ Recall@10 (overlap of the returned top 10 with the true top 10), QPS on all cores (`query_batch`) and bytes per vector for Flat, IVF, IVF + PQ (ADC and SDC, m ∈ {8, 16, 32}) and IVF + PQ + re-ranking (`k_factor` ∈ {4, 16}), sweeping `nprobe` from 1 to 128. Training uses each dataset's `learn` set; `nlist ≈ √N`. Each configuration's QPS is the best of 3 full passes, so cold caches after a build don't distort the first measurement. Writes `benchmarks/results/<dataset>.csv` and a recall-vs-QPS plot.
208
+
209
+ ```bash
210
+ PYTHONPATH=build python3 benchmarks/bench_indexes.py # siftsmall (10k)
211
+ PYTHONPATH=build python3 benchmarks/bench_indexes.py --check benchmarks/baseline_siftsmall.json
212
+ PYTHONPATH=build python3 benchmarks/bench_indexes.py --threads 1 # single-threaded (default 0 = all cores)
213
+ python3 benchmarks/compare.py [ref] [--dataset sift] [--rounds N] # ref (default HEAD) vs working tree
214
+ PYTHONPATH=build python3 benchmarks/bench_indexes.py --dataset sift # SIFT1M
215
+ PYTHONPATH=build python3 benchmarks/bench_indexes.py --max-queries 1000
216
+ ```
217
+
218
+ - **Recall gate**: `--seed N --write-baseline FILE` records recall for every row; `--check FILE` reruns with the baseline's dataset and seed and exits 1 if any row drops more than `--tolerance` (default 0.02). `benchmarks/baseline_siftsmall.json` (seed 1) is the committed baseline. Regenerate it when a change is *meant* to move recall.
219
+ - **`compare.py`**: builds `ref` in a temporary copy (reusing `build/_deps` sources), runs each side's own `bench_indexes.py` alternately for `--rounds`, keeps the best QPS and fastest build per row, and prints recall, QPS ratios and build times. Speed is compared locally rather than in CI because shared CI runners vary 10 to 20% run to run.
220
+
221
+ Datasets come from `ftp://ftp.irisa.fr/local/texmex/corpus/` (`siftsmall.tar.gz`, `sift.tar.gz`), extracted into `benchmarks/data/`.
222
+
223
+ **SIFT1M findings** (Apple M3, 4 performance + 4 efficiency cores, `nlist = 1000`, 10k queries, seed 1, idle machine):
224
+
225
+ Thread scaling at `nprobe = 16` (`set_num_threads`):
226
+
227
+ | threads | Plain IVF (recall 0.93) | IVF + PQ16 (0.56) | IVF + PQ16 + R16 (0.93) |
228
+ |---|---|---|---|
229
+ | 1 | 3.3k QPS | 5.7k | 3.6k |
230
+ | 2 | 6.4k (1.9×) | 11.0k (1.9×) | 7.4k (2.0×) |
231
+ | 4 | 7.6k (2.3×) | 19.2k (3.4×) | 13.3k (3.7×) |
232
+ | 8 | 10.6k (3.2×) | 24.5k (4.3×) | 18.0k (4.9×) |
233
+
234
+ - Plain IVF stalls between 2 and 4 threads while PQ keeps scaling: IVF streams 512-byte raw vectors and is limited by memory bandwidth, PQ scans 16-byte codes. The efficiency cores add 25 to 40% on top of 4 threads.
235
+ - On all cores, IVF + PQ16 + re-ranking is 1.7× faster than plain IVF at the same 0.93 recall. Single-threaded, plain IVF was faster; PQ's smaller scan wins once bandwidth is the bottleneck.
236
+ - Plain IVF on all cores: 10.2k QPS at 0.93 recall (`nprobe = 16`), 5.3k at 0.98 (`nprobe = 32`), 2.6k at 0.995 (`nprobe = 64`); Flat is 196 QPS.
237
+ - PQ alone plateaus at 0.38 / 0.57 / 0.73 recall for m = 8 / 16 / 32; re-ranking lifts m = 16 to 0.99 and m = 32 to 0.999.
238
+ - SDC recall *falls* as `nprobe` grows, because more candidates expose its query-side quantization error.
239
+
240
+ **Build times** on all cores: IVF trains in 2.5 s and adds 1M vectors in 1.3 s (`add_batch`); IVF + PQ16 trains in 3.4 s (8.6 s on one thread) and adds in 3.0 s. Before `assign_nearest`, IVF training took 37 s and IVF + PQ16 55 s.
241
+
242
+ The PQ rows of `benchmarks/results/sift.csv` from this run are ~20% low: other jobs loaded the machine during its second half. The thread-scaling numbers above come from a separate clean run.
243
+
244
+ ---
245
+
246
+ ## Python Bindings (`bindings/python_bindings.cpp`)
247
+
248
+ Inputs are float32 C-contiguous NumPy arrays, passed zero-copy as `nb::ndarray` (`FloatMatrix` 2D, `FloatVector` 1D, `ByteVector` uint8 1D). Wrappers validate shapes and raise `ValueError` so mismatched input never reaches C++ in Release builds, where asserts are off.
249
+
250
+ ```python
251
+ import numpy as np, qann
252
+
253
+ x = np.random.rand(10_000, 128).astype(np.float32)
254
+
255
+ ivf = qann.IVFIndex(128, nlist=100, nprobe=10)
256
+ ivf.enable_pq(16) # optional; before train/add
257
+ ivf.train(x)
258
+ refine = qann.RefineIndex(ivf, k_factor=10) # optional; wrap before adding
259
+ refine.add(x) # adds to ivf too
260
+ ivf.nprobe = 32 # tune after training
261
+ ivf.pq_distance = qann.PQDistance.ADC # default ADC
262
+ ids, dists = refine.query(x[:1], 10)
263
+ ids, dists = refine.batch_query(x[:5], 10) # (5, 10) arrays
264
+
265
+ pq = qann.PQCodebook(128, 16)
266
+ pq.train(x)
267
+ code = pq.encode(x[0]) # (16,) uint8
268
+ table = pq.compute_adc_table(x[1]) # (16 * 256,) float32
269
+ pq.distance_adc(table, code)
270
+ ```
271
+
272
+ `RefineIndex` holds its base by reference; `keep_alive` keeps the Python base object alive as long as the wrapper exists.
273
+
274
+ Also exposed: `FlatIndex`, `IndexType`, `IndexOptions`, `make_index`, `set_num_threads(n)` / `num_threads()`, and row-wise `l2_distance(a, b)` / `cosine_distance(a, b)`.
275
+
276
+ Results are NumPy arrays that take ownership of the C++ buffer (no copy): IDs are `int64`, distances `float32`, codes `uint8`. `query` returns 1D arrays with one entry per result. `batch_query` returns `(num_queries, k)` arrays; a row with fewer than `k` results is padded with ID `-1` and distance `inf`.
277
+
278
+ ---
279
+
280
+ ## Build, Test, Benchmark
281
+
282
+ ```bash
283
+ pip install . # builds the Python module into a wheel (no C++ tests)
284
+ ./build.sh # Release build in ./build/, then ctest
285
+ ./build.sh build Debug # Debug build (asserts on)
286
+ ./build/test_vecengine "[pq]" # run one tag: [l2] [cosine] [index] [ivf] [pq] [adc] [sdc] [refine] ...
287
+ ./build/bench_vecengine
288
+ ```
289
+
290
+ - Dependencies come from FetchContent: Catch2 v3.5.3, Google Benchmark v1.8.3, nanobind v1.9.2.
291
+ - BLAS for k-means (`VECENGINE_USE_BLAS`): on Apple, `vecengine_core` links `-framework Accelerate` (PUBLIC, so tests and the Python module link it too). Elsewhere `find_package(BLAS)` plus a `cblas.h` search (also under `include/openblas`); on Debian/Ubuntu `apt install libopenblas-dev`. Configure prints which one it used, or that it fell back to the per-point scan. OpenBLAS runs its own threads inside `sgemm`; if PQ training (which calls it from every `parallel_for` worker) oversubscribes, set `OPENBLAS_NUM_THREADS=1`.
292
+ - Packaging: `pyproject.toml` uses scikit-build-core, builds only `vecengine_py` with `VECENGINE_BUILD_TESTS=OFF` (skips fetching Catch2 and Google Benchmark), and installs the module at the wheel root. Runtime dependency: NumPy.
293
+ - `-mavx2 -mfma` are only added on x86 so arm64 builds never see unsupported flags.
294
+ - Google Benchmark's own tests are disabled via cache variables to avoid its stale GoogleTest download.
295
+ - Link-time optimization was tried and left off: no query speedup, and k-means training got ~35% slower.
296
+ - `CMAKE_EXPORT_COMPILE_COMMANDS=ON` writes `build/compile_commands.json` for clangd.
297
+ - Targets: `vecengine_core` (static lib), `test_vecengine`, `bench_vecengine`, `vecengine_py` (Python module).
298
+ - CI (`.github/workflows/ci.yml`, on pushes to `main` and on PRs): builds and runs `ctest` on `ubuntu-latest` (with OpenBLAS) and `macos-14`, then runs the siftsmall recall gate on macOS only, since the baseline was recorded with Accelerate. siftsmall is cached between runs.
299
+
300
+ ---
301
+
302
+ ## Not Yet Implemented
303
+
304
+ - **IVF-PQ precomputed tables**: FAISS-style decomposition to avoid building an ADC table per probed list.
305
+ - **HNSW**: graph index; slot reserved in `IndexType`.
306
+ - **Persistence**: no save/load of trained indexes.
@@ -0,0 +1,174 @@
1
+ cmake_minimum_required(VERSION 3.20)
2
+ project(vecengine VERSION 0.1.0 LANGUAGES CXX)
3
+
4
+ set(CMAKE_CXX_STANDARD 20)
5
+ set(CMAKE_CXX_STANDARD_REQUIRED ON)
6
+ set(CMAKE_CXX_EXTENSIONS OFF)
7
+
8
+ # ---------------------------------------------------------------------------
9
+ # Dependencies via FetchContent
10
+ # ---------------------------------------------------------------------------
11
+ include(FetchContent)
12
+
13
+ # Off for Python wheel builds (pyproject.toml) so they skip fetching and
14
+ # compiling Catch2 and Google Benchmark.
15
+ option(VECENGINE_BUILD_TESTS "Build C++ tests and microbenchmarks" ON)
16
+ # Per-section query timers (src/profile.hpp) and the profile_query tool.
17
+ # Adds timer calls to the query path, so keep it off for normal builds.
18
+ option(VECENGINE_PROFILE "Build with query section timers and profile_query" OFF)
19
+
20
+ FetchContent_Declare(
21
+ Catch2
22
+ GIT_REPOSITORY https://github.com/catchorg/Catch2.git
23
+ GIT_TAG v3.5.3
24
+ )
25
+
26
+ FetchContent_Declare(
27
+ benchmark
28
+ GIT_REPOSITORY https://github.com/google/benchmark.git
29
+ GIT_TAG v1.8.3
30
+ )
31
+
32
+ # Must be set before FetchContent_MakeAvailable to suppress GoogleTest download
33
+ set(BENCHMARK_ENABLE_TESTING OFF CACHE BOOL "" FORCE)
34
+ set(BENCHMARK_ENABLE_GTEST_TESTS OFF CACHE BOOL "" FORCE)
35
+ set(BENCHMARK_INSTALL_DOCS OFF CACHE BOOL "" FORCE)
36
+
37
+ FetchContent_Declare(
38
+ nanobind
39
+ GIT_REPOSITORY https://github.com/wjakob/nanobind.git
40
+ GIT_TAG v1.9.2
41
+ )
42
+
43
+ if(VECENGINE_BUILD_TESTS)
44
+ FetchContent_MakeAvailable(Catch2 benchmark nanobind)
45
+ else()
46
+ FetchContent_MakeAvailable(nanobind)
47
+ endif()
48
+
49
+ # ---------------------------------------------------------------------------
50
+ # Compiler flags
51
+ # ---------------------------------------------------------------------------
52
+ add_compile_options(
53
+ $<$<CONFIG:Release>:-O3>
54
+ $<$<CXX_COMPILER_ID:GNU,Clang>:-Wall>
55
+ $<$<CXX_COMPILER_ID:GNU,Clang>:-Wextra>
56
+ )
57
+
58
+ # AVX2+FMA are x86-only; skip silently on arm64 (Apple Silicon uses NEON).
59
+ if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64|i686")
60
+ add_compile_options(
61
+ $<$<CONFIG:Release>:-mavx2>
62
+ $<$<CONFIG:Release>:-mfma>
63
+ )
64
+ endif()
65
+
66
+ # ---------------------------------------------------------------------------
67
+ # Core static library
68
+ # ---------------------------------------------------------------------------
69
+ add_library(vecengine_core STATIC
70
+ src/distance.cpp
71
+ src/index.cpp
72
+ src/threads.cpp
73
+ src/flat_index.cpp
74
+ src/ivf_index.cpp
75
+ src/pq.cpp
76
+ src/kmeans.cpp
77
+ src/refine_index.cpp
78
+ src/index_factory.cpp
79
+ )
80
+
81
+ target_include_directories(vecengine_core PUBLIC
82
+ $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
83
+ $<INSTALL_INTERFACE:include>
84
+ )
85
+
86
+ target_compile_features(vecengine_core PUBLIC cxx_std_20)
87
+
88
+ # The static lib is linked into the Python shared module, which Linux requires
89
+ # to be position-independent code (macOS builds PIC by default).
90
+ set_target_properties(vecengine_core PROPERTIES POSITION_INDEPENDENT_CODE ON)
91
+
92
+ if(VECENGINE_PROFILE)
93
+ target_compile_definitions(vecengine_core PUBLIC VECENGINE_PROFILE)
94
+ add_executable(profile_query benchmarks/profile_query.cpp)
95
+ target_include_directories(profile_query PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
96
+ target_link_libraries(profile_query PRIVATE vecengine_core)
97
+ endif()
98
+
99
+ # k-means assignment uses cblas_sgemm when available: Accelerate on Apple (runs
100
+ # on the matrix unit), otherwise any BLAS that ships cblas.h (e.g. OpenBLAS).
101
+ # Without one it falls back to a scalar scan. PUBLIC link: anything linking the
102
+ # static lib needs it too.
103
+ if(APPLE)
104
+ target_link_libraries(vecengine_core PUBLIC "-framework Accelerate")
105
+ target_compile_definitions(vecengine_core PRIVATE VECENGINE_USE_BLAS ACCELERATE_NEW_LAPACK)
106
+ else()
107
+ find_package(BLAS)
108
+ find_path(CBLAS_INCLUDE_DIR cblas.h PATH_SUFFIXES openblas)
109
+ if(BLAS_FOUND AND CBLAS_INCLUDE_DIR)
110
+ message(STATUS "vecengine: using BLAS (${BLAS_LIBRARIES}) with ${CBLAS_INCLUDE_DIR}/cblas.h")
111
+ target_link_libraries(vecengine_core PUBLIC BLAS::BLAS)
112
+ target_include_directories(vecengine_core PRIVATE ${CBLAS_INCLUDE_DIR})
113
+ target_compile_definitions(vecengine_core PRIVATE VECENGINE_USE_BLAS)
114
+ else()
115
+ message(STATUS "vecengine: no BLAS with cblas.h found; k-means assignment uses the scalar fallback")
116
+ endif()
117
+ endif()
118
+
119
+ if(VECENGINE_BUILD_TESTS)
120
+ # ---------------------------------------------------------------------------
121
+ # Unit tests (Catch2 v3)
122
+ # ---------------------------------------------------------------------------
123
+ add_executable(test_vecengine
124
+ tests/test_distance.cpp
125
+ )
126
+
127
+ target_link_libraries(test_vecengine PRIVATE
128
+ vecengine_core
129
+ Catch2::Catch2WithMain
130
+ )
131
+
132
+ # Lets tests reach internal headers like parallel.hpp.
133
+ target_include_directories(test_vecengine PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
134
+
135
+ include(CTest)
136
+ include(Catch)
137
+ catch_discover_tests(test_vecengine)
138
+
139
+ # ---------------------------------------------------------------------------
140
+ # Microbenchmarks (Google Benchmark)
141
+ # ---------------------------------------------------------------------------
142
+ add_executable(bench_vecengine
143
+ benchmarks/bench_distance.cpp
144
+ )
145
+
146
+ target_link_libraries(bench_vecengine PRIVATE
147
+ vecengine_core
148
+ benchmark::benchmark
149
+ )
150
+ endif()
151
+
152
+ # ---------------------------------------------------------------------------
153
+ # Python bindings (nanobind)
154
+ # ---------------------------------------------------------------------------
155
+ find_package(Python 3.9 REQUIRED COMPONENTS Interpreter Development.Module)
156
+
157
+ nanobind_add_module(vecengine_py
158
+ bindings/python_bindings.cpp
159
+ )
160
+
161
+ target_link_libraries(vecengine_py PRIVATE vecengine_core)
162
+
163
+ set_target_properties(vecengine_py PROPERTIES
164
+ OUTPUT_NAME "qann"
165
+ )
166
+
167
+ # scikit-build-core sets SKBUILD; the module goes at the wheel root so
168
+ # `import qann` finds it.
169
+ if(SKBUILD)
170
+ install(TARGETS vecengine_py LIBRARY DESTINATION .)
171
+ target_compile_definitions(vecengine_py PRIVATE QANN_VERSION="${SKBUILD_PROJECT_VERSION_FULL}")
172
+ else()
173
+ target_compile_definitions(vecengine_py PRIVATE QANN_VERSION="dev")
174
+ endif()
qann-0.1.0a1/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Maximilian Miller
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.