diff --git a/CMakeLists.txt b/CMakeLists.txt index 2649d66d..3d8afc4e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -84,6 +84,7 @@ endfunction() set(RABITQ_CORE_SOURCES src/core/api.cpp + src/core/blob.cpp src/core/fast_scale_table.cpp src/core/provisional.cpp src/utils/cpu_features.cpp @@ -247,6 +248,7 @@ if(RABITQ_BUILD_FASTSCAN) endif() if(RABITQ_BUILD_INDEX) add_library(rabitq_index STATIC + src/index/format2.cpp src/index/hnsw.cpp src/index/ivf.cpp src/index/ivf_initializer.cpp @@ -261,6 +263,10 @@ if(RABITQ_BUILD_INDEX) ) add_library(rabitqlib::index ALIAS rabitq_index) target_link_libraries(rabitq_index PUBLIC rabitq_fastscan) + # Index files record the library version that wrote them. + target_compile_definitions(rabitq_index PRIVATE + RABITQ_LIBRARY_VERSION="${PROJECT_VERSION}" + ) set_target_properties(rabitq_index PROPERTIES EXPORT_NAME index POSITION_INDEPENDENT_CODE ON diff --git a/benchmarks/hotpaths/load.cpp b/benchmarks/hotpaths/load.cpp index 8e4d0de2..dc1c6c84 100644 --- a/benchmarks/hotpaths/load.cpp +++ b/benchmarks/hotpaths/load.cpp @@ -37,8 +37,11 @@ void run(int argc, char** argv) { if (!std::filesystem::exists(rabitqlib::io_impl::filesystem_path(path + ".hnsw"))) { rabitqlib::ivf::HNSWInitializer index(dim, count); index.add_vectors(data.data(), 1); - std::ofstream unused; - index.save(unused, path.c_str()); + std::ofstream output( + rabitqlib::io_impl::filesystem_path(path + ".hnsw"), std::ios::binary + ); + output.exceptions(std::ios::failbit | std::ios::badbit); + index.write(output); return; } } @@ -46,9 +49,19 @@ void run(int argc, char** argv) { uint64_t expected = 0; for (size_t rep = 0; rep <= reps; ++rep) { rabitqlib::ivf::HNSWInitializer index(dim, count); - std::ifstream unused; const auto start = bench::Clock::now(); - index.load(unused, path.c_str()); + // Open the graph file inside the timing, as loading an index does. + std::ifstream input( + rabitqlib::io_impl::filesystem_path(path + ".hnsw"), + std::ios::binary | std::ios::ate + ); + if (!input) { + throw std::runtime_error("Cannot open the HNSW centroid graph file"); + } + const auto bytes = + static_cast(static_cast(input.tellg())); + input.seekg(0); + index.read(input, bytes); const double seconds = std::chrono::duration(bench::Clock::now() - start).count(); uint64_t h = 14695981039346656037ULL; diff --git a/docs/development.md b/docs/development.md index 5800fe56..3977d182 100644 --- a/docs/development.md +++ b/docs/development.md @@ -505,12 +505,13 @@ in place of extra-bit codes. #### IVF formats and initialization -New IVF files have a magic/version header, an explicit padded dimension, and a resolved initializer type (v2); v1 files infer the historical -Flat/HNSW choice at 20,000 clusters, independently of current auto selection. -New IVF auto routing uses Flat below 5,000 clusters, Flat RaBitQ below 60,000, -and HNSW otherwise. Flat RaBitQ persists its mean and codes in the main file; -HNSW retains its sidecar. Legacy quantized files and raw v1 files infer -64-dimension padding. +New IVF files use index format 2 (`docs/formats.md`): the resolved initializer type +is in `ivf_params`, and the routing data, including the HNSW graph, is a section of +the one file. Earlier `RABQIDX1` v2 files store the initializer type; v1 files infer +the historical Flat/HNSW choice at 20,000 clusters, independently of current auto +selection. Those layouts keep HNSW routing in the `.hnsw` sidecar. New IVF auto +routing uses Flat below 5,000 clusters, Flat RaBitQ below 60,000, and HNSW otherwise. +Legacy quantized files and raw v1 files infer 64-dimension padding. #### Search defaults @@ -532,7 +533,7 @@ and buffer changes must also preserve the [ID marker](#search-buffer-id-marker) Filesystem path strings are UTF-8 on Windows and native path bytes on POSIX. Use `io_impl::filesystem_path` from `utils/path.hpp` before opening files or querying -filesystem metadata, including IVF's HNSW sidecar. It converts UTF-8 to native Unicode +filesystem metadata, including the HNSW sidecar of earlier IVF files. It converts UTF-8 to native Unicode on Windows and preserves existing path bytes on Linux/macOS. For path-only changes, check path conversion and the affected I/O operations; the @@ -550,15 +551,20 @@ keep their provenance and hashes, and do not regenerate them with the current wr Index files are trusted input (see the usage guide): validate structure and sizes so that corrupt or mismatched files are reported before allocation, without treating a crafted file as a supported input. +Readers throw `rabitqlib::FormatError` (`include/rabitqlib/index/format_error.hpp`) for +any truncated, corrupt, mismatched or unsupported content, including errors raised by +shared validators while loading; only a file that cannot be opened throws a plain +`std::runtime_error`. Python maps it to `rabitqlib.FormatError`; rejection tests assert it. Do not silently reinterpret an old file. Add a magic/version discriminator, use fixed-width -serialized fields for new formats (the current IVF, HNSW and QG bodies still store several -`size_t` fields; change those only with a format version bump), validate sizes before -allocation, check every read, and retain a -compatibility test fixture or an explicit rejection path. Review each affected format; -changes to shared serialization helpers require auditing every format using them. -IVF-RaBitQ, HNSW-RaBitQ, and SymphonyQG have separate formats. SymphonyQG includes a versioned quantized format plus a legacy -raw-format fallback; preserve both unless a breaking change is explicitly requested. +serialized fields for new formats (earlier layouts store several `size_t` fields; keep their readers unchanged), validate sizes before allocation, check every +read, and retain a compatibility test fixture or an explicit rejection path. Review each affected +format; changes to shared serialization helpers require auditing every format using them. +Every index writes index format 2 ([`docs/formats.md`](formats.md#index-format-2), container in +`src/index/format2.*`) and still reads its earlier layouts: IVF-RaBitQ's `RABQIDX1` versions and +unversioned files, HNSW-RaBitQ's unversioned layout, and SymphonyQG's versioned quantized layout +and raw fallback. Preserve those readers unless a breaking change is explicitly requested. When a writer moves to +a new format, freeze the last files of the old writer as release-produced fixtures first. ### Change Python bindings @@ -840,8 +846,8 @@ Centroids, data, and queries compared by the same estimator must be in the same Rotator state is part of persisted index state. Loading an index must restore the exact rotation used at construction; generating a new random rotator produces silently incorrect distances. -Honor stored padded dimensions instead of recomputing them from current defaults. IVF `RABQIDX1` -v1 and v2 files explicitly store this field; preserve it for both versions. Only legacy unversioned +Honor stored padded dimensions instead of recomputing them from current defaults. Index format 2 +and IVF `RABQIDX1` v1 and v2 files explicitly store this field; preserve it. Only legacy unversioned quantized IVF files and old raw v1 files omit it and infer historical 64-dimension padding. Unreleased pad16 files whose stored dimension is not divisible by 32 are rejected. @@ -875,9 +881,9 @@ updates or persistence. These markers affect live search as well as saved files. IVF-RaBitQ `remove` stores a removed point as `f_add = +inf` inside the existing batch data, so it needs no discriminator. No index written before `remove` existed has an infinite `f_add`, so nothing in an old file is reinterpreted. In the legacy file formats readable by release 0.5.0, that reader also -excludes removed points. This compatibility of the removal marker does not extend to the current -writer's `RABQIDX1` v2 format: older releases cannot read current saves, even when resaving a loaded -legacy index. Keep this convention when touching IVF-RaBitQ code: +excludes removed points. This compatibility of the removal marker does not extend to index +format 2, which the current writer uses: releases up to 0.8 cannot read current saves, even when +resaving a loaded legacy index. Keep this convention when touching IVF-RaBitQ code: - Every estimator and scan path must turn `f_add = +inf` into an infinite estimated distance and an infinite lower bound, and never into NaN. A NaN distance would enter the result buffer. diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index a846c58d..a5fe8e53 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -82,7 +82,7 @@ available CPU, so **default builds are not reproducible**: pass `num_threads=1` still change results. Loading restores the saved rotation. Results use squared L2 or `1 - dot` estimates, ordered from smallest to largest. -Python fills missing results with ID `2**32 - 1` and distance `+inf`; C++ returns +Python returns `int64` IDs and fills missing results with ID `-1` and distance `+inf`; C++ returns only the available pairs. See [result handling](../quick_start.md#read-the-results) and the shared [input, count and concurrency contracts](usage.md). @@ -271,8 +271,8 @@ removed = index.remove(new_ids[:10]) # IDs in [0, num_stored_points) A removed point keeps its codes, its links, and its place in the graph. Search still walks through it and `add` may link new points to it, which keeps the graph connected, as in hnswlib. Search never returns it, so a query can get fewer than -`k` results once points are removed: in Python the missing slots hold `kPidMax` -(`2**32 - 1`) with an infinite distance. Removed points still count in +`k` results once points are removed: in Python the missing slots hold ID `-1` +with an infinite distance. Removed points still count in `num_stored_points()` and cannot be restored. Removal requires the same exclusive access as other updates. @@ -316,6 +316,10 @@ through 0.5.1 check every cluster ID on `load` and reject the file as invalid instead of returning removed points. Releases before 0.3.7 do not check cluster IDs, so do not open such a file with them. +`save` writes [index format 2](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/docs/formats.md#index-format-2), which also records the construction seed, so +`seed` survives a reload. Releases up to 0.8 cannot read these files; every earlier +file still loads, without a recorded seed, and saving it writes format 2. + New HNSW-RaBitQ indexes pad dimensions to multiples of 32 for all supported bit widths. Loading preserves the saved padded dimension and rotation, so older indexes using multiples of 64 remain readable. Older library versions cannot load new diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 86e11790..669892ce 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -52,11 +52,11 @@ See [C++ integration](../cpp_integration.md) for build setup and | `seed` | Rotation seed in `[0, 2**32 - 1]` and the initial `cp.seed`. Python's default `None` draws one at random; `index.seed` reports it, and passing it back reproduces the index. | | `metric` | Python `"l2"` or `"ip"`; defaults to `"l2"`. | | `initializer` | Defaults to `"auto"`; see [centroid routing](#centroid-routing). | -| `k`, `nprobe` | Positive result count and cluster probe count. `k` may exceed the stored point count; missing results are padded with `kPidMax` (`2**32 - 1` in Python) and `+inf`. Python's `nprobe` defaults to `num_clusters / 16`, rounded and at least 1, when omitted, `None` or 0; C++ requires a positive `nprobe`. | +| `k`, `nprobe` | Positive result count and cluster probe count. `k` may exceed the stored point count; missing results are padded with `kPidMax` (`-1` in Python) and `+inf`. Python's `nprobe` defaults to `num_clusters / 16`, rounded and at least 1, when omitted, `None` or 0; C++ requires a positive `nprobe`. | | `high_accuracy` | Python defaults to `None`: HACC for 4–9 bits, standard FastScan otherwise. | Results contain IDs and squared L2 or `1 - dot` distances, ordered from smallest -to largest. Quantized distances are estimates. Unfilled slots use ID `2**32 - 1` +to largest. Quantized distances are estimates. IDs are `int64` in Python. Unfilled slots use ID `-1` and distance `+inf`; see [result handling](../quick_start.md#read-the-results). See the shared [input and concurrency contracts](usage.md). @@ -450,15 +450,15 @@ enters the result buffer and never triggers reranking. Any code that estimates IVF-RaBitQ distances must keep this behavior and must never turn `+inf` into NaN. Search may return fewer than `k` results once points are removed: the missing slots -hold `kPidMax` (`2**32 - 1` in Python) with an infinite distance, as when the probed +hold `kPidMax` (`-1` in Python) with an infinite distance, as when the probed clusters hold fewer than `k` points. The removal marker lives in the ordinary batch data and needs no additional format change; removal survives `save` and `load`. No file written before `remove` existed has an infinite `f_add`, so nothing in an old file is reinterpreted. Release 0.5.0 also excludes removed points when reading its -supported legacy formats. However, current saves always use `RABQIDX1` v2, -including when resaving a loaded legacy index. Older releases cannot read these +supported legacy formats. However, current saves always use index format 2, +including when resaving a loaded legacy index. Releases up to 0.8 cannot read these files, regardless of whether they contain removals. ## Persistence @@ -468,19 +468,23 @@ Save a constructed index with: ```c++ ivf.save(output_index_file); ``` -New files use the `RABQIDX1` magic, version 2, storage flags, an explicit -padded dimension, and a uint32 initializer type (1 = Flat, 2 = Flat RaBitQ, -3 = HNSW) before the index metadata. `Auto` is resolved before saving. Loading -restores the routing type, storage mode, rotation state, and padded dimension. -Raw files also contain the original vectors, so querying does not require an external dataset. - -Version-1 `RABQIDX1` files remain readable with their stored padding. Legacy -unversioned quantized files and raw v1 files retain their original 64-dimension -padding. All three historical formats infer routing with the original 20,000 -cluster threshold. HNSW routing still requires the accompanying `.hnsw` file. -Flat RaBitQ loads its saved codes without random re-quantization. Saving a loaded -legacy index writes the new header without changing its padded dimension or encoded data. Older library -versions cannot read the new format. +Files are written in +[index format 2](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/docs/formats.md#index-format-2): +one file with a header, a section table, and sections for the rotation, the routing +data, the codes and the IDs. HNSW routing is a section of that file, so there is no +separate sidecar. `Auto` is resolved before saving. Loading restores the routing type, +storage mode, rotation state, and padded dimension. Raw files also contain the original +vectors, so querying does not require an external dataset. Releases up to 0.8 cannot +read format 2. Indexes of more than 65,536 dimensions cannot be saved. + +Files from earlier releases remain readable. `RABQIDX1` version 2 files store their +routing type; version-1 `RABQIDX1` files keep their stored padding; legacy unversioned +quantized files and raw v1 files retain their original 64-dimension padding. These three +older layouts infer routing with the original 20,000 cluster threshold. All of these +layouts keep HNSW routing in the accompanying `.hnsw` file, which loading reads. +Flat RaBitQ loads its saved codes without random re-quantization. Saving a loaded file +writes format 2 without changing its padded dimension or encoded data; a `.hnsw` file +left beside it is not removed and is no longer used. ## Implementation details diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index cbab008b..8f73bb51 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -77,7 +77,7 @@ qg.save("qg_example.index"); | `init` | Python keyword: `"pipnn"` (default) or `"random"`. | | `metric` | Python keyword: `"l2"` (default) or `"ip"`. | | `seed` | Python keyword: rotation and graph-construction seed in `[0, 2**32 - 1]`. The default `None` draws one at random; `index.seed` reports it, and passing it back with `num_threads=1` rebuilds the same graph from the same data. | -| `k`, `ef` | `k` is at least 1 and may exceed the point count; slots the graph cannot fill are padded with ID `2**32 - 1` and distance `+inf`. Omitted or zero `ef` uses `max(k, 64)` in Python, and smaller values are raised to `k`, as for HNSW-RaBitQ. | +| `k`, `ef` | `k` is at least 1 and may exceed the point count; slots the graph cannot fill are padded with ID `-1` in Python (`kPidMax` in C++) and distance `+inf`. Omitted or zero `ef` uses `max(k, 64)` in Python, and smaller values are raised to `k`, as for HNSW-RaBitQ. | | `num_threads` | Python `add` defaults to every available logical CPU, capped by `OMP_NUM_THREADS`, and `search` to 1; `None` selects these defaults. | Python search returns `(ids, distances)` arrays. Distances are squared @@ -209,6 +209,10 @@ Python counterparts for complete examples. ## Persistence +`save` writes [index format 2](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/docs/formats.md#index-format-2), which also records the seed, `ef_construction` +and `init`, so they survive a reload. Releases up to 0.8 cannot read these files; every +earlier file still loads, without these settings, and saving it writes format 2. + All three modes pad the rotated domain to a multiple of 32. Loading an older index preserves its saved padded dimension, including historical multiples of 64. Older library versions that enforce 64-dimension padding cannot load new files diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md index f96c5d2f..bdfadbaa 100644 --- a/docs/docs/index/usage.md +++ b/docs/docs/index/usage.md @@ -126,7 +126,7 @@ A conflicting call immediately raises It does not wait or queue an update. Access protection includes Python input conversion and is released when the operation returns or raises an exception. Different index objects can operate independently. Concurrent saves must use -different output paths, including any sidecar files. +different output paths. Search parameters such as `ef`, `nprobe`, precision, and `num_threads` belong to each call. When using a Python thread pool, consider `num_threads=1` to avoid @@ -148,7 +148,8 @@ native call, or pass `num_threads=1`, in such processes. clustering classes default to `0`. `search` defaults to one thread (`None` or `1`), as do the deprecated `build`, which requires an integer, and `search_batch`. Multithreaded HNSW-RaBitQ `add` links points in a nondeterministic order, so pass -`num_threads=1` when a build must be reproducible. +`num_threads=1` when a build must be reproducible. The same holds for the centroid graph +of IVF-RaBitQ HNSW routing, which `train`, `set_centroids` or the first `add` builds. On Linux, CPU detection accounts for affinity and OpenMP binding. Apply binding before starting the program. @@ -156,25 +157,34 @@ C++ paths use UTF-8 on Windows and native path bytes on POSIX. Python `save`/`load` accept strings, `pathlib.Path` and `os.PathLike` objects; existing `bytes` and `bytearray` paths remain supported. POSIX paths preserve native bytes, but non-UTF-8 filenames require filesystem support. Windows byte paths use -UTF-8. Embedded null characters are rejected before file access. IVF HNSW sidecars -use the same path conversion. +UTF-8. Embedded null characters are rejected before file access. The HNSW sidecars +of IVF files from earlier releases use the same path conversion. `save` writes a temporary file beside the destination and renames it into place, so a failed save leaves an existing index unchanged. A replaced file keeps its permission bits, but not its owner or other hard links. Symbolic links and other non-regular -destinations are written in place. IVF renames its HNSW sidecar before the index file; -a failure between the two renames can leave a new sidecar beside the old index. +destinations are written in place. IVF-RaBitQ writes one file that includes its HNSW +routing graph. ## Index file compatibility -The compatibility suite loads frozen files written by v0.5.2, compares search -IDs and distances, and verifies save/load migration and rejection of truncated -or unsupported versions. See the +The compatibility suite loads frozen files written by v0.5.2 and v0.8.3, compares +search IDs and distances, and verifies save/load migration and rejection of truncated +or unsupported versions. Every index saves in +[index format 2](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/docs/formats.md#index-format-2), +which releases up to 0.8 cannot read, and which records the construction seed. See the [fixture provenance and coverage](https://github.com/VectorDB-NTU/RaBitQ-Library/tree/main/tests/fixtures/persistence). Compatibility is format-specific; preserving old readers' ability to open new files is not guaranteed. Consult each index's format documentation before downgrading the library. +A truncated, corrupt, mismatched or unsupported file raises `rabitqlib.FormatError`, +a subclass of both `ValueError` and `RuntimeError`, so `except RuntimeError` +handlers still catch it. A file that cannot be opened raises a plain `RuntimeError`; +a read that fails after opening looks like truncation and raises `FormatError`. +C++ loaders throw `rabitqlib::FormatError` from `rabitqlib/index/format_error.hpp`, +a subclass of `std::runtime_error`. + Index files are trusted input. Loaders validate structure and reject truncated, inconsistent, or unsupported files so that mistakes are reported, but they are not a security boundary: a crafted file can make `load` allocate its stored diff --git a/docs/docs/python_migration.md b/docs/docs/python_migration.md index 52fad117..85d2d0c3 100644 --- a/docs/docs/python_migration.md +++ b/docs/docs/python_migration.md @@ -54,7 +54,8 @@ work if you want to run k-means yourself. - **`k` above the point count.** 0.8.0 and 0.8.1 raised `ValueError` when `k` exceeded `num_stored_points` (HNSW-RaBitQ) or the point count (IVF-RaBitQ, SymphonyQG). Since 0.8.2, every index accepts any `k >= 1` and fills missing - slots with ID `2**32 - 1` and distance `inf`; test `np.isfinite(distances)`. + slots with ID `2**32 - 1` (`-1` since 0.9) and distance `inf`; test + `np.isfinite(distances)`, which works in both. - **HNSW-RaBitQ capacity grows.** `add` resizes the index, at least doubling it, instead of raising `ValueError` when the batch does not fit. This also applies to `add` on indexes created in the old form, and to the deprecated `build`, diff --git a/docs/docs/quick_start.md b/docs/docs/quick_start.md index 99032928..fae22f68 100644 --- a/docs/docs/quick_start.md +++ b/docs/docs/quick_start.md @@ -66,7 +66,7 @@ not exact distances to the original vectors. A search may return fewer than `k` valid neighbors, for example when `k` exceeds the number of points, after removals, or when the probed IVF-RaBitQ clusters contain too -few points. Unfilled slots contain ID `2**32 - 1` and distance `+inf`. Filter them with +few points. IDs are `int64`; unfilled slots contain ID `-1` and distance `+inf`. Filter them with `valid = np.isfinite(distances)` before using IDs to index your data. The `metric` argument accepts `"l2"` and `"ip"` (also `"innerproduct"`). For @@ -122,9 +122,9 @@ with TemporaryDirectory() as directory: np.testing.assert_allclose(loaded_distances, distances) ``` -The loaded index owns its data and restores the saved rotation. IVF-RaBitQ indexes using -HNSW centroid routing also need the accompanying `.hnsw` file; keep it with the -main index. This small example uses Flat routing and has no sidecar. +The loaded index owns its data and restores the saved rotation. IVF-RaBitQ saves one +file, including its HNSW centroid routing; files from releases up to 0.8 with HNSW +routing also need the accompanying `.hnsw` file. See [file compatibility](index/usage.md#index-file-compatibility) before downgrading. ## Next steps diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md index 12df7922..0716cc8c 100644 --- a/docs/docs/rabitq/quant_python.md +++ b/docs/docs/rabitq/quant_python.md @@ -1,7 +1,7 @@ # Python quantization building blocks !!! warning "Provisional" - `rabitqlib.quant` is provisional in the sense of [PEP 411](https://peps.python.org/pep-0411/): any minor release may change or remove part of its API, with a note in the release notes, until it is declared stable. Byte formats are versioned from the first release, and bytes of an unknown version raise `quant.FormatError` instead of giving wrong results. + `rabitqlib.quant` is provisional in the sense of [PEP 411](https://peps.python.org/pep-0411/): any minor release may change or remove part of its API, with a note in the release notes, until it is declared stable. Byte formats are versioned from the first release, and bytes of an unknown version raise `quant.FormatError` instead of giving wrong results. It is `rabitqlib.FormatError`, which index loads also raise, a subclass of both `ValueError` and `RuntimeError`. `rabitqlib.quant` exposes the pieces the indexes are built from: rotation, quantization, correction factors, packed codes and distance estimation. Use it to store RaBitQ codes in your own database or file format and to score them with your own selection logic. For a ready-made index, use [IVF-RaBitQ](../index/ivf.md), [HNSW-RaBitQ](../index/hnsw.md) or [SymphonyQG](../index/qg.md) instead. diff --git a/docs/docs/releases/unreleased.md b/docs/docs/releases/unreleased.md new file mode 100644 index 00000000..27b31180 --- /dev/null +++ b/docs/docs/releases/unreleased.md @@ -0,0 +1,44 @@ +# Unreleased + +## Index files + +- IVF-RaBitQ, HNSW-RaBitQ and SymphonyQG `save` write [index format 2](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/docs/formats.md#index-format-2), a single file with a header, a section table and checksummed metadata. +- Releases up to 0.8 cannot read these files. Keep a copy saved by 0.8 if you may downgrade. +- Files written by every earlier release still load, including the `.hnsw` sidecars of HNSW-routed IVF files. Saving one writes format 2 with the same rotation, codes and results. +- IVF-RaBitQ stores HNSW centroid routing inside the file, so `save` no longer writes a `.hnsw` sidecar. A sidecar left by an earlier save is not removed. +- The saved bytes now include that routing graph. It is reproducible only when `train`, `set_centroids` or the first `add` builds it with `num_threads=1`. Saving the same index twice still gives identical files. +- `save` rejects IVF-RaBitQ indexes of more than 65,536 dimensions. +- A format-2 file that uses read-only-compatible features unknown to this release loads for search only: operations that change or save the index raise. Rebuilding with IVF-RaBitQ `construct`, or a new SymphonyQG build, makes it writable again. + +## Python changes + +- Point IDs are `int64` at the Python boundary (#211): + - `search` returns `int64` IDs, and slots without a neighbor hold `-1` instead of `2**32 - 1`; their distance is still `inf`. Code that compares IDs with `2**32 - 1`, `0xFFFFFFFF` or `np.iinfo(np.uint32).max` must change; `np.isfinite(distances)`, recommended since 0.8.2, works before and after. + - `add` returns `int64` IDs. `remove` already took any integer IDs and still rejects `-1` and other out-of-range values without removing anything. + - IDs remain the row numbers the index assigns, and C++ point IDs remain 32-bit `PID` values with `kPidMax` for a missing result. + - txtai 9.14.0 and earlier filter results with `uid != 0xFFFFFFFF`, so with this release they return padded slots as ID `-1`; see #213 for the first compatible txtai release. + +- Loading a format-2 file keeps the `seed` of every index, and `ef_construction` and `init` of SymphonyQG. They were `None` after `load`, and still are for files from releases up to 0.8. +- With a recorded seed, `cp.seed` after `load` follows it as for a new index: the seed for IVF-RaBitQ, its low 32 bits for HNSW-RaBitQ. It was the k-means default. +- A new exception, `rabitqlib.FormatError`, subclasses both `ValueError` and `RuntimeError`: + - Every index `load` raises it for a truncated, corrupt, mismatched or unsupported file, in format 2 and in every earlier layout. These were `RuntimeError`, which still catches them. A file that cannot be opened still raises a plain `RuntimeError`. + - `quant.FormatError` is now the same class. It was a subclass of `ValueError` only, which still catches it. +- Messages for damaged or mismatched files change: + - A format-2 file names the problem, for example `IVF index file is 6916 bytes, but its sections end at 6852` for trailing bytes. + - Errors in an HNSW routing graph name the graph rather than the sidecar: `Truncated HNSW centroid graph` and `HNSW centroid graph changed while loading` replace their `sidecar` forms. + - A missing sidecar of an earlier file still raises `Cannot open HNSW centroid sidecar`. + - A failed HNSW-RaBitQ write raises `HNSW index file cannot be written` instead of `HNSW failed to write index file`. + +## C++ changes + +- Loading an index throws `rabitqlib::FormatError`, a subclass of `std::runtime_error` in the new header `rabitqlib/index/format_error.hpp`, for a damaged, mismatched or unsupported file. Failing to open the file still throws `std::runtime_error`. Files of earlier layouts that threw other exceptions for invalid contents now throw `FormatError` too: an IVF-RaBitQ file with an invalid metric or a dimension its rotation does not support, which threw `std::invalid_argument`, and a SymphonyQG file with invalid settings or dimensions, which threw `std::invalid_argument` or `std::length_error`. + +- `ivf::Initializer` and its subclasses: + - `save(std::ofstream&, const char*)` becomes `write(std::ostream&)`, and `load(std::ifstream&, const char*)` becomes `read(std::istream&, uint64_t bytes)`. + - The HNSW initializer no longer writes a sidecar itself. +- New accessors report what format-2 files record, or no value when unknown: + - `IVF::seed()`; + - `HierarchicalNSW::seed()`; + - `QuantizedGraph::seed()`, `ef_construction()` and `initialization()`. +- `symqg::QGInitialization` moves from `qg_builder.hpp` to `qg.hpp`, which `qg_builder.hpp` includes. +- The vendored hnswlib adds `saveIndex(std::ostream&)` and `loadIndex(std::istream&, size_t bytes, ...)`; the path overloads are unchanged. diff --git a/docs/formats.md b/docs/formats.md index 7aa60efb..9634b590 100644 --- a/docs/formats.md +++ b/docs/formats.md @@ -1,6 +1,6 @@ # Byte formats -This is the specification of the byte formats `rabitqlib.quant` reads and writes. It is shared with the planned index format revision, which will add its own blob kinds. While `rabitqlib.quant` is provisional, a release may stop reading an old format version; reading it then raises `quant.FormatError` instead of returning wrong results. A promise that every released version stays readable comes with stabilisation. +This is the specification of the byte formats `rabitqlib.quant` reads and writes. It is shared with the planned index format revision, which will add its own blob kinds. While `rabitqlib.quant` is provisional, a release may stop reading an old format version; reading it then raises `rabitqlib.FormatError` (also `quant.FormatError`) instead of returning wrong results. A promise that every released version stays readable comes with stabilisation. Every format has a version. A change to any layout is a new version, never an edit of an existing one. The golden files in `tests/python/fixtures/quant/` freeze each version and are never regenerated. @@ -92,3 +92,287 @@ The body is what `Rotator::save()` writes, and `body_bytes` must equal it | Matrix | The rotation matrix as row-major float32 | `4 * dim * padded_dim` | The body defines the rotation; `seed` only records where it came from. A seeded FHT/Kac rotation is identical on every platform, while a matrix drawn from a seed is reproducible only with the same standard library and build, which is why the matrix itself is stored. + + + +## Index files, format 2 (draft) + +> **Status: draft for 0.9 (#211, #212).** The container (`src/index/format2.*`) is implemented, and IVF-RaBitQ, HNSW-RaBitQ and SymphonyQG read and write format 2; no index stores the cosine metric yet, and IVF-SSD and companion files are not implemented. Field numbers, section kinds and the open questions at the end may still change until a release writes format 2. + +Format 2 is one container for IVF-RaBitQ, HNSW-RaBitQ and SymphonyQG, and later for the SSD-resident IVF index from RaBitQ-SSD. It replaces the per-index layouts: IVF-RaBitQ `RABQIDX1` versions 1 and 2, the unversioned HNSW-RaBitQ layout, SymphonyQG's quantized version 1 and its raw layout, and the unversioned hnswlib sidecar `.hnsw` of IVF-RaBitQ. Every released layout stays readable (see [Earlier layouts](#earlier-layouts)); `save` always writes format 2. + +The number continues the name used in #212. An index file is told apart from earlier layouts by its magic, not by the version, so the unrelated IVF-RaBitQ `RABQIDX1` version 2 does not clash. + +Unlike the provisional `rabitqlib.quant` formats above, index files carry the library's full compatibility promise: every format-2 file written by a release stays readable by every later release. + +### Conventions + +- Multi-byte fields are little-endian; types are fixed-width (`u32`, `u64`, `i32`, `f32`, `f64`). No field is `size_t`. +- Offsets are absolute byte offsets from the start of the file. +- The section table starts at a multiple of 8, and every section starts at a multiple of 64, so a reader can map the file and read sections in place. Writers fill gaps with zero bytes, and readers do not inspect them. +- Reserved fields and undefined flag bits are written as zero. A reader rejects a nonzero value unless a rule below says otherwise. +- Enumerations start at 1, so zero always means unset; their values are never reused or renumbered. +- Checksums are CRC-32 as `zlib.crc32` computes it, the same as in [blobs](#blobs). + +### Layout + +| Part | Position | +| --- | --- | +| Header | Offset 0, `header_bytes` long | +| Section table | `section_table_offset`, `section_count` entries of 32 bytes | +| Sections | Anywhere after the header, each at a multiple of 64 | + +Every section starts at or after `header_bytes`. Non-empty sections do not overlap the table or each other; empty sections take no bytes and may share an offset. Table entries are in non-decreasing offset order. The file ends exactly where the table or the last section ends, whichever is later; an empty section still has an aligned offset, so writers pad up to it. + +### Header + +`header_bytes` is 128 in this version. It is a multiple of 8 and at least 128; a later version may append fields, which a reader that knows only these 128 bytes skips (see [Compatibility rules](#compatibility-rules)). + +| Offset | Size | Type | Field | Rule | +| --- | --- | --- | --- | --- | +| 0 | 8 | ASCII | `magic` = `"RABQINDX"` | Anything else is not a format-2 file; the reader tries [earlier layouts](#earlier-layouts) | +| 8 | 4 | u32 | `version` = 2 | Any other value is rejected | +| 12 | 4 | u32 | `header_bytes` | See above | +| 16 | 8 | u64 | `incompat_features` | See [Feature flags](#feature-flags) | +| 24 | 8 | u64 | `ro_compat_features` | | +| 32 | 8 | u64 | `compat_features` | | +| 40 | 4 | u32 | `index_kind` | 1 = IVF-RaBitQ, 2 = HNSW-RaBitQ, 3 = SymphonyQG; 4 = IVF-SSD (reserved). Unknown values are rejected | +| 44 | 4 | u32 | `metric` | 1 = L2, 2 = inner product, 3 = cosine. See [Metrics](#metrics) | +| 48 | 8 | u64 | `dim` | Original dimension, at least 1 | +| 56 | 8 | u64 | `padded_dim` | Working dimension; a multiple of 32, at least `dim` and at most 65,536 (the rotator's limit). Readers use the stored value and never recompute it | +| 64 | 4 | u32 | `code_bits` | Bits per coordinate of the stored quantized code, 1–9; 0 when the index stores no quantized code of its points | +| 68 | 4 | u32 | `raw_type` | Element type of stored raw vectors: 0 = none, 1 = float32, 2 = uint8 | +| 72 | 4 | u32 | `id_bits` | Width of the point IDs stored in the file: 32 or 64 | +| 76 | 4 | u32 | `flags` | Bit 0, `seed_known`: `seed` is valid | +| 80 | 8 | u64 | `seed` | The index's construction seed. Zero when `seed_known` is clear, which a file migrated from an earlier layout has. IVF-RaBitQ: the rotation and training seed, at most 2^32 − 1. HNSW-RaBitQ: the graph-level and training seed, whose low 32 bits seeded the rotation. SymphonyQG: the rotation seed, at most 2^32 − 1, which Python also uses for graph construction. The rotator blob records the rotation seed again | +| 88 | 8 | u64 | `num_points` | Stored points, including removed ones | +| 96 | 8 | u64 | `section_table_offset` | A multiple of 8, at least `header_bytes` | +| 104 | 4 | u32 | `section_count` | At least 1 | +| 108 | 4 | u32 | `header_crc32` | Over the header with this field taken as zero, followed by the section table | +| 112 | 16 | ASCII | `writer_version` | Library version that wrote the file, printable ASCII, NUL-padded. For error messages only; readers never branch on it | + +The Python `nbits` maps to these fields as follows. IVF-RaBitQ `nbits` 1–9 is `code_bits`. IVF-RaBitQ `nbits = 32` is `code_bits = 1` with `raw_type = float32`, because it keeps 1-bit FastScan codes and reranks with raw vectors. SymphonyQG `nbits = 32` is `code_bits = 0` with `raw_type = float32`. + +### Section table + +| Offset | Size | Type | Field | Rule | +| --- | --- | --- | --- | --- | +| 0 | 4 | u32 | `kind` | See [Section kinds](#section-kinds) | +| 4 | 4 | u32 | `flags` | Bit 0, `required`; bit 1, `checksummed` | +| 8 | 8 | u64 | `offset` | A multiple of 64 | +| 16 | 8 | u64 | `size` | Bytes; may be 0 | +| 24 | 4 | u32 | `crc32` | Over the section's bytes when `checksummed` is set; zero otherwise | +| 28 | 4 | — | reserved | | + +Kind 0 is reserved, and a kind appears at most once. Section flags other than the two above are rejected; a later version cannot add a section flag without a new `version`. Writers checksum the header-like sections (parameters, cluster sizes, companions), whose values size later allocations, and leave bulk data unchecksummed, so that loading a large index does not read it twice. Per-point arrays such as `hnsw_levels` count as bulk data: the header fixes their size, and readers check every value against its limit. + +A section the reader does not know is handled by its `required` bit: a required one rejects the file, and an optional one is skipped without reading it, so its checksum is not verified. A writer marks a section optional only if dropping it leaves a correct index, for example a cache that can be recomputed. `save` drops optional sections it does not know. + +### Feature flags + +The three masks follow the ext4 superblock: they say what an older reader may do with a file that uses something newer than itself. + +| Mask | A reader that does not know a set bit | +| --- | --- | +| `incompat_features` | Rejects the file; reading it would give wrong results | +| `ro_compat_features` | May load and search; operations that change or save the index (`add`, `reserve`, `remove`, `save`) raise, because they would leave the unknown data stale. Rebuilding replaces that data and makes the index writable again: IVF-RaBitQ `construct`, or a new SymphonyQG build; HNSW-RaBitQ cannot be rebuilt in place | +| `compat_features` | Ignores the bit; `save` clears it | + +The error names the unknown bits and `writer_version`, such as "index file uses incompatible features 0x4 that this release does not know; written by rabitqlib 0.10.0". + +A writer sets only the bits it knows. Resaving a file therefore drops the unknown `compat` bits, together with the unknown optional sections and header fields past the writer's `header_bytes`. A `compat` bit may describe only data whose loss on such a resave is harmless. + +No bit is defined in version 2. Bits are assigned when a feature needs one. Examples: + +| Possible feature | Mask | Why | +| --- | --- | --- | +| Caller-provided IDs (`add(x, ids=...)`), an external-ID section | `incompat` | An older reader would return row numbers instead of the caller's IDs | +| A per-cluster statistics section that `add` updates | `ro_compat` | An older reader can search, but its `add` would not update the section | +| A new code layout of an existing section | `incompat`, or a new section kind | The bytes mean something else | + +### Metrics + +L2 and inner product are as in [Distance conventions](development.md#distance-conventions). Cosine stores the codes and factors of inner product over L2-normalized vectors. The index normalizes data on `add` and queries on `search`, and reports `1 - cos` as the distance. It is a stored metric, not a flag of the caller, so a reloaded index keeps normalizing. + +### Point IDs + +Python `search`, `add` and `remove` use `int64` IDs, and a missing search result is `-1` (#211). `id_bits` records what the file stores. The in-memory indexes keep 32-bit row numbers internally and write `id_bits = 32`; their upper limit stays below 2^31 points (the search-buffer marker bit). IVF-SSD writes `id_bits = 64`. A reader rejects an `id_bits` its index kind does not support. + +### Section kinds + +Kinds 1–255 are shared, 256–511 are IVF-RaBitQ, 512–767 HNSW-RaBitQ, 768–1023 SymphonyQG and 1024–1279 IVF-SSD. Every section listed here is `required`. + +| Kind | Name | Used by | Content | +| --- | --- | --- | --- | +| 1 | `rotator` | All | A [rotator blob](#rotator-kind-1-version-1), byte for byte what `Rotator.to_bytes()` writes. Its `dim` and `padded_dim` must equal the header's. Not checksummed by the container: the blob carries its own CRC-32 | +| 2 | `companions` | IVF-SSD | The [companion table](#companion-files) | +| 256 | `ivf_params` | IVF | `num_clusters` u64; `initializer` u32 (1 = flat, 2 = FlatRaBitQ, 3 = HNSW); reserved u32 | +| 257 | `ivf_cluster_sizes` | IVF | `num_clusters` × u64; the sum equals `num_points` | +| 258 | `ivf_flat_routing` | IVF, flat | The rotated centroids, `num_clusters × padded_dim` f32 | +| 259 | `ivf_flat_rabitq_routing` | IVF, FlatRaBitQ | The rotated centroids as above, the global centroid (their mean, `padded_dim` f32), then `ceil(num_clusters / 32)` [FastScan blocks](#fastscan-blocks-version-1) of the centroids, quantized against the global centroid under the index metric | +| 260 | `ivf_hnsw_routing` | IVF, HNSW | The [HNSW routing graph](#hnsw-routing-graph), byte for byte the `.hnsw` sidecar of earlier releases | +| 261 | `ivf_codes` | IVF | Per cluster in cluster order, `ceil(size / 32)` [FastScan blocks](#fastscan-blocks-version-1). Removed points have `add = +inf` | +| 262 | `ivf_rerank` | IVF, `code_bits > 1` or raw | Per point in cluster order: a [single-code record](#single-code-record) with `b = code_bits - 1`, or `dim` raw values of `raw_type` | +| 263 | `ivf_ids` | IVF | Per point in cluster order, `id_bits`-wide IDs | +| 512 | `hnsw_params` | HNSW | `max_elements`, `num_clusters`, `M`, `upper_max_degree`, `base_max_degree`, `ef_construction` as u64; `level_multiplier` f64; `max_level` i32; reserved u32; `entry_point` u64. Record sizes are derived, not stored (see [HNSW-RaBitQ records](#hnsw-rabitq-records)) | +| 513 | `hnsw_centroids` | HNSW | `num_clusters × padded_dim` f32 | +| 514 | `hnsw_records` | HNSW | `num_points` [base-layer records](#hnsw-rabitq-records) | +| 515 | `hnsw_levels` | HNSW | `num_points` × u32 levels | +| 516 | `hnsw_upper_links` | HNSW | The upper-layer link lists of points with level > 0 in point order; each is `level × upper_links_bytes` | +| 768 | `qg_params` | SymphonyQG | `max_degree` u64; `entry_point` u64; `ef_construction` u64 (0 = unknown); `init` u32 (0 = unknown, 1 = PiPNN, 2 = random); reserved u32 | +| 769 | `qg_centroid` | SymphonyQG, `code_bits > 0` | `padded_dim` f32 | +| 770 | `qg_records` | SymphonyQG | `num_points` [records](#symphonyqg-records) | + +An IVF file has exactly one routing section, the one that `ivf_params` names; a rerank section exactly when `code_bits > 1` or raw vectors are stored. + +The HNSW per-point length prefix of the current layout becomes `hnsw_levels`, so a reader knows every size before reading the link lists. + +### Record layouts + +These are the in-memory layouts the indexes use today, written as is, so that `load` copies sections without converting them. Record sizes follow from the header and parameter sections; readers derive them and compare them with the section sizes. Floats inside records are not aligned. Reserved bytes inside bulk data are written as zero and ignored by readers, which do not scan bulk data for them. + +The tables use `id_bits = 32`, so an ID takes 4 bytes; this is the only value the in-memory indexes write. + +#### Single-code record + +One vector's `b`-bit code with its two factors. IVF-RaBitQ and HNSW-RaBitQ use it for the extra bits (`b = code_bits - 1`), SymphonyQG for its whole code (`b = code_bits`). + +| Offset | Size | Content | +| --- | --- | --- | +| 0 | `padded_dim * b / 8` | The codes as [packed codes](#packed-codes-version-1) of width `b` | +| `padded_dim * b / 8` | 4 | `add`, f32 | +| `padded_dim * b / 8 + 4` | 4 | `rescale`, f32 | + +The record has no `error` factor. It should also become a raw-array format of `rabitqlib.quant`, so that both share one definition. + +#### HNSW sign code + +HNSW-RaBitQ stores the 1-bit code in a different bit order from the `rabitqlib.quant` bitset. Split the `padded_dim` coordinates into groups of 64, where the last group has 32 coordinates when `padded_dim` is not a multiple of 64. A group of `G` bytes, read as a little-endian integer, stores coordinate `t` of the group in bit `8G - 1 - t`. This is what `pack_binary_to_bytes` writes (checked against it for `padded_dim` 32, 64, 96, 160 and 1024); the implementation adds a golden test that freezes it. + +#### HNSW-RaBitQ records + +Let `L = 4 * (base_max_degree + 1)` and `U = 4 * (upper_max_degree + 1)`. + +| Offset | Size | Content | +| --- | --- | --- | +| 0 | 2 | Base-layer neighbor count, u16 | +| 2 | 2 | Reserved | +| 4 | `4 * base_max_degree` | Base-layer neighbor IDs; entries past the count are unspecified | +| `L` | 4 | Cluster ID, u32; bit 31 is the removal mark (see [HNSW removal](development.md#hnsw-removal)) | +| `L + 4` | 4 | Row ID of the point, u32 | +| `L + 8` | `padded_dim / 8` | [Sign code](#hnsw-sign-code) | +| `L + 8 + padded_dim / 8` | 12 | `add`, `rescale` and `error`, f32 | +| `L + 20 + padded_dim / 8` | | When `code_bits > 1`, a [single-code record](#single-code-record) with `b = code_bits - 1` | + +A point at level `l > 0` has `l` upper-layer link lists in `hnsw_upper_links`, for levels 1 to `l`, each `U` bytes: a u16 count, 2 reserved bytes, then `upper_max_degree` neighbor IDs. + +#### HNSW routing graph + +The centroid graph of IVF-RaBitQ's HNSW routing is the layout hnswlib's `saveIndex` writes. Unlike the rest of format 2 it is not defined field by field by this library, so it is fixed here as written on the supported 64-bit little-endian platforms (`size_t` is 8 bytes, `int` and `unsigned` 4): + +| Offset | Size | Type | Field | Rule | +| --- | --- | --- | --- | --- | +| 0 | 8 | u64 | `offset_level0` | 0 | +| 8 | 8 | u64 | `max_elements` | At least `count` | +| 16 | 8 | u64 | `count` | `num_clusters` | +| 24 | 8 | u64 | `record_bytes` | `label_offset + 8` | +| 32 | 8 | u64 | `label_offset` | `data_offset + 4 * padded_dim` | +| 40 | 8 | u64 | `data_offset` | `4 + 4 * max_m0` | +| 48 | 4 | i32 | `max_level` | The entry point's level, at least 0 | +| 52 | 4 | u32 | `entry_point` | Below `count` | +| 56 | 8 | u64 | `max_m` | `m` | +| 64 | 8 | u64 | `max_m0` | `2 * m` | +| 72 | 8 | u64 | `m` | 2 to 10,000 | +| 80 | 8 | f64 | `level_multiplier` | Finite and positive | +| 88 | 8 | u64 | `ef_construction` | At least `m` | + +Then `count` base records of `record_bytes`: a u16 neighbor count of at most `max_m0`, a flags byte whose bit 0 (hnswlib's deletion mark) must be clear, a reserved byte, `max_m0` u32 neighbor IDs, the rotated centroid (`padded_dim` f32) and its label (u64). Then, per point in order, a u32 byte length followed by that many bytes of upper-layer link lists: `level` lists of `4 + 4 * max_m` bytes each, laid out as in the base record. Every neighbor ID is below `count` and has at least the level of the list that holds it. + +#### SymphonyQG records + +Every vertex has `max_degree` neighbor slots, a multiple of 32. Let `V` be the size of the vector part: `4 * dim` for raw vectors, or the [single-code record](#single-code-record) size with `b = code_bits` otherwise. + +| Offset | Size | Content | +| --- | --- | --- | +| 0 | `V` | The vector: `dim` float32 values (`raw_type = float32`, `code_bits = 0`), or a single-code record | +| `V` | `(max_degree / 32) * (padded_dim * 4 + 256)` | The neighbors' 1-bit codes, quantized against this vertex, as [FastScan blocks](#fastscan-blocks-version-1) without the `error` factor: sign codes, then 32 `add` and 32 `rescale` values | +| `V + (max_degree / 32) * (padded_dim * 4 + 256)` | `4 * max_degree` | Neighbor IDs | + +Unused lanes of a partial neighbor block have zero factors. Their IDs are left over from construction, but like every ID slot they are below `num_points`, which readers check. + +### IVF-RaBitQ files + +Besides the rules above, an IVF-RaBitQ reader requires: + +- `num_points` from 1 to 2^31, and every stored ID below it; +- `code_bits` from 1 to 9, with `raw_type = float32` only when `code_bits = 1`, and `raw_type` none or float32; +- `id_bits = 32`; +- in `ivf_params`, `num_clusters` from 1 to 2^31, `initializer` 1 to 3, and the reserved field zero; +- cluster sizes that add up to `num_points`; +- a seed, when known, of at most 2^32 − 1. + +This release also rejects the cosine metric, which no IVF-RaBitQ file stores yet. + +### HNSW-RaBitQ files + +An HNSW-RaBitQ reader requires `num_points` from 1 to 2^31, `code_bits` from 1 to 9, `raw_type` none, `id_bits = 32`, `dim` at least 64, and an FHT/Kac rotator. In `hnsw_params`: `max_elements` from `num_points` to 2^31; `num_clusters` from 1 to 2^31; `M` from 2 to 10,000, with `upper_max_degree = M` and `base_max_degree = 2 * M`; `ef_construction` at least `M`; a finite positive `level_multiplier`; `max_level` at least 0; `entry_point` below `num_points`; and the reserved field zero. Every level in `hnsw_levels` is at most `max_level`, and the entry point's level equals it. Every record's cluster ID (without the removal mark) is below `num_clusters`, its labels are distinct, its base-layer neighbor count is at most `base_max_degree` and each upper-layer count at most `M`, and every neighbor ID is below `num_points` with at least the level of the list that holds it. This release rejects the cosine metric. + +### SymphonyQG files + +A SymphonyQG reader requires `id_bits = 32` and either raw vectors (`code_bits = 0`, `raw_type = float32`) or `code_bits` 4 or 8 with `raw_type` none; a seed, when known, of at most 2^32 − 1; `num_points` up to 2^31; in `qg_params`, a `max_degree` that is a positive multiple of 32 and below `num_points`, `entry_point` below `num_points`, `init` 0 to 2 and the reserved field zero; `qg_centroid` exactly when the codes are quantized; and every neighbor ID below `num_points`. This release rejects the cosine metric. + +### Companion files + +Only IVF-SSD uses them. `.ssd` holds the page-packed inverted lists, read with `O_DIRECT` and too large to keep in the main file. The `companions` section lists each one: + +| Offset | Size | Type | Field | +| --- | --- | --- | --- | +| 0 | 4 | u32 | `role`: 1 = SSD inverted lists | +| 4 | 4 | u32 | reserved | +| 8 | 32 | ASCII | `suffix` appended to the main file's path, such as `".ssd"`, NUL-padded; no path separators | +| 40 | 8 | u64 | `size` in bytes | +| 48 | 8 | u64 | `fingerprint` | + +`fingerprint` is the 64-bit FNV-1a hash that RaBitQ-SSD already uses (offset basis 1469598103934665603, prime 1099511628211) over the rotated centroids as f32 bytes. It is deterministic, so saving the same index twice gives identical files. + +A companion starts with a 4,096-byte header page: magic `"RABQSSD\0"`, version, `fingerprint`, `num_points` and `code_bits`, then zeros, with a CRC-32 over the page. Data pages follow at offset 4,096, keeping the 4,096-byte alignment that `O_DIRECT` needs. Loading rejects a missing companion, or one whose size, fingerprint or header does not match the main file. + +RaBitQ-SSD's swappable coarse quantizer is not a companion: the base does not pin one. A coarse-quantizer file is its own file and records the base's fingerprint, so the pairing is checked from the coarse side. Its layout is specified when IVF-SSD moves into the library. + +### Reading + +A reader checks, in order, and validates every size before allocating. The loader of each index knows the index kind it expects, so a file of another kind is reported as such before any section is checked: + +1. The file is at least 128 bytes long, and the magic, `version` and `header_bytes` are valid. +2. The section table is non-empty and lies within the file, and `header_crc32` matches. +3. Feature flags, as in the table above. +4. `index_kind` is the expected one; the fields that every kind shares are valid; then the expecting index's own header rules hold. +5. Each table entry, in order: its flags and reserved field, alignment, bounds, order, overlap and repetition, then its kind. An unknown required kind rejects the file here. +6. The file ends where the table or the last section ends. +7. Each known checksummed section matches its `crc32`. +8. The index's required sections are present, and each section's size equals what the header and parameter sections imply, before its content is read. + +Steps 1–7 belong to the container; step 8 and the index's header rules in step 4 belong to each index. Readers read through the stream buffer, so the stream's exception mask does not matter, and need seekable input. + +A file that fails any of these checks, or a damaged file of an [earlier layout](#earlier-layouts), raises `rabitqlib::FormatError` in C++, a subclass of `std::runtime_error`, and `rabitqlib.FormatError` in Python, a subclass of both `ValueError` and `RuntimeError`. `rabitqlib.quant` raises the same class for blobs. A file that cannot be opened raises a plain `std::runtime_error` (Python `RuntimeError`). A read that fails after opening cannot be told from a truncated file, so it raises `FormatError`. + +Index files remain trusted input (see [Index file compatibility](docs/index/usage.md#index-file-compatibility)). These checks report truncated, corrupt or mismatched files; they do not make crafted files safe. + +### Writing + +A writer writes the file once, front to back, so it does not need seekable output (`save` writes devices and FIFOs in place): it buffers the checksummed sections, which are small, to know their checksums before the table. `save` writes to a temporary file and renames it into place, as today. With the routing graph inside the main file, IVF-RaBitQ no longer has a second rename that can leave a new sidecar beside an old index. Resaving a file loaded from an earlier layout writes format 2. IVF-SSD companion files are written before the main file, so a failed save never leaves a main file that points to a missing companion. + +### Earlier layouts + +Readers dispatch on the leading bytes: `"RABQINDX"`, then each index's legacy magic (IVF `RABQIDX1`, SymphonyQG quantized), then the unversioned fallbacks the readers accept today. Legacy files have no cosine metric, seed or 64-bit IDs; they load as `seed_known = false` with 32-bit IDs, as they do in 0.8. The frozen fixtures in `tests/fixtures/persistence/` keep covering every layout. + +### Compatibility rules + +- A new release reads every earlier released index file. +- An older release reads a newer format-2 file or rejects it with an error naming the problem; it never misreads it. Older releases reject `version` values other than 2, unknown `index_kind`, `metric` or `id_bits` values, unknown required sections and unknown `incompat` bits. +- An additive change appends header fields (raising `header_bytes`), adds sections, or adds feature bits, chosen by the table in [Feature flags](#feature-flags). A change that an older reader cannot detect from these is not allowed; it needs `version = 3`. +- Every released change adds a frozen fixture. Synthetic "future" fixtures check the rules above: an unknown optional section loads; an unknown required section, an `incompat` bit or a larger `version` is rejected with its message; a `ro_compat` bit loads read-only; a longer header loads. + +### Open questions + +- Whether the header should be a [blob](#blobs) kind instead of its own magic. Proposed: its own magic, as above. A blob checksums its whole body, which is too costly for large index files, so index files get their own container and nest blobs in sections. diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index a1dfd332..7b40b1af 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -14,6 +14,7 @@ nav: - C++ Integration: cpp_integration.md # Group by minor version: nested groups stay collapsed except the current page's. - Release notes: + - Unreleased: releases/unreleased.md - "0.8": - "0.8.3": releases/0.8.3.md - "0.8.2": releases/0.8.2.md diff --git a/include/rabitqlib/index/format_error.hpp b/include/rabitqlib/index/format_error.hpp new file mode 100644 index 00000000..104a7374 --- /dev/null +++ b/include/rabitqlib/index/format_error.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include + +namespace rabitqlib { + +/** + * @brief An index file whose contents are invalid, or which this release cannot read. + * + * Loading any index throws it for a damaged, mismatched or unsupported file in every + * layout; failing to open the file is a plain std::runtime_error. Python raises it as + * rabitqlib.FormatError, a subclass of both ValueError and RuntimeError. + */ +class FormatError : public std::runtime_error { + public: + using std::runtime_error::runtime_error; +}; + +} // namespace rabitqlib diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 001ffc65..1a3cbdd2 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -7,9 +7,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -70,6 +72,9 @@ class HierarchicalNSW { [[nodiscard]] size_t M() const { return M_; } [[nodiscard]] size_t ef_construction() const { return ef_construction_; } [[nodiscard]] MetricType metric_type() const { return metric_type_; } + // The construction seed, or no value for an index loaded from a file that does not + // record it (files from releases up to 0.8). + [[nodiscard]] std::optional seed() const { return seed_; } // Logical capacity, nonremoved points, and stored records (including removals). [[nodiscard]] size_t max_elements() const { return max_elements_; } [[nodiscard]] size_t num_points() const { @@ -284,6 +289,10 @@ class HierarchicalNSW { std::unordered_map label_lookup_; std::default_random_engine level_generator_; + std::optional seed_; // the construction seed, when known + // Loaded from a file with read-only-compatible features this release does not know: + // search works, but add, resize, remove and save refuse the index. + bool read_only_ = false; std::unique_ptr visited_set_pool_{nullptr}; @@ -381,9 +390,20 @@ class HierarchicalNSW { int random_level(double multiplier); - // Write the index file at exactly this path. + // Write the index file in format 2 at exactly this path. void write_file(const char* filename) const; + // Read a file into this freshly created index: format 2, or the earlier layout. + void read_format2(std::istream& input); + void read_legacy(std::ifstream& input); + + // Check the loaded records and link lists, and rebuild what they imply: the label + // lookup and the removed-point count. + void validate_graph(); + + // Refuse an operation that would modify or save a read-only index. + void check_writable(const char* operation) const; + // A link list is a count header followed by neighbor IDs. [[nodiscard]] PID* upper_layer_links(PID internal_id, int level) { return reinterpret_cast( diff --git a/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp b/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp index 32bde339..080a42ef 100644 --- a/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp +++ b/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp @@ -5,37 +5,40 @@ #include #include #include -#include #include +#include #include #include -#include +#include "rabitqlib/index/format_error.hpp" #include "rabitqlib/third/hnswlib/hnswlib.h" -#include "rabitqlib/utils/path.hpp" namespace rabitqlib::ivf::detail { -// The HNSW initializer stores its graph at the index path plus this suffix. Saving, -// loading and atomic replacement must all agree on it. +// Releases before index format 2 stored the HNSW routing graph at the index path plus +// this suffix; loading those files reads it from there. inline constexpr char kHnswSidecarSuffix[] = ".hnsw"; -// hnswlib's native sidecar has no version tag. Check the existing writer's -// header and record lengths before its loader uses offsets or allocates links. -// As with the main index file, callers must not mutate the file during loading. -inline void validate_hnsw_sidecar_file( - const std::string& path, size_t expected_count, size_t vector_bytes +// hnswlib's native layout has no version tag. Check the header and record lengths of +// the `bytes` bytes at the current position of `input` before hnswlib's loader uses +// offsets or allocates links, then restore the position. As with the index file, +// callers must not mutate the data during loading. +inline void validate_hnsw_graph( + std::istream& input, uint64_t bytes, size_t expected_count, size_t vector_bytes ) { - std::ifstream input(io_impl::filesystem_path(path), std::ios::binary | std::ios::ate); - if (!input) { - throw std::runtime_error("Cannot open HNSW centroid sidecar"); + const std::streampos begin = input.tellg(); + if (begin == std::streampos(std::streamoff(-1))) { + throw FormatError("Cannot read HNSW centroid graph"); } - const auto end = input.tellg(); - input.seekg(0); + // Stay within the graph's bytes, as a budget rather than by querying the position: + // in an index file, other data follows them. + auto remaining = static_cast(bytes); const auto read = [&](auto& value) { - if (!input.read(reinterpret_cast(&value), sizeof(value))) { - throw std::runtime_error("Truncated HNSW centroid sidecar"); + if (remaining < static_cast(sizeof(value)) || + !input.read(reinterpret_cast(&value), sizeof(value))) { + throw FormatError("Truncated HNSW centroid graph"); } + remaining -= static_cast(sizeof(value)); }; size_t level0_offset = 0, capacity = 0, count = 0, level0_record_bytes = 0; size_t label_offset = 0, data_offset = 0; @@ -58,12 +61,12 @@ inline void validate_hnsw_sidecar_file( read(ef_construction); if (count != expected_count || count == 0 || capacity < count || count > std::numeric_limits::max()) { - throw std::runtime_error("HNSW centroid count does not match IVF cluster count"); + throw FormatError("HNSW centroid count does not match IVF cluster count"); } constexpr size_t kMaxM = 10000; // hnswlib caps M at this value. if (m < 2 || m > kMaxM || max_m != m || max_m0 != 2 * m || !std::isfinite(mult) || mult <= 0 || ef_construction < m) { - throw std::runtime_error("Invalid HNSW centroid graph parameters"); + throw FormatError("Invalid HNSW centroid graph parameters"); } const size_t upper_layer_links_bytes = sizeof(hnswlib::linklistsizeint) + max_m * sizeof(hnswlib::tableint); @@ -73,14 +76,13 @@ inline void validate_hnsw_sidecar_file( data_offset > label_offset || label_offset > level0_record_bytes || label_offset - data_offset != vector_bytes || level0_record_bytes - label_offset != sizeof(hnswlib::labeltype)) { - throw std::runtime_error("HNSW centroid layout does not match IVF dimension"); + throw FormatError("HNSW centroid layout does not match IVF dimension"); } if (max_level < 0 || entry >= count) { - throw std::runtime_error("Invalid HNSW centroid entry point or maximum level"); + throw FormatError("Invalid HNSW centroid entry point or maximum level"); } - auto remaining = end - input.tellg(); - if (remaining < 0 || count > static_cast(remaining) / level0_record_bytes) { - throw std::runtime_error("Truncated HNSW centroid base records"); + if (count > static_cast(remaining) / level0_record_bytes) { + throw FormatError("Truncated HNSW centroid base records"); } const auto level0_records_bytes = static_cast(count * level0_record_bytes); @@ -89,14 +91,10 @@ inline void validate_hnsw_sidecar_file( size_t highest_level = 0, entry_level = 0; for (size_t i = 0; i < count; ++i) { unsigned int link_bytes = 0; // Native type written by hnswlib::saveIndex. - if (remaining < static_cast(sizeof(link_bytes))) { - throw std::runtime_error("Truncated HNSW centroid sidecar"); - } read(link_bytes); - remaining -= static_cast(sizeof(link_bytes)); if (link_bytes % upper_layer_links_bytes != 0 || link_bytes > static_cast(remaining)) { - throw std::runtime_error("Invalid HNSW centroid upper-link length"); + throw FormatError("Invalid HNSW centroid upper-link length"); } const size_t levels = link_bytes / upper_layer_links_bytes; highest_level = std::max(highest_level, levels); @@ -112,8 +110,10 @@ inline void validate_hnsw_sidecar_file( } if (!input || remaining != 0 || highest_level != static_cast(max_level) || entry_level != highest_level) { - throw std::runtime_error("Invalid HNSW centroid level structure"); + throw FormatError("Invalid HNSW centroid level structure"); } + input.clear(); + input.seekg(begin); } inline void validate_hnsw_centroid_graph(const hnswlib::HierarchicalNSW& graph) { @@ -121,11 +121,11 @@ inline void validate_hnsw_centroid_graph(const hnswlib::HierarchicalNSW& if (graph.enterpoint_node_ >= count || graph.maxlevel_ < 0 || graph.element_levels_[graph.enterpoint_node_] != graph.maxlevel_ || graph.num_deleted_ != 0) { - throw std::runtime_error("Invalid HNSW centroid entry point or deletion state"); + throw FormatError("Invalid HNSW centroid entry point or deletion state"); } for (size_t i = 0; i < count; ++i) { if (graph.element_levels_[i] < 0 || graph.element_levels_[i] > graph.maxlevel_) { - throw std::runtime_error("Invalid HNSW centroid node level"); + throw FormatError("Invalid HNSW centroid node level"); } for (int level = 0; level <= graph.element_levels_[i]; ++level) { const char* links = @@ -136,7 +136,7 @@ inline void validate_hnsw_centroid_graph(const hnswlib::HierarchicalNSW& uint16_t degree = 0; // The high two bytes hold flags, not the degree. std::memcpy(°ree, links, sizeof(degree)); if (degree > (level == 0 ? graph.maxM0_ : graph.maxM_)) { - throw std::runtime_error("Invalid HNSW centroid neighbor count"); + throw FormatError("Invalid HNSW centroid neighbor count"); } for (size_t j = 0; j < degree; ++j) { hnswlib::tableint neighbor = 0; @@ -146,9 +146,7 @@ inline void validate_hnsw_centroid_graph(const hnswlib::HierarchicalNSW& sizeof(neighbor) ); if (neighbor >= count || graph.element_levels_[neighbor] < level) { - throw std::runtime_error( - "Invalid HNSW centroid neighbor or target level" - ); + throw FormatError("Invalid HNSW centroid neighbor or target level"); } } } diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index 618d6731..1104965a 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -3,10 +3,12 @@ #include #include #include +#include #include -#include #include +#include #include +#include #include #include #include @@ -15,6 +17,7 @@ #include "rabitqlib/core/metric.hpp" #include "rabitqlib/index/detail/parallel_for.hpp" +#include "rabitqlib/index/format_error.hpp" #include "rabitqlib/index/ivf/detail/hnsw_sidecar.hpp" #include "rabitqlib/index/ivf/types.hpp" // IWYU pragma: export #include "rabitqlib/index/types.hpp" @@ -61,8 +64,11 @@ class Initializer { virtual void centroids_distances( const float* query, size_t nprobe, std::vector>& candidates ) const = 0; - virtual void load(std::ifstream& input, const char* filename) = 0; - virtual void save(std::ofstream& output, const char* filename) const = 0; + // Write the routing data that read() restores. + virtual void write(std::ostream& output) const = 0; + // Read `bytes` bytes that write() wrote, from the current position of `input`. + // Data that does not fit this initializer raises FormatError. + virtual void read(std::istream& input, uint64_t bytes) = 0; }; inline Initializer::~Initializer() = default; @@ -113,8 +119,8 @@ class FlatInitializer : public Initializer { std::copy_n(centroid_candidates.begin(), num_probes, candidates.begin()); } - // The flat initializer stores its centroids inline in the index file stream. - void save(std::ofstream& output, const char* /*filename*/) const override { + // The flat initializer stores only its centroids. + void write(std::ostream& output) const override { output.write( reinterpret_cast(centroids_.data()), static_cast(sizeof(float) * dim_ * num_cluster_) @@ -124,13 +130,15 @@ class FlatInitializer : public Initializer { } } - void load(std::ifstream& input, const char* /*filename*/) override { + void read(std::istream& input, uint64_t bytes) override { + if (bytes != sizeof(float) * dim_ * num_cluster_) { + throw FormatError("IVF Flat initializer data has the wrong size"); + } input.read( - reinterpret_cast(centroids_.data()), - static_cast(sizeof(float) * dim_ * num_cluster_) + reinterpret_cast(centroids_.data()), static_cast(bytes) ); if (!input) { - throw std::runtime_error("IVF Flat initializer cannot be read"); + throw FormatError("IVF Flat initializer cannot be read"); } } }; @@ -153,8 +161,8 @@ class FlatRaBitQInitializer : public Initializer { void centroids_distances( const float* query, size_t nprobe, std::vector>& candidates ) const override; - void save(std::ofstream& output, const char* filename) const override; - void load(std::ifstream& input, const char* filename) override; + void write(std::ostream& output) const override; + void read(std::istream& input, uint64_t bytes) override; }; // Keep centroid routing on the same runtime-selected distance kernels as flat IVF. @@ -291,28 +299,30 @@ class HNSWInitializer : public Initializer { } } - // The HNSW initializer stores its graph in a separate hnswlib sidecar file. - void save(std::ofstream& /*output*/, const char* filename) const override { - alg_hnsw_->saveIndex(std::string(filename) + detail::kHnswSidecarSuffix); - } + // The HNSW initializer stores its graph in hnswlib's own layout, which earlier + // releases wrote to the sidecar file .hnsw. + void write(std::ostream& output) const override { alg_hnsw_->saveIndex(output); } - void load(std::ifstream& /*input*/, const char* filename) override { - const std::string sidecar = std::string(filename) + detail::kHnswSidecarSuffix; - detail::validate_hnsw_sidecar_file(sidecar, num_cluster_, space_->get_data_size()); - alg_hnsw_->loadIndex(sidecar, space_, num_cluster_); + void read(std::istream& input, uint64_t bytes) override { + detail::validate_hnsw_graph(input, bytes, num_cluster_, space_->get_data_size()); + alg_hnsw_->loadIndex(input, static_cast(bytes), space_, num_cluster_); const auto& graph = *alg_hnsw_; - // validate_hnsw_sidecar_file already checked these header fields, and loadIndex - // reads them again from a second open of the file. Recheck the loaded values so - // a sidecar replaced in between cannot give centroid() an out-of-record offset. + // validate_hnsw_graph already checked these header fields, and loadIndex reads + // them again. Recheck the loaded values so that data changed in between cannot + // give centroid() an out-of-record offset. if (graph.offsetData_ > graph.label_offset_ || graph.label_offset_ > graph.size_data_per_element_ || graph.label_offset_ - graph.offsetData_ != space_->get_data_size() || graph.size_data_per_element_ - graph.label_offset_ != sizeof(hnswlib::labeltype)) { - throw std::runtime_error("HNSW centroid sidecar changed while loading"); + throw FormatError("HNSW centroid graph changed while loading"); } detail::validate_hnsw_centroid_graph(graph); - rebuild_label_lookup(); + try { + rebuild_label_lookup(); + } catch (const std::runtime_error& error) { + throw FormatError(error.what()); // stored labels that do not fit + } } ~HNSWInitializer() override = default; diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index 745ed983..05829655 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -90,7 +91,11 @@ class IVF { size_t num_clusters_ = 0; bool raw_reranking_ = false; // raw vectors replace extra-bit codes bool ready_ = false; // centroids set, construction or loading completed - size_t ex_bits_ = 0; // total bits = ex_bits_ + 1; 0 with raw reranking + // Loaded from a file with read-only-compatible features this release does not know: + // search works, but add, reserve, remove and save refuse the index. + bool read_only_ = false; + std::optional seed_; // the construction seed, when known + size_t ex_bits_ = 0; // total bits = ex_bits_ + 1; 0 with raw reranking RotatorType rotator_type_ = RotatorType::FhtKacRotator; std::unique_ptr> rotator_; std::vector clusters_; @@ -190,9 +195,22 @@ class IVF { const float* query, size_t k, size_t nprobe, const PID* result_ids ) const; - // Write the index file, and any initializer sidecar, at exactly this path. + // Write the index file in format 2 at exactly this path. void write_file(const char* filename) const; + // Load a format-2 file into this freshly created index. + void read_format2(std::istream& input); + + // Bytes of Flat or FlatRaBitQ routing data; HNSW routing has no fixed size. + [[nodiscard]] size_t flat_routing_bytes() const; + + // Check the stored IDs, create the cluster views, count removed points and mark the + // index ready: the end of loading in every format. + void complete_load(const std::vector& cluster_sizes); + + // Refuse an operation that would modify or save a read-only index. + void check_writable(const char* operation) const; + public: explicit IVF() noexcept; // dim is the original input dimension; total_bits is 1-9 for quantized @@ -247,6 +265,9 @@ class IVF { [[nodiscard]] size_t dimension() const { return dim_; } [[nodiscard]] size_t nbits() const { return raw_reranking_ ? 32 : ex_bits_ + 1; } [[nodiscard]] MetricType metric_type() const { return metric_type_; } + // The rotation seed of the seeded constructor, or no value when it is unknown: the + // point-count constructor draws one, and files from releases up to 0.8 omit it. + [[nodiscard]] std::optional seed() const { return seed_; } [[nodiscard]] RotatorType rotator_type() const { return rotator_type_; } [[nodiscard]] InitializerType initializer_type() const { return initializer_type_; } // Centroids are set, by train, set_centroids, construct or load. diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 1e82f6d0..b72b35e2 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -2,7 +2,9 @@ #include #include +#include #include +#include #include #include @@ -26,6 +28,9 @@ class SearchBuffer; namespace rabitqlib::symqg { +// How a builder initializes the graph before refinement. +enum class QGInitialization { PiPNN, Random }; + class QuantizedQuery { private: const float* rotated_query_; @@ -114,6 +119,13 @@ class QuantizedGraph { size_t default_ef_search_ = 0; bool ready_ = false; uint32_t seed_ = 42; + bool seed_known_ = false; // false for files that do not record the seed + // Construction settings a builder records; unknown (0 and no value) otherwise. + size_t ef_construction_ = 0; + std::optional initialization_; + // Loaded from a file with read-only-compatible features this release does not know: + // search works, but save refuses the graph until a builder rebuilds it. + bool read_only_ = false; [[nodiscard]] static size_t checked_add(size_t lhs, size_t rhs); @@ -127,9 +139,12 @@ class QuantizedGraph { void initialize(bool randomize_rotation = true); - // Write the graph file at exactly this path. + // Write the graph file in index format 2 at exactly this path. void write_file(const char* filename) const; + // Load a format-2 file into this freshly created graph. + void read_format2(std::istream& input); + void copy_vectors(const float* data, size_t num_workers); void set_quantization_centroid(const float* centroid); @@ -279,6 +294,14 @@ class QuantizedGraph { [[nodiscard]] size_t quantization_bits() const; + // The rotation seed, or no value for a graph loaded from a file that does not + // record it (files from releases up to 0.8). + [[nodiscard]] std::optional seed() const; + + // The settings the builder used, or no value when they are unknown. + [[nodiscard]] std::optional ef_construction() const; + [[nodiscard]] std::optional initialization() const; + [[nodiscard]] bool is_quantized() const; void set_ep(PID entry); diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index e2bf47e7..70e5f4b3 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -14,8 +14,6 @@ namespace rabitqlib::symqg { inline constexpr size_t kMaxBsIter = 5; // max iter for binary search of pruning bar using CandidateList = std::vector>; -enum class QGInitialization { PiPNN, Random }; - // Owns temporary state for SymphonyQG initialization and refinement. class QGBuilder { friend struct QGConstructionTestAccess; diff --git a/include/rabitqlib/third/hnswlib/hnswalg.h b/include/rabitqlib/third/hnswlib/hnswalg.h index c8ff85d1..d6ccf26a 100644 --- a/include/rabitqlib/third/hnswlib/hnswalg.h +++ b/include/rabitqlib/third/hnswlib/hnswalg.h @@ -688,8 +688,14 @@ class HierarchicalNSW : public AlgorithmInterface { if (!output.is_open()) throw std::runtime_error("Cannot open HNSW index file for writing"); output.exceptions(std::ios::failbit | std::ios::badbit); - std::streampos position; + saveIndex(output); + output.close(); + } + + // RaBitQ: the stream overloads let an index embed the graph in a section of its own + // file. The bytes are those of the path overload, which delegates to them. + void saveIndex(std::ostream &output) { writeBinaryPOD(output, offsetLevel0_); writeBinaryPOD(output, max_elements_); writeBinaryPOD(output, cur_element_count); @@ -713,7 +719,6 @@ class HierarchicalNSW : public AlgorithmInterface { if (linkListSize) output.write(linkLists_[i], linkListSize); } - output.close(); } @@ -723,11 +728,20 @@ class HierarchicalNSW : public AlgorithmInterface { if (!input.is_open()) throw std::runtime_error("Cannot open file"); - clear(); // get file size: input.seekg(0, input.end); - std::streampos total_filesize = input.tellg(); + const std::streamoff total_filesize = input.tellg(); input.seekg(0, input.beg); + loadIndex(input, static_cast(total_filesize), s, max_elements_i); + input.close(); + } + + + // RaBitQ: read a graph of `total_bytes` bytes that starts at the current position. + void loadIndex(std::istream &input, size_t total_bytes, SpaceInterface *s, size_t max_elements_i = 0) { + clear(); + const std::streampos begin = input.tellg(); + const std::streampos end = begin + static_cast(total_bytes); readBinaryPOD(input, offsetLevel0_); readBinaryPOD(input, max_elements_); @@ -758,7 +772,7 @@ class HierarchicalNSW : public AlgorithmInterface { /// Optional - check if index is ok: input.seekg(cur_element_count * size_data_per_element_, input.cur); for (size_t i = 0; i < cur_element_count; i++) { - if (input.tellg() < 0 || input.tellg() >= total_filesize) { + if (input.tellg() < begin || input.tellg() >= end) { throw std::runtime_error("Index seems to be corrupted or unsupported"); } @@ -770,13 +784,13 @@ class HierarchicalNSW : public AlgorithmInterface { } // throw exception if it either corrupted or old index - if (input.tellg() != total_filesize) + if (input.tellg() != end) throw std::runtime_error("Index seems to be corrupted or unsupported"); input.clear(); /// Optional check end - input.seekg(pos, input.beg); + input.seekg(pos); data_level0_memory_ = (char *) malloc(max_elements * size_data_per_element_); if (data_level0_memory_ == nullptr) @@ -819,10 +833,6 @@ class HierarchicalNSW : public AlgorithmInterface { if (allow_replace_deleted_) deleted_elements.insert(i); } } - - input.close(); - - return; } diff --git a/python_bindings/__init__.py b/python_bindings/__init__.py index b84c57a5..1c29aea8 100644 --- a/python_bindings/__init__.py +++ b/python_bindings/__init__.py @@ -9,6 +9,7 @@ from ._rabitqlib import ( FinalAssignmentMode, + FormatError, HnswIndex, IvfIndex, QGKMeansIterationStats, @@ -627,6 +628,7 @@ class RaBitQKMeans(_KMeans): __all__ = [ "FinalAssignmentMode", + "FormatError", "HnswIndex", "IvfIndex", "QGKMeans", diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index 62d4f0b5..ef53ad8a 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -15,6 +15,7 @@ __version__: str __all__ = [ "FinalAssignmentMode", + "FormatError", "HnswIndex", "IvfIndex", "QGKMeans", @@ -34,6 +35,9 @@ _QGInit: TypeAlias = Literal["pipnn", "random"] # Rank, ID ranges and Python bool (an int subtype) are validated at runtime. _IntegerIds: TypeAlias = npt.NDArray[np.integer[Any]] | Sequence[int | np.integer[Any]] +class FormatError(ValueError, RuntimeError): + """A damaged or mismatched index file or serialized byte string, or one this release cannot read. A subclass of both ValueError and RuntimeError.""" + # pybind11 enums are not subclasses of enum.Enum. class FinalAssignmentMode: """Final centroid assignment policy. @@ -208,7 +212,7 @@ class IvfIndex: @property def seed(self) -> int | None: - """Rotation seed, or None for an index loaded from disk or created in the 0.7 form.""" + """Rotation seed, or None for an index loaded from a file from a release up to 0.8 or created in the 0.7 form.""" @property def cp(self) -> RaBitQKMeansParameters: @@ -312,7 +316,7 @@ class IvfIndex: cluster_ids: _IntegerIds | None = None, num_threads: SupportsIndex | None = None, fast_quantization: bool | np.bool_ = False, - ) -> npt.NDArray[np.uint32]: + ) -> npt.NDArray[np.int64]: """Add vectors, training the index first if needed. Parameters @@ -332,7 +336,7 @@ class IvfIndex: Returns ------- - ids : ndarray of uint32, shape (n,) + ids : ndarray of int64, shape (n,) Assigned point IDs. Notes @@ -358,7 +362,7 @@ class IvfIndex: cluster_ids: _IntegerIds | None, num_threads: SupportsIndex = ..., fast_quantization: bool | np.bool_ = ..., - ) -> npt.NDArray[np.uint32]: ... + ) -> npt.NDArray[np.int64]: ... def remove(self, ids: _IntegerIds) -> int: """Exclude point IDs from subsequent search results. @@ -385,7 +389,7 @@ class IvfIndex: nprobe: SupportsIndex | None = None, high_accuracy: bool | np.bool_ | None = None, num_threads: SupportsIndex | None = None, - ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: + ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float32]]: """Search a batch of query vectors. Parameters @@ -408,7 +412,7 @@ class IvfIndex: Returns ------- - ids : ndarray of uint32, shape (n_queries, k) + ids : ndarray of int64, shape (n_queries, k) Point IDs of the selected neighbors. distances : ndarray of float32, shape (n_queries, k) Lower values are better; squared L2 estimates for l2 and inner-product @@ -417,7 +421,7 @@ class IvfIndex: Notes ----- Requires a trained or loaded index. Returns (ids, distances), in that - order. Missing neighbors have ID 4294967295 and distance inf; test + order. Missing neighbors have ID -1 and distance inf; test np.isfinite(distances) to find them. Array shapes and value ranges are checked at runtime, not by static type checkers. """ @@ -434,7 +438,7 @@ class IvfIndex: Notes ----- A trained index without vectors raises RuntimeError; add vectors first. - IVF with HNSW centroid routing also writes its required sidecar file. + The one file includes HNSW centroid routing; releases up to 0.8 cannot read it. """ @staticmethod @@ -451,6 +455,11 @@ class IvfIndex: ------- index : IvfIndex A trained index ready for search. + + Notes + ----- + Files from releases up to 0.8 with HNSW centroid routing also need the + accompanying ``.hnsw`` file. """ @property @@ -636,7 +645,7 @@ class HnswIndex: @property def seed(self) -> int | None: - """Graph and rotation seed, or None for an index loaded from disk.""" + """Graph and rotation seed, or None for an index loaded from a file from a release up to 0.8.""" @property def cp(self) -> RaBitQKMeansParameters: @@ -747,7 +756,7 @@ class HnswIndex: cluster_ids: _IntegerIds | None = None, num_threads: SupportsIndex | None = None, fast_quantization: bool | np.bool_ = False, - ) -> npt.NDArray[np.uint32]: + ) -> npt.NDArray[np.int64]: """Add vectors, training the index first if needed. Parameters @@ -767,7 +776,7 @@ class HnswIndex: Returns ------- - ids : ndarray of uint32, shape (n,) + ids : ndarray of int64, shape (n,) Assigned point IDs. Notes @@ -795,7 +804,7 @@ class HnswIndex: cluster_ids: _IntegerIds | None, num_threads: SupportsIndex = ..., fast_quantization: bool | np.bool_ = ..., - ) -> npt.NDArray[np.uint32]: ... + ) -> npt.NDArray[np.int64]: ... def remove(self, ids: _IntegerIds) -> int: """Exclude point IDs from subsequent search results. @@ -821,7 +830,7 @@ class HnswIndex: k: SupportsIndex, ef: SupportsIndex | None = None, num_threads: SupportsIndex | None = None, - ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: + ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float32]]: """Search a batch of query vectors. Parameters @@ -841,7 +850,7 @@ class HnswIndex: Returns ------- - ids : ndarray of uint32, shape (n_queries, k) + ids : ndarray of int64, shape (n_queries, k) Point IDs of the selected neighbors. distances : ndarray of float32, shape (n_queries, k) Lower values are better; squared L2 estimates for l2 and inner-product @@ -850,7 +859,7 @@ class HnswIndex: Notes ----- Requires a trained or loaded index. Returns (ids, distances), in that - order. Missing neighbors have ID 4294967295 and distance inf; test + order. Missing neighbors have ID -1 and distance inf; test np.isfinite(distances) to find them. Array shapes and value ranges are checked at runtime, not by static type checkers. """ @@ -867,7 +876,8 @@ class HnswIndex: Notes ----- A trained index without vectors raises RuntimeError; add vectors first. - The saved state includes the rotation used to build the index. + The saved state includes the rotation used to build the index and its seed; + releases up to 0.8 cannot read the file. """ @staticmethod @@ -1063,11 +1073,11 @@ class SymqgIndex: @property def ef_construction(self) -> int | None: - """Construction search breadth of the current graph or the next add, or None for an index loaded from disk.""" + """Construction search breadth of the current graph or the next add, or None for an index loaded from a file from a release up to 0.8.""" @property def init(self) -> _QGInit | None: - """Graph initialization method of the current graph or the next add, or None for an index loaded from disk.""" + """Graph initialization method of the current graph or the next add, or None for an index loaded from a file from a release up to 0.8.""" @property def nbits(self) -> int: @@ -1083,7 +1093,7 @@ class SymqgIndex: @property def seed(self) -> int | None: - """Rotation and graph-construction seed, or None for an index loaded from disk.""" + """Rotation and graph-construction seed, or None for an index loaded from a file from a release up to 0.8.""" def train( self, x: npt.ArrayLike, *, num_threads: SupportsIndex | None = None @@ -1108,7 +1118,7 @@ class SymqgIndex: def add( self, data: npt.ArrayLike, *, num_threads: SupportsIndex | None = None - ) -> npt.NDArray[np.uint32]: + ) -> npt.NDArray[np.int64]: """Build the graph from all vectors. Parameters @@ -1122,7 +1132,7 @@ class SymqgIndex: Returns ------- - ids : ndarray of uint32, shape (n,) + ids : ndarray of int64, shape (n,) Point IDs, the row numbers of data. Notes @@ -1170,7 +1180,7 @@ class SymqgIndex: k: SupportsIndex, ef: SupportsIndex | None = None, num_threads: SupportsIndex | None = None, - ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: + ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float32]]: """Search a batch of query vectors. Parameters @@ -1190,7 +1200,7 @@ class SymqgIndex: Returns ------- - ids : ndarray of uint32, shape (n_queries, k) + ids : ndarray of int64, shape (n_queries, k) Point IDs of the selected neighbors. distances : ndarray of float32, shape (n_queries, k) Lower values are better; squared L2 estimates for l2 and inner-product @@ -1199,7 +1209,7 @@ class SymqgIndex: Notes ----- Returns (ids, distances), in that order. When the graph reaches fewer than k - points, the remaining neighbors have ID 4294967295 and distance inf; test + points, the remaining neighbors have ID -1 and distance inf; test np.isfinite(distances) to find them. Array shapes and value ranges are checked at runtime, not by static type checkers. """ @@ -1211,7 +1221,7 @@ class SymqgIndex: k: SupportsIndex, ef: SupportsIndex, num_threads: SupportsIndex = 1, - ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: + ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float32]]: """Deprecated since 0.8: the same as search(); use search().""" def save(self, path: _Path) -> None: @@ -1225,7 +1235,8 @@ class SymqgIndex: Notes ----- - The saved state includes the rotation used to build the index. + The saved state includes the rotation used to build the index and its seed; + releases up to 0.8 cannot read the file. """ @staticmethod diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index f88da6fc..03538ff9 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -353,13 +353,30 @@ RowView checked_output_rows( return make_row_view(std::move(array)); } +// Point IDs as Python sees them: int64, with -1 where the index reports kPidMax. +inline py::array_t python_ids(const PID* ids, std::vector shape) { + py::array_t result(std::move(shape)); + int64_t* out = result.mutable_data(); + for (py::ssize_t i = 0; i < result.size(); ++i) { + out[i] = ids[i] == kPidMax ? int64_t{-1} : static_cast(ids[i]); + } + return result; +} + +// (num_queries, k) outputs. The index writes row numbers into ids; slots without a +// neighbor keep kPidMax and +inf. Searches pad slots without a result, so k may exceed +// the number of points. struct SearchResults { - py::array_t ids; + std::vector shape; + std::vector ids; py::array_t distances; + + // (ids, distances) for Python, with int64 IDs and -1 for a missing neighbor. + [[nodiscard]] py::tuple to_python() const { + return py::make_tuple(python_ids(ids.data(), shape), distances); + } }; -// (num_queries, k) outputs. Slots without a neighbor keep kPidMax and +inf. -// Searches pad slots without a result, so k may exceed the number of points. inline void check_k(size_t k) { if (k == 0 || k > buffer::kSearchBufferMaxPointCount) { throw std::invalid_argument("k must be between 1 and 2**31"); @@ -369,10 +386,8 @@ inline void check_k(size_t k) { inline SearchResults unfilled_search_results(size_t num_queries, size_t k) { const auto shape = std::vector{ static_cast(num_queries), static_cast(k)}; - SearchResults results{py::array_t(shape), py::array_t(shape)}; - std::fill( - results.ids.mutable_data(), results.ids.mutable_data() + results.ids.size(), kPidMax - ); + SearchResults results{ + shape, std::vector(num_queries * k, kPidMax), py::array_t(shape)}; std::fill( results.distances.mutable_data(), results.distances.mutable_data() + results.distances.size(), diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 18297530..e2830b13 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -177,7 +177,7 @@ class HnswIndex { num_clusters_ = num_clusters; } - py::array_t add( + py::array_t add( py::handle data, const py::object& cluster_ids, std::optional num_threads, @@ -203,7 +203,7 @@ class HnswIndex { } // Allocated before the index is touched. - auto ids = py::array_t(static_cast(rows)); + auto ids = py::array_t(static_cast(rows)); const auto* data_ptr = data_array.data(); const size_t threads = num_threads.value_or(0); if (trains) { @@ -287,7 +287,6 @@ class HnswIndex { const size_t ef_search = ef.value_or(0) == 0 ? std::max(k, 64) : *ef; auto output = unfilled_search_results(static_cast(query_array.shape(0)), k); - auto result_ids_view = output.ids.mutable_unchecked<2>(); auto distances_view = output.distances.mutable_unchecked<2>(); const auto* queries_data = query_array.data(); @@ -305,13 +304,13 @@ class HnswIndex { k, results[static_cast(i)].size() )); ++j) { - result_ids_view(i, j) = + output.ids[(static_cast(i) * k) + static_cast(j)] = results[static_cast(i)][static_cast(j)].second; distances_view(i, j) = results[static_cast(i)][static_cast(j)].first; } } - return py::make_tuple(output.ids, output.distances); + return output.to_python(); } void save(py::handle path) const { @@ -337,8 +336,15 @@ class HnswIndex { wrapper.num_clusters_ = wrapper.index_->is_trained() ? wrapper.index_->num_clusters() : kDefaultNumClusters; + // Format-2 files record the seed; files from releases up to 0.8 do not. + const std::optional seed = wrapper.index_->seed(); + if (seed.has_value()) { + wrapper.seed_ = static_cast(*seed); + } wrapper.cp_ = default_training_parameters( - wrapper.index_->metric_type(), rabitqkmeans::RaBitQKMeansParameters{}.seed + wrapper.index_->metric_type(), + seed.has_value() ? static_cast(*seed) + : rabitqkmeans::RaBitQKMeansParameters{}.seed ); return wrapper; } @@ -447,7 +453,7 @@ class HnswIndex { std::unique_ptr index_; rabitqkmeans::RaBitQKMeansParameters cp_; size_t num_clusters_ = 0; // Centroid count train and set_centroids expect. - std::optional seed_; // Unknown for loaded indexes. + std::optional seed_; // Unknown for files from releases up to 0.8. }; } // namespace rabitqlib::python_bindings @@ -526,7 +532,7 @@ fast_quantization : bool, optional Returns ------- -ids : ndarray of uint32, shape (n,) +ids : ndarray of int64, shape (n,) Assigned point IDs. Notes @@ -652,7 +658,7 @@ num_threads : int, optional Returns ------- -ids : ndarray of uint32, shape (n_queries, k) +ids : ndarray of int64, shape (n_queries, k) Point IDs of the selected neighbors. distances : ndarray of float32, shape (n_queries, k) Lower values are better; squared L2 estimates for l2 and inner-product @@ -661,7 +667,7 @@ distances : ndarray of float32, shape (n_queries, k) Notes ----- Requires a trained or loaded index. Returns (ids, distances), in that -order. Missing neighbors have ID 4294967295 and distance inf; test +order. Missing neighbors have ID -1 and distance inf; test np.isfinite(distances) to find them. Array shapes and value ranges are checked at runtime, not by static type checkers.)doc" ) @@ -722,7 +728,8 @@ path : str, bytes, bytearray, or os.PathLike Notes ----- A trained index without vectors raises RuntimeError; add vectors first. -The saved state includes the rotation used to build the index.)doc" +The saved state includes the rotation used to build the index and its seed; +releases up to 0.8 cannot read the file.)doc" ) .def_static( "load", @@ -793,7 +800,7 @@ index : HnswIndex .def_property_readonly( "seed", &HnswIndex::seed, - R"doc(Graph and rotation seed, or None for an index loaded from disk.)doc" + R"doc(Graph and rotation seed, or None for an index loaded from a file from a release up to 0.8.)doc" ) .def_property( "cp", diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 779eae5d..3f384ee7 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -203,7 +203,7 @@ class IvfIndex { } } - py::array_t add( + py::array_t add( py::handle data, const py::object& cluster_ids, std::optional num_threads, @@ -237,10 +237,10 @@ class IvfIndex { const size_t first = index_->num_stored_points(); // Allocate the returned IDs before committing any changes to the index. - auto ids = py::array_t(static_cast(rows)); + auto ids = py::array_t(static_cast(rows)); auto* ids_data = ids.mutable_data(); for (size_t i = 0; i < rows; ++i) { - ids_data[i] = static_cast(first + i); + ids_data[i] = static_cast(first + i); } const auto* data_ptr = data_array.data(); @@ -322,7 +322,7 @@ class IvfIndex { const size_t num_queries = static_cast(query_array.shape(0)); auto output = unfilled_search_results(num_queries, k); - auto* ids_data = output.ids.mutable_data(); + auto* ids_data = output.ids.data(); auto* distances_data = output.distances.mutable_data(); const auto* queries_data = query_array.data(); @@ -340,7 +340,7 @@ class IvfIndex { ); } - return py::make_tuple(output.ids, output.distances); + return output.to_python(); } void save(py::handle path) const { @@ -362,8 +362,12 @@ class IvfIndex { py::gil_scoped_release release; wrapper.index_->load(filename.c_str()); } - wrapper.cp_ = untrainable_parameters(wrapper.index_->metric_type()); wrapper.builds_by_construct_ = true; + wrapper.seed_ = wrapper.index_->seed(); // recorded by format-2 files + wrapper.cp_ = default_training_parameters( + wrapper.index_->metric_type(), + wrapper.seed_.value_or(rabitqkmeans::RaBitQKMeansParameters{}.seed) + ); return wrapper; } @@ -461,7 +465,7 @@ class IvfIndex { mutable IndexAccess access_; std::unique_ptr index_; rabitqkmeans::RaBitQKMeansParameters cp_; - std::optional seed_; // Unknown for point-count and loaded indexes. + std::optional seed_; // Unknown for point-count indexes and 0.8 files. // Point-count and loaded indexes build by construct, never by training. bool builds_by_construct_ = false; }; @@ -535,7 +539,7 @@ fast_quantization : bool, optional Returns ------- -ids : ndarray of uint32, shape (n,) +ids : ndarray of int64, shape (n,) Assigned point IDs. Notes @@ -719,7 +723,7 @@ num_threads : int, optional Returns ------- -ids : ndarray of uint32, shape (n_queries, k) +ids : ndarray of int64, shape (n_queries, k) Point IDs of the selected neighbors. distances : ndarray of float32, shape (n_queries, k) Lower values are better; squared L2 estimates for l2 and inner-product @@ -728,7 +732,7 @@ distances : ndarray of float32, shape (n_queries, k) Notes ----- Requires a trained or loaded index. Returns (ids, distances), in that -order. Missing neighbors have ID 4294967295 and distance inf; test +order. Missing neighbors have ID -1 and distance inf; test np.isfinite(distances) to find them. Array shapes and value ranges are checked at runtime, not by static type checkers.)doc" ) @@ -747,7 +751,7 @@ path : str, bytes, bytearray, or os.PathLike Notes ----- A trained index without vectors raises RuntimeError; add vectors first. -IVF with HNSW centroid routing also writes its required sidecar file.)doc" +The one file includes HNSW centroid routing; releases up to 0.8 cannot read it.)doc" ) .def_static( "load", @@ -764,7 +768,12 @@ path : str, bytes, bytearray, or os.PathLike Returns ------- index : IvfIndex - A trained index ready for search.)doc" + A trained index ready for search. + +Notes +----- +Files from releases up to 0.8 with HNSW centroid routing also need the +accompanying ``.hnsw`` file.)doc" ) .def( "__repr__", @@ -808,7 +817,7 @@ index : IvfIndex .def_property_readonly( "seed", &IvfIndex::seed, - R"doc(Rotation seed, or None for an index loaded from disk or created in the 0.7 form.)doc" + R"doc(Rotation seed, or None for an index loaded from a file from a release up to 0.8 or created in the 0.7 form.)doc" ) .def_property( "cp", diff --git a/python_bindings/quant.py b/python_bindings/quant.py index c31cac2b..f11ff978 100644 --- a/python_bindings/quant.py +++ b/python_bindings/quant.py @@ -50,8 +50,9 @@ FASTSCAN_BLOCKS_VERSION = 1 """Version of the pack_fastscan_blocks layout.""" -# Classes report this module, so pickle and help() find them here. -for _public in (FormatError, Query, Rotator): +# Classes report this module, so pickle and help() find them here. FormatError is +# rabitqlib.FormatError, shared with index loading. +for _public in (Query, Rotator): _public.__module__ = __name__ del _public diff --git a/python_bindings/quant.pyi b/python_bindings/quant.pyi index 9cdbf2f6..b50851b6 100644 --- a/python_bindings/quant.pyi +++ b/python_bindings/quant.pyi @@ -8,6 +8,7 @@ from typing import Any, Literal, overload import numpy as np import numpy.typing as npt +from rabitqlib import FormatError as FormatError __all__ = [ "ESTIMATE_DTYPE", @@ -37,9 +38,6 @@ PACKED_CODES_VERSION: int FACTORS_VERSION: int FASTSCAN_BLOCKS_VERSION: int -class FormatError(ValueError): - """Corrupt or incompatible serialized bytes. A subclass of ValueError.""" - class Rotator: """A seeded rotation from the original dimension to the padded working space.""" diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp index a713042d..a364262c 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -50,11 +50,14 @@ namespace py = pybind11; namespace rabitqlib::python_bindings { namespace { -// Raised for corrupt or incompatible serialized bytes; a ValueError in Python. +// Raised for corrupt or incompatible serialized bytes; rabitqlib.FormatError in Python. struct FormatError : std::invalid_argument { using std::invalid_argument::invalid_argument; }; +// rabitqlib.FormatError, owned by the parent module. +PyObject* quant_format_error_type = nullptr; + // The blob format and both rotation methods share the FHT/Kac limit. constexpr size_t kMaxRotatorDim = rotator_impl::FhtKacRotator::kMaxDim; constexpr size_t kMaxTotalBits = 9; @@ -1823,11 +1826,20 @@ void register_quant(py::module_& m) { "_quant", "Native backend of rabitqlib.quant. Import rabitqlib.quant instead of this module." ); - py::register_exception( - quant, "FormatError", py::module_::import("builtins").attr("ValueError") - ); - quant.attr("FormatError").attr("__doc__") = - "Corrupt or incompatible serialized bytes. A subclass of ValueError."; + // Rejected bytes raise rabitqlib.FormatError, which the parent module created. + quant.attr("FormatError") = m.attr("FormatError"); + quant_format_error_type = m.attr("FormatError").ptr(); + // pybind11 translators take std::exception_ptr by value. + // NOLINTNEXTLINE(performance-unnecessary-value-param) + py::register_exception_translator([](std::exception_ptr exception) { + try { + if (exception) { + std::rethrow_exception(exception); + } + } catch (const FormatError& error) { + PyErr_SetString(quant_format_error_type, error.what()); + } + }); quant.attr("FACTORS_DTYPE") = factors_dtype(); quant.attr("ESTIMATE_DTYPE") = estimate_dtype(); diff --git a/python_bindings/rabitq_bindings.cpp b/python_bindings/rabitq_bindings.cpp index c609ca06..bb4984fb 100644 --- a/python_bindings/rabitq_bindings.cpp +++ b/python_bindings/rabitq_bindings.cpp @@ -3,10 +3,20 @@ #include "rabitqlib/core/metric.hpp" #include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/index/format_error.hpp" #include "rabitqlib/utils/tools.hpp" namespace py = pybind11; +namespace { +// rabitqlib.FormatError; the module keeps it alive for the interpreter's lifetime. +PyObject* format_error_type = nullptr; + +constexpr const char* kFormatErrorDoc = + "A damaged or mismatched index file or serialized byte string, or one this release " + "cannot read. A subclass of both ValueError and RuntimeError."; +} // namespace + // Forward declarations of registration functions implemented in other cpp files void register_hnsw(py::module_& m); void register_ivf(py::module_& m); @@ -18,6 +28,33 @@ PYBIND11_MODULE(_rabitqlib, m) { m.doc() = "RabitQ Python bindings combined module"; m.attr("__version__") = RABITQ_VERSION; + + // One exception for every rejected file or byte string, registered before the + // modules that raise it. + { + const py::object builtins = py::module_::import("builtins"); + const py::tuple bases = + py::make_tuple(builtins.attr("ValueError"), builtins.attr("RuntimeError")); + PyObject* type = PyErr_NewExceptionWithDoc( + "rabitqlib.FormatError", kFormatErrorDoc, bases.ptr(), nullptr + ); + if (type == nullptr) { + throw py::error_already_set(); + } + m.attr("FormatError") = py::reinterpret_steal(type); + format_error_type = type; + } + // pybind11 translators take std::exception_ptr by value. + // NOLINTNEXTLINE(performance-unnecessary-value-param) + py::register_exception_translator([](std::exception_ptr exception) { + try { + if (exception) { + std::rethrow_exception(exception); + } + } catch (const rabitqlib::FormatError& error) { + PyErr_SetString(format_error_type, error.what()); + } + }); m.def("_available_cpu_count", &rabitqlib::total_threads); m.def("_default_thread_count", &rabitqlib::default_threads); diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index 3dc00723..0430494a 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -77,7 +77,7 @@ class SymqgIndex { check_coordinates(x_ptr, count, dim_, "x"); } - py::array_t add(py::handle data, std::optional num_threads) { + py::array_t add(py::handle data, std::optional num_threads) { const auto operation = access_.write(); if (built()) { throw std::runtime_error( @@ -89,8 +89,8 @@ class SymqgIndex { build_graph(data, *ef_construction_, num_threads.value_or(0), *init_, *seed_); // Point IDs are the row numbers of data. const size_t rows = index_->num_points(); - auto ids = py::array_t(static_cast(rows)); - std::iota(ids.mutable_data(), ids.mutable_data() + rows, PID{0}); + auto ids = py::array_t(static_cast(rows)); + std::iota(ids.mutable_data(), ids.mutable_data() + rows, int64_t{0}); return ids; } @@ -124,12 +124,9 @@ class SymqgIndex { ensure_dimension(query_array, dim_, "query dimension does not match index dim"); const size_t num_queries = static_cast(query_array.shape(0)); - const auto shape = std::vector{ - static_cast(num_queries), static_cast(k)}; - auto result_ids = py::array_t(shape); - auto distances = py::array_t(shape); - auto* ids_data = result_ids.mutable_data(); - auto* distances_data = distances.mutable_data(); + auto output = unfilled_search_results(num_queries, k); + auto* ids_data = output.ids.data(); + auto* distances_data = output.distances.mutable_data(); const auto* queries_data = query_array.data(); { py::gil_scoped_release release; @@ -144,7 +141,7 @@ class SymqgIndex { ); } - return py::make_tuple(result_ids, distances); + return output.to_python(); } void save(py::handle path) const { @@ -167,9 +164,14 @@ class SymqgIndex { wrapper.max_degree_ = wrapper.index_->degree_bound(); wrapper.metric_type_ = wrapper.index_->metric_type(); wrapper.quantization_bits_ = wrapper.index_->quantization_bits(); - // Files store neither the seed nor the construction settings. - wrapper.ef_construction_.reset(); + // Format-2 files record the seed and the construction settings; earlier files + // record neither. + wrapper.seed_ = wrapper.index_->seed(); + wrapper.ef_construction_ = wrapper.index_->ef_construction(); wrapper.init_.reset(); + if (const auto init = wrapper.index_->initialization(); init.has_value()) { + wrapper.init_ = *init == symqg::QGInitialization::PiPNN ? "pipnn" : "random"; + } return wrapper; } @@ -314,7 +316,7 @@ class SymqgIndex { // Settings of the current graph, or of the next add; unknown after load. std::optional ef_construction_ = kDefaultEfConstruction; std::optional init_ = "pipnn"; - std::optional seed_; // Unknown for loaded indexes. + std::optional seed_; // Unknown for files from releases up to 0.8. bool seed_given_ = false; // The user chose the seed. mutable IndexAccess access_; std::unique_ptr> index_; @@ -394,7 +396,7 @@ num_threads : int, optional Returns ------- -ids : ndarray of uint32, shape (n_queries, k) +ids : ndarray of int64, shape (n_queries, k) Point IDs of the selected neighbors. distances : ndarray of float32, shape (n_queries, k) Lower values are better; squared L2 estimates for l2 and inner-product @@ -403,7 +405,7 @@ distances : ndarray of float32, shape (n_queries, k) Notes ----- Returns (ids, distances), in that order. When the graph reaches fewer than k -points, the remaining neighbors have ID 4294967295 and distance inf; test +points, the remaining neighbors have ID -1 and distance inf; test np.isfinite(distances) to find them. Array shapes and value ranges are checked at runtime, not by static type checkers.)doc"; @@ -469,7 +471,7 @@ num_threads : int, optional Returns ------- -ids : ndarray of uint32, shape (n,) +ids : ndarray of int64, shape (n,) Point IDs, the row numbers of data. Notes @@ -510,7 +512,8 @@ path : str, bytes, bytearray, or os.PathLike Notes ----- -The saved state includes the rotation used to build the index.)doc" +The saved state includes the rotation used to build the index and its seed; +releases up to 0.8 cannot read the file.)doc" ) .def_static( "load", @@ -548,12 +551,12 @@ index : SymqgIndex .def_property_readonly( "ef_construction", &SymqgIndex::ef_construction, - R"doc(Construction search breadth of the current graph or the next add, or None for an index loaded from disk.)doc" + R"doc(Construction search breadth of the current graph or the next add, or None for an index loaded from a file from a release up to 0.8.)doc" ) .def_property_readonly( "init", &SymqgIndex::init, - R"doc(Graph initialization method of the current graph or the next add, or None for an index loaded from disk.)doc" + R"doc(Graph initialization method of the current graph or the next add, or None for an index loaded from a file from a release up to 0.8.)doc" ) .def_property_readonly( "num_points", @@ -585,6 +588,6 @@ index : SymqgIndex .def_property_readonly( "seed", &SymqgIndex::seed, - R"doc(Rotation and graph-construction seed, or None for an index loaded from disk.)doc" + R"doc(Rotation and graph-construction seed, or None for an index loaded from a file from a release up to 0.8.)doc" ); } diff --git a/src/core/blob.cpp b/src/core/blob.cpp new file mode 100644 index 00000000..3d8132aa --- /dev/null +++ b/src/core/blob.cpp @@ -0,0 +1,232 @@ +#include "blob.hpp" + +#include +#include +#include +#include + +#include "crc32.hpp" + +namespace rabitqlib::detail { +namespace { + +constexpr std::array kBlobMagic = {'R', 'A', 'B', 'Q', 'B', 'L', 'O', 'B'}; +constexpr size_t kCommonHeaderBytes = 40; +constexpr size_t kChecksumOffset = 32; +constexpr uint32_t kRotatorKind = 1; +constexpr uint32_t kRotatorVersion = 1; +constexpr uint32_t kSeedKnown = 1; +constexpr uint32_t kMethodFhtKac = 1; +constexpr uint32_t kMethodMatrix = 2; +constexpr uint32_t kScalarFloat32 = 1; +constexpr uint32_t kScalarFloat64 = 2; +constexpr size_t kMaxRotatorDim = rotator_impl::FhtKacRotator::kMaxDim; + +constexpr std::array make_crc_table() { + std::array table{}; + for (uint32_t i = 0; i < 256; ++i) { + uint32_t value = i; + for (int bit = 0; bit < 8; ++bit) { + value = (value & 1U) != 0 ? 0xEDB88320U ^ (value >> 1) : value >> 1; + } + table[i] = value; + } + return table; +} + +constexpr std::array kCrcTable = make_crc_table(); + +void put_u32(uint8_t* out, uint32_t value) { + for (size_t i = 0; i < 4; ++i) { + out[i] = static_cast(value >> (8 * i)); + } +} + +void put_u64(uint8_t* out, uint64_t value) { + for (size_t i = 0; i < 8; ++i) { + out[i] = static_cast(value >> (8 * i)); + } +} + +uint32_t get_u32(const uint8_t* in) { + uint32_t value = 0; + for (size_t i = 0; i < 4; ++i) { + value |= static_cast(in[i]) << (8 * i); + } + return value; +} + +uint64_t get_u64(const uint8_t* in) { + uint64_t value = 0; + for (size_t i = 0; i < 8; ++i) { + value |= static_cast(in[i]) << (8 * i); + } + return value; +} + +// The blob checksum: the whole blob with its crc32 field taken as zero. +uint32_t blob_checksum(const uint8_t* blob, size_t size) { + constexpr std::array kZero{}; + uint32_t crc = crc32(blob, kChecksumOffset); + crc = crc32(kZero.data(), kZero.size(), crc); + return crc32(blob + kChecksumOffset + 4, size - kChecksumOffset - 4, crc); +} + +// The same rules and messages as rabitqlib.quant's Rotator. +void check_dimensions(RotatorType type, size_t dim, size_t padded_dim) { + if (dim == 0 || dim > kMaxRotatorDim) { + throw std::invalid_argument( + "dim must be in [1, " + std::to_string(kMaxRotatorDim) + "]; got " + + std::to_string(dim) + ); + } + if (type == RotatorType::FhtKacRotator && dim < 64) { + throw std::invalid_argument( + "dim must be at least 64 for method 'fht_kac'; got " + std::to_string(dim) + ); + } + if (padded_dim % 32 != 0 || padded_dim < dim || padded_dim > kMaxRotatorDim) { + throw std::invalid_argument( + "padded_dim must be a multiple of 32 in [dim, " + + std::to_string(kMaxRotatorDim) + "]; got " + std::to_string(padded_dim) + ); + } +} + +size_t body_bytes(RotatorType type, size_t dim, size_t padded_dim) { + return type == RotatorType::FhtKacRotator ? padded_dim / 2 + : sizeof(float) * dim * padded_dim; +} + +} // namespace + +uint32_t crc32(const void* data, size_t size, uint32_t crc) { + const auto* bytes = static_cast(data); + crc = ~crc; + for (size_t i = 0; i < size; ++i) { + crc = kCrcTable[(crc ^ bytes[i]) & 0xFFU] ^ (crc >> 8); + } + return ~crc; +} + +size_t rotator_blob_bytes(RotatorType type, size_t dim, size_t padded_dim) { + return kRotatorBlobHeaderBytes + body_bytes(type, dim, padded_dim); +} + +void encode_rotator_blob( + const Rotator& rotator, + RotatorType type, + size_t dim, + std::optional seed, + uint8_t* out +) { + const size_t padded_dim = rotator.size(); + check_dimensions(type, dim, padded_dim); + const size_t body = body_bytes(type, dim, padded_dim); + if (rotator.dump_bytes() != body) { + throw std::invalid_argument("rotator does not match its type and dimensions"); + } + const size_t size = kRotatorBlobHeaderBytes + body; + std::memset(out, 0, kRotatorBlobHeaderBytes); + std::memcpy(out, kBlobMagic.data(), kBlobMagic.size()); + put_u32(out + 8, kRotatorKind); + put_u32(out + 12, kRotatorVersion); + put_u32(out + 16, static_cast(kRotatorBlobHeaderBytes)); + put_u32(out + 20, seed.has_value() ? kSeedKnown : 0); + put_u64(out + 24, body); + put_u32(out + 40, type == RotatorType::FhtKacRotator ? kMethodFhtKac : kMethodMatrix); + put_u32(out + 44, kScalarFloat32); + put_u32(out + 48, static_cast(dim)); + put_u32(out + 52, static_cast(padded_dim)); + put_u32(out + 56, seed.value_or(0)); + rotator.save(reinterpret_cast(out + kRotatorBlobHeaderBytes)); + put_u32(out + kChecksumOffset, blob_checksum(out, size)); +} + +DecodedRotator decode_rotator_blob(const uint8_t* data, size_t size) { + if (size < kCommonHeaderBytes) { + throw std::invalid_argument("data is too short for a rabitqlib blob"); + } + if (std::memcmp(data, kBlobMagic.data(), kBlobMagic.size()) != 0) { + throw std::invalid_argument("data is not a rabitqlib blob (bad magic)"); + } + const uint32_t kind = get_u32(data + 8); + if (kind != kRotatorKind) { + throw std::invalid_argument( + "data is not a rotator blob; got kind " + std::to_string(kind) + ); + } + const uint32_t version = get_u32(data + 12); + if (version != kRotatorVersion) { + throw std::invalid_argument( + "data has unsupported rotator format version " + std::to_string(version) + + "; this release reads version " + std::to_string(kRotatorVersion) + ); + } + if (get_u32(data + 16) != kRotatorBlobHeaderBytes) { + throw std::invalid_argument( + "data has an invalid header size for rotator version 1" + ); + } + const uint64_t body = get_u64(data + 24); + if (size < kRotatorBlobHeaderBytes || body != size - kRotatorBlobHeaderBytes) { + throw std::invalid_argument( + "data length does not match its header (truncated or padded)" + ); + } + if (get_u32(data + kChecksumOffset) != blob_checksum(data, size)) { + throw std::invalid_argument("data checksum mismatch (corrupt blob)"); + } + const uint32_t flags = get_u32(data + 20); + if ((flags & ~kSeedKnown) != 0) { + throw std::invalid_argument("data has unknown flags"); + } + if (get_u32(data + 36) != 0 || get_u32(data + 60) != 0) { + throw std::invalid_argument("data has non-zero reserved fields"); + } + const uint32_t method = get_u32(data + 40); + if (method != kMethodFhtKac && method != kMethodMatrix) { + throw std::invalid_argument( + "data has an unknown rotator method " + std::to_string(method) + ); + } + const uint32_t scalar = get_u32(data + 44); + if (scalar == kScalarFloat64) { + throw std::invalid_argument( + "data stores a float64 rotator, which version 1 cannot load" + ); + } + if (scalar != kScalarFloat32) { + throw std::invalid_argument( + "data has an unknown scalar type " + std::to_string(scalar) + ); + } + + DecodedRotator result; + result.type = + method == kMethodFhtKac ? RotatorType::FhtKacRotator : RotatorType::MatrixRotator; + result.dim = get_u32(data + 48); + result.padded_dim = get_u32(data + 52); + try { + check_dimensions(result.type, result.dim, result.padded_dim); + } catch (const std::invalid_argument& error) { + throw std::invalid_argument( + std::string("data has invalid rotator dimensions: ") + error.what() + ); + } + const uint32_t seed = get_u32(data + 56); + if ((flags & kSeedKnown) != 0) { + result.seed = seed; + } else if (seed != 0) { + throw std::invalid_argument("data has a seed without the seed_known flag"); + } + if (body != body_bytes(result.type, result.dim, result.padded_dim)) { + throw std::invalid_argument("data body size does not match the rotator parameters"); + } + result.rotator = + rotator_impl::make_rotator_for_loading(result.dim, result.type, result.padded_dim); + result.rotator->load(reinterpret_cast(data + kRotatorBlobHeaderBytes)); + return result; +} + +} // namespace rabitqlib::detail diff --git a/src/core/blob.hpp b/src/core/blob.hpp new file mode 100644 index 00000000..c4c514d2 --- /dev/null +++ b/src/core/blob.hpp @@ -0,0 +1,45 @@ +#pragma once + +#include +#include +#include +#include + +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/core/types.hpp" + +// The self-describing blobs of docs/formats.md ("Blobs"), for C++ readers and writers +// such as index files; an implementation detail, not installed. rabitqlib.quant keeps +// its own rotator blob code, which uses zlib's faster checksum; both are tested against +// the same golden files. +namespace rabitqlib::detail { + +inline constexpr size_t kRotatorBlobHeaderBytes = 64; + +// The size of the rotator blob for these parameters, which must be valid. +size_t rotator_blob_bytes(RotatorType type, size_t dim, size_t padded_dim); + +// Write the blob of `rotator`, which has the given type and dimensions, to `out`, +// which holds rotator_blob_bytes(type, dim, rotator.size()) bytes. Dimensions the +// format does not allow raise std::invalid_argument. +void encode_rotator_blob( + const Rotator& rotator, + RotatorType type, + size_t dim, + std::optional seed, + uint8_t* out +); + +struct DecodedRotator { + RotatorType type = RotatorType::FhtKacRotator; + size_t dim = 0; + size_t padded_dim = 0; + std::optional seed; + std::unique_ptr> rotator; +}; + +// Validate and decode a rotator blob of `size` bytes. Invalid data raises +// std::invalid_argument with the messages rabitqlib.quant uses, which start with "data". +DecodedRotator decode_rotator_blob(const uint8_t* data, size_t size); + +} // namespace rabitqlib::detail diff --git a/src/core/crc32.hpp b/src/core/crc32.hpp new file mode 100644 index 00000000..69ad4975 --- /dev/null +++ b/src/core/crc32.hpp @@ -0,0 +1,12 @@ +#pragma once + +#include +#include + +namespace rabitqlib::detail { + +// CRC-32 as zlib.crc32 computes it; pass an earlier result as `crc` to continue it. +// Byte-wise: meant for headers and metadata, not bulk data. +uint32_t crc32(const void* data, size_t size, uint32_t crc = 0); + +} // namespace rabitqlib::detail diff --git a/src/index/format2.cpp b/src/index/format2.cpp new file mode 100644 index 00000000..c66e0159 --- /dev/null +++ b/src/index/format2.cpp @@ -0,0 +1,568 @@ +#include "format2.hpp" + +#include +#include +#include +#include +#include +#include +#include + +#include "../core/crc32.hpp" + +namespace rabitqlib::index_detail::format2 { +namespace { + +using rabitqlib::detail::crc32; + +constexpr uint32_t kHeaderFlagSeedKnown = 1U << 0; +constexpr uint32_t kKnownSectionFlags = kSectionRequired | kSectionChecksummed; +constexpr size_t kChecksumOffset = 108; +constexpr size_t kChunkBytes = size_t{1} << 16; + +// Round up to a multiple of kSectionAlignment; the caller has bounded `value`. +uint64_t align_section(uint64_t value) { + return (value + kSectionAlignment - 1) / kSectionAlignment * kSectionAlignment; +} + +std::string hex(uint64_t value) { + std::ostringstream out; + out << "0x" << std::hex << value; + return out.str(); +} + +// The error-message component of each index, as in docs/development.md. +const char* component(IndexKind kind) { + switch (kind) { + case IndexKind::Ivf: + return "IVF"; + case IndexKind::Hnsw: + return "HNSW"; + case IndexKind::SymphonyQg: + return "QuantizedGraph"; + } + return "Index"; +} + +const char* kind_name(uint32_t kind) { + switch (kind) { + case 1: + return "an IVF-RaBitQ"; + case 2: + return "an HNSW-RaBitQ"; + case 3: + return "a SymphonyQG"; + case 4: + return "an IVF-SSD"; + default: + return nullptr; + } +} + +// The header rules that hold for every index kind, as the text after "index file", +// or an empty string. The index kind itself is checked by the caller. +std::string header_error(const Header& header) { + const auto metric = static_cast(header.metric); + if (metric < 1 || metric > 3) { + return "has unknown metric " + std::to_string(metric); + } + if (header.dim == 0) { + return "has dimension 0"; + } + if (header.padded_dim % 32 != 0 || header.padded_dim < header.dim || + header.padded_dim > kMaxPaddedDim) { + return "has padded dimension " + std::to_string(header.padded_dim) + + ", which must be a multiple of 32 from the dimension " + + std::to_string(header.dim) + " to 65536"; + } + if (header.code_bits > 9) { + return "has " + std::to_string(header.code_bits) + + " code bits; at most 9 are supported"; + } + const auto raw_type = static_cast(header.raw_type); + if (raw_type > 2) { + return "has unknown raw element type " + std::to_string(raw_type); + } + if (header.id_bits != 32 && header.id_bits != 64) { + return "has " + std::to_string(header.id_bits) + + "-bit IDs; only 32 and 64 are valid"; + } + return {}; +} + +void encode_header( + const Header& header, + uint64_t section_count, + std::string_view writer_version, + uint8_t* out +) { + std::memset(out, 0, kHeaderBytes); + std::memcpy(out, kMagic.data(), kMagic.size()); + put_u32(out + 8, kVersion); + put_u32(out + 12, kHeaderBytes); + put_u64(out + 16, header.incompat_features); + put_u64(out + 24, header.ro_compat_features); + put_u64(out + 32, header.compat_features); + put_u32(out + 40, static_cast(header.index_kind)); + put_u32(out + 44, static_cast(header.metric)); + put_u64(out + 48, header.dim); + put_u64(out + 56, header.padded_dim); + put_u32(out + 64, header.code_bits); + put_u32(out + 68, static_cast(header.raw_type)); + put_u32(out + 72, header.id_bits); + put_u32(out + 76, header.seed.has_value() ? kHeaderFlagSeedKnown : 0); + put_u64(out + 80, header.seed.value_or(0)); + put_u64(out + 88, header.num_points); + put_u64(out + 96, kHeaderBytes); + put_u32(out + 104, static_cast(section_count)); + std::memcpy(out + 112, writer_version.data(), writer_version.size()); +} + +void encode_section(const Section& section, uint8_t* out) { + std::memset(out, 0, kSectionEntryBytes); + put_u32(out, section.kind); + put_u32(out + 4, section.flags); + put_u64(out + 8, section.offset); + put_u64(out + 16, section.size); + put_u32(out + 24, section.crc32); +} + +// The header checksum: the header with its checksum field taken as zero, then the table. +uint32_t header_checksum( + const std::vector& header_bytes, const std::vector& table_bytes +) { + constexpr std::array kZero{}; + uint32_t crc = crc32(header_bytes.data(), kChecksumOffset); + crc = crc32(kZero.data(), kZero.size(), crc); + crc = crc32( + header_bytes.data() + kChecksumOffset + 4, + header_bytes.size() - kChecksumOffset - 4, + crc + ); + return crc32(table_bytes.data(), table_bytes.size(), crc); +} + +void write_zeros(std::ostream& output, uint64_t count) { + static const std::array kZeros{}; + while (count > 0) { + const auto chunk = static_cast(std::min(count, kZeros.size())); + output.write(kZeros.data(), static_cast(chunk)); + count -= chunk; + } +} + +} // namespace + +void SectionSink::write(const void* data, size_t size) { + if (size > remaining_) { + throw std::invalid_argument("index section writer wrote more bytes than declared"); + } + output_.write(static_cast(data), static_cast(size)); + remaining_ -= size; +} + +void write_file( + std::ostream& output, + const Header& header, + const std::vector& sections, + std::string_view writer_version +) { + const std::string name = component(header.index_kind); + const auto kind = static_cast(header.index_kind); + if (kind < 1 || kind > 3) { + throw std::invalid_argument( + name + " index file cannot have index kind " + std::to_string(kind) + ); + } + if (const std::string error = header_error(header); !error.empty()) { + throw std::invalid_argument(name + " index file " + error); + } + if ((header.incompat_features & ~kKnownIncompatFeatures) != 0 || + (header.ro_compat_features & ~kKnownRoCompatFeatures) != 0 || + (header.compat_features & ~kKnownCompatFeatures) != 0) { + throw std::invalid_argument(name + " index file cannot use unknown feature bits"); + } + if (sections.empty()) { + throw std::invalid_argument(name + " index file needs at least one section"); + } + const bool printable = + std::all_of(writer_version.begin(), writer_version.end(), [](char c) { + return c >= 0x20 && c <= 0x7E; + }); + if (writer_version.size() > kWriterVersionBytes || !printable) { + throw std::invalid_argument( + name + + " index file writer version must be at most 16 printable ASCII characters" + ); + } + + std::vector
entries(sections.size()); + std::unordered_set kinds; + const uint64_t table_end = kHeaderBytes + (kSectionEntryBytes * sections.size()); + uint64_t position = table_end; + for (size_t i = 0; i < sections.size(); ++i) { + const SectionWrite& section = sections[i]; + if (section.kind == 0) { + throw std::invalid_argument(name + " index file section kind 0 is reserved"); + } + if ((section.flags & ~kKnownSectionFlags) != 0) { + throw std::invalid_argument(name + " index file section has unknown flags"); + } + if (!kinds.insert(section.kind).second) { + throw std::invalid_argument( + name + " index file section kind " + std::to_string(section.kind) + + " is repeated" + ); + } + position = align_section(position); + if (section.size > + std::numeric_limits::max() - kSectionAlignment - position) { + throw std::invalid_argument( + name + " index file sections exceed the offset range" + ); + } + entries[i].kind = section.kind; + entries[i].flags = section.flags; + entries[i].offset = position; + entries[i].size = section.size; + position += section.size; + } + + // Run a section's writer and check that it wrote exactly the declared size. + const auto write_section = [&name](const SectionWrite& section, SectionSink& sink) { + if (section.write) { + section.write(sink); + } + if (sink.remaining() != 0) { + throw std::invalid_argument( + name + " index file section kind " + std::to_string(section.kind) + + " writer wrote fewer bytes than declared" + ); + } + }; + + // Checksummed sections are small: buffer them, so that the header and the table, + // which hold their checksums, can be written first. + std::vector buffered(sections.size()); + for (size_t i = 0; i < sections.size(); ++i) { + if ((entries[i].flags & kSectionChecksummed) == 0) { + continue; + } + std::ostringstream buffer; + SectionSink sink(buffer, entries[i].size); + write_section(sections[i], sink); + buffered[i] = buffer.str(); + entries[i].crc32 = crc32(buffered[i].data(), buffered[i].size()); + } + + std::vector header_bytes(kHeaderBytes); + encode_header(header, sections.size(), writer_version, header_bytes.data()); + std::vector table_bytes(kSectionEntryBytes * sections.size()); + for (size_t i = 0; i < entries.size(); ++i) { + encode_section(entries[i], table_bytes.data() + (i * kSectionEntryBytes)); + } + put_u32( + header_bytes.data() + kChecksumOffset, header_checksum(header_bytes, table_bytes) + ); + + output.write( + reinterpret_cast(header_bytes.data()), + static_cast(header_bytes.size()) + ); + output.write( + reinterpret_cast(table_bytes.data()), + static_cast(table_bytes.size()) + ); + uint64_t written = table_end; + for (size_t i = 0; i < sections.size(); ++i) { + write_zeros(output, entries[i].offset - written); + if ((entries[i].flags & kSectionChecksummed) != 0) { + output.write( + buffered[i].data(), static_cast(buffered[i].size()) + ); + } else { + SectionSink sink(output, entries[i].size); + write_section(sections[i], sink); + } + written = entries[i].offset + entries[i].size; + } + output.flush(); + if (!output) { + throw std::runtime_error(name + " index file cannot be written"); + } +} + +bool has_magic(std::istream& input, IndexKind kind) { + std::streambuf& buffer = *input.rdbuf(); + const auto start = buffer.pubseekoff(0, std::ios::cur, std::ios::in); + if (start == std::streampos(std::streamoff(-1))) { + throw FormatError( + std::string(component(kind)) + " index file input must be seekable" + ); + } + std::array magic{}; + const std::streamsize count = buffer.sgetn(magic.data(), magic.size()); + buffer.pubseekpos(start, std::ios::in); + return count == static_cast(magic.size()) && magic == kMagic; +} + +Reader::Reader( + std::istream& input, + IndexKind expected_kind, + const std::vector& known_kinds, + const std::function& check_header +) + : input_(*input.rdbuf()) + , component_(std::string(component(expected_kind)) + " index file") { + const auto end = input_.pubseekoff(0, std::ios::end, std::ios::in); + if (end == std::streampos(std::streamoff(-1))) { + throw FormatError(component_ + " input must be seekable"); + } + const auto file_size = static_cast(static_cast(end)); + + // 1. Magic, version and header size. + if (file_size < kHeaderBytes) { + throw FormatError(component_ + " is too short for a format-2 header"); + } + std::vector header_bytes(kHeaderBytes); + read_at(0, header_bytes.data(), kHeaderBytes, "header"); + if (std::memcmp(header_bytes.data(), kMagic.data(), kMagic.size()) != 0) { + throw FormatError(component_ + " is not in format 2 (bad magic)"); + } + const uint32_t version = get_u32(header_bytes.data() + 8); + if (version != kVersion) { + throw FormatError( + component_ + " has format version " + std::to_string(version) + + "; this release reads version 2" + ); + } + const uint32_t header_size = get_u32(header_bytes.data() + 12); + if (header_size < kHeaderBytes || header_size % 8 != 0 || header_size > file_size) { + throw FormatError( + component_ + " has an invalid header size " + std::to_string(header_size) + ); + } + header_bytes.resize(header_size); + read_at( + kHeaderBytes, + header_bytes.data() + kHeaderBytes, + header_size - kHeaderBytes, + "header" + ); + + // 2. The section table lies within the file, and the header checksum matches. + const uint8_t* fields = header_bytes.data(); + const uint64_t table_offset = get_u64(fields + 96); + const uint32_t section_count = get_u32(fields + 104); + if (section_count == 0) { + throw FormatError(component_ + " has no sections"); + } + if (table_offset % 8 != 0 || table_offset < header_size || table_offset > file_size || + section_count > (file_size - table_offset) / kSectionEntryBytes) { + throw FormatError(component_ + " section table lies outside the file"); + } + const uint64_t table_end = table_offset + (kSectionEntryBytes * section_count); + std::vector table_bytes(kSectionEntryBytes * section_count); + read_at(table_offset, table_bytes.data(), table_bytes.size(), "section table"); + if (get_u32(fields + kChecksumOffset) != header_checksum(header_bytes, table_bytes)) { + throw FormatError(component_ + " header checksum mismatch (corrupt file)"); + } + + // 3. Feature flags. + const char* const writer = reinterpret_cast(fields + 112); + writer_version_.assign(writer, std::find(writer, writer + kWriterVersionBytes, '\0')); + header_.incompat_features = get_u64(fields + 16); + header_.ro_compat_features = get_u64(fields + 24); + header_.compat_features = get_u64(fields + 32); + const uint64_t unknown_incompat = header_.incompat_features & ~kKnownIncompatFeatures; + if (unknown_incompat != 0) { + throw FormatError( + component_ + " uses incompatible features " + hex(unknown_incompat) + + " that this release does not know; written by rabitqlib " + writer_version_ + ); + } + read_only_ = (header_.ro_compat_features & ~kKnownRoCompatFeatures) != 0; + + // 4. The index kind, the header rules of every kind, then the index's own. + const uint32_t kind = get_u32(fields + 40); + if (kind != static_cast(expected_kind)) { + const char* name = kind_name(kind); + if (name == nullptr) { + throw FormatError( + component_ + " has unknown index kind " + std::to_string(kind) + ); + } + throw FormatError(component_ + " holds " + name + " index"); + } + header_.index_kind = expected_kind; + header_.metric = static_cast(get_u32(fields + 44)); + header_.dim = get_u64(fields + 48); + header_.padded_dim = get_u64(fields + 56); + header_.code_bits = get_u32(fields + 64); + header_.raw_type = static_cast(get_u32(fields + 68)); + header_.id_bits = get_u32(fields + 72); + const uint32_t header_flags = get_u32(fields + 76); + const uint64_t seed = get_u64(fields + 80); + header_.num_points = get_u64(fields + 88); + if ((header_flags & ~kHeaderFlagSeedKnown) != 0) { + throw FormatError(component_ + " has unknown header flags " + hex(header_flags)); + } + if ((header_flags & kHeaderFlagSeedKnown) != 0) { + header_.seed = seed; + } else if (seed != 0) { + throw FormatError(component_ + " has a seed without the seed_known flag"); + } + if (const std::string error = header_error(header_); !error.empty()) { + throw FormatError(component_ + " " + error); + } + if (check_header) { + check_header(header_); + } + + // 5. Each entry in table order: its fields, alignment, bounds, order, overlap and + // repetition, then whether its kind is known. + std::unordered_set kinds; + uint64_t previous_end = 0; + uint64_t content_end = table_end; + for (uint32_t i = 0; i < section_count; ++i) { + const uint8_t* entry = table_bytes.data() + (i * kSectionEntryBytes); + Section section; + section.kind = get_u32(entry); + section.flags = get_u32(entry + 4); + section.offset = get_u64(entry + 8); + section.size = get_u64(entry + 16); + section.crc32 = get_u32(entry + 24); + const std::string name = + component_ + " section kind " + std::to_string(section.kind); + if (section.kind == 0) { + throw FormatError(name + " is reserved"); + } + if ((section.flags & ~kKnownSectionFlags) != 0) { + throw FormatError(name + " has unknown flags " + hex(section.flags)); + } + if (get_u32(entry + 28) != 0) { + throw FormatError(name + " has a non-zero reserved field"); + } + if ((section.flags & kSectionChecksummed) == 0 && section.crc32 != 0) { + throw FormatError(name + " has a checksum but is not checksummed"); + } + if (section.offset % kSectionAlignment != 0) { + throw FormatError(name + " is not aligned to 64 bytes"); + } + if (section.offset > file_size || section.size > file_size - section.offset) { + throw FormatError(name + " lies outside the file"); + } + const uint64_t end_offset = section.offset + section.size; + if (section.offset < previous_end) { + throw FormatError(name + " overlaps the previous section or is out of order"); + } + const bool overlaps_table = + section.size != 0 && section.offset < table_end && table_offset < end_offset; + if (section.offset < header_size || overlaps_table) { + throw FormatError(name + " overlaps the header or the section table"); + } + if (!kinds.insert(section.kind).second) { + throw FormatError(name + " is repeated"); + } + previous_end = end_offset; + content_end = std::max(content_end, end_offset); + + if (std::find(known_kinds.begin(), known_kinds.end(), section.kind) == + known_kinds.end()) { + if ((section.flags & kSectionRequired) != 0) { + throw FormatError( + name + + " is required but unknown to this release; written by rabitqlib " + + writer_version_ + ); + } + continue; + } + sections_.push_back(section); + } + + // 6. The file ends where the table or the last section ends. + if (content_end != file_size) { + throw FormatError( + component_ + " is " + std::to_string(file_size) + + " bytes, but its sections end at " + std::to_string(content_end) + ); + } + + // 7. Known checksummed sections, streamed without holding them in memory. + std::vector chunk; + for (const Section& section : sections_) { + if ((section.flags & kSectionChecksummed) == 0) { + continue; + } + chunk.resize(static_cast(std::min(section.size, kChunkBytes))); + uint32_t crc = 0; + for (uint64_t done = 0; done < section.size;) { + const auto count = + static_cast(std::min(section.size - done, kChunkBytes)); + read(section, done, chunk.data(), count); + crc = crc32(chunk.data(), count, crc); + done += count; + } + if (crc != section.crc32) { + throw FormatError( + component_ + " section kind " + std::to_string(section.kind) + + " checksum mismatch (corrupt file)" + ); + } + } +} + +const Section* Reader::find(uint32_t kind) const { + for (const Section& section : sections_) { + if (section.kind == kind) { + return §ion; + } + } + return nullptr; +} + +void Reader::read(const Section& section, uint64_t offset, void* out, size_t size) { + if (offset > section.size || size > section.size - offset) { + throw std::out_of_range(component_ + " read past the end of a section"); + } + read_at( + section.offset + offset, out, size, "section kind " + std::to_string(section.kind) + ); +} + +SectionInput Reader::open(const Section& section) { + const std::string name = component_ + " section kind " + std::to_string(section.kind); + const auto offset = static_cast(section.offset); + if (input_.pubseekpos(offset, std::ios::in) != std::streampos(offset)) { + throw FormatError(name + " cannot be read"); + } + return SectionInput(input_, name, section.size); +} + +void SectionInput::read(void* out, size_t size) { + if (size > remaining_) { + throw std::out_of_range(name_ + " read past its end"); + } + if (size != 0 && + input_.sgetn(static_cast(out), static_cast(size)) != + static_cast(size)) { + throw FormatError(name_ + " cannot be read"); + } + remaining_ -= size; +} + +void Reader::read_at(uint64_t offset, void* out, size_t size, const std::string& what) { + if (size == 0) { + return; + } + const auto position = + input_.pubseekpos(static_cast(offset), std::ios::in); + if (position != std::streampos(static_cast(offset)) || + input_.sgetn(static_cast(out), static_cast(size)) != + static_cast(size)) { + throw FormatError(component_ + " " + what + " cannot be read"); + } +} + +} // namespace rabitqlib::index_detail::format2 diff --git a/src/index/format2.hpp b/src/index/format2.hpp new file mode 100644 index 00000000..7bf2df90 --- /dev/null +++ b/src/index/format2.hpp @@ -0,0 +1,262 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/index/format_error.hpp" + +// The index container of format 2, specified in docs/formats.md ("Index files, +// format 2"). This layer reads and writes the header, the section table and section +// bytes; each index defines its own sections and validates their contents. +namespace rabitqlib::index_detail::format2 { + +// A file that is not a valid format-2 container, or uses something this release +// does not support: the library-wide index file error. +using FormatError = ::rabitqlib::FormatError; + +inline constexpr std::array kMagic = {'R', 'A', 'B', 'Q', 'I', 'N', 'D', 'X'}; +inline constexpr uint32_t kVersion = 2; +inline constexpr uint32_t kHeaderBytes = 128; +inline constexpr uint64_t kSectionEntryBytes = 32; +inline constexpr uint64_t kSectionAlignment = 64; +inline constexpr size_t kWriterVersionBytes = 16; +// The rotator blob's limit, which bounds every index's working dimension. +inline constexpr uint64_t kMaxPaddedDim = 65536; + +// Feature bits this release knows. Version 2 defines none. +inline constexpr uint64_t kKnownIncompatFeatures = 0; +inline constexpr uint64_t kKnownRoCompatFeatures = 0; +inline constexpr uint64_t kKnownCompatFeatures = 0; + +enum class IndexKind : uint32_t { Ivf = 1, Hnsw = 2, SymphonyQg = 3 }; +enum class Metric : uint32_t { L2 = 1, InnerProduct = 2, Cosine = 3 }; +enum class RawType : uint32_t { None = 0, Float32 = 1, Uint8 = 2 }; + +inline constexpr uint32_t kSectionRequired = 1U << 0; +inline constexpr uint32_t kSectionChecksummed = 1U << 1; + +// Little-endian fields of fixed width, for section contents. +inline void put_u32(uint8_t* out, uint32_t value) { + for (size_t i = 0; i < 4; ++i) { + out[i] = static_cast(value >> (8 * i)); + } +} + +inline void put_u64(uint8_t* out, uint64_t value) { + for (size_t i = 0; i < 8; ++i) { + out[i] = static_cast(value >> (8 * i)); + } +} + +[[nodiscard]] inline uint32_t get_u32(const uint8_t* in) { + uint32_t value = 0; + for (size_t i = 0; i < 4; ++i) { + value |= static_cast(in[i]) << (8 * i); + } + return value; +} + +[[nodiscard]] inline uint64_t get_u64(const uint8_t* in) { + uint64_t value = 0; + for (size_t i = 0; i < 8; ++i) { + value |= static_cast(in[i]) << (8 * i); + } + return value; +} + +// The section kind every index shares: its rotator blob. +inline constexpr uint32_t kRotatorSection = 1; + +struct Header { + uint64_t incompat_features = 0; + uint64_t ro_compat_features = 0; + uint64_t compat_features = 0; + IndexKind index_kind = IndexKind::Ivf; + Metric metric = Metric::L2; + uint64_t dim = 0; + uint64_t padded_dim = 0; + uint32_t code_bits = 0; + RawType raw_type = RawType::None; + uint32_t id_bits = 32; + std::optional seed; + uint64_t num_points = 0; +}; + +// A section table entry. +struct Section { + uint32_t kind = 0; + uint32_t flags = 0; + uint64_t offset = 0; + uint64_t size = 0; + uint32_t crc32 = 0; +}; + +class SectionSink; + +// One section to write: `write` must pass exactly `size` bytes to the sink. Writers +// mark only small sections checksummed: those are buffered in memory before writing. +struct SectionWrite { + uint32_t kind = 0; + uint32_t flags = 0; + uint64_t size = 0; + std::function write; +}; + +// Write a complete file from the current position of `output`, which must be the +// start of the file: the header, the section table right after it, then the sections +// in the given order, each at a multiple of 64. The output is written once, front to +// back, so it need not be seekable. Invalid arguments raise std::invalid_argument and +// write nothing; a section writer that passes the wrong number of bytes raises +// std::invalid_argument after part of the file is written. +void write_file( + std::ostream& output, + const Header& header, + const std::vector& sections, + std::string_view writer_version +); + +// Receives one section's bytes during write_file; it rejects more bytes than declared. +class SectionSink { + public: + void write(const void* data, size_t size); + [[nodiscard]] uint64_t remaining() const { return remaining_; } + + private: + friend void write_file( + std::ostream&, const Header&, const std::vector&, std::string_view + ); + SectionSink(std::ostream& output, uint64_t size) : output_(output), remaining_(size) {} + + std::ostream& output_; + uint64_t remaining_; +}; + +// An output stream that passes its bytes to a section sink, for section writers that +// take a std::ostream. A sink error propagates as an exception. +class SinkStream : public std::ostream { + public: + explicit SinkStream(SectionSink& sink) : std::ostream(nullptr), buffer_(sink) { + rdbuf(&buffer_); + exceptions(std::ios::badbit); + } + + private: + class Buffer : public std::streambuf { + public: + explicit Buffer(SectionSink& sink) : sink_(sink) {} + + private: + int_type overflow(int_type c) override { + if (!traits_type::eq_int_type(c, traits_type::eof())) { + const char value = traits_type::to_char_type(c); + sink_.write(&value, 1); + } + return traits_type::not_eof(c); + } + std::streamsize xsputn(const char* data, std::streamsize count) override { + sink_.write(data, static_cast(count)); + return count; + } + + SectionSink& sink_; + }; + + Buffer buffer_; +}; + +// An output stream that only counts its bytes, to size a section before writing it. +class CountingStream : public std::ostream { + public: + CountingStream() : std::ostream(nullptr) { rdbuf(&buffer_); } + [[nodiscard]] uint64_t count() const { return buffer_.count; } + + private: + class Buffer : public std::streambuf { + public: + uint64_t count = 0; + + private: + int_type overflow(int_type c) override { + count += traits_type::eq_int_type(c, traits_type::eof()) ? 0 : 1; + return traits_type::not_eof(c); + } + std::streamsize xsputn(const char* /*data*/, std::streamsize n) override { + count += static_cast(n); + return n; + } + }; + + Buffer buffer_; +}; + +// Whether the input starts with the format-2 magic. It reads through the stream +// buffer, so the stream's state and exception mask do not matter, and restores the +// position. Input that cannot seek raises FormatError, named for `kind`. +bool has_magic(std::istream& input, IndexKind kind); + +// Reads one section front to back in pieces, seeking only once. Other reads of the +// same input move its position, so finish one section before reading another. +class SectionInput { + public: + // Read the next `size` bytes of the section. + void read(void* out, size_t size); + + private: + friend class Reader; + SectionInput(std::streambuf& input, const std::string& name, uint64_t size) + : input_(input), name_(name), remaining_(size) {} + + std::streambuf& input_; + std::string name_; + uint64_t remaining_; +}; + +class Reader { + public: + // Read and validate the header and the section table, in the order of "Reading" + // in the specification: the generic header rules, then `expected_kind`, then + // `check_header` (the index's own header rules, which throw FormatError), then + // the sections. Sections whose kind is not in `known_kinds` are rejected when + // required and skipped otherwise. The reader reads through the stream buffer of + // `input`, which must outlive it, and ignores the stream's state. + Reader( + std::istream& input, + IndexKind expected_kind, + const std::vector& known_kinds, + const std::function& check_header = {} + ); + + [[nodiscard]] const Header& header() const { return header_; } + // True when the file uses read-only-compatible features this release does not know: + // it may be searched, but add, remove and save must refuse it. + [[nodiscard]] bool read_only() const { return read_only_; } + [[nodiscard]] const std::string& writer_version() const { return writer_version_; } + // The entry of a known section kind, or nullptr when the file has none. + [[nodiscard]] const Section* find(uint32_t kind) const; + // Read `size` bytes starting `offset` bytes into `section`. + void read(const Section& section, uint64_t offset, void* out, size_t size); + // Read `section` from its start in consecutive pieces; for large sections. + SectionInput open(const Section& section); + + private: + void read_at(uint64_t offset, void* out, size_t size, const std::string& what); + + std::streambuf& input_; + std::string component_; + Header header_; + bool read_only_ = false; + std::string writer_version_; + std::vector
sections_; +}; + +} // namespace rabitqlib::index_detail::format2 diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index 72dc4fa8..a1756bec 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -1,6 +1,7 @@ #include "rabitqlib/index/hnsw/hnsw.hpp" #include +#include #include #include #include @@ -14,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -22,7 +24,9 @@ #include #include +#include "../core/blob.hpp" #include "file_replacement.hpp" +#include "format2.hpp" #include "rabitqlib/clustering/rabitqkmeans.hpp" #include "rabitqlib/core/data_layout.hpp" #include "rabitqlib/core/encoder.hpp" @@ -35,6 +39,7 @@ #include "rabitqlib/index/detail/hnsw_dispatch.hpp" #include "rabitqlib/index/detail/parallel_for.hpp" #include "rabitqlib/index/detail/visited_pool.hpp" +#include "rabitqlib/index/format_error.hpp" #include "rabitqlib/index/hnsw/hnsw_quant.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/input_checks.hpp" @@ -45,6 +50,16 @@ namespace rabitqlib::hnsw { +namespace { +// Index format 2 sections of HNSW-RaBitQ (docs/formats.md, "Section kinds"). +constexpr uint32_t kParamsSection = 512; +constexpr uint32_t kCentroidsSection = 513; +constexpr uint32_t kRecordsSection = 514; +constexpr uint32_t kLevelsSection = 515; +constexpr uint32_t kUpperLinksSection = 516; +constexpr size_t kParamsBytes = 72; +} // namespace + HierarchicalNSW::HierarchicalNSW() = default; HierarchicalNSW::~HierarchicalNSW() = default; void HierarchicalNSW::AlignedDeleter::operator()(char* pointer) const noexcept { @@ -118,6 +133,7 @@ HierarchicalNSW::HierarchicalNSW( ); level_generator_.seed(random_seed); + seed_ = random_seed; visited_set_pool_ = std::make_unique(1, max_elements_); @@ -149,69 +165,107 @@ void HierarchicalNSW::save(const char* filename) const { if (filename == nullptr || filename[0] == '\0') { throw std::invalid_argument("HNSW index filename must not be empty"); } + check_writable("save"); index_detail::save_replacing(filename, {}, [this](const std::string& path) { write_file(path.c_str()); }); } void HierarchicalNSW::write_file(const char* filename) const { + namespace f2 = index_detail::format2; std::ofstream output(io_impl::filesystem_path(filename), std::ios::binary); if (!output.is_open()) { throw std::runtime_error("HNSW cannot open index file for writing"); } - - // Serialize one count value rather than the bytes of the atomic object. - const size_t num_stored_points = num_stored_points_.load(); - output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); - output.write(reinterpret_cast(&num_stored_points), sizeof(size_t)); - - output.write(reinterpret_cast(&dim_), sizeof(size_t)); - output.write(reinterpret_cast(&padded_dim_), sizeof(size_t)); - output.write(reinterpret_cast(&num_clusters_), sizeof(size_t)); - output.write(reinterpret_cast(&ex_bits_), sizeof(size_t)); - output.write(reinterpret_cast(&metric_type_), sizeof(metric_type_)); - - output.write(reinterpret_cast(&bin_data_bytes_), sizeof(size_t)); - output.write(reinterpret_cast(&ex_data_bytes_), sizeof(size_t)); - output.write(reinterpret_cast(&level0_links_bytes_), sizeof(size_t)); - output.write(reinterpret_cast(&bin_data_offset_), sizeof(size_t)); - output.write(reinterpret_cast(&ex_data_offset_), sizeof(size_t)); - output.write(reinterpret_cast(&label_offset_), sizeof(PID)); - output.write(reinterpret_cast(&level0_record_bytes_), sizeof(size_t)); - output.write(reinterpret_cast(&upper_layer_links_bytes_), sizeof(size_t)); - - output.write(reinterpret_cast(&max_level_), sizeof(int)); - output.write(reinterpret_cast(&entry_point_id_), sizeof(PID)); - - output.write(reinterpret_cast(&M_), sizeof(size_t)); - output.write(reinterpret_cast(&upper_layer_max_degree_), sizeof(size_t)); - output.write(reinterpret_cast(&base_layer_max_degree_), sizeof(size_t)); - output.write(reinterpret_cast(&level_multiplier_), sizeof(double)); - output.write(reinterpret_cast(&ef_construction_), sizeof(size_t)); - - output.write( - reinterpret_cast(centroids_memory_.get()), - static_cast(num_clusters_ * padded_dim_ * sizeof(float)) + const size_t element_count = num_stored_points_.load(); + + f2::Header header; + header.index_kind = f2::IndexKind::Hnsw; + header.metric = metric_type_ == METRIC_IP ? f2::Metric::InnerProduct : f2::Metric::L2; + header.dim = dim_; + header.padded_dim = padded_dim_; + header.code_bits = static_cast(ex_bits_ + 1); + header.id_bits = 32; + header.seed = seed_; + header.num_points = element_count; + + // The rotation was drawn from the low 32 bits of the seed. + std::vector rotator( + rabitqlib::detail::rotator_blob_bytes(RotatorType::FhtKacRotator, dim_, padded_dim_) ); - - output.write( - reinterpret_cast(data_level0_memory_.get()), - static_cast(num_stored_points * level0_record_bytes_) + rabitqlib::detail::encode_rotator_blob( + *rotator_, + RotatorType::FhtKacRotator, + dim_, + seed_.has_value() ? std::optional(static_cast(*seed_)) + : std::nullopt, + rotator.data() ); - - for (size_t i = 0; i < num_stored_points; ++i) { - uint32_t link_list_bytes = - element_levels_[i] > 0 ? upper_layer_links_bytes_ * element_levels_[i] : 0; - output.write(reinterpret_cast(&link_list_bytes), sizeof(uint32_t)); - if (link_list_bytes != 0) { - output.write(link_lists_[i].get(), link_list_bytes); - } + std::array params{}; + f2::put_u64(params.data(), max_elements_); + f2::put_u64(params.data() + 8, num_clusters_); + f2::put_u64(params.data() + 16, M_); + f2::put_u64(params.data() + 24, upper_layer_max_degree_); + f2::put_u64(params.data() + 32, base_layer_max_degree_); + f2::put_u64(params.data() + 40, ef_construction_); + uint64_t multiplier_bits = 0; + std::memcpy(&multiplier_bits, &level_multiplier_, sizeof(multiplier_bits)); + f2::put_u64(params.data() + 48, multiplier_bits); + f2::put_u32( + params.data() + 56, static_cast(static_cast(max_level_)) + ); + f2::put_u64(params.data() + 64, entry_point_id_); + std::vector levels(sizeof(uint32_t) * element_count); + uint64_t upper_bytes = 0; + for (size_t i = 0; i < element_count; ++i) { + f2::put_u32(levels.data() + (sizeof(uint32_t) * i), element_levels_[i]); + upper_bytes += uint64_t(element_levels_[i]) * upper_layer_links_bytes_; } - rotator_->save(output); + const auto bytes = [](const auto& buffer) { + return + [&buffer](f2::SectionSink& sink) { sink.write(buffer.data(), buffer.size()); }; + }; + const std::vector sections = { + {f2::kRotatorSection, f2::kSectionRequired, rotator.size(), bytes(rotator)}, + {kParamsSection, + f2::kSectionRequired | f2::kSectionChecksummed, + params.size(), + bytes(params)}, + {kCentroidsSection, + f2::kSectionRequired, + sizeof(float) * num_clusters_ * padded_dim_, + [this](f2::SectionSink& sink) { + sink.write( + centroids_memory_.get(), sizeof(float) * num_clusters_ * padded_dim_ + ); + }}, + {kRecordsSection, + f2::kSectionRequired, + element_count * level0_record_bytes_, + [&](f2::SectionSink& sink) { + sink.write(data_level0_memory_.get(), element_count * level0_record_bytes_); + }}, + {kLevelsSection, f2::kSectionRequired, levels.size(), bytes(levels)}, + {kUpperLinksSection, + f2::kSectionRequired, + upper_bytes, + [&](f2::SectionSink& sink) { + for (size_t i = 0; i < element_count; ++i) { + if (element_levels_[i] > 0) { + sink.write( + link_lists_[i].get(), + static_cast(element_levels_[i]) * upper_layer_links_bytes_ + ); + } + } + }}, + }; + // write_file checks the stream after flushing; close reports a late failure. + f2::write_file(output, header, sections, RABITQ_LIBRARY_VERSION); output.close(); if (!output) { - throw std::runtime_error("HNSW failed to write index file"); + throw std::runtime_error("HNSW index file cannot be written"); } } @@ -224,8 +278,68 @@ void HierarchicalNSW::load(const char* filename) { throw std::runtime_error("HNSW cannot open index file"); } + // Keep the current index intact until the complete file has been validated. + HierarchicalNSW loaded; + if (index_detail::format2::has_magic(input, index_detail::format2::IndexKind::Hnsw)) { + loaded.read_format2(input); + } else { + loaded.read_legacy(input); + } + + loaded.visited_set_pool_ = std::make_unique(1, loaded.max_elements_); + loaded.raw_dist_func_ = + (loaded.metric_type_ == METRIC_IP) ? dot_product_dis : euclidean_sqr; + loaded.ip_func_ = select_excode_ipfunc(loaded.ex_bits_); + loaded.query_config_ = + quant::faster_config(loaded.padded_dim_, SplitSingleQuery::kNumBits); + + // All swaps below are non-allocating. The temporary releases the old state. + using std::swap; + swap(construction_data_, loaded.construction_data_); + swap(max_elements_, loaded.max_elements_); + const size_t old_count = num_stored_points_.load(); + num_stored_points_ = loaded.num_stored_points_.load(); + loaded.num_stored_points_ = old_count; + swap(num_removed_, loaded.num_removed_); + swap(level0_record_bytes_, loaded.level0_record_bytes_); + swap(upper_layer_links_bytes_, loaded.upper_layer_links_bytes_); + swap(M_, loaded.M_); + swap(upper_layer_max_degree_, loaded.upper_layer_max_degree_); + swap(base_layer_max_degree_, loaded.base_layer_max_degree_); + swap(ef_construction_, loaded.ef_construction_); + swap(metric_type_, loaded.metric_type_); + swap(level_multiplier_, loaded.level_multiplier_); + swap(max_level_, loaded.max_level_); + link_list_locks_.swap(loaded.link_list_locks_); + swap(entry_point_id_, loaded.entry_point_id_); + swap(level0_links_bytes_, loaded.level0_links_bytes_); + swap(bin_data_offset_, loaded.bin_data_offset_); + swap(ex_data_offset_, loaded.ex_data_offset_); + swap(label_offset_, loaded.label_offset_); + swap(bin_data_bytes_, loaded.bin_data_bytes_); + swap(ex_data_bytes_, loaded.ex_data_bytes_); + swap(ex_bits_, loaded.ex_bits_); + swap(data_level0_memory_, loaded.data_level0_memory_); + swap(link_lists_, loaded.link_lists_); + element_levels_.swap(loaded.element_levels_); + swap(num_clusters_, loaded.num_clusters_); + swap(dim_, loaded.dim_); + swap(padded_dim_, loaded.padded_dim_); + swap(centroids_memory_, loaded.centroids_memory_); + label_lookup_.swap(loaded.label_lookup_); + swap(visited_set_pool_, loaded.visited_set_pool_); + swap(ip_func_, loaded.ip_func_); + swap(rotator_, loaded.rotator_); + swap(query_config_, loaded.query_config_); + swap(raw_dist_func_, loaded.raw_dist_func_); + swap(seed_, loaded.seed_); + swap(read_only_, loaded.read_only_); +} + +// The unversioned layout of releases up to 0.8. +void HierarchicalNSW::read_legacy(std::ifstream& input) { auto invalid_file = []() -> void { - throw std::runtime_error("HNSW index file is invalid or truncated"); + throw FormatError("HNSW index file is invalid or truncated"); }; auto read_header = [&](auto& value) { input.read(reinterpret_cast(&value), sizeof(value)); @@ -240,74 +354,65 @@ void HierarchicalNSW::load(const char* filename) { return left * right; }; - // Keep the current index intact until the complete file has been validated. - HierarchicalNSW loaded; size_t element_count = 0; PID label_offset = 0; - read_header(loaded.max_elements_); + read_header(max_elements_); read_header(element_count); - read_header(loaded.dim_); - read_header(loaded.padded_dim_); - read_header(loaded.num_clusters_); - read_header(loaded.ex_bits_); - read_header(loaded.metric_type_); - read_header(loaded.bin_data_bytes_); - read_header(loaded.ex_data_bytes_); - read_header(loaded.level0_links_bytes_); - read_header(loaded.bin_data_offset_); - read_header(loaded.ex_data_offset_); + read_header(dim_); + read_header(padded_dim_); + read_header(num_clusters_); + read_header(ex_bits_); + read_header(metric_type_); + read_header(bin_data_bytes_); + read_header(ex_data_bytes_); + read_header(level0_links_bytes_); + read_header(bin_data_offset_); + read_header(ex_data_offset_); read_header(label_offset); - loaded.label_offset_ = label_offset; - read_header(loaded.level0_record_bytes_); - read_header(loaded.upper_layer_links_bytes_); - read_header(loaded.max_level_); - read_header(loaded.entry_point_id_); - read_header(loaded.M_); - read_header(loaded.upper_layer_max_degree_); - read_header(loaded.base_layer_max_degree_); - read_header(loaded.level_multiplier_); - read_header(loaded.ef_construction_); - - if (loaded.max_elements_ == 0 || - loaded.max_elements_ > buffer::kSearchBufferMaxPointCount || - element_count > loaded.max_elements_ || - loaded.num_clusters_ > buffer::kSearchBufferMaxPointCount || - (element_count != 0 && loaded.num_clusters_ == 0) || loaded.dim_ < 64 || - loaded.dim_ > rotator_impl::FhtKacRotator::kMaxDim || - loaded.padded_dim_ < loaded.dim_ || loaded.padded_dim_ % 32 != 0 || - loaded.padded_dim_ > rotator_impl::FhtKacRotator::kMaxDim || loaded.ex_bits_ > 8 || - (loaded.metric_type_ != METRIC_L2 && loaded.metric_type_ != METRIC_IP) || - loaded.M_ < 2 || loaded.M_ > kMaxM || loaded.upper_layer_max_degree_ != loaded.M_ || - loaded.base_layer_max_degree_ != 2 * loaded.M_ || - loaded.ef_construction_ < loaded.M_ || - (element_count == 0 && - (loaded.max_level_ != -1 || loaded.entry_point_id_ != kPidMax)) || - (element_count != 0 && - (loaded.max_level_ < 0 || loaded.entry_point_id_ >= element_count))) { + label_offset_ = label_offset; + read_header(level0_record_bytes_); + read_header(upper_layer_links_bytes_); + read_header(max_level_); + read_header(entry_point_id_); + read_header(M_); + read_header(upper_layer_max_degree_); + read_header(base_layer_max_degree_); + read_header(level_multiplier_); + read_header(ef_construction_); + + if (max_elements_ == 0 || max_elements_ > buffer::kSearchBufferMaxPointCount || + element_count > max_elements_ || + num_clusters_ > buffer::kSearchBufferMaxPointCount || + (element_count != 0 && num_clusters_ == 0) || dim_ < 64 || + dim_ > rotator_impl::FhtKacRotator::kMaxDim || padded_dim_ < dim_ || + padded_dim_ % 32 != 0 || padded_dim_ > rotator_impl::FhtKacRotator::kMaxDim || + ex_bits_ > 8 || (metric_type_ != METRIC_L2 && metric_type_ != METRIC_IP) || + M_ < 2 || M_ > kMaxM || upper_layer_max_degree_ != M_ || + base_layer_max_degree_ != 2 * M_ || ef_construction_ < M_ || + (element_count == 0 && (max_level_ != -1 || entry_point_id_ != kPidMax)) || + (element_count != 0 && (max_level_ < 0 || entry_point_id_ >= element_count))) { invalid_file(); } // The stored layout must be the one the validated parameters derive. - const auto stored_layout = loaded.record_layout(); - loaded.set_record_layout(); - if (loaded.record_layout() != stored_layout || - !std::isfinite(loaded.level_multiplier_) || loaded.level_multiplier_ <= 0.0) { + const auto stored_layout = record_layout(); + set_record_layout(); + if (record_layout() != stored_layout || !std::isfinite(level_multiplier_) || + level_multiplier_ <= 0.0) { invalid_file(); } - const size_t centroids_bytes = checked_product( - checked_product(loaded.num_clusters_, loaded.padded_dim_), sizeof(float) - ); - const size_t data_capacity_bytes = - checked_product(loaded.max_elements_, loaded.level0_record_bytes_); - const size_t data_bytes = checked_product(element_count, loaded.level0_record_bytes_); + const size_t centroids_bytes = + checked_product(checked_product(num_clusters_, padded_dim_), sizeof(float)); + const size_t data_capacity_bytes = checked_product(max_elements_, level0_record_bytes_); + const size_t data_bytes = checked_product(element_count, level0_record_bytes_); // Validate pointer-table capacity before allocating it. - (void)checked_product(loaded.max_elements_, sizeof(LinkStorage)); + (void)checked_product(max_elements_, sizeof(LinkStorage)); const size_t link_headers_bytes = checked_product(element_count, sizeof(uint32_t)); - loaded.rotator_ = rotator_impl::make_rotator_for_loading( - loaded.dim_, RotatorType::FhtKacRotator, loaded.padded_dim_ + rotator_ = rotator_impl::make_rotator_for_loading( + dim_, RotatorType::FhtKacRotator, padded_dim_ ); - const size_t rotator_bytes = loaded.rotator_->dump_bytes(); + const size_t rotator_bytes = rotator_->dump_bytes(); const auto payload_start = input.tellg(); input.seekg(0, std::ios::end); @@ -337,85 +442,268 @@ void HierarchicalNSW::load(const char* filename) { }; if (centroids_bytes != 0) { - loaded.centroids_memory_.reset(memory::huge_page_allocate(centroids_bytes)); - read_payload(loaded.centroids_memory_.get(), centroids_bytes); + centroids_memory_.reset(memory::huge_page_allocate(centroids_bytes)); + read_payload(centroids_memory_.get(), centroids_bytes); } - loaded.data_level0_memory_.reset(memory::huge_page_allocate(data_capacity_bytes)); - read_payload(loaded.data_level0_memory_.get(), data_bytes); + data_level0_memory_.reset(memory::huge_page_allocate(data_capacity_bytes)); + read_payload(data_level0_memory_.get(), data_bytes); - loaded.element_levels_ = std::vector(loaded.max_elements_); - loaded.link_list_locks_ = std::vector(loaded.max_elements_); - loaded.link_lists_.reserve(loaded.max_elements_); - loaded.link_lists_.resize(element_count); + element_levels_ = std::vector(max_elements_); + link_list_locks_ = std::vector(max_elements_); + link_lists_.reserve(max_elements_); + link_lists_.resize(element_count); for (size_t i = 0; i < element_count; ++i) { - const char* row = - loaded.data_level0_memory_.get() + i * loaded.level0_record_bytes_; + uint32_t link_list_bytes = 0; + read_payload(reinterpret_cast(&link_list_bytes), sizeof(link_list_bytes)); + if (link_list_bytes != 0) { + if (link_list_bytes % upper_layer_links_bytes_ != 0 || + link_list_bytes / upper_layer_links_bytes_ > + static_cast(max_level_) || + link_list_bytes > remaining) { + invalid_file(); + } + link_lists_[i].reset(static_cast(std::malloc(link_list_bytes))); + if (link_lists_[i] == nullptr) { + throw std::bad_alloc(); + } + element_levels_[i] = + static_cast(link_list_bytes / upper_layer_links_bytes_); + } + num_stored_points_ = i + 1; + if (link_list_bytes != 0) { + read_payload(link_lists_[i].get(), link_list_bytes); + } + } + validate_graph(); + + if (rotator_bytes > remaining) { + invalid_file(); + } + rotator_->load(input); + if (!input) { + invalid_file(); + } + remaining -= rotator_bytes; + if (remaining != 0) { + invalid_file(); + } +} + +void HierarchicalNSW::read_format2(std::istream& input) { + namespace f2 = index_detail::format2; + const std::vector known_kinds = { + f2::kRotatorSection, + kParamsSection, + kCentroidsSection, + kRecordsSection, + kLevelsSection, + kUpperLinksSection}; + f2::Reader reader( + input, f2::IndexKind::Hnsw, known_kinds, [](const f2::Header& header) { + if (header.metric == f2::Metric::Cosine) { + throw f2::FormatError( + "HNSW index file uses the cosine metric, which this release does not " + "support" + ); + } + if (header.id_bits != 32 || header.raw_type != f2::RawType::None) { + throw f2::FormatError("HNSW index file stores 64-bit IDs or raw vectors"); + } + if (header.code_bits == 0 || header.dim < 64) { + throw f2::FormatError( + "HNSW index file has an invalid code width or dimension" + ); + } + if (header.num_points == 0 || + header.num_points > buffer::kSearchBufferMaxPointCount) { + throw f2::FormatError( + "HNSW index file has an invalid point count " + + std::to_string(header.num_points) + ); + } + } + ); + const f2::Header& header = reader.header(); + metric_type_ = header.metric == f2::Metric::InnerProduct ? METRIC_IP : METRIC_L2; + dim_ = static_cast(header.dim); + padded_dim_ = static_cast(header.padded_dim); + ex_bits_ = header.code_bits - 1; + seed_ = header.seed; + const auto element_count = static_cast(header.num_points); + + const auto require = [&reader](uint32_t kind, std::optional size) { + const f2::Section* section = reader.find(kind); + if (section == nullptr) { + throw f2::FormatError( + "HNSW index file has no section kind " + std::to_string(kind) + ); + } + if (size.has_value() && section->size != *size) { + throw f2::FormatError( + "HNSW index file section kind " + std::to_string(kind) + " has " + + std::to_string(section->size) + " bytes; expected " + std::to_string(*size) + ); + } + return section; + }; + const auto read_all = [&reader](const f2::Section& section) { + std::vector bytes(static_cast(section.size)); + reader.read(section, 0, bytes.data(), bytes.size()); + return bytes; + }; + + const std::vector params = read_all(*require(kParamsSection, kParamsBytes)); + max_elements_ = static_cast(f2::get_u64(params.data())); + num_clusters_ = static_cast(f2::get_u64(params.data() + 8)); + M_ = static_cast(f2::get_u64(params.data() + 16)); + upper_layer_max_degree_ = static_cast(f2::get_u64(params.data() + 24)); + base_layer_max_degree_ = static_cast(f2::get_u64(params.data() + 32)); + ef_construction_ = static_cast(f2::get_u64(params.data() + 40)); + const uint64_t multiplier_bits = f2::get_u64(params.data() + 48); + std::memcpy(&level_multiplier_, &multiplier_bits, sizeof(level_multiplier_)); + max_level_ = static_cast(static_cast(f2::get_u32(params.data() + 56))); + const uint64_t entry_point = f2::get_u64(params.data() + 64); + if (max_elements_ < element_count || + max_elements_ > buffer::kSearchBufferMaxPointCount || num_clusters_ == 0 || + num_clusters_ > buffer::kSearchBufferMaxPointCount || M_ < 2 || M_ > kMaxM || + upper_layer_max_degree_ != M_ || base_layer_max_degree_ != 2 * M_ || + ef_construction_ < M_ || !std::isfinite(level_multiplier_) || + level_multiplier_ <= 0.0 || max_level_ < 0 || entry_point >= element_count || + f2::get_u32(params.data() + 60) != 0) { + throw f2::FormatError("HNSW index file has invalid graph parameters"); + } + entry_point_id_ = static_cast(entry_point); + set_record_layout(); + + const std::vector blob = read_all(*require(f2::kRotatorSection, std::nullopt)); + rabitqlib::detail::DecodedRotator rotator; + try { + rotator = rabitqlib::detail::decode_rotator_blob(blob.data(), blob.size()); + } catch (const std::invalid_argument& error) { + throw f2::FormatError( + std::string("HNSW index file has an invalid rotator (") + error.what() + ")" + ); + } + if (rotator.type != RotatorType::FhtKacRotator || rotator.dim != dim_ || + rotator.padded_dim != padded_dim_) { + throw f2::FormatError("HNSW index file rotator does not match its dimensions"); + } + rotator_ = std::move(rotator.rotator); + + const size_t centroids_bytes = sizeof(float) * num_clusters_ * padded_dim_; + const f2::Section* centroids = require(kCentroidsSection, centroids_bytes); + const f2::Section* records = + require(kRecordsSection, element_count * level0_record_bytes_); + const std::vector levels = + read_all(*require(kLevelsSection, sizeof(uint32_t) * element_count)); + uint64_t upper_bytes = 0; + element_levels_ = std::vector(max_elements_); + for (size_t i = 0; i < element_count; ++i) { + const uint32_t level = f2::get_u32(levels.data() + (sizeof(uint32_t) * i)); + if (level > static_cast(max_level_)) { + throw f2::FormatError("HNSW index file has a point above the maximum level"); + } + element_levels_[i] = static_cast(level); + // Guard the sum, so that a wrapped total cannot match the section size. + const uint64_t bytes = uint64_t{level} * upper_layer_links_bytes_; + if (bytes > std::numeric_limits::max() - upper_bytes) { + throw f2::FormatError("HNSW index file has an invalid level structure"); + } + upper_bytes += bytes; + } + const f2::Section* upper = require(kUpperLinksSection, upper_bytes); + + centroids_memory_.reset(memory::huge_page_allocate(centroids_bytes)); + f2::SectionInput centroids_input = reader.open(*centroids); + centroids_input.read(centroids_memory_.get(), centroids_bytes); + data_level0_memory_.reset( + memory::huge_page_allocate(max_elements_ * level0_record_bytes_) + ); + f2::SectionInput records_input = reader.open(*records); + records_input.read(data_level0_memory_.get(), element_count * level0_record_bytes_); + link_list_locks_ = std::vector(max_elements_); + link_lists_.reserve(max_elements_); + link_lists_.resize(element_count); + f2::SectionInput upper_input = reader.open(*upper); + for (size_t i = 0; i < element_count; ++i) { + const size_t bytes = + static_cast(element_levels_[i]) * upper_layer_links_bytes_; + if (bytes == 0) { + continue; + } + link_lists_[i].reset(static_cast(std::malloc(bytes))); + if (link_lists_[i] == nullptr) { + throw std::bad_alloc(); + } + upper_input.read(link_lists_[i].get(), bytes); + } + num_stored_points_ = element_count; + validate_graph(); + read_only_ = reader.read_only(); +} + +void HierarchicalNSW::check_writable(const char* operation) const { + if (read_only_) { + throw std::logic_error( + std::string( + "HNSW index was loaded from a file with features this release " + "cannot update, so it cannot " + ) + + operation + ); + } +} + +void HierarchicalNSW::validate_graph() { + const auto invalid = []() -> void { + throw FormatError("HNSW index file is invalid or truncated"); + }; + const size_t element_count = num_stored_points_; + for (size_t i = 0; i < element_count; ++i) { + const char* row = data_level0_memory_.get() + i * level0_record_bytes_; PID cluster_id = 0; PID label = 0; uint16_t level0_degree = 0; std::memcpy(&level0_degree, row, sizeof(level0_degree)); - std::memcpy(&cluster_id, row + loaded.level0_links_bytes_, sizeof(cluster_id)); - std::memcpy(&label, row + loaded.label_offset_, sizeof(label)); - if ((cluster_id & ~kRemovedMask) >= loaded.num_clusters_ || - level0_degree > loaded.base_layer_max_degree_) { - invalid_file(); + std::memcpy(&cluster_id, row + level0_links_bytes_, sizeof(cluster_id)); + std::memcpy(&label, row + label_offset_, sizeof(label)); + if ((cluster_id & ~kRemovedMask) >= num_clusters_ || + level0_degree > base_layer_max_degree_) { + invalid(); } - loaded.num_removed_ += (cluster_id & kRemovedMask) != 0; + num_removed_ += (cluster_id & kRemovedMask) != 0; for (size_t edge = 0; edge < level0_degree; ++edge) { PID neighbor = 0; std::memcpy(&neighbor, row + sizeof(PID) + edge * sizeof(PID), sizeof(PID)); if (neighbor >= element_count) { - invalid_file(); + invalid(); } } - if (!loaded.label_lookup_.emplace(label, static_cast(i)).second) { - invalid_file(); + if (!label_lookup_.emplace(label, static_cast(i)).second) { + invalid(); } - - uint32_t link_list_bytes = 0; - read_payload(reinterpret_cast(&link_list_bytes), sizeof(link_list_bytes)); - if (link_list_bytes != 0) { - if (link_list_bytes % loaded.upper_layer_links_bytes_ != 0 || - link_list_bytes / loaded.upper_layer_links_bytes_ > - static_cast(loaded.max_level_) || - link_list_bytes > remaining) { - invalid_file(); - } - loaded.link_lists_[i].reset(static_cast(std::malloc(link_list_bytes))); - if (loaded.link_lists_[i] == nullptr) { - throw std::bad_alloc(); + for (int level = 0; level < element_levels_[i]; ++level) { + const char* link = link_lists_[i].get() + level * upper_layer_links_bytes_; + uint16_t degree = 0; + std::memcpy(°ree, link, sizeof(degree)); + if (degree > upper_layer_max_degree_) { + invalid(); } - loaded.element_levels_[i] = - static_cast(link_list_bytes / loaded.upper_layer_links_bytes_); - } - loaded.num_stored_points_ = i + 1; - if (link_list_bytes != 0) { - read_payload(loaded.link_lists_[i].get(), link_list_bytes); - for (int level = 0; level < loaded.element_levels_[i]; ++level) { - const char* link = - loaded.link_lists_[i].get() + level * loaded.upper_layer_links_bytes_; - uint16_t degree = 0; - std::memcpy(°ree, link, sizeof(degree)); - if (degree > loaded.upper_layer_max_degree_) { - invalid_file(); - } - for (size_t edge = 0; edge < degree; ++edge) { - PID neighbor = 0; - std::memcpy( - &neighbor, link + sizeof(PID) + edge * sizeof(PID), sizeof(PID) - ); - if (neighbor >= element_count) { - invalid_file(); - } + for (size_t edge = 0; edge < degree; ++edge) { + PID neighbor = 0; + std::memcpy( + &neighbor, link + sizeof(PID) + edge * sizeof(PID), sizeof(PID) + ); + if (neighbor >= element_count) { + invalid(); } } } } for (size_t i = 0; i < element_count; ++i) { - for (int level = 1; level <= loaded.element_levels_[i]; ++level) { - const char* link = - loaded.link_lists_[i].get() + - (static_cast(level) - 1) * loaded.upper_layer_links_bytes_; + for (int level = 1; level <= element_levels_[i]; ++level) { + const char* link = link_lists_[i].get() + + (static_cast(level) - 1) * upper_layer_links_bytes_; uint16_t degree = 0; std::memcpy(°ree, link, sizeof(degree)); for (size_t edge = 0; edge < degree; ++edge) { @@ -423,75 +711,15 @@ void HierarchicalNSW::load(const char* filename) { std::memcpy( &neighbor, link + sizeof(PID) + edge * sizeof(PID), sizeof(PID) ); - if (loaded.element_levels_[neighbor] < level) { - invalid_file(); + if (element_levels_[neighbor] < level) { + invalid(); } } } } - if (element_count != 0 && - loaded.element_levels_[loaded.entry_point_id_] != loaded.max_level_) { - invalid_file(); + if (element_count != 0 && element_levels_[entry_point_id_] != max_level_) { + invalid(); } - - if (rotator_bytes > remaining) { - invalid_file(); - } - loaded.rotator_->load(input); - if (!input) { - invalid_file(); - } - remaining -= rotator_bytes; - if (remaining != 0) { - invalid_file(); - } - - loaded.visited_set_pool_ = std::make_unique(1, loaded.max_elements_); - loaded.raw_dist_func_ = - (loaded.metric_type_ == METRIC_IP) ? dot_product_dis : euclidean_sqr; - loaded.ip_func_ = select_excode_ipfunc(loaded.ex_bits_); - loaded.query_config_ = - quant::faster_config(loaded.padded_dim_, SplitSingleQuery::kNumBits); - - // All swaps below are non-allocating. The temporary releases the old state. - using std::swap; - swap(construction_data_, loaded.construction_data_); - swap(max_elements_, loaded.max_elements_); - const size_t old_count = num_stored_points_.load(); - num_stored_points_ = loaded.num_stored_points_.load(); - loaded.num_stored_points_ = old_count; - swap(num_removed_, loaded.num_removed_); - swap(level0_record_bytes_, loaded.level0_record_bytes_); - swap(upper_layer_links_bytes_, loaded.upper_layer_links_bytes_); - swap(M_, loaded.M_); - swap(upper_layer_max_degree_, loaded.upper_layer_max_degree_); - swap(base_layer_max_degree_, loaded.base_layer_max_degree_); - swap(ef_construction_, loaded.ef_construction_); - swap(metric_type_, loaded.metric_type_); - swap(level_multiplier_, loaded.level_multiplier_); - swap(max_level_, loaded.max_level_); - link_list_locks_.swap(loaded.link_list_locks_); - swap(entry_point_id_, loaded.entry_point_id_); - swap(level0_links_bytes_, loaded.level0_links_bytes_); - swap(bin_data_offset_, loaded.bin_data_offset_); - swap(ex_data_offset_, loaded.ex_data_offset_); - swap(label_offset_, loaded.label_offset_); - swap(bin_data_bytes_, loaded.bin_data_bytes_); - swap(ex_data_bytes_, loaded.ex_data_bytes_); - swap(ex_bits_, loaded.ex_bits_); - swap(data_level0_memory_, loaded.data_level0_memory_); - swap(link_lists_, loaded.link_lists_); - element_levels_.swap(loaded.element_levels_); - swap(num_clusters_, loaded.num_clusters_); - swap(dim_, loaded.dim_); - swap(padded_dim_, loaded.padded_dim_); - swap(centroids_memory_, loaded.centroids_memory_); - label_lookup_.swap(loaded.label_lookup_); - swap(visited_set_pool_, loaded.visited_set_pool_); - swap(ip_func_, loaded.ip_func_); - swap(rotator_, loaded.rotator_); - swap(query_config_, loaded.query_config_); - swap(raw_dist_func_, loaded.raw_dist_func_); } void HierarchicalNSW::construct( @@ -1112,6 +1340,7 @@ size_t HierarchicalNSW::remove(const PID* labels, size_t num_points) { rotator_ == nullptr || num_clusters_ == 0) { throw std::logic_error("HNSW index must be constructed or loaded before remove"); } + check_writable("remove"); if (num_points == 0) { return 0; } @@ -1319,6 +1548,7 @@ void HierarchicalNSW::resize(size_t new_max_elements) { if (level0_record_bytes_ == 0) { throw std::logic_error("HNSW index must be initialized before resize"); } + check_writable("resize"); if (new_max_elements == 0) { throw std::invalid_argument("HNSW resize capacity must be greater than zero"); } @@ -1379,6 +1609,7 @@ std::vector HierarchicalNSW::add( "HNSW index must be trained, constructed or loaded before add" ); } + check_writable("add"); if (num_points == 0) { return {}; } diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index abc91769..32e8e5c2 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -20,7 +20,9 @@ #include #include +#include "../core/blob.hpp" #include "file_replacement.hpp" +#include "format2.hpp" #include "ivf_search.hpp" #include "rabitqlib/clustering/detail/kmeans.hpp" #include "rabitqlib/clustering/qgkmeans.hpp" @@ -38,6 +40,7 @@ #include "rabitqlib/index/detail/buffer.hpp" #include "rabitqlib/index/detail/parallel_exception.hpp" #include "rabitqlib/index/detail/parallel_for.hpp" +#include "rabitqlib/index/format_error.hpp" #include "rabitqlib/index/ivf/cluster.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/training.hpp" @@ -50,6 +53,11 @@ namespace rabitqlib::ivf { namespace { +using index_detail::format2::get_u32; +using index_detail::format2::get_u64; +using index_detail::format2::put_u32; +using index_detail::format2::put_u64; + // remove() marks a point by this f_add; save and load keep the marker. constexpr float kRemovedFAdd = std::numeric_limits::infinity(); // Files without routing metadata used HNSW routing from this many clusters. @@ -59,6 +67,31 @@ constexpr const char* kPointCountConstructOnly = "IVF index created with a point count must be built by construct"; // Scratch for the rows that flat routing in add rotates at a time. constexpr size_t kRoutingBlockBytes = size_t{128} << 20; + +// Index format 2 sections of IVF-RaBitQ (docs/formats.md, "Section kinds"). +constexpr uint32_t kParamsSection = 256; +constexpr uint32_t kClusterSizesSection = 257; +constexpr uint32_t kFlatRoutingSection = 258; +constexpr uint32_t kFlatRaBitQRoutingSection = 259; +constexpr uint32_t kHnswRoutingSection = 260; +constexpr uint32_t kCodesSection = 261; +constexpr uint32_t kRerankSection = 262; +constexpr uint32_t kIdsSection = 263; +constexpr size_t kParamsBytes = 16; +// The rotator blob stores dimensions in 32 bits, up to the FHT/Kac limit. +constexpr size_t kMaxSavedDim = rotator_impl::FhtKacRotator::kMaxDim; + +uint32_t routing_section(InitializerType type) { + switch (type) { + case InitializerType::Flat: + return kFlatRoutingSection; + case InitializerType::FlatRaBitQ: + return kFlatRaBitQRoutingSection; + default: + return kHnswRoutingSection; + } +} + } // namespace IVF::IVF() noexcept = default; @@ -147,6 +180,8 @@ void IVF::swap(IVF& other) noexcept { swap(num_clusters_, other.num_clusters_); swap(raw_reranking_, other.raw_reranking_); swap(ready_, other.ready_); + swap(read_only_, other.read_only_); + swap(seed_, other.seed_); swap(ex_bits_, other.ex_bits_); swap(rotator_type_, other.rotator_type_); swap(rotator_, other.rotator_); @@ -197,6 +232,7 @@ IVF::IVF( , padded_dim_(dim) , num_clusters_(num_clusters) , raw_reranking_(total_bits == 32) + , seed_(seed) , ex_bits_(total_bits == 32 ? 0 : total_bits - 1) , rotator_type_(rotator_type) , metric_type_(metric_type) { @@ -393,6 +429,7 @@ void IVF::construct( void IVF::allocate_storage(const std::vector& cluster_sizes) { ready_ = false; + read_only_ = false; num_removed_ = 0; free_storage(); clusters_.clear(); @@ -477,6 +514,7 @@ void IVF::add( "IVF index must be trained, constructed or loaded before add" ); } + check_writable("add"); if (num_points == 0) { return; } @@ -781,6 +819,7 @@ void IVF::reserve(PID cluster_id, size_t num_points) { if (!ready_) { throw std::logic_error("IVF index must be constructed or loaded before reserve"); } + check_writable("reserve"); if (cluster_id >= num_clusters_) { throw std::invalid_argument("IVF cluster ID is out of range"); } @@ -819,6 +858,7 @@ size_t IVF::remove(const PID* ids_to_remove, size_t num_points) { if (!ready_) { throw std::logic_error("IVF index must be constructed or loaded before remove"); } + check_writable("remove"); if (num_points == 0) { return 0; } @@ -869,82 +909,117 @@ void IVF::save(const char* filename) const { if (filename == nullptr || filename[0] == '\0') { throw std::invalid_argument("IVF save filename must not be empty"); } - - // An HNSW initializer writes its sidecar next to the path it is given. - const auto write = [this](const std::string& path) { write_file(path.c_str()); }; - if (initializer_type_ == InitializerType::HNSW) { - index_detail::save_replacing(filename, {detail::kHnswSidecarSuffix}, write); - } else { - index_detail::save_replacing(filename, {}, write); + check_writable("save"); + if (dim_ > kMaxSavedDim) { + throw std::invalid_argument( + "IVF cannot save an index of more than " + std::to_string(kMaxSavedDim) + + " dimensions" + ); } + index_detail::save_replacing(filename, {}, [this](const std::string& path) { + write_file(path.c_str()); + }); } void IVF::write_file(const char* filename) const { + namespace f2 = index_detail::format2; std::ofstream output(io_impl::filesystem_path(filename), std::ios::binary); if (!output.is_open()) { throw std::runtime_error("IVF cannot open index file for writing"); } output.exceptions(std::ios::failbit | std::ios::badbit); - const uint32_t flags = raw_reranking_ ? 1 : 0; - const uint64_t stored_padded_dim = padded_dim_; - output.write(reinterpret_cast(&kFormatMagic), sizeof(kFormatMagic)); - output.write(reinterpret_cast(&kFormatVersion), sizeof(kFormatVersion)); - output.write(reinterpret_cast(&flags), sizeof(flags)); - output.write( - reinterpret_cast(&stored_padded_dim), sizeof(stored_padded_dim) - ); - const auto stored_initializer = static_cast(initializer_type_); - output.write( - reinterpret_cast(&stored_initializer), sizeof(stored_initializer) + f2::Header header; + header.index_kind = f2::IndexKind::Ivf; + header.metric = metric_type_ == METRIC_IP ? f2::Metric::InnerProduct : f2::Metric::L2; + header.dim = dim_; + header.padded_dim = padded_dim_; + header.code_bits = raw_reranking_ ? 1 : static_cast(ex_bits_ + 1); + header.raw_type = raw_reranking_ ? f2::RawType::Float32 : f2::RawType::None; + header.id_bits = 32; + header.seed = seed_; + header.num_points = max_elements_; + + std::vector rotator( + rabitqlib::detail::rotator_blob_bytes(rotator_type_, dim_, padded_dim_) ); - - // Save metadata - output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); - output.write(reinterpret_cast(&dim_), sizeof(size_t)); - output.write(reinterpret_cast(&num_clusters_), sizeof(size_t)); - output.write(reinterpret_cast(&ex_bits_), sizeof(size_t)); - output.write(reinterpret_cast(&rotator_type_), sizeof(rotator_type_)); - output.write(reinterpret_cast(&metric_type_), sizeof(metric_type_)); - - // Save the number of points in each cluster - std::vector cluster_sizes; - cluster_sizes.reserve(num_clusters_); - for (const auto& cluster : clusters_) { - cluster_sizes.push_back(cluster.num()); - } - output.write( - reinterpret_cast(cluster_sizes.data()), - static_cast(sizeof(size_t) * num_clusters_) + rabitqlib::detail::encode_rotator_blob( + *rotator_, rotator_type_, dim_, seed_, rotator.data() ); - - // Save rotator - rotator_->save(output); - - // Save data: initializer payload, then all batch data, all ex data and all IDs, - // each in cluster order - initializer_->save(output, filename); + std::array params{}; + put_u64(params.data(), num_clusters_); + put_u32(params.data() + 8, static_cast(initializer_type_)); + std::vector cluster_sizes(num_clusters_); + std::vector sizes(sizeof(uint64_t) * num_clusters_); + uint64_t codes_bytes = 0; + uint64_t rerank_bytes = 0; for (size_t i = 0; i < num_clusters_; ++i) { - output.write( - cluster_storage_[i].batch_data(), - static_cast(batch_data_bytes(cluster_sizes[i])) - ); + cluster_sizes[i] = clusters_[i].num(); + put_u64(sizes.data() + (sizeof(uint64_t) * i), cluster_sizes[i]); + codes_bytes += batch_data_bytes(cluster_sizes[i]); + rerank_bytes += has_rerank_data() ? rerank_data_bytes(cluster_sizes[i]) : 0; } + f2::CountingStream routing_size; + initializer_->write(routing_size); + + const auto bytes = [](const auto& buffer) { + return + [&buffer](f2::SectionSink& sink) { sink.write(buffer.data(), buffer.size()); }; + }; + std::vector sections = { + {f2::kRotatorSection, f2::kSectionRequired, rotator.size(), bytes(rotator)}, + {kParamsSection, + f2::kSectionRequired | f2::kSectionChecksummed, + params.size(), + bytes(params)}, + {kClusterSizesSection, + f2::kSectionRequired | f2::kSectionChecksummed, + sizes.size(), + bytes(sizes)}, + {routing_section(initializer_type_), + f2::kSectionRequired, + routing_size.count(), + [this](f2::SectionSink& sink) { + f2::SinkStream stream(sink); + initializer_->write(stream); + }}, + {kCodesSection, + f2::kSectionRequired, + codes_bytes, + [&](f2::SectionSink& sink) { + for (size_t i = 0; i < num_clusters_; ++i) { + sink.write( + cluster_storage_[i].batch_data(), batch_data_bytes(cluster_sizes[i]) + ); + } + }}, + }; if (has_rerank_data()) { - for (size_t i = 0; i < num_clusters_; ++i) { - output.write( - cluster_storage_[i].rerank_data(), - static_cast(rerank_data_bytes(cluster_sizes[i])) - ); - } - } - for (size_t i = 0; i < num_clusters_; ++i) { - output.write( - reinterpret_cast(cluster_storage_[i].ids()), - static_cast(ids_bytes(cluster_sizes[i])) + sections.push_back( + {kRerankSection, + f2::kSectionRequired, + rerank_bytes, + [&](f2::SectionSink& sink) { + for (size_t i = 0; i < num_clusters_; ++i) { + sink.write( + cluster_storage_[i].rerank_data(), + rerank_data_bytes(cluster_sizes[i]) + ); + } + }} ); } - + sections.push_back( + {kIdsSection, + f2::kSectionRequired, + ids_bytes(max_elements_), + [&](f2::SectionSink& sink) { + for (size_t i = 0; i < num_clusters_; ++i) { + sink.write(cluster_storage_[i].ids(), ids_bytes(cluster_sizes[i])); + } + }} + ); + f2::write_file(output, header, sections, RABITQ_LIBRARY_VERSION); output.close(); } @@ -956,14 +1031,27 @@ void IVF::load(const char* filename) { if (!input.is_open()) { throw std::runtime_error("Cannot open IVF index file"); } + if (index_detail::format2::has_magic(input, index_detail::format2::IndexKind::Ivf)) { + IVF loaded; + loaded.read_format2(input); + swap(loaded); + return; + } + // An earlier layout: RABQIDX1, the raw-vector v1 layout or the unversioned one. IVF loaded; - input.exceptions(std::ios::failbit | std::ios::badbit); + // Every group of reads is checked before its values are used. + const auto check_read = [&input] { + if (!input) { + throw FormatError("IVF index file is truncated"); + } + }; input.seekg(0, std::ios::end); const auto file_bytes = static_cast(input.tellg()); input.seekg(0); uint64_t magic = 0; input.read(reinterpret_cast(&magic), sizeof(magic)); + check_read(); const bool explicit_padding = magic == kFormatMagic; if (explicit_padding) { uint32_t version = 0, flags = 0; @@ -971,20 +1059,22 @@ void IVF::load(const char* filename) { input.read(reinterpret_cast(&version), sizeof(version)); input.read(reinterpret_cast(&flags), sizeof(flags)); input.read(reinterpret_cast(&stored_padded_dim), sizeof(stored_padded_dim)); + check_read(); if (version != 1 && version != kFormatVersion) { - throw std::runtime_error("Unsupported IVF index version"); + throw FormatError("Unsupported IVF index version"); } if (flags > 1 || stored_padded_dim > std::numeric_limits::max()) { - throw std::runtime_error("Invalid IVF index metadata"); + throw FormatError("Invalid IVF index metadata"); } loaded.raw_reranking_ = flags == 1; loaded.padded_dim_ = static_cast(stored_padded_dim); if (version >= 2) { uint32_t initializer = 0; input.read(reinterpret_cast(&initializer), sizeof(initializer)); + check_read(); if (initializer < static_cast(InitializerType::Flat) || initializer > static_cast(InitializerType::HNSW)) { - throw std::runtime_error("Invalid IVF initializer type in index file"); + throw FormatError("Invalid IVF initializer type in index file"); } loaded.initializer_type_ = static_cast(initializer); } @@ -992,8 +1082,9 @@ void IVF::load(const char* filename) { loaded.raw_reranking_ = true; uint32_t version = 0; input.read(reinterpret_cast(&version), sizeof(version)); + check_read(); if (version != kRawFormatVersion) { - throw std::runtime_error("Unsupported raw IVF index version"); + throw FormatError("Unsupported raw IVF index version"); } } else { input.seekg(0); @@ -1008,7 +1099,7 @@ void IVF::load(const char* filename) { reinterpret_cast(&loaded.rotator_type_), sizeof(loaded.rotator_type_) ); input.read(reinterpret_cast(&loaded.metric_type_), sizeof(loaded.metric_type_)); - validate_metric_type(loaded.metric_type_); + check_read(); // Missing routing metadata means a historical file. Its payload layout must // follow the original threshold, independently of the current Auto policy. if (loaded.initializer_type_ == InitializerType::Auto) { @@ -1016,13 +1107,14 @@ void IVF::load(const char* filename) { ? InitializerType::Flat : InitializerType::HNSW; } - if (loaded.max_elements_ == 0 || + if ((loaded.metric_type_ != METRIC_L2 && loaded.metric_type_ != METRIC_IP) || + loaded.max_elements_ == 0 || loaded.max_elements_ > buffer::kSearchBufferMaxPointCount || loaded.dim_ == 0 || loaded.num_clusters_ == 0 || loaded.num_clusters_ > buffer::kSearchBufferMaxPointCount || loaded.ex_bits_ > 8 || (loaded.raw_reranking_ && loaded.ex_bits_ != 0) || loaded.dim_ > (std::numeric_limits::max() / 32) - 64) { - throw std::runtime_error("Invalid IVF index metadata"); + throw FormatError("Invalid IVF index metadata"); } if (!explicit_padding) { // Both historical formats omit padded_dim and always used 64-coordinate blocks. @@ -1030,13 +1122,13 @@ void IVF::load(const char* filename) { } if (loaded.padded_dim_ < loaded.dim_ || loaded.padded_dim_ % 32 != 0 || loaded.padded_dim_ > (std::numeric_limits::max() / 32) - 64) { - throw std::runtime_error("Invalid padded dimension in IVF index file"); + throw FormatError("Invalid padded dimension in IVF index file"); } // Bound every payload by the actual file before allocating or multiplying sizes. size_t remaining = file_bytes - static_cast(input.tellg()); const auto consume = [&remaining](size_t count, size_t bytes) { if (bytes != 0 && count > remaining / bytes) { - throw std::runtime_error("Invalid or truncated IVF index payload"); + throw FormatError("Invalid or truncated IVF index payload"); } remaining -= count * bytes; }; @@ -1048,7 +1140,7 @@ void IVF::load(const char* filename) { } else if (loaded.rotator_type_ == RotatorType::FhtKacRotator) { consume(1, loaded.padded_dim_ / 2); } else { - throw std::runtime_error("Invalid IVF rotator type"); + throw FormatError("Invalid IVF rotator type"); } if (loaded.initializer_type_ != InitializerType::HNSW) { consume(loaded.num_clusters_, sizeof(float) * loaded.padded_dim_); @@ -1066,29 +1158,51 @@ void IVF::load(const char* filename) { reinterpret_cast(cluster_sizes.data()), static_cast(sizeof(size_t) * loaded.num_clusters_) ); + check_read(); size_t total = 0; for (size_t size : cluster_sizes) { if (size > loaded.max_elements_ - total) { - throw std::runtime_error("Invalid cluster counts in IVF index file"); + throw FormatError("Invalid cluster counts in IVF index file"); } total += size; consume(div_round_up(size, fastscan::kBatchSize), loaded.batch_bytes()); } if (total != loaded.max_elements_) { - throw std::runtime_error("Invalid cluster counts in IVF index file"); + throw FormatError("Invalid cluster counts in IVF index file"); + } + try { + loaded.rotator_ = rotator_impl::make_rotator_for_loading( + loaded.dim_, loaded.rotator_type_, loaded.padded_dim_ + ); + } catch (const std::invalid_argument& error) { + // Dimensions the stored rotation does not support, such as FHT/Kac below 64. + throw FormatError(std::string("Invalid IVF rotator dimensions: ") + error.what()); } - loaded.rotator_ = rotator_impl::make_rotator_for_loading( - loaded.dim_, loaded.rotator_type_, loaded.padded_dim_ - ); // Load rotator loaded.rotator_->load(input); + check_read(); // Load data: initializer payload, then all batch data, all ex data and all IDs, // each in cluster order loaded.allocate_storage(cluster_sizes); - loaded.initializer_->load(input, filename); + if (loaded.initializer_type_ == InitializerType::HNSW) { + // These layouts kept the routing graph in a sidecar file beside the index. + const std::string sidecar = std::string(filename) + detail::kHnswSidecarSuffix; + std::ifstream graph( + io_impl::filesystem_path(sidecar), std::ios::binary | std::ios::ate + ); + if (!graph) { + throw std::runtime_error("Cannot open HNSW centroid sidecar"); + } + const auto graph_bytes = + static_cast(static_cast(graph.tellg())); + graph.seekg(0); + loaded.initializer_->read(graph, graph_bytes); + } else { + loaded.initializer_->read(input, loaded.flat_routing_bytes()); + } for (size_t i = 0; i < loaded.num_clusters_; ++i) { input.read( loaded.cluster_storage_[i].batch_data(), @@ -1109,37 +1223,247 @@ void IVF::load(const char* filename) { static_cast(ids_bytes(cluster_sizes[i])) ); } - for (size_t i = 0; i < loaded.num_clusters_; ++i) { - for (PID id : loaded.cluster_storage_[i].id_storage) { - if (id >= loaded.max_elements_ || id >= buffer::kSearchBufferCheckedMask) { - throw std::runtime_error("Invalid point ID in IVF index file"); + check_read(); + if (remaining != 0) { + throw FormatError("IVF index file has unexpected trailing data"); + } + loaded.complete_load(cluster_sizes); + + input.close(); + swap(loaded); +} + +void IVF::read_format2(std::istream& input) { + namespace f2 = index_detail::format2; + const std::vector known_kinds = { + f2::kRotatorSection, + kParamsSection, + kClusterSizesSection, + kFlatRoutingSection, + kFlatRaBitQRoutingSection, + kHnswRoutingSection, + kCodesSection, + kRerankSection, + kIdsSection}; + f2::Reader reader(input, f2::IndexKind::Ivf, known_kinds, [](const f2::Header& header) { + if (header.metric == f2::Metric::Cosine) { + throw f2::FormatError( + "IVF index file uses the cosine metric, which this release does not support" + ); + } + if (header.id_bits != 32) { + throw f2::FormatError( + "IVF index file has 64-bit IDs; IVF-RaBitQ stores 32 bits" + ); + } + if (header.raw_type == f2::RawType::Uint8) { + throw f2::FormatError( + "IVF index file has uint8 raw vectors, which IVF-RaBitQ lacks" + ); + } + if (header.code_bits == 0 || + (header.raw_type == f2::RawType::Float32 && header.code_bits != 1)) { + throw f2::FormatError( + "IVF index file has " + std::to_string(header.code_bits) + + " code bits, which its raw vectors do not allow" + ); + } + if (header.num_points == 0 || + header.num_points > buffer::kSearchBufferMaxPointCount) { + throw f2::FormatError( + "IVF index file has an invalid point count " + + std::to_string(header.num_points) + ); + } + if (header.seed.has_value() && + *header.seed > std::numeric_limits::max()) { + throw f2::FormatError("IVF index file has a seed above 32 bits"); + } + }); + const f2::Header& header = reader.header(); + metric_type_ = header.metric == f2::Metric::InnerProduct ? METRIC_IP : METRIC_L2; + dim_ = static_cast(header.dim); + padded_dim_ = static_cast(header.padded_dim); + raw_reranking_ = header.raw_type == f2::RawType::Float32; + ex_bits_ = raw_reranking_ ? 0 : header.code_bits - 1; + max_elements_ = static_cast(header.num_points); + if (header.seed.has_value()) { + seed_ = static_cast(*header.seed); + } + + // A section of this kind, whose size, when given, must match before it is read. + const auto require = [&reader](uint32_t kind, std::optional size) { + const f2::Section* section = reader.find(kind); + if (section == nullptr) { + throw f2::FormatError( + "IVF index file has no section kind " + std::to_string(kind) + ); + } + if (size.has_value() && section->size != *size) { + throw f2::FormatError( + "IVF index file section kind " + std::to_string(kind) + " has " + + std::to_string(section->size) + " bytes; expected " + std::to_string(*size) + ); + } + return section; + }; + const auto read_all = [&reader](const f2::Section& section) { + std::vector bytes(static_cast(section.size)); + reader.read(section, 0, bytes.data(), bytes.size()); + return bytes; + }; + + const std::vector params = read_all(*require(kParamsSection, kParamsBytes)); + const uint64_t num_clusters = get_u64(params.data()); + const uint32_t initializer = get_u32(params.data() + 8); + if (num_clusters == 0 || num_clusters > buffer::kSearchBufferMaxPointCount) { + throw f2::FormatError( + "IVF index file has an invalid cluster count " + std::to_string(num_clusters) + ); + } + if (initializer < static_cast(InitializerType::Flat) || + initializer > static_cast(InitializerType::HNSW) || + get_u32(params.data() + 12) != 0) { + throw f2::FormatError("IVF index file has invalid routing parameters"); + } + num_clusters_ = static_cast(num_clusters); + initializer_type_ = static_cast(initializer); + + const std::vector sizes = + read_all(*require(kClusterSizesSection, sizeof(uint64_t) * num_clusters_)); + std::vector cluster_sizes(num_clusters_); + uint64_t total = 0; + uint64_t codes_bytes = 0; + uint64_t rerank_bytes = 0; + for (size_t i = 0; i < num_clusters_; ++i) { + const uint64_t size = get_u64(sizes.data() + (sizeof(uint64_t) * i)); + if (size > max_elements_ - total) { + throw f2::FormatError("IVF index file cluster sizes exceed its point count"); + } + total += size; + cluster_sizes[i] = static_cast(size); + codes_bytes += batch_data_bytes(cluster_sizes[i]); + rerank_bytes += has_rerank_data() ? rerank_data_bytes(cluster_sizes[i]) : 0; + } + if (total != max_elements_) { + throw f2::FormatError( + "IVF index file cluster sizes do not add up to its point count" + ); + } + + const std::vector blob = read_all(*require(f2::kRotatorSection, std::nullopt)); + rabitqlib::detail::DecodedRotator rotator; + try { + rotator = rabitqlib::detail::decode_rotator_blob(blob.data(), blob.size()); + } catch (const std::invalid_argument& error) { + throw f2::FormatError( + std::string("IVF index file has an invalid rotator (") + error.what() + ")" + ); + } + if (rotator.dim != dim_ || rotator.padded_dim != padded_dim_) { + throw f2::FormatError("IVF index file rotator does not match its dimensions"); + } + rotator_type_ = rotator.type; + rotator_ = std::move(rotator.rotator); + + const uint32_t routing_kind = routing_section(initializer_type_); + for (const uint32_t kind : + {kFlatRoutingSection, kFlatRaBitQRoutingSection, kHnswRoutingSection}) { + if (kind != routing_kind && reader.find(kind) != nullptr) { + throw f2::FormatError("IVF index file has routing data of another initializer"); + } + } + const f2::Section* routing = require( + routing_kind, + initializer_type_ == InitializerType::HNSW + ? std::nullopt + : std::optional(flat_routing_bytes()) + ); + const f2::Section* codes = require(kCodesSection, codes_bytes); + const f2::Section* rerank = nullptr; + if (has_rerank_data()) { + rerank = require(kRerankSection, rerank_bytes); + } else if (reader.find(kRerankSection) != nullptr) { + throw f2::FormatError( + "IVF index file has rerank data that its code bits do not use" + ); + } + const f2::Section* ids = require(kIdsSection, ids_bytes(max_elements_)); + + allocate_storage(cluster_sizes); + std::istream routing_input(input.rdbuf()); + routing_input.seekg(static_cast(routing->offset)); + initializer_->read(routing_input, routing->size); + // Read each bulk section front to back: a seek per cluster would discard the file + // buffer every time. + f2::SectionInput codes_input = reader.open(*codes); + for (size_t i = 0; i < num_clusters_; ++i) { + codes_input.read( + cluster_storage_[i].batch_data(), batch_data_bytes(cluster_sizes[i]) + ); + } + if (rerank != nullptr) { + f2::SectionInput rerank_input = reader.open(*rerank); + for (size_t i = 0; i < num_clusters_; ++i) { + rerank_input.read( + cluster_storage_[i].rerank_data(), rerank_data_bytes(cluster_sizes[i]) + ); + } + } + f2::SectionInput ids_input = reader.open(*ids); + for (size_t i = 0; i < num_clusters_; ++i) { + ids_input.read(cluster_storage_[i].ids(), ids_bytes(cluster_sizes[i])); + } + complete_load(cluster_sizes); + read_only_ = reader.read_only(); +} + +size_t IVF::flat_routing_bytes() const { + const size_t centroids = sizeof(float) * num_clusters_ * padded_dim_; + if (initializer_type_ == InitializerType::FlatRaBitQ) { + return centroids + (sizeof(float) * padded_dim_) + + (div_round_up(num_clusters_, fastscan::kBatchSize) * batch_bytes()); + } + return centroids; +} + +void IVF::complete_load(const std::vector& cluster_sizes) { + for (size_t i = 0; i < num_clusters_; ++i) { + for (PID id : cluster_storage_[i].id_storage) { + if (id >= max_elements_ || id >= buffer::kSearchBufferCheckedMask) { + throw FormatError("Invalid point ID in IVF index file"); } } } - loaded.create_cluster_views(cluster_sizes); + create_cluster_views(cluster_sizes); // Removal markers are already persisted, including in historical files. // Count only logical rows: FastScan padding lanes are not stored points. - for (const auto& cluster : loaded.clusters_) { + for (const auto& cluster : clusters_) { for (size_t first = 0; first < cluster.num(); first += fastscan::kBatchSize) { ConstBatchDataMap batch( - cluster.batch_data() + - (first / fastscan::kBatchSize) * loaded.batch_bytes(), - loaded.padded_dim_ + cluster.batch_data() + (first / fastscan::kBatchSize) * batch_bytes(), + padded_dim_ ); const size_t batch_size = std::min(fastscan::kBatchSize, cluster.num() - first); for (size_t lane = 0; lane < batch_size; ++lane) { - loaded.num_removed_ += batch.f_add()[lane] == kRemovedFAdd; + num_removed_ += batch.f_add()[lane] == kRemovedFAdd; } } } - if (remaining != 0) { - throw std::runtime_error("IVF index file has unexpected trailing data"); - } - loaded.ready_ = true; + ready_ = true; +} - input.close(); - swap(loaded); +void IVF::check_writable(const char* operation) const { + if (read_only_) { + throw std::logic_error( + std::string( + "IVF index was loaded from a file with features this release " + "cannot update, so it cannot " + ) + + operation + ); + } } void IVF::search( diff --git a/src/index/ivf_initializer.cpp b/src/index/ivf_initializer.cpp index cbaf7e00..4fb51a55 100644 --- a/src/index/ivf_initializer.cpp +++ b/src/index/ivf_initializer.cpp @@ -2,10 +2,12 @@ #include #include #include +#include #include -#include #include +#include #include +#include #include #include @@ -17,6 +19,7 @@ #include "rabitqlib/fastscan/query.hpp" #include "rabitqlib/index/detail/buffer.hpp" #include "rabitqlib/index/detail/parallel_for.hpp" +#include "rabitqlib/index/format_error.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/space.hpp" @@ -163,7 +166,7 @@ void FlatRaBitQInitializer::centroids_distances( } } -void FlatRaBitQInitializer::save(std::ofstream& output, const char* /*filename*/) const { +void FlatRaBitQInitializer::write(std::ostream& output) const { output.write( reinterpret_cast(centroids_.data()), static_cast(centroids_.size() * sizeof(float)) @@ -181,7 +184,11 @@ void FlatRaBitQInitializer::save(std::ofstream& output, const char* /*filename*/ } } -void FlatRaBitQInitializer::load(std::ifstream& input, const char* /*filename*/) { +void FlatRaBitQInitializer::read(std::istream& input, uint64_t bytes) { + if (bytes != sizeof(float) * (centroids_.size() + global_centroid_.size()) + + batch_data_.size()) { + throw FormatError("IVF Flat RaBitQ initializer data has the wrong size"); + } input.read( reinterpret_cast(centroids_.data()), static_cast(centroids_.size() * sizeof(float)) @@ -195,7 +202,7 @@ void FlatRaBitQInitializer::load(std::ifstream& input, const char* /*filename*/) static_cast(batch_data_.size()) ); if (!input) { - throw std::runtime_error("IVF Flat RaBitQ initializer cannot be read"); + throw FormatError("IVF Flat RaBitQ initializer cannot be read"); } } } // namespace rabitqlib::ivf diff --git a/src/index/qg.cpp b/src/index/qg.cpp index 02da4bed..cff2ab74 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -1,6 +1,7 @@ #include "rabitqlib/index/symqg/qg.hpp" #include +#include #include #include #include @@ -15,7 +16,9 @@ #include #include +#include "../core/blob.hpp" #include "file_replacement.hpp" +#include "format2.hpp" #include "rabitqlib/core/data_layout.hpp" #include "rabitqlib/core/encoder.hpp" #include "rabitqlib/core/metric.hpp" @@ -29,6 +32,7 @@ #include "rabitqlib/index/detail/buffer.hpp" #include "rabitqlib/index/detail/parallel_exception.hpp" #include "rabitqlib/index/detail/visited_set.hpp" +#include "rabitqlib/index/format_error.hpp" #include "rabitqlib/index/symqg/detail/search_scratch.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/bitops.hpp" @@ -45,6 +49,15 @@ namespace { // Native-endian numeric tag; little-endian files contain the bytes "QTIBARGQ". constexpr uint64_t kFormatMagic = 0x5147524142495451ULL; constexpr uint32_t kFormatVersion = 1; + +// Index format 2 sections of SymphonyQG (docs/formats.md, "Section kinds"). +constexpr uint32_t kParamsSection = 768; +constexpr uint32_t kCentroidSection = 769; +constexpr uint32_t kRecordsSection = 770; +constexpr size_t kParamsBytes = 32; +// The stored initialization: 0 is unknown. +constexpr uint32_t kInitPiPNN = 1; +constexpr uint32_t kInitRandom = 2; } // namespace QuantizedQuery::QuantizedQuery( @@ -139,7 +152,8 @@ QuantizedGraph::QuantizedGraph( , metric_type_(metric_type) , rotator_type_(rotator_type) , quantization_bits_(quantization_bits) - , seed_(seed) { + , seed_(seed) + , seed_known_(true) { validate_configuration(); padded_dim_ = padded_dimension(dim_); initialize(); @@ -262,54 +276,233 @@ void QuantizedGraph::save(const char* filename) const { if (filename == nullptr || filename[0] == '\0') { throw std::invalid_argument("QuantizedGraph save filename must not be empty"); } + if (read_only_) { + throw std::logic_error( + "QuantizedGraph was loaded from a file with features this release cannot " + "update, so it cannot save" + ); + } index_detail::save_replacing(filename, {}, [this](const std::string& path) { write_file(path.c_str()); }); } void QuantizedGraph::write_file(const char* filename) const { + namespace f2 = index_detail::format2; std::ofstream output(io_impl::filesystem_path(filename), std::ios::binary); if (!output.is_open()) { throw std::runtime_error("QuantizedGraph cannot open file for writing"); } output.exceptions(std::ios::badbit | std::ios::failbit); + f2::Header header; + header.index_kind = f2::IndexKind::SymphonyQg; + header.metric = metric_type_ == METRIC_IP ? f2::Metric::InnerProduct : f2::Metric::L2; + header.dim = dim_; + header.padded_dim = padded_dim_; + header.code_bits = static_cast(quantization_bits_); + header.raw_type = quantization_bits_ == 0 ? f2::RawType::Float32 : f2::RawType::None; + header.id_bits = 32; + if (seed_known_) { + header.seed = seed_; + } + header.num_points = max_elements_; + + std::vector rotator( + rabitqlib::detail::rotator_blob_bytes(rotator_type_, dim_, padded_dim_) + ); + rabitqlib::detail::encode_rotator_blob( + *rotator_, rotator_type_, dim_, seed(), rotator.data() + ); + std::array params{}; + f2::put_u64(params.data(), max_degree_); + f2::put_u64(params.data() + 8, entry_point_id_); + f2::put_u64(params.data() + 16, ef_construction_); + uint32_t init = 0; + if (initialization_.has_value()) { + init = *initialization_ == QGInitialization::PiPNN ? kInitPiPNN : kInitRandom; + } + f2::put_u32(params.data() + 24, init); + + const auto bytes = [](const auto& buffer) { + return + [&buffer](f2::SectionSink& sink) { sink.write(buffer.data(), buffer.size()); }; + }; + std::vector sections = { + {f2::kRotatorSection, f2::kSectionRequired, rotator.size(), bytes(rotator)}, + {kParamsSection, + f2::kSectionRequired | f2::kSectionChecksummed, + params.size(), + bytes(params)}, + }; if (quantization_bits_ != 0) { - output.write(reinterpret_cast(&kFormatMagic), sizeof(kFormatMagic)); - output.write( - reinterpret_cast(&kFormatVersion), sizeof(kFormatVersion) + sections.push_back( + {kCentroidSection, + f2::kSectionRequired, + padded_dim_ * sizeof(float), + [this](f2::SectionSink& sink) { + sink.write(centroid_.data(), padded_dim_ * sizeof(float)); + }} ); } + const size_t data_bytes = checked_multiply(max_elements_, record_bytes_); + sections.push_back( + {kRecordsSection, + f2::kSectionRequired, + data_bytes, + [this, data_bytes](f2::SectionSink& sink) { + sink.write(get_row_data(0), data_bytes); + }} + ); + f2::write_file(output, header, sections, RABITQ_LIBRARY_VERSION); + output.close(); +} - // Basic variants - output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); - output.write(reinterpret_cast(&max_degree_), sizeof(size_t)); - output.write(reinterpret_cast(&dim_), sizeof(size_t)); - output.write(reinterpret_cast(&padded_dim_), sizeof(size_t)); - output.write(reinterpret_cast(&entry_point_id_), sizeof(PID)); - output.write(reinterpret_cast(&rotator_type_), sizeof(RotatorType)); - output.write(reinterpret_cast(&metric_type_), sizeof(MetricType)); - if (quantization_bits_ != 0) { - output.write( - reinterpret_cast(&quantization_bits_), sizeof(quantization_bits_) +void QuantizedGraph::read_format2(std::istream& input) { + namespace f2 = index_detail::format2; + const std::vector known_kinds = { + f2::kRotatorSection, kParamsSection, kCentroidSection, kRecordsSection}; + f2::Reader reader( + input, f2::IndexKind::SymphonyQg, known_kinds, [](const f2::Header& header) { + if (header.metric == f2::Metric::Cosine) { + throw f2::FormatError( + "QuantizedGraph index file uses the cosine metric, which this release " + "does not support" + ); + } + const bool raw = + header.code_bits == 0 && header.raw_type == f2::RawType::Float32; + const bool quantized = (header.code_bits == 4 || header.code_bits == 8) && + header.raw_type == f2::RawType::None; + if (header.id_bits != 32 || (!raw && !quantized)) { + throw f2::FormatError( + "QuantizedGraph index file stores neither raw vectors nor 4- or 8-bit " + "codes with 32-bit IDs" + ); + } + if (header.seed.has_value() && + *header.seed > std::numeric_limits::max()) { + throw f2::FormatError("QuantizedGraph index file has a seed above 32 bits"); + } + } + ); + const f2::Header& header = reader.header(); + metric_type_ = header.metric == f2::Metric::InnerProduct ? METRIC_IP : METRIC_L2; + dim_ = static_cast(header.dim); + padded_dim_ = static_cast(header.padded_dim); + quantization_bits_ = header.code_bits; + max_elements_ = static_cast(header.num_points); + seed_known_ = header.seed.has_value(); + seed_ = static_cast(header.seed.value_or(0)); + + const auto require = [&reader](uint32_t kind, std::optional size) { + const f2::Section* section = reader.find(kind); + if (section == nullptr) { + throw f2::FormatError( + "QuantizedGraph index file has no section kind " + std::to_string(kind) + ); + } + if (size.has_value() && section->size != *size) { + throw f2::FormatError( + "QuantizedGraph index file section kind " + std::to_string(kind) + " has " + + std::to_string(section->size) + " bytes; expected " + std::to_string(*size) + ); + } + return section; + }; + const auto read_all = [&reader](const f2::Section& section) { + std::vector bytes(static_cast(section.size)); + reader.read(section, 0, bytes.data(), bytes.size()); + return bytes; + }; + + const std::vector blob = read_all(*require(f2::kRotatorSection, std::nullopt)); + rabitqlib::detail::DecodedRotator rotator; + try { + rotator = rabitqlib::detail::decode_rotator_blob(blob.data(), blob.size()); + } catch (const std::invalid_argument& error) { + throw f2::FormatError( + std::string("QuantizedGraph index file has an invalid rotator (") + + error.what() + ")" ); - output.write( - reinterpret_cast(centroid_.data()), - static_cast(padded_dim_ * sizeof(float)) + } + if (rotator.dim != dim_ || rotator.padded_dim != padded_dim_) { + throw f2::FormatError( + "QuantizedGraph index file rotator does not match its dimensions" + ); + } + rotator_type_ = rotator.type; + + const std::vector params = read_all(*require(kParamsSection, kParamsBytes)); + const uint64_t max_degree = f2::get_u64(params.data()); + const uint64_t entry_point = f2::get_u64(params.data() + 8); + const uint32_t init = f2::get_u32(params.data() + 24); + if (max_degree > max_elements_ || entry_point >= max_elements_ || init > kInitRandom || + f2::get_u32(params.data() + 28) != 0) { + throw f2::FormatError("QuantizedGraph index file has invalid graph parameters"); + } + max_degree_ = static_cast(max_degree); + entry_point_id_ = static_cast(entry_point); + ef_construction_ = static_cast(f2::get_u64(params.data() + 16)); + if (init != 0) { + initialization_ = + init == kInitPiPNN ? QGInitialization::PiPNN : QGInitialization::Random; + } + try { + validate_configuration(); + } catch (const std::invalid_argument& error) { + throw f2::FormatError( + std::string("QuantizedGraph index file is invalid: ") + error.what() ); } - // Data - output.write( - get_row_data(0), - static_cast(checked_multiply(max_elements_, record_bytes_)) - ); + size_t data_bytes = 0; + try { + initialize_layout(); + data_bytes = checked_multiply(max_elements_, record_bytes_); + } catch (const std::length_error&) { + throw f2::FormatError( + "QuantizedGraph index file describes records too large to address" + ); + } + const f2::Section* centroid = nullptr; + if (quantization_bits_ != 0) { + centroid = require(kCentroidSection, padded_dim_ * sizeof(float)); + } else if (reader.find(kCentroidSection) != nullptr) { + throw f2::FormatError( + "QuantizedGraph index file has a centroid but stores raw vectors" + ); + } + const f2::Section* records = require(kRecordsSection, data_bytes); + + initialize(false); + rotator_ = std::move(rotator.rotator); + if (centroid != nullptr) { + centroid_.resize(padded_dim_); + reader.open(*centroid).read(centroid_.data(), padded_dim_ * sizeof(float)); + } + reader.open(*records).read(get_row_data(0), data_bytes); + read_only_ = reader.read_only(); + for (PID source = 0; source < max_elements_; ++source) { + const auto neighbors = get_neighbors(source); + for (size_t i = 0; i < max_degree_; ++i) { + if (neighbors[i] >= max_elements_) { + throw FormatError("QuantizedGraph neighbor ID is invalid"); + } + } + } +} - // Rotator - rotator_->save(output); +std::optional QuantizedGraph::seed() const { + return seed_known_ ? std::optional(seed_) : std::nullopt; +} - output.flush(); - output.close(); +std::optional QuantizedGraph::ef_construction() const { + return ef_construction_ != 0 ? std::optional(ef_construction_) : std::nullopt; +} + +std::optional QuantizedGraph::initialization() const { + return initialization_; } void QuantizedGraph::load(const char* filename) { @@ -324,16 +517,27 @@ void QuantizedGraph::load(const char* filename) { if (!input.is_open()) { throw std::runtime_error("QuantizedGraph cannot open file"); } + if (index_detail::format2::has_magic( + input, index_detail::format2::IndexKind::SymphonyQg + )) { + QuantizedGraph loaded; + loaded.read_format2(input); + // ef_search is a runtime search setting rather than persisted index state. + loaded.default_ef_search_ = default_ef_search_; + loaded.ready_ = true; + *this = std::move(loaded); + return; + } auto read_exact = [&](void* destination, size_t bytes, const char* field) { if (bytes > static_cast(std::numeric_limits::max())) { - throw std::runtime_error("QuantizedGraph field is too large to read"); + throw FormatError("QuantizedGraph field is too large to read"); } input.read( reinterpret_cast(destination), static_cast(bytes) ); if (!input) { - throw std::runtime_error( + throw FormatError( std::string("QuantizedGraph file is truncated while reading ") + field ); } @@ -348,7 +552,7 @@ void QuantizedGraph::load(const char* filename) { uint32_t version = 0; read_value(version, "format version"); if (version != kFormatVersion) { - throw std::runtime_error("QuantizedGraph file version is unsupported"); + throw FormatError("QuantizedGraph file version is unsupported"); } } else { // Files produced before qg-quant have no header and always contain raw vectors. @@ -371,27 +575,37 @@ void QuantizedGraph::load(const char* filename) { loaded.quantization_bits_ = 0; } - loaded.validate_configuration(); + try { + loaded.validate_configuration(); + } catch (const std::invalid_argument& error) { + throw FormatError(error.what()); // stored settings that do not fit + } if (stored_padded_dim < loaded.dim_ || stored_padded_dim % 32 != 0) { - throw std::runtime_error("QuantizedGraph file has an invalid padded dimension"); + throw FormatError("QuantizedGraph file has an invalid padded dimension"); } // Older files use 64-coordinate padding. Their saved rotation and codes must // retain that dimension, including when initialize() allocates the row storage. loaded.padded_dim_ = stored_padded_dim; - loaded.initialize_layout(); - - const size_t centroid_bytes = loaded.quantization_bits_ == 0 - ? 0 - : checked_multiply(loaded.padded_dim_, sizeof(float)); - const size_t data_bytes = checked_multiply(loaded.max_elements_, loaded.record_bytes_); - const size_t rotator_bytes = - loaded.rotator_type_ == RotatorType::MatrixRotator - ? checked_multiply( - checked_multiply(sizeof(float), loaded.dim_), loaded.padded_dim_ - ) - : checked_multiply(loaded.padded_dim_, size_t{4}) / 8; - const size_t expected_payload_bytes = - checked_add(checked_add(centroid_bytes, data_bytes), rotator_bytes); + size_t centroid_bytes = 0; + size_t data_bytes = 0; + size_t expected_payload_bytes = 0; + try { + loaded.initialize_layout(); + centroid_bytes = loaded.quantization_bits_ == 0 + ? 0 + : checked_multiply(loaded.padded_dim_, sizeof(float)); + data_bytes = checked_multiply(loaded.max_elements_, loaded.record_bytes_); + const size_t rotator_bytes = + loaded.rotator_type_ == RotatorType::MatrixRotator + ? checked_multiply( + checked_multiply(sizeof(float), loaded.dim_), loaded.padded_dim_ + ) + : checked_multiply(loaded.padded_dim_, size_t{4}) / 8; + expected_payload_bytes = + checked_add(checked_add(centroid_bytes, data_bytes), rotator_bytes); + } catch (const std::length_error&) { + throw FormatError("QuantizedGraph file describes records too large to address"); + } const auto payload_position = input.tellg(); if (payload_position < 0) { @@ -401,10 +615,15 @@ void QuantizedGraph::load(const char* filename) { const size_t payload_offset = static_cast(payload_position); if (payload_offset > file_size || file_size - payload_offset != expected_payload_bytes) { - throw std::runtime_error("QuantizedGraph payload size is invalid"); + throw FormatError("QuantizedGraph payload size is invalid"); } - loaded.initialize(false); + try { + loaded.initialize(false); + } catch (const std::invalid_argument& error) { + // Dimensions the stored rotation does not support, such as FHT/Kac below 64. + throw FormatError(std::string("QuantizedGraph file is invalid: ") + error.what()); + } if (loaded.quantization_bits_ != 0) { loaded.centroid_.resize(loaded.padded_dim_); @@ -417,14 +636,14 @@ void QuantizedGraph::load(const char* filename) { const auto neighbors = loaded.get_neighbors(source); for (size_t i = 0; i < loaded.max_degree_; ++i) { if (neighbors[i] >= loaded.max_elements_) { - throw std::runtime_error("QuantizedGraph neighbor ID is invalid"); + throw FormatError("QuantizedGraph neighbor ID is invalid"); } } } loaded.rotator_->load(input); if (!input) { - throw std::runtime_error("QuantizedGraph file is truncated while reading rotator"); + throw FormatError("QuantizedGraph file is truncated while reading rotator"); } input.close(); diff --git a/src/index/qg_builder.cpp b/src/index/qg_builder.cpp index 470a2f78..4163f474 100644 --- a/src/index/qg_builder.cpp +++ b/src/index/qg_builder.cpp @@ -79,6 +79,11 @@ QGBuilder::QGBuilder( throw std::invalid_argument("QuantizedGraph initialization is invalid"); } initialization_ = init; + // The graph records its construction settings for its index file. Rebuilding replaces + // whatever unknown features a loaded file described. + qg_.ef_construction_ = ef_construction; + qg_.initialization_ = init; + qg_.read_only_ = false; reset(data); } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 879b224f..62ff0125 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -40,9 +40,16 @@ add_executable(rabitq_tests ${UNIT_TESTS} ) target_include_directories(rabitq_tests PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/common) +target_include_directories(rabitq_tests PRIVATE ${PROJECT_SOURCE_DIR}/src/core) +target_compile_definitions(rabitq_tests PRIVATE + RABITQ_TEST_LIBRARY_VERSION="${PROJECT_VERSION}" +) if(RABITQ_BUILD_INDEX) target_include_directories(rabitq_tests PRIVATE ${PROJECT_SOURCE_DIR}/src/index) endif() +set_source_files_properties(unit/rabitqlib/utils/blob_test.cpp + PROPERTIES COMPILE_DEFINITIONS + "RABITQ_QUANT_FIXTURE_ROOT=\"${CMAKE_CURRENT_SOURCE_DIR}/python/fixtures/quant\"") set_source_files_properties(unit/rabitqlib/index/historical_persistence_test.cpp PROPERTIES COMPILE_DEFINITIONS "RABITQ_TEST_FIXTURE_ROOT=\"${CMAKE_CURRENT_SOURCE_DIR}/fixtures/persistence\"") diff --git a/tests/common/format2_file.hpp b/tests/common/format2_file.hpp new file mode 100644 index 00000000..5c432532 --- /dev/null +++ b/tests/common/format2_file.hpp @@ -0,0 +1,160 @@ +#ifndef RABITQ_TEST_FORMAT2_FILE_HPP +#define RABITQ_TEST_FORMAT2_FILE_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "crc32.hpp" +#include "format2.hpp" + +namespace rabitq_test { +namespace format2 = rabitqlib::index_detail::format2; + +// Every section kind of each index, for reading back a file in full. +inline const std::vector kIvfSections = { + 1, 256, 257, 258, 259, 260, 261, 262, 263}; +inline const std::vector kHnswSections = {1, 512, 513, 514, 515, 516}; +inline const std::vector kQgSections = {1, 768, 769, 770}; + +// A whole format-2 file: the header and each known section with its bytes, in file order. +struct Format2Contents { + format2::Header header; + std::string writer_version; + std::vector> sections; + + // The bytes of a section; an error if the file has none. + std::string& bytes(uint32_t kind) { + for (auto& [section, data] : sections) { + if (section.kind == kind) { + return data; + } + } + throw std::runtime_error("test file has no such section"); + } + + void remove(uint32_t kind) { + sections.erase( + std::remove_if( + sections.begin(), + sections.end(), + [kind](const auto& entry) { return entry.first.kind == kind; } + ), + sections.end() + ); + } +}; + +inline Format2Contents read_format2_file( + const std::string& path, + format2::IndexKind kind, + const std::vector& known_kinds +) { + Format2Contents contents; + std::ifstream input(path, std::ios::binary); + format2::Reader reader(input, kind, known_kinds); + contents.header = reader.header(); + contents.writer_version = reader.writer_version(); + for (const uint32_t known : known_kinds) { + if (const format2::Section* section = reader.find(known)) { + std::string bytes(section->size, '\0'); + reader.read(*section, 0, bytes.data(), bytes.size()); + contents.sections.emplace_back(*section, std::move(bytes)); + } + } + std::sort( + contents.sections.begin(), + contents.sections.end(), + [](const auto& a, const auto& b) { return a.first.offset < b.first.offset; } + ); + return contents; +} + +// Write `contents` as a structurally valid file, so that a test can damage what a +// section holds rather than the container. +inline void write_format2_file(const std::string& path, const Format2Contents& contents) { + std::vector writes; + for (const auto& [section, bytes] : contents.sections) { + writes.push_back( + {section.kind, + section.flags, + bytes.size(), + [&data = bytes](format2::SectionSink& sink) { + sink.write(data.data(), data.size()); + }} + ); + } + std::ofstream output(path, std::ios::binary | std::ios::trunc); + output.exceptions(std::ios::failbit | std::ios::badbit); + format2::write_file(output, contents.header, writes, contents.writer_version); +} + +// The bytes of one section of a format-2 file. +inline std::string read_format2_section( + const std::string& path, + format2::IndexKind kind, + const std::vector& known_kinds, + uint32_t section_kind +) { + return read_format2_file(path, kind, known_kinds).bytes(section_kind); +} + +// Rewrite a format-2 file with one section's bytes replaced and every other part kept. +inline void replace_format2_section( + const std::string& path, + format2::IndexKind kind, + const std::vector& known_kinds, + uint32_t section_kind, + const std::string& replacement +) { + Format2Contents contents = read_format2_file(path, kind, known_kinds); + contents.bytes(section_kind) = replacement; + write_format2_file(path, contents); +} + +// Overwrite a little-endian header field of `size` bytes and recompute the header +// checksum; the writer refuses some values tests need, such as unknown feature bits. +inline void patch_format2_header( + const std::string& path, size_t offset, uint64_t value, size_t size +) { + std::string bytes; + { + std::ifstream input(path, std::ios::binary); + bytes.assign(std::istreambuf_iterator(input), {}); + } + for (size_t i = 0; i < size; ++i) { + bytes[offset + i] = static_cast(value >> (8 * i)); + } + const auto u32 = [&bytes](size_t at) { + uint32_t result = 0; + for (size_t i = 0; i < 4; ++i) { + result |= static_cast(static_cast(bytes[at + i])) << (8 * i); + } + return result; + }; + uint64_t table_offset = 0; + for (size_t i = 0; i < 8; ++i) { + table_offset |= static_cast(static_cast(bytes[96 + i])) + << (8 * i); + } + std::string header = bytes.substr(0, u32(12)); + header.replace(108, 4, 4, '\0'); + uint32_t crc = rabitqlib::detail::crc32(header.data(), header.size()); + crc = rabitqlib::detail::crc32(bytes.data() + table_offset, size_t{32} * u32(104), crc); + for (size_t i = 0; i < 4; ++i) { + bytes[108 + i] = static_cast(crc >> (8 * i)); + } + std::ofstream output(path, std::ios::binary | std::ios::trunc); + output.write(bytes.data(), static_cast(bytes.size())); +} + +} // namespace rabitq_test + +#endif // RABITQ_TEST_FORMAT2_FILE_HPP diff --git a/tests/fixtures/persistence/.gitattributes b/tests/fixtures/persistence/.gitattributes index df5ad8ed..afe6b1ab 100644 --- a/tests/fixtures/persistence/.gitattributes +++ b/tests/fixtures/persistence/.gitattributes @@ -1,3 +1,4 @@ *.index -text +*.hnsw -text *.expected text eol=lf manifest.json text eol=lf diff --git a/tests/fixtures/persistence/README.md b/tests/fixtures/persistence/README.md index ee820cf2..1273ab8f 100644 --- a/tests/fixtures/persistence/README.md +++ b/tests/fixtures/persistence/README.md @@ -28,9 +28,12 @@ interchange format for arbitrary architectures. Python also tests synthetic IVF- v1 headers around the frozen historical payload, unknown versions, and invalid padding. That header test is distinct from release-produced fixtures. -The small IVF-RaBitQ fixtures use Flat centroid routing. Current round-trip and failure -tests exercise explicit HNSW routing and its sidecar; historical auto-selected -HNSW routing at 20,000 clusters is not represented by these small fixtures. +The small IVF-RaBitQ fixtures use Flat centroid routing. The [v0.8.3](v0.8.3/README.md) +set adds the last layouts before index format 2: IVF-RaBitQ with Flat RaBitQ and HNSW +routing (with its `.hnsw` sidecar), HNSW-RaBitQ at 1, 5 and 9 bits, and SymphonyQG raw, +4- and 8-bit with FHT/Kac and matrix rotation. Historical auto-selected HNSW routing at 20,000 clusters is +not represented by fixtures; Python tests cover it with headers built around the +v0.8.3 layout. ## Deliberate regeneration diff --git a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/README.md b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/README.md index 3e7c742d..89b1ea05 100644 --- a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/README.md +++ b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/README.md @@ -13,10 +13,12 @@ single-point cluster, a full FastScan batch and a partial tail batch. Point 0 is removed. The `.expected` files hold reference results of a search that probes every cluster. -The tests load each file, compare search results, and require that saving it -again reproduces the original bytes. Because the clusters differ in size, a -reader or writer that orders the batch data, extra-bit data and IDs differently -from the shared-buffer layout fails these checks. +The tests load each file and compare search results. Until index format 2, they also +required that saving it again reproduces the original bytes; saving now writes format +2, and the tests require that the migrated file gives the same results and that saving +it again reproduces it. Because the clusters differ in size, a reader that orders the +batch data, extra-bit data and IDs differently from the shared-buffer layout fails +these checks. ## Deliberate regeneration diff --git a/tests/fixtures/persistence/v0.8.3/CMakeLists.txt b/tests/fixtures/persistence/v0.8.3/CMakeLists.txt new file mode 100644 index 00000000..4fd4b3e5 --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/CMakeLists.txt @@ -0,0 +1,15 @@ +cmake_minimum_required(VERSION 3.20) +project(rabitq_v083_fixtures LANGUAGES CXX) + +if(NOT IS_DIRECTORY "${RABITQ_HISTORICAL_SOURCE}") + message(FATAL_ERROR "Set RABITQ_HISTORICAL_SOURCE to the pinned v0.8.3 checkout") +endif() +set(RABITQ_BUILD_TESTS OFF CACHE BOOL "" FORCE) +set(RABITQ_BUILD_SAMPLES OFF CACHE BOOL "" FORCE) +set(RABITQ_BUILD_PYTHON_BINDINGS OFF CACHE BOOL "" FORCE) +set(RABITQ_ENABLE_NATIVE_OPTIMIZATION OFF CACHE BOOL "" FORCE) +add_subdirectory("${RABITQ_HISTORICAL_SOURCE}" historical) +add_executable(generate_fixtures generate.cpp) +target_link_libraries(generate_fixtures PRIVATE rabitqlib::index) +add_executable(generate_graph_fixtures generate_graphs.cpp) +target_link_libraries(generate_graph_fixtures PRIVATE rabitqlib::index) diff --git a/tests/fixtures/persistence/v0.8.3/README.md b/tests/fixtures/persistence/v0.8.3/README.md new file mode 100644 index 00000000..74091fa0 --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/README.md @@ -0,0 +1,57 @@ +# Index fixtures from v0.8.3 + +These files were written by release **v0.8.3**, commit +`dabfdb8c200a6f155df4edd0b9c825e6728b14a1`, using [generate.cpp](generate.cpp) for +IVF-RaBitQ and [generate_graphs.cpp](generate_graphs.cpp) for HNSW-RaBitQ and SymphonyQG. +[manifest.json](manifest.json) records the producer and SHA-256 hashes. +Keep these files unchanged; tests never regenerate them with the current writer. + +They freeze the last IVF-RaBitQ layout before index format 2: `RABQIDX1` version 2, +which stores the centroid-routing type, and the hnswlib sidecar `.hnsw` that +HNSW routing writes beside the index. The five cases are: + +| File | Bits | Metric | Rotator | Routing | +| --- | --- | --- | --- | --- | +| `ivf4_flat_l2` | 4 | L2 | FHT/Kac | Flat | +| `ivf7_flatrabitq_ip` | 7 | inner product | FHT/Kac | FlatRaBitQ | +| `ivf4_hnsw_l2` | 4 | L2 | FHT/Kac | HNSW, with sidecar | +| `ivf32_hnsw_ip` | raw (`nbits=32`) | inner product | FHT/Kac | HNSW, with sidecar | +| `ivf1_flat_matrix_l2` | 1 | L2 | matrix | Flat | + +The graph generator, run separately, adds the last HNSW-RaBitQ and SymphonyQG layouts +before index format 2: + +| File | Bits | Metric | Rotator | Notes | +| --- | --- | --- | --- | --- | +| `hnsw1_l2` | 1 | L2 | FHT/Kac | M = 8, `ef_construction` 64, seed 42, point 0 removed | +| `hnsw5_ip` | 5 | inner product | FHT/Kac | As above | +| `hnsw9_l2` | 9 | L2 | FHT/Kac | As above | +| `qg0_ip` | raw | inner product | FHT/Kac | Degree 32, PiPNN, `ef_construction` 64, seed 42 | +| `qg4_l2` | 4 | L2 | FHT/Kac | As above | +| `qg8_ip` | 8 | inner product | FHT/Kac | As above | +| `qg4_matrix_l2` | 4 | L2 | matrix | As above | + +Each index holds 130 points in 65 original dimensions (96 padded); the IVF-RaBitQ and +HNSW-RaBitQ files use four clusters of 0, 1, 32 and 97 points, with point 0 removed, as in +[pre-ivf-per-cluster-storage](../pre-ivf-per-cluster-storage/README.md). The +`.expected` files hold reference results: IVF-RaBitQ probes every cluster, and the graph +indexes search with `ef` 64. + +## Deliberate regeneration + +Build the generator only against an isolated copy of the release: + +```bash +mkdir -p build/fixtures/v0.8.3 +git archive v0.8.3 | tar -x -C build/fixtures/v0.8.3 +cmake -S tests/fixtures/persistence/v0.8.3 -B build/fixtures/v0.8.3-generator \ + -DRABITQ_HISTORICAL_SOURCE="$PWD/build/fixtures/v0.8.3" \ + -DCMAKE_BUILD_TYPE=Release +cmake --build build/fixtures/v0.8.3-generator --target generate_fixtures generate_graph_fixtures +``` + +Run either generator with a new output directory and `OMP_NUM_THREADS=1`. +IVF-RaBitQ rotation is seeded from `random_device`, so those bytes differ on every run; +the graph indexes use seed 42. Each index, its sidecar and its `.expected` file form a +fixed set. Record new hashes in the manifest +before adopting them. diff --git a/tests/fixtures/persistence/v0.8.3/generate.cpp b/tests/fixtures/persistence/v0.8.3/generate.cpp new file mode 100644 index 00000000..ff2441d5 --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/generate.cpp @@ -0,0 +1,111 @@ +// Build only against the pinned v0.8.3 source, never the current checkout. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/types.hpp" +#include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/ivf/types.hpp" +#include "rabitqlib/index/types.hpp" + +int main(int argc, char** argv) { + if (argc != 2) { + throw std::invalid_argument("Pass the fixture output directory"); + } + using namespace rabitqlib; + // Cluster sizes 0, 1, 32 and 97 cover empty, single, full-batch and tail clusters. + constexpr size_t kCount = 130, kDim = 65, kClusters = 4, kQueries = 3, kTopK = 5; + const std::filesystem::path directory(argv[1]); + std::filesystem::create_directories(directory); + uint32_t state = 713; + const auto next = [&state] { + state = state * 1664525U + 1013904223U; + return static_cast(static_cast((state >> 8) % 2001) - 1000) / 256.0F; + }; + std::vector data(kCount * kDim), queries(kQueries * kDim); + std::vector centroids(kClusters * kDim); + for (auto& value : data) { + value = next(); + } + for (auto& value : queries) { + value = next(); + } + for (auto& value : centroids) { + value = next(); + } + std::vector clusters(kCount, 3); + clusters[0] = 1; + for (size_t i = 1; i < 33; ++i) { + clusters[i] = 2; + } + + const struct { + const char* name; + size_t total_bits; + MetricType metric; + RotatorType rotator; + ivf::InitializerType initializer; + } cases[] = { + {"ivf4_flat_l2", + 4, + METRIC_L2, + RotatorType::FhtKacRotator, + ivf::InitializerType::Flat}, + {"ivf7_flatrabitq_ip", + 7, + METRIC_IP, + RotatorType::FhtKacRotator, + ivf::InitializerType::FlatRaBitQ}, + {"ivf4_hnsw_l2", + 4, + METRIC_L2, + RotatorType::FhtKacRotator, + ivf::InitializerType::HNSW}, + {"ivf32_hnsw_ip", + 32, + METRIC_IP, + RotatorType::FhtKacRotator, + ivf::InitializerType::HNSW}, + {"ivf1_flat_matrix_l2", + 1, + METRIC_L2, + RotatorType::MatrixRotator, + ivf::InitializerType::Flat}, + }; + for (const auto& c : cases) { + const std::string name = c.name; + const auto path = (directory / (name + ".index")).string(); + std::vector ids(kQueries * kTopK); + std::vector distances(ids.size()); + ivf::IVF index( + kCount, kDim, kClusters, c.total_bits, c.metric, c.rotator, c.initializer + ); + index.construct(data.data(), centroids.data(), clusters.data(), false, 1); + const PID removed = 0; + index.remove(&removed, 1); + index.save(path.c_str()); + index.search_batch( + queries.data(), kQueries, kTopK, kClusters, ids.data(), distances.data() + ); + std::ofstream expected(directory / (name + ".expected"), std::ios::binary); + expected.exceptions(std::ios::failbit | std::ios::badbit); + expected << std::setprecision(std::numeric_limits::max_digits10); + expected << kCount << ' ' << kDim << ' ' << kQueries << ' ' << kTopK << '\n'; + for (size_t i = 0; i < queries.size(); ++i) { + expected << queries[i] << (i + 1 == queries.size() ? '\n' : ' '); + } + for (size_t i = 0; i < ids.size(); ++i) { + expected << ids[i] << ' ' << distances[i] << '\n'; + } + expected.close(); + } + return 0; +} diff --git a/tests/fixtures/persistence/v0.8.3/generate_graphs.cpp b/tests/fixtures/persistence/v0.8.3/generate_graphs.cpp new file mode 100644 index 00000000..abe912cb --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/generate_graphs.cpp @@ -0,0 +1,129 @@ +// Build only against the pinned v0.8.3 source, never the current checkout. Writes the +// HNSW-RaBitQ and SymphonyQG fixtures; generate.cpp writes the IVF-RaBitQ ones. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/types.hpp" +#include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/symqg/qg.hpp" +#include "rabitqlib/index/symqg/qg_builder.hpp" +#include "rabitqlib/index/types.hpp" + +int main(int argc, char** argv) { + if (argc != 2) { + throw std::invalid_argument("Pass the fixture output directory"); + } + using namespace rabitqlib; + // The IVF fixtures' data: cluster sizes 0, 1, 32 and 97. + constexpr size_t kCount = 130, kDim = 65, kClusters = 4, kQueries = 3, kTopK = 5; + const std::filesystem::path directory(argv[1]); + std::filesystem::create_directories(directory); + uint32_t state = 713; + const auto next = [&state] { + state = state * 1664525U + 1013904223U; + return static_cast(static_cast((state >> 8) % 2001) - 1000) / 256.0F; + }; + std::vector data(kCount * kDim), queries(kQueries * kDim); + std::vector centroids(kClusters * kDim); + for (auto& value : data) { + value = next(); + } + for (auto& value : queries) { + value = next(); + } + for (auto& value : centroids) { + value = next(); + } + std::vector clusters(kCount, 3); + clusters[0] = 1; + for (size_t i = 1; i < 33; ++i) { + clusters[i] = 2; + } + + const auto write_expected = [&](const std::string& name, + const std::vector& ids, + const std::vector& distances) { + std::ofstream expected(directory / (name + ".expected"), std::ios::binary); + expected.exceptions(std::ios::failbit | std::ios::badbit); + expected << std::setprecision(std::numeric_limits::max_digits10); + expected << kCount << ' ' << kDim << ' ' << kQueries << ' ' << kTopK << '\n'; + for (size_t i = 0; i < queries.size(); ++i) { + expected << queries[i] << (i + 1 == queries.size() ? '\n' : ' '); + } + for (size_t i = 0; i < ids.size(); ++i) { + expected << ids[i] << ' ' << distances[i] << '\n'; + } + expected.close(); + }; + + const struct { + const char* name; + size_t total_bits; + MetricType metric; + } hnsw_cases[] = { + {"hnsw1_l2", 1, METRIC_L2}, + {"hnsw5_ip", 5, METRIC_IP}, + {"hnsw9_l2", 9, METRIC_L2}, + }; + for (const auto& c : hnsw_cases) { + const std::string name = c.name; + hnsw::HierarchicalNSW index(kCount, kDim, c.total_bits, 8, 64, 42, c.metric); + std::vector assigned = clusters; + index.construct( + kClusters, centroids.data(), kCount, data.data(), assigned.data(), 1 + ); + const PID removed = 0; + index.remove(&removed, 1); + index.save((directory / (name + ".index")).string().c_str()); + std::vector ids(kQueries * kTopK); + std::vector distances(ids.size()); + const auto result = index.search(queries.data(), kQueries, kTopK, 64, 1); + for (size_t i = 0; i < kQueries; ++i) { + for (size_t j = 0; j < kTopK; ++j) { + ids[i * kTopK + j] = result[i][j].second; + distances[i * kTopK + j] = result[i][j].first; + } + } + write_expected(name, ids, distances); + } + + const struct { + const char* name; + size_t bits; // 0 stores raw vectors + MetricType metric; + RotatorType rotator; + } qg_cases[] = { + {"qg0_ip", 0, METRIC_IP, RotatorType::FhtKacRotator}, + {"qg4_l2", 4, METRIC_L2, RotatorType::FhtKacRotator}, + {"qg8_ip", 8, METRIC_IP, RotatorType::FhtKacRotator}, + {"qg4_matrix_l2", 4, METRIC_L2, RotatorType::MatrixRotator}, + }; + for (const auto& c : qg_cases) { + const std::string name = c.name; + symqg::QuantizedGraph index( + kCount, kDim, 32, c.metric, c.rotator, c.bits, 42 + ); + symqg::QGBuilder builder( + index, 64, data.data(), 1, symqg::QGInitialization::PiPNN, 42 + ); + builder.build(); + index.save((directory / (name + ".index")).string().c_str()); + std::vector ids(kQueries * kTopK); + std::vector distances(ids.size()); + index.set_ef(64); + index.search_batch( + queries.data(), kQueries, kTopK, ids.data(), distances.data(), 1 + ); + write_expected(name, ids, distances); + } + return 0; +} diff --git a/tests/fixtures/persistence/v0.8.3/hnsw1_l2.expected b/tests/fixtures/persistence/v0.8.3/hnsw1_l2.expected new file mode 100644 index 00000000..96a62dad --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/hnsw1_l2.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +59 457.198364 +94 473.985107 +92 476.936432 +13 481.009399 +1 482.431854 +71 283.250702 +120 364.552856 +102 403.540588 +50 411.44928 +28 413.655273 +119 332.883972 +48 336.432617 +12 339.079712 +100 404.918793 +105 407.847778 diff --git a/tests/fixtures/persistence/v0.8.3/hnsw1_l2.index b/tests/fixtures/persistence/v0.8.3/hnsw1_l2.index new file mode 100644 index 00000000..fae1c705 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/hnsw1_l2.index differ diff --git a/tests/fixtures/persistence/v0.8.3/hnsw5_ip.expected b/tests/fixtures/persistence/v0.8.3/hnsw5_ip.expected new file mode 100644 index 00000000..193d8d0d --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/hnsw5_ip.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +128 -94.5418472 +46 -83.159584 +92 -80.4502106 +50 -75.6382751 +13 -72.0860596 +128 -75.1309814 +11 -73.6424179 +112 -73.1209106 +31 -65.0838242 +14 -64.2643051 +32 -69.4516983 +106 -66.4129791 +3 -62.6622734 +87 -57.6153526 +127 -55.5868149 diff --git a/tests/fixtures/persistence/v0.8.3/hnsw5_ip.index b/tests/fixtures/persistence/v0.8.3/hnsw5_ip.index new file mode 100644 index 00000000..2e69c115 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/hnsw5_ip.index differ diff --git a/tests/fixtures/persistence/v0.8.3/hnsw9_l2.expected b/tests/fixtures/persistence/v0.8.3/hnsw9_l2.expected new file mode 100644 index 00000000..38c4ce7a --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/hnsw9_l2.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +92 420.362671 +46 471.622772 +128 481.912445 +43 495.599579 +3 497.779449 +31 412.638641 +112 416.755554 +129 460.20517 +71 463.54837 +128 464.295135 +32 423.867615 +3 448.949097 +87 462.59491 +12 477.650421 +48 481.020172 diff --git a/tests/fixtures/persistence/v0.8.3/hnsw9_l2.index b/tests/fixtures/persistence/v0.8.3/hnsw9_l2.index new file mode 100644 index 00000000..194e5d64 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/hnsw9_l2.index differ diff --git a/tests/fixtures/persistence/v0.8.3/ivf1_flat_matrix_l2.expected b/tests/fixtures/persistence/v0.8.3/ivf1_flat_matrix_l2.expected new file mode 100644 index 00000000..c7b534ad --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/ivf1_flat_matrix_l2.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +83 361.387726 +36 367.132172 +46 399.009552 +50 400.204742 +65 405.736511 +31 288.149902 +112 353.039429 +129 425.505188 +40 444.491913 +95 452.522827 +84 289.657135 +72 334.618195 +87 337.608826 +50 342.41745 +100 356.274933 diff --git a/tests/fixtures/persistence/v0.8.3/ivf1_flat_matrix_l2.index b/tests/fixtures/persistence/v0.8.3/ivf1_flat_matrix_l2.index new file mode 100644 index 00000000..269d4966 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/ivf1_flat_matrix_l2.index differ diff --git a/tests/fixtures/persistence/v0.8.3/ivf32_hnsw_ip.expected b/tests/fixtures/persistence/v0.8.3/ivf32_hnsw_ip.expected new file mode 100644 index 00000000..d0f55dd7 --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/ivf32_hnsw_ip.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +128 -94.8744507 +92 -82.2540436 +46 -80.8800812 +50 -71.6611023 +5 -69.8279419 +112 -73.9705811 +128 -73.9433899 +11 -71.4328156 +14 -64.1144104 +31 -59.65448 +32 -72.5804443 +106 -66.686676 +3 -59.7108917 +87 -57.427536 +127 -53.9758759 diff --git a/tests/fixtures/persistence/v0.8.3/ivf32_hnsw_ip.index b/tests/fixtures/persistence/v0.8.3/ivf32_hnsw_ip.index new file mode 100644 index 00000000..2929eafe Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/ivf32_hnsw_ip.index differ diff --git a/tests/fixtures/persistence/v0.8.3/ivf32_hnsw_ip.index.hnsw b/tests/fixtures/persistence/v0.8.3/ivf32_hnsw_ip.index.hnsw new file mode 100644 index 00000000..59f96bf2 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/ivf32_hnsw_ip.index.hnsw differ diff --git a/tests/fixtures/persistence/v0.8.3/ivf4_flat_l2.expected b/tests/fixtures/persistence/v0.8.3/ivf4_flat_l2.expected new file mode 100644 index 00000000..2adbbb05 --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/ivf4_flat_l2.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +92 421.434753 +46 460.21283 +128 473.936279 +13 483.826019 +3 495.024872 +31 406.883026 +112 410.993774 +129 454.92041 +71 455.203857 +128 456.861572 +32 436.845459 +3 452.321411 +87 459.564636 +12 475.835114 +48 481.350098 diff --git a/tests/fixtures/persistence/v0.8.3/ivf4_flat_l2.index b/tests/fixtures/persistence/v0.8.3/ivf4_flat_l2.index new file mode 100644 index 00000000..581ded95 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/ivf4_flat_l2.index differ diff --git a/tests/fixtures/persistence/v0.8.3/ivf4_hnsw_l2.expected b/tests/fixtures/persistence/v0.8.3/ivf4_hnsw_l2.expected new file mode 100644 index 00000000..e91a5bbe --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/ivf4_hnsw_l2.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +92 431.944397 +128 470.540405 +46 476.507202 +3 490.621582 +13 493.916656 +112 412.643494 +31 416.188904 +128 453.381653 +129 459.386047 +90 459.440277 +32 424.608551 +3 457.105835 +87 463.820557 +48 464.434967 +12 477.980652 diff --git a/tests/fixtures/persistence/v0.8.3/ivf4_hnsw_l2.index b/tests/fixtures/persistence/v0.8.3/ivf4_hnsw_l2.index new file mode 100644 index 00000000..8e1426b2 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/ivf4_hnsw_l2.index differ diff --git a/tests/fixtures/persistence/v0.8.3/ivf4_hnsw_l2.index.hnsw b/tests/fixtures/persistence/v0.8.3/ivf4_hnsw_l2.index.hnsw new file mode 100644 index 00000000..06c069f5 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/ivf4_hnsw_l2.index.hnsw differ diff --git a/tests/fixtures/persistence/v0.8.3/ivf7_flatrabitq_ip.expected b/tests/fixtures/persistence/v0.8.3/ivf7_flatrabitq_ip.expected new file mode 100644 index 00000000..3f2607f1 --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/ivf7_flatrabitq_ip.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +128 -94.6611481 +92 -82.9637375 +46 -81.7006149 +13 -72.5498047 +50 -71.9551773 +112 -74.3031998 +128 -74.0310822 +11 -72.1495819 +14 -64.1889343 +71 -59.4696617 +32 -72.5847778 +106 -67.2913208 +3 -59.6324234 +87 -56.3743286 +127 -54.9878922 diff --git a/tests/fixtures/persistence/v0.8.3/ivf7_flatrabitq_ip.index b/tests/fixtures/persistence/v0.8.3/ivf7_flatrabitq_ip.index new file mode 100644 index 00000000..da3ca3b9 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/ivf7_flatrabitq_ip.index differ diff --git a/tests/fixtures/persistence/v0.8.3/manifest.json b/tests/fixtures/persistence/v0.8.3/manifest.json new file mode 100644 index 00000000..15790eb6 --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/manifest.json @@ -0,0 +1,36 @@ +{ + "source_tag": "v0.8.3", + "source_commit": "dabfdb8c200a6f155df4edd0b9c825e6728b14a1", + "writer": "tests/fixtures/persistence/v0.8.3/generate.cpp (IVF-RaBitQ) and generate_graphs.cpp (HNSW-RaBitQ, SymphonyQG)", + "producer": "Linux x86-64, g++ 14.3.0, Release, one thread, native optimization OFF", + "byte_order": "little", + "size_t_bytes": 8, + "sha256": { + "v0.8.3/hnsw1_l2.expected": "00af920380168a49f0739b9f180a764c3fac3bbc3a7b31b6acbbf0ecd621f62a", + "v0.8.3/hnsw1_l2.index": "cb89a17e02c04233b599d9077d81b030586eba9bc32fe1ead9ea30271be0b7e1", + "v0.8.3/hnsw5_ip.expected": "199106ce08b2667df5b4cbc47bc06309564637e6790cf74448324c27e377b2f2", + "v0.8.3/hnsw5_ip.index": "5780fb6ad6843df24cfa0b9a545e5251505e50040e3ac73593e1f897d29b5967", + "v0.8.3/hnsw9_l2.expected": "ef0fa91553a6effcdbba13bf0fc5263fa75a3af5ac37acc948809ef051d93ab3", + "v0.8.3/hnsw9_l2.index": "116d2496789bfb392f0cb03aca6d0163985298c71aa0a8259f4e3092efc8a47f", + "v0.8.3/ivf1_flat_matrix_l2.expected": "f1616ce0a6f59a57984e62eda1a16b92a54f603434e163ce5a2f33c336a0a252", + "v0.8.3/ivf1_flat_matrix_l2.index": "c864c367f4eaac449984694402d42ba81706391c61405f3fcdcf233a7f4d92f4", + "v0.8.3/ivf32_hnsw_ip.expected": "8280cf8c79e396c94ddb0a3bf77d69f45f7261b1ea71e90e5522395286376f89", + "v0.8.3/ivf32_hnsw_ip.index": "61935ec5f4a5c3b2aeb990ccb9bafa25e7323658c65a20f0344e79e0c8df7194", + "v0.8.3/ivf32_hnsw_ip.index.hnsw": "320c644a727fa94f7e5ced3821d8fcca6562c22904acfbd158fc5b3f6c3c5eb4", + "v0.8.3/ivf4_flat_l2.expected": "c2d41f4b7c18332d94f01ebdf493658b8f468b5c517bc9961cbfea2ffdb33dbe", + "v0.8.3/ivf4_flat_l2.index": "d324a9df508c32fee0a862868d3f7b2fcd7acc999db4e7bd5761b4f6cd1d07be", + "v0.8.3/ivf4_hnsw_l2.expected": "4235ed49a3d008e34b646358ec0de0649c63122b3f5f60e813df8a0d494c3a88", + "v0.8.3/ivf4_hnsw_l2.index": "88b5b7d1b7e5be99506421cd36ad6bb1785490446f55788d5b39a1069d5b304c", + "v0.8.3/ivf4_hnsw_l2.index.hnsw": "1797edb3eb1fb16a20684d8873cbc2c0713bdf1623ff31f366e8f016f5e616cf", + "v0.8.3/ivf7_flatrabitq_ip.expected": "c602bd98139b597532067094b001ff992d4498ea07ad7ddaf6762b8261f67d83", + "v0.8.3/ivf7_flatrabitq_ip.index": "54c11bcb813395eed099c8e8f65cdd96d97ca1ed684703e66abe27fd06887929", + "v0.8.3/qg0_ip.expected": "3ecb6604dce2108556e5e48130d4e9902d7ef7f0ef81aba39a3caa40b62230fe", + "v0.8.3/qg0_ip.index": "5b48546797a13825d27f0296b3c21c9c79fcf80ce44121fa5e9540358bb5762f", + "v0.8.3/qg4_l2.expected": "e973d22674ccde5c5afe4cf916e7c9271a79e8b6230c2da11bbc304218f488cf", + "v0.8.3/qg4_l2.index": "9f61d0570b75e67f101bc8e50e9af9db35194fa51187317d92a831db1a00d1e1", + "v0.8.3/qg4_matrix_l2.expected": "8f7d74551eb36226a89b195d23b69eac2d7194fbfe8f47b73caee8364757730e", + "v0.8.3/qg4_matrix_l2.index": "f60b01bb2e6b670eb63a0ed108346cc5c3cde0614b5178462b9a117b3c563484", + "v0.8.3/qg8_ip.expected": "46b4b3bd3e0b6ad1ec7e00df843e679977a7092fd3dcca78eb85a5239293cfe5", + "v0.8.3/qg8_ip.index": "5519fecdbde73515940de91d213ca4ab7ad3c486cd97a773fe02927c121ce879" + } +} diff --git a/tests/fixtures/persistence/v0.8.3/qg0_ip.expected b/tests/fixtures/persistence/v0.8.3/qg0_ip.expected new file mode 100644 index 00000000..750cfd52 --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/qg0_ip.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +128 -94.8744507 +92 -82.2540436 +46 -80.8800812 +13 -71.9051971 +50 -71.6611023 +112 -73.9705811 +128 -73.9433899 +11 -71.4328156 +14 -64.1144104 +31 -59.65448 +32 -72.5804443 +106 -66.686676 +3 -59.7108917 +87 -57.427536 +127 -53.9758759 diff --git a/tests/fixtures/persistence/v0.8.3/qg0_ip.index b/tests/fixtures/persistence/v0.8.3/qg0_ip.index new file mode 100644 index 00000000..c25390c8 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/qg0_ip.index differ diff --git a/tests/fixtures/persistence/v0.8.3/qg4_l2.expected b/tests/fixtures/persistence/v0.8.3/qg4_l2.expected new file mode 100644 index 00000000..24704b89 --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/qg4_l2.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +92 426.139954 +46 474.039948 +128 487.332336 +43 495.476349 +3 501.042847 +31 412.692932 +112 417.631104 +71 458.311768 +129 460.74765 +14 467.755432 +32 428.727905 +87 450.430389 +3 451.144958 +12 467.519897 +0 474.705017 diff --git a/tests/fixtures/persistence/v0.8.3/qg4_l2.index b/tests/fixtures/persistence/v0.8.3/qg4_l2.index new file mode 100644 index 00000000..74b0b7b6 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/qg4_l2.index differ diff --git a/tests/fixtures/persistence/v0.8.3/qg4_matrix_l2.expected b/tests/fixtures/persistence/v0.8.3/qg4_matrix_l2.expected new file mode 100644 index 00000000..8529804a --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/qg4_matrix_l2.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +92 422.084595 +46 478.977295 +128 487.529175 +43 493.764648 +3 502.006531 +31 412.601593 +112 417.63501 +128 451.895691 +71 463.59549 +129 465.134125 +32 430.10022 +3 447.459839 +87 459.170197 +48 465.867981 +106 471.233765 diff --git a/tests/fixtures/persistence/v0.8.3/qg4_matrix_l2.index b/tests/fixtures/persistence/v0.8.3/qg4_matrix_l2.index new file mode 100644 index 00000000..3e1736ad Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/qg4_matrix_l2.index differ diff --git a/tests/fixtures/persistence/v0.8.3/qg8_ip.expected b/tests/fixtures/persistence/v0.8.3/qg8_ip.expected new file mode 100644 index 00000000..328daa94 --- /dev/null +++ b/tests/fixtures/persistence/v0.8.3/qg8_ip.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +128 -94.953392 +92 -82.2493134 +46 -81.0288467 +13 -71.8557434 +50 -71.6811447 +112 -74.0330429 +128 -73.8238602 +11 -71.6039658 +14 -64.0551224 +31 -59.4229736 +32 -72.4512024 +106 -66.6851273 +3 -59.6906853 +87 -57.5236473 +127 -53.835289 diff --git a/tests/fixtures/persistence/v0.8.3/qg8_ip.index b/tests/fixtures/persistence/v0.8.3/qg8_ip.index new file mode 100644 index 00000000..09e1ffa2 Binary files /dev/null and b/tests/fixtures/persistence/v0.8.3/qg8_ip.index differ diff --git a/tests/python/format2_file.py b/tests/python/format2_file.py new file mode 100644 index 00000000..8887df5e --- /dev/null +++ b/tests/python/format2_file.py @@ -0,0 +1,93 @@ +"""Read, patch and convert index format 2 files in tests (docs/formats.md).""" + +from __future__ import annotations + +import struct +import zlib + +MAGIC = b"RABQINDX" +ROTATOR = 1 +IVF_PARAMS = 256 +IVF_CLUSTER_SIZES = 257 +IVF_ROUTING = {1: 258, 2: 259, 3: 260} # by stored initializer: flat, flat_rabitq, hnsw +IVF_HNSW_ROUTING = 260 +IVF_CODES = 261 +IVF_RERANK = 262 +IVF_IDS = 263 +HNSW_PARAMS = 512 +QG_PARAMS = 768 +QG_CENTROID = 769 +QG_RECORDS = 770 + + +def sections(data: bytes) -> dict[int, tuple[int, int]]: + """Map each section kind of a format-2 file to its (offset, size).""" + table, count = struct.unpack_from(" bytes: + offset, size = sections(data)[kind] + return bytes(data[offset : offset + size]) + + +def padded_dim(data: bytes) -> int: + return struct.unpack_from(" None: + """Recompute section and header checksums after a same-size edit in place.""" + table, count = struct.unpack_from(" tuple[bytes, bytes | None]: + """Convert an IVF-RaBitQ format-2 file to the layout releases up to 0.8 wrote. + + Returns the RABQIDX1 version 2 index and, for HNSW routing, its ``.hnsw`` + sidecar. Tests use it to build earlier layouts that no fixture covers. + """ + assert data[:8] == MAGIC + present = sections(data) + metric = struct.unpack_from("= 0][:3]) == 3 + assert len(index) == 7 + + @pytest.mark.parametrize("kind", KINDS) def test_len_counts_live_points(kind, base_data): index = new_index(kind) @@ -496,9 +514,11 @@ def make(): return index implicit = make() - implicit.add(base_data, num_threads=1) + implicit.add(base_data, num_threads=1) # trains with the add's single thread explicit = make() - explicit.train(base_data) + # The saved file includes HNSW routing, whose graph is reproducible only when + # it is built with one thread. + explicit.train(base_data, num_threads=1) explicit.add(base_data, num_threads=1) assert saved_bytes(implicit, tmp_path / "implicit.index") == saved_bytes( explicit, tmp_path / "explicit.index" @@ -541,3 +561,33 @@ def test_empty_first_add_requires_training(kind, base_data): ): index.add(base_data[:0]) assert not index.is_trained + + +@pytest.mark.parametrize("kind", KINDS + ["qg"]) +def test_unknown_read_only_features_allow_search_only(kind, base_data, tmp_path): + if kind == "qg": + index = SymqgIndex(DIM, ef_construction=40, seed=2) + index.add(base_data) + else: + index = new_index(kind) + index.add(base_data, num_threads=1) + path = tmp_path / "read_only.index" + index.save(path) + payload = bytearray(path.read_bytes()) + payload[24:32] = (1).to_bytes(8, "little") # an unknown ro_compat feature bit + format2_file.reseal(payload) + path.write_bytes(payload) + + loaded = type(index).load(path) + np.testing.assert_array_equal( + search(loaded, base_data[:5], 3)[0], search(index, base_data[:5], 3)[0] + ) + message = "cannot update, so it cannot save$" + with pytest.raises(RuntimeError, match=message): + loaded.save(tmp_path / "copy.index") + if kind != "qg": + # HNSW-RaBitQ grows its capacity before adding, which is refused first. + with pytest.raises( + RuntimeError, match="cannot update, so it cannot (add|resize)$" + ): + loaded.add(base_data[:1]) diff --git a/tests/python/test_txtai_compat.py b/tests/python/test_txtai_compat.py index 91257889..089fba84 100644 --- a/tests/python/test_txtai_compat.py +++ b/tests/python/test_txtai_compat.py @@ -13,6 +13,9 @@ import pytest from rabitqlib import HnswIndex, IvfIndex +# txtai's filter for missing results. Since 0.9 a missing result is -1, which this no +# longer matches; these searches never leave slots unfilled. Follow txtai's filter once +# its first compatible release exists (#213). SENTINEL = 0xFFFFFFFF diff --git a/tests/python/typing/public_api.py b/tests/python/typing/public_api.py index a0b27263..f835bbdb 100644 --- a/tests/python/typing/public_api.py +++ b/tests/python/typing/public_api.py @@ -36,16 +36,16 @@ def check(x: npt.NDArray[np.float32], labels: npt.NDArray[np.uint32]) -> None: index.set_centroids(x[:2]) assert_type(index.is_trained, bool) assert_type(index.seed, int | None) - assert_type(index.add(x, cluster_ids=labels), npt.NDArray[np.uint32]) + assert_type(index.add(x, cluster_ids=labels), npt.NDArray[np.int64]) assert_type(index.remove([0, 1]), int) assert_type(index.dim, int) assert_type(len(index), int) index.save("index.bin") ids, distances = ivf.search(x, k=5, nprobe=2) - assert_type(ids, npt.NDArray[np.uint32]) + assert_type(ids, npt.NDArray[np.int64]) assert_type(distances, npt.NDArray[np.float32]) assert_type( - hnsw.search(x, 5), tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]] + hnsw.search(x, 5), tuple[npt.NDArray[np.int64], npt.NDArray[np.float32]] ) graph.train(x) graph.add(x, num_threads=1) @@ -54,7 +54,7 @@ def check(x: npt.NDArray[np.float32], labels: npt.NDArray[np.uint32]) -> None: assert_type(len(graph), int) assert_type( graph.search(x, 5), - tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]], + tuple[npt.NDArray[np.int64], npt.NDArray[np.float32]], ) assert_type(IvfIndex.load("index.bin"), IvfIndex) assert_type(HnswIndex.load("index.bin"), HnswIndex) @@ -135,7 +135,7 @@ def check_numpy_scalars() -> None: ids, distances = ivf.search( x[:2], k=np.int64(3), nprobe=np.int64(1), high_accuracy=np.bool_(True) ) - assert_type(ids, npt.NDArray[np.uint32]) + assert_type(ids, npt.NDArray[np.int64]) assert_type(distances, npt.NDArray[np.float32]) hnsw.resize(np.int64(100)) hnsw.search(x[:2], np.int64(3), ef=np.int64(10)) diff --git a/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp b/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp index 17daef67..6d0a007a 100644 --- a/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp +++ b/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp @@ -5,16 +5,17 @@ #include #include #include -#include +#include #include -#include #include +#include #include #include #include #include #include +#include "format2_file.hpp" #include "rabitqlib/core/metric.hpp" #include "rabitqlib/core/rotator.hpp" #include "rabitqlib/index/ivf/initializer.hpp" @@ -127,11 +128,15 @@ TEST(IvfInitializerTest, ReloadChangesRoutingAndFailedLoadPreservesExistingIndex EXPECT_EQ(actual_ids, expected_ids); EXPECT_EQ(actual_distances, expected_distances); { - std::fstream file(path, std::ios::binary | std::ios::in | std::ios::out); - // A persisted v2 file must contain a concrete initializer, never Auto. - file.seekp(24); - const uint32_t invalid = 0; - file.write(reinterpret_cast(&invalid), sizeof(invalid)); + // A persisted file must contain a concrete initializer, never Auto. + using rabitq_test::format2::IndexKind; + std::string params = rabitq_test::read_format2_section( + path, IndexKind::Ivf, rabitq_test::kIvfSections, 256 + ); + std::memset(params.data() + 8, 0, sizeof(uint32_t)); + rabitq_test::replace_format2_section( + path, IndexKind::Ivf, rabitq_test::kIvfSections, 256, params + ); } EXPECT_THROW(target.load(path.c_str()), std::runtime_error); EXPECT_EQ(target.initializer_type(), type); @@ -223,19 +228,18 @@ TEST(IvfFlatRaBitQInitializerTest, PersistsCodesAndSupportsConcurrentSearch) { std::vector centroids(kDim * kCount), queries(8 * kDim); std::generate(centroids.begin(), centroids.end(), [&] { return normal(generator); }); std::generate(queries.begin(), queries.end(), [&] { return normal(generator); }); - const rabitq_test::TempFile path_file("flat_rabitq_initializer.bin"); - const std::string& path = path_file.path(); for (auto metric : {METRIC_L2, METRIC_IP}) { FlatRaBitQInitializer index(kDim, kCount, metric), loaded(kDim, kCount, metric); index.add_vectors(centroids.data(), 2); + std::stringstream stored; + index.write(stored); + const std::string bytes = stored.str(); { - std::ofstream out(path, std::ios::binary); - index.save(out, nullptr); - } - { - std::ifstream input(path, std::ios::binary); - loaded.load(input, nullptr); + FlatRaBitQInitializer wrong_size(kDim, kCount, metric); + std::stringstream copy(bytes); + EXPECT_THROW(wrong_size.read(copy, bytes.size() - 1), std::runtime_error); } + loaded.read(stored, bytes.size()); for (size_t i = 0; i < kCount; ++i) { EXPECT_TRUE( std::equal(index.centroid(i), index.centroid(i) + kDim, loaded.centroid(i)) @@ -259,15 +263,9 @@ TEST(IvfFlatRaBitQInitializerTest, PersistsCodesAndSupportsConcurrentSearch) { EXPECT_FLOAT_EQ(actual[i][j].distance, expected[i][j].distance); } } - std::ifstream before(path, std::ios::binary); - const std::string bytes{std::istreambuf_iterator(before), {}}; - before.close(); - { - std::ofstream out(path, std::ios::binary); - loaded.save(out, nullptr); - } - std::ifstream after(path, std::ios::binary); - EXPECT_EQ((std::string{std::istreambuf_iterator(after), {}}), bytes); + std::ostringstream resaved; + loaded.write(resaved); + EXPECT_EQ(resaved.str(), bytes); } } } // namespace diff --git a/tests/unit/rabitqlib/index/format2_test.cpp b/tests/unit/rabitqlib/index/format2_test.cpp new file mode 100644 index 00000000..362b7b3b --- /dev/null +++ b/tests/unit/rabitqlib/index/format2_test.cpp @@ -0,0 +1,607 @@ +#include "format2.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "crc32.hpp" + +namespace { +namespace f2 = rabitqlib::index_detail::format2; +using rabitqlib::detail::crc32; + +constexpr uint32_t kParams = 256; +constexpr uint32_t kBulk = 261; +constexpr uint32_t kEmpty = 263; +const std::vector kKnown = {kParams, kBulk, kEmpty}; + +uint32_t get_u32(const std::string& bytes, size_t offset) { + uint32_t value = 0; + for (size_t i = 0; i < 4; ++i) { + value |= static_cast(static_cast(bytes[offset + i])) << (8 * i); + } + return value; +} + +uint64_t get_u64(const std::string& bytes, size_t offset) { + uint64_t value = 0; + for (size_t i = 0; i < 8; ++i) { + value |= static_cast(static_cast(bytes[offset + i])) << (8 * i); + } + return value; +} + +void put_u32(std::string& bytes, size_t offset, uint32_t value) { + for (size_t i = 0; i < 4; ++i) { + bytes[offset + i] = static_cast(value >> (8 * i)); + } +} + +void put_u64(std::string& bytes, size_t offset, uint64_t value) { + for (size_t i = 0; i < 8; ++i) { + bytes[offset + i] = static_cast(value >> (8 * i)); + } +} + +// Recompute the header checksum after a test edits the header or the table, so that +// the check under test fires rather than the checksum. +void reseal(std::string& bytes) { + const uint32_t header_bytes = get_u32(bytes, 12); + const uint64_t table_offset = get_u64(bytes, 96); + const uint32_t count = get_u32(bytes, 104); + std::string header = bytes.substr(0, header_bytes); + put_u32(header, 108, 0); + uint32_t crc = crc32(header.data(), header.size()); + crc = crc32(bytes.data() + table_offset, size_t{32} * count, crc); + put_u32(bytes, 108, crc); +} + +f2::Header test_header() { + f2::Header header; + header.index_kind = f2::IndexKind::Ivf; + header.metric = f2::Metric::InnerProduct; + header.dim = 100; + header.padded_dim = 128; + header.code_bits = 7; + header.raw_type = f2::RawType::None; + header.id_bits = 32; + header.seed = 42; + header.num_points = 3; + return header; +} + +const std::string kParamBytes = "params!!"; + +std::string bulk_bytes() { + std::string bulk(1000, '\0'); + for (size_t i = 0; i < bulk.size(); ++i) { + bulk[i] = static_cast(i * 7); + } + return bulk; +} + +// A checksummed parameter section, an unchecksummed bulk section written in two +// pieces, and an empty section. Layout: header 0-127, table 128-223, parameters at +// 256, bulk at 320 (1000 bytes), the empty section at 1344, which ends the file. +std::vector test_sections() { + const std::string bulk = bulk_bytes(); + return { + {kParams, + f2::kSectionRequired | f2::kSectionChecksummed, + kParamBytes.size(), + [](f2::SectionSink& sink) { sink.write(kParamBytes.data(), kParamBytes.size()); }}, + {kBulk, + f2::kSectionRequired, + bulk.size(), + [bulk](f2::SectionSink& sink) { + sink.write(bulk.data(), 400); + sink.write(bulk.data() + 400, bulk.size() - 400); + }}, + {kEmpty, f2::kSectionRequired, 0, nullptr}}; +} + +std::string write(const f2::Header& header, const std::vector& sections) { + std::ostringstream out; + f2::write_file(out, header, sections, "0.9.0"); + return out.str(); +} + +std::string valid_file() { return write(test_header(), test_sections()); } + +// Byte offset of section table entry `i` when the table follows a 128-byte header. +size_t entry(size_t i) { return 128 + (32 * i); } + +f2::Reader open_reader(std::istream& stream, f2::IndexKind kind = f2::IndexKind::Ivf) { + return f2::Reader(stream, kind, kKnown); +} + +void expect_format_error( + const std::string& bytes, + const std::string& message, + f2::IndexKind kind = f2::IndexKind::Ivf +) { + std::istringstream stream(bytes); + try { + f2::Reader reader(stream, kind, kKnown); + ADD_FAILURE() << "expected FormatError containing: " << message; + } catch (const f2::FormatError& error) { + EXPECT_NE(std::string(error.what()).find(message), std::string::npos) + << "message: " << error.what(); + } +} + +// Input and output buffers that cannot seek, like pipes. +class SequentialInput : public std::streambuf { + public: + explicit SequentialInput(std::string bytes) : bytes_(std::move(bytes)) { + setg(bytes_.data(), bytes_.data(), bytes_.data() + bytes_.size()); + } + + private: + std::string bytes_; +}; + +class SequentialOutput : public std::streambuf { + public: + std::string bytes; + + private: + int_type overflow(int_type c) override { + if (!traits_type::eq_int_type(c, traits_type::eof())) { + bytes.push_back(traits_type::to_char_type(c)); + } + return traits_type::not_eof(c); + } + std::streamsize xsputn(const char* data, std::streamsize count) override { + bytes.append(data, static_cast(count)); + return count; + } +}; + +TEST(Format2ContainerTest, RoundTripsHeaderAndSections) { + const std::string bytes = valid_file(); + std::istringstream stream(bytes); + f2::Reader reader = open_reader(stream); + + const f2::Header& header = reader.header(); + EXPECT_EQ(header.index_kind, f2::IndexKind::Ivf); + EXPECT_EQ(header.metric, f2::Metric::InnerProduct); + EXPECT_EQ(header.dim, 100U); + EXPECT_EQ(header.padded_dim, 128U); + EXPECT_EQ(header.code_bits, 7U); + EXPECT_EQ(header.raw_type, f2::RawType::None); + EXPECT_EQ(header.id_bits, 32U); + ASSERT_TRUE(header.seed.has_value()); + EXPECT_EQ(*header.seed, 42U); + EXPECT_EQ(header.num_points, 3U); + EXPECT_FALSE(reader.read_only()); + EXPECT_EQ(reader.writer_version(), "0.9.0"); + + const f2::Section* params = reader.find(kParams); + const f2::Section* bulk = reader.find(kBulk); + const f2::Section* empty = reader.find(kEmpty); + ASSERT_NE(params, nullptr); + ASSERT_NE(bulk, nullptr); + ASSERT_NE(empty, nullptr); + EXPECT_EQ(reader.find(999), nullptr); + EXPECT_EQ(params->offset, 256U); + EXPECT_EQ(bulk->offset, 320U); + EXPECT_EQ(bulk->crc32, 0U); + EXPECT_EQ(empty->offset, 1344U); // 320 + 1000, padded to 64 + EXPECT_EQ(bytes.size(), 1344U); // the empty section ends the file + + std::string read_params(params->size, '\0'); + reader.read(*params, 0, read_params.data(), read_params.size()); + EXPECT_EQ(read_params, kParamBytes); + std::string read_bulk(bulk->size - 10, '\0'); + reader.read(*bulk, 10, read_bulk.data(), read_bulk.size()); + EXPECT_EQ(read_bulk, bulk_bytes().substr(10)); + EXPECT_THROW( + reader.read(*bulk, 10, read_bulk.data(), read_bulk.size() + 1), std::out_of_range + ); +} + +TEST(Format2ContainerTest, WritesIdenticalBytesTwice) { + EXPECT_EQ(valid_file(), valid_file()); +} + +TEST(Format2ContainerTest, WritesAnUnknownSeed) { + f2::Header header = test_header(); + header.seed.reset(); + const std::string bytes = write(header, {{kEmpty, f2::kSectionRequired, 0, nullptr}}); + EXPECT_EQ(get_u32(bytes, 76), 0U); + EXPECT_EQ(get_u64(bytes, 80), 0U); + std::istringstream stream(bytes); + EXPECT_FALSE(open_reader(stream).header().seed.has_value()); +} + +TEST(Format2ContainerTest, WritesAndReadsWithoutSeeking) { + SequentialOutput output_buffer; + std::ostream output(&output_buffer); + f2::write_file(output, test_header(), test_sections(), "0.9.0"); + EXPECT_EQ(output_buffer.bytes, valid_file()); + + SequentialInput input_buffer(valid_file()); + std::istream input(&input_buffer); + EXPECT_THROW(f2::has_magic(input, f2::IndexKind::Ivf), f2::FormatError); + try { + f2::Reader reader(input, f2::IndexKind::Ivf, kKnown); + ADD_FAILURE() << "expected FormatError"; + } catch (const f2::FormatError& error) { + EXPECT_STREQ(error.what(), "IVF index file input must be seekable"); + } +} + +TEST(Format2ContainerTest, DetectsTheMagicAndRestoresThePosition) { + std::istringstream good(valid_file()); + EXPECT_TRUE(f2::has_magic(good, f2::IndexKind::Ivf)); + EXPECT_EQ(good.tellg(), std::streampos(0)); + std::istringstream legacy(std::string("RABQIDX1") + std::string(200, '\0')); + EXPECT_FALSE(f2::has_magic(legacy, f2::IndexKind::Ivf)); + + std::istringstream short_file("RABQ"); + short_file.exceptions(std::ios::failbit | std::ios::badbit); + EXPECT_FALSE(f2::has_magic(short_file, f2::IndexKind::Ivf)); + EXPECT_EQ(short_file.tellg(), std::streampos(0)); +} + +TEST(Format2ContainerTest, IgnoresTheStreamExceptionMask) { + // Truncated in the padding before the last section: the header and the table are + // read, then a section check fails. + std::istringstream stream(valid_file().substr(0, 1336)); + stream.exceptions(std::ios::failbit | std::ios::badbit | std::ios::eofbit); + EXPECT_THROW(open_reader(stream), f2::FormatError); + std::istringstream valid(valid_file()); + valid.exceptions(std::ios::failbit | std::ios::badbit | std::ios::eofbit); + EXPECT_NO_THROW(open_reader(valid)); +} + +TEST(Format2ContainerTest, WriterRejectsInvalidArgumentsBeforeWriting) { + const f2::SectionWrite empty{kEmpty, f2::kSectionRequired, 0, nullptr}; + const auto expect_rejected = [](const f2::Header& header, + const std::vector& sections, + const char* version = "0.9.0") { + std::ostringstream out; + EXPECT_THROW(f2::write_file(out, header, sections, version), std::invalid_argument); + EXPECT_TRUE(out.str().empty()); + }; + + f2::Header header = test_header(); + header.padded_dim = 100; + expect_rejected(header, {empty}); + header = test_header(); + header.index_kind = static_cast(4); + expect_rejected(header, {empty}); + for (const int mask : {0, 1, 2}) { + header = test_header(); + (mask == 0 ? header.incompat_features + : mask == 1 ? header.ro_compat_features + : header.compat_features) = 1; + expect_rejected(header, {empty}); + } + expect_rejected(test_header(), {}); + expect_rejected(test_header(), {empty, empty}); + expect_rejected(test_header(), {{0, f2::kSectionRequired, 0, nullptr}}); + expect_rejected(test_header(), {{kEmpty, 4, 0, nullptr}}); + expect_rejected( + test_header(), + {{kBulk, 0, 1, nullptr}, + {kEmpty, 0, std::numeric_limits::max() - 10, nullptr}} + ); + expect_rejected(test_header(), {empty}, "a version longer than 16"); + expect_rejected(test_header(), {empty}, "0.9\n"); +} + +TEST(Format2ContainerTest, WriterChecksSectionSizesAndPropagatesFailures) { + const auto two_bytes = [](f2::SectionSink& sink) { sink.write("ab", 2); }; + for (const uint32_t flags : {0U, f2::kSectionChecksummed}) { + EXPECT_THROW( + write(test_header(), {{kBulk, flags, 1, two_bytes}}), std::invalid_argument + ); + EXPECT_THROW( + write(test_header(), {{kBulk, flags, 3, two_bytes}}), std::invalid_argument + ); + } + const auto fails = [](f2::SectionSink&) { throw std::runtime_error("section failed"); }; + EXPECT_THROW(write(test_header(), {{kBulk, 0, 1, fails}}), std::runtime_error); + + std::ostringstream broken; + broken.setstate(std::ios::badbit); + try { + f2::write_file(broken, test_header(), test_sections(), "0.9.0"); + ADD_FAILURE() << "expected runtime_error"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "IVF index file cannot be written"); + } +} + +TEST(Format2ContainerTest, RejectsBrokenHeaders) { + const std::string good = valid_file(); + expect_format_error(good.substr(0, 127), "IVF index file is too short"); + + std::string bytes = good; + bytes[0] = 'X'; + expect_format_error(bytes, "IVF index file is not in format 2 (bad magic)"); + + bytes = good; + put_u32(bytes, 8, 3); + reseal(bytes); + expect_format_error(bytes, "format version 3; this release reads version 2"); + + for (const uint32_t size : {120U, 132U, 4096U}) { + bytes = good; + put_u32(bytes, 12, size); + expect_format_error(bytes, "invalid header size " + std::to_string(size)); + } + + bytes = good; + bytes[50] ^= 1; // dim, without resealing + expect_format_error(bytes, "header checksum mismatch"); + + bytes = good; + put_u32(bytes, 104, 0); + reseal(bytes); + expect_format_error(bytes, "has no sections"); + + for (const uint64_t table_offset : {uint64_t{132}, uint64_t{120}, uint64_t{1344}}) { + bytes = good; + put_u64(bytes, 96, table_offset); + expect_format_error(bytes, "section table lies outside the file"); + } + bytes = good; + put_u32(bytes, 104, 1000); + expect_format_error(bytes, "section table lies outside the file"); +} + +TEST(Format2ContainerTest, ChecksTheIndexKindBeforeTheSections) { + std::string bytes = valid_file(); + put_u32(bytes, entry(1), 777); // an unknown required section + reseal(bytes); + expect_format_error( + bytes, "HNSW index file holds an IVF-RaBitQ index", f2::IndexKind::Hnsw + ); + expect_format_error( + bytes, + "QuantizedGraph index file holds an IVF-RaBitQ index", + f2::IndexKind::SymphonyQg + ); + + const struct { + uint32_t kind; + const char* message; + } cases[] = { + {4, "IVF index file holds an IVF-SSD index"}, + {9, "IVF index file has unknown index kind 9"}, + {0, "IVF index file has unknown index kind 0"}, + }; + for (const auto& c : cases) { + bytes = valid_file(); + put_u32(bytes, 40, c.kind); + reseal(bytes); + expect_format_error(bytes, c.message); + } + + // The index's own header rules also run before the sections. + bytes = valid_file(); + put_u32(bytes, entry(1), 777); + reseal(bytes); + std::istringstream stream(bytes); + try { + f2::Reader reader(stream, f2::IndexKind::Ivf, kKnown, [](const f2::Header& header) { + if (header.id_bits != 64) { + throw f2::FormatError("IVF index file needs 64-bit IDs"); + } + }); + ADD_FAILURE() << "expected FormatError"; + } catch (const f2::FormatError& error) { + EXPECT_STREQ(error.what(), "IVF index file needs 64-bit IDs"); + } +} + +TEST(Format2ContainerTest, RejectsInvalidHeaderFields) { + const std::string good = valid_file(); + const struct { + size_t offset; + uint64_t value; + const char* message; + } cases[] = { + {44, 7, "IVF index file has unknown metric 7"}, + {44, 0, "IVF index file has unknown metric 0"}, + {48, 0, "IVF index file has dimension 0"}, + {56, 100, "has padded dimension 100, which must be a multiple of 32"}, + {56, + 96, + "has padded dimension 96, which must be a multiple of 32 from the dimension 100"}, + {56, 65568, "has padded dimension 65568, which must be a multiple of 32"}, + {64, 10, "has 10 code bits; at most 9 are supported"}, + {68, 5, "has unknown raw element type 5"}, + {72, 16, "has 16-bit IDs; only 32 and 64 are valid"}, + {76, 3, "has unknown header flags 0x3"}, + }; + for (const auto& c : cases) { + std::string bytes = good; + if (c.offset == 48 || c.offset == 56) { + put_u64(bytes, c.offset, c.value); + } else { + put_u32(bytes, c.offset, static_cast(c.value)); + } + reseal(bytes); + expect_format_error(bytes, c.message); + } + + std::string bytes = good; + put_u32(bytes, 76, 0); // seed 42 without seed_known + reseal(bytes); + expect_format_error(bytes, "has a seed without the seed_known flag"); +} + +TEST(Format2ContainerTest, AppliesFeatureFlags) { + const std::string good = valid_file(); + + std::string bytes = good; + put_u64(bytes, 16, 4); + reseal(bytes); + expect_format_error( + bytes, + "IVF index file uses incompatible features 0x4 that this release does not know; " + "written by rabitqlib 0.9.0" + ); + + bytes = good; + put_u64(bytes, 24, 1); + reseal(bytes); + std::istringstream read_only(bytes); + EXPECT_TRUE(open_reader(read_only).read_only()); + + bytes = good; + put_u64(bytes, 32, 0x80); + reseal(bytes); + std::istringstream compat(bytes); + const f2::Reader reader = open_reader(compat); + EXPECT_FALSE(reader.read_only()); + EXPECT_EQ(reader.header().compat_features, 0x80U); +} + +TEST(Format2ContainerTest, SkipsUnknownOptionalSectionsAndRejectsUnknownRequiredOnes) { + std::string bytes = valid_file(); + put_u32(bytes, entry(1), 777); + put_u32(bytes, entry(1) + 4, 0); // optional + reseal(bytes); + std::istringstream optional(bytes); + const f2::Reader reader = open_reader(optional); + EXPECT_EQ(reader.find(kBulk), nullptr); + EXPECT_EQ(reader.find(777), nullptr); + EXPECT_NE(reader.find(kParams), nullptr); + + put_u32(bytes, entry(1) + 4, f2::kSectionRequired); + reseal(bytes); + expect_format_error( + bytes, + "IVF index file section kind 777 is required but unknown to this release; written " + "by " + "rabitqlib 0.9.0" + ); + + // An unknown optional section is not read, so its checksum is not verified. + bytes = valid_file(); + put_u32(bytes, entry(0), 778); + put_u32(bytes, entry(0) + 4, f2::kSectionChecksummed); + reseal(bytes); + bytes[256] ^= 1; + std::istringstream unverified(bytes); + EXPECT_EQ(open_reader(unverified).find(kParams), nullptr); +} + +// The valid file with a 136-byte header whose extra bytes come from a later version. +std::string longer_header_file() { + const std::string good = valid_file(); + std::string bytes = good.substr(0, 128) + std::string(8, '\x5A'); + put_u32(bytes, 12, 136); + put_u64(bytes, 96, 136); + bytes += good.substr(128, 96); // the table + bytes.resize(256, '\0'); // sections start at 256 in both layouts + bytes += good.substr(256); + reseal(bytes); + return bytes; +} + +TEST(Format2ContainerTest, LoadsALongerHeaderAndChecksumsItsExtension) { + std::string bytes = longer_header_file(); + std::istringstream stream(bytes); + const f2::Reader reader = open_reader(stream); + EXPECT_EQ(reader.header().dim, 100U); + EXPECT_NE(reader.find(kBulk), nullptr); + + bytes[130] ^= 1; + expect_format_error(bytes, "header checksum mismatch"); +} + +TEST(Format2ContainerTest, LoadsATableAfterTheSections) { + const std::string good = valid_file(); + std::string bytes = good.substr(0, 128) + std::string(128, '\0') + good.substr(256); + put_u64(bytes, 96, bytes.size()); // 1344, a multiple of 8 + bytes += good.substr(128, 96); + reseal(bytes); + std::istringstream stream(bytes); + EXPECT_NE(open_reader(stream).find(kBulk), nullptr); + + // The empty last section now grows into the table. + put_u64(bytes, 1344 + 64 + 16, 50); + reseal(bytes); + expect_format_error(bytes, "section kind 263 overlaps the header or the section table"); +} + +TEST(Format2ContainerTest, RejectsBrokenSectionEntries) { + const std::string good = valid_file(); + const struct { + size_t offset; + uint64_t value; + bool wide; + const char* message; + } cases[] = { + {entry(0), 0, false, "IVF index file section kind 0 is reserved"}, + {entry(0) + 4, 0x7, false, "section kind 256 has unknown flags 0x7"}, + {entry(0) + 28, 1, false, "section kind 256 has a non-zero reserved field"}, + {entry(1) + 24, 5, false, "section kind 261 has a checksum but is not checksummed"}, + {entry(1) + 8, 328, true, "section kind 261 is not aligned to 64 bytes"}, + {entry(1) + 16, 1344, true, "section kind 261 lies outside the file"}, + {entry(1) + 16, + std::numeric_limits::max(), + true, + "261 lies outside the file"}, + {entry(1) + 8, + std::numeric_limits::max() - 63, + true, + "261 lies outside the file"}, + {entry(1) + 8, 256, true, "section kind 261 overlaps the previous section"}, + {entry(0) + 8, + 128, + true, + "section kind 256 overlaps the header or the section table"}, + {entry(0) + 8, + 0, + true, + "section kind 256 overlaps the header or the section table"}, + {entry(2), kBulk, false, "section kind 261 is repeated"}, + }; + for (const auto& c : cases) { + std::string bytes = good; + if (c.wide) { + put_u64(bytes, c.offset, c.value); + } else { + put_u32(bytes, c.offset, static_cast(c.value)); + } + reseal(bytes); + expect_format_error(bytes, c.message); + } +} + +TEST(Format2ContainerTest, RejectsTruncatedTrailingAndCorruptFiles) { + const std::string good = valid_file(); + expect_format_error( + good + std::string(64, '\0'), + "IVF index file is 1408 bytes, but its sections end at 1344" + ); + // Truncated inside the padding before the empty last section. + expect_format_error(good.substr(0, 1336), "section kind 263 lies outside the file"); + + std::string bytes = good; + bytes[256 + 2] ^= 1; // inside the checksummed parameters + expect_format_error(bytes, "section kind 256 checksum mismatch (corrupt file)"); + + bytes = good; + bytes[320 + 2] ^= 1; // bulk data is not checksummed + std::istringstream stream(bytes); + EXPECT_NO_THROW(open_reader(stream)); +} + +} // namespace diff --git a/tests/unit/rabitqlib/index/historical_persistence_test.cpp b/tests/unit/rabitqlib/index/historical_persistence_test.cpp index e8d8156b..4514c80a 100644 --- a/tests/unit/rabitqlib/index/historical_persistence_test.cpp +++ b/tests/unit/rabitqlib/index/historical_persistence_test.cpp @@ -3,16 +3,20 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include +#include #include #include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/format_error.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/ivf/ivf.hpp" @@ -34,7 +38,7 @@ struct HistoricalFixtureCase { size_t padded_dim; size_t count = 33; size_t ivf_nprobe = 1; - bool resave_identical = false; // the current writer must reproduce the file + ivf::InitializerType initializer = ivf::InitializerType::Flat; }; void PrintTo(const HistoricalFixtureCase& value, std::ostream* output) { @@ -82,6 +86,8 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail const auto exercise = [&](auto& index, const auto& search) { const size_t removed = name.find("qg") == 0 ? 0 : 1; const auto check_counts = [&] { + // Earlier layouts do not record the seed, and migrating them invents none. + EXPECT_FALSE(index.seed().has_value()); EXPECT_EQ(index.num_stored_points(), count); EXPECT_EQ(index.num_points(), count - removed); EXPECT_GE(index.max_elements(), index.num_stored_points()); @@ -94,15 +100,18 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail search(); compare(); index.save(output.string().c_str()); - if (GetParam().resave_identical) { - EXPECT_TRUE(read_bytes(output) == read_bytes(path)); - } + // Every index saves format 2. + EXPECT_EQ(read_bytes(output).substr(0, 8), "RABQINDX"); index.load(output.string().c_str()); check_counts(); search(); compare(); + // Reading format 2 back and saving again reproduces the file. + const std::string saved = read_bytes(output); + index.save(output.string().c_str()); + EXPECT_TRUE(read_bytes(output) == saved); std::filesystem::resize_file(output, std::filesystem::file_size(output) - 1); - EXPECT_THROW(index.load(output.string().c_str()), std::runtime_error); + EXPECT_THROW(index.load(output.string().c_str()), FormatError); check_counts(); search(); compare(); @@ -112,7 +121,7 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail ivf::IVF index; exercise(index, [&] { EXPECT_EQ(index.padded_dim(), expected_padded_dim); - EXPECT_EQ(index.initializer_type(), ivf::InitializerType::Flat); + EXPECT_EQ(index.initializer_type(), GetParam().initializer); index.search_batch( queries.data(), nq, k, GetParam().ivf_nprobe, ids.data(), distances.data() ); @@ -133,6 +142,9 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail symqg::QuantizedGraph index; exercise(index, [&] { EXPECT_EQ(index.padded_dim(), expected_padded_dim); + // Earlier layouts do not record the construction settings either. + EXPECT_FALSE(index.ef_construction().has_value()); + EXPECT_FALSE(index.initialization().has_value()); index.set_ef(64); index.search_batch(queries.data(), nq, k, ids.data(), distances.data(), 1); }); @@ -164,10 +176,274 @@ INSTANTIATE_TEST_SUITE_P( PreIvfPerClusterStorage, HistoricalPersistenceTest, ::testing::Values( - HistoricalFixtureCase{"pre-ivf-per-cluster-storage", "ivf4_l2", 96, 130, 4, true}, - HistoricalFixtureCase{"pre-ivf-per-cluster-storage", "ivf32_l2", 96, 130, 4, true} + HistoricalFixtureCase{"pre-ivf-per-cluster-storage", "ivf4_l2", 96, 130, 4}, + HistoricalFixtureCase{"pre-ivf-per-cluster-storage", "ivf32_l2", 96, 130, 4} + ), + [](const auto& info) { return info.param.name; } +); + +// The last IVF layout before format 2: RABQIDX1 version 2 and the HNSW sidecar. +INSTANTIATE_TEST_SUITE_P( + V083, + HistoricalPersistenceTest, + ::testing::Values( + HistoricalFixtureCase{"v0.8.3", "ivf4_flat_l2", 96, 130, 4}, + HistoricalFixtureCase{ + "v0.8.3", "ivf7_flatrabitq_ip", 96, 130, 4, ivf::InitializerType::FlatRaBitQ}, + HistoricalFixtureCase{ + "v0.8.3", "ivf4_hnsw_l2", 96, 130, 4, ivf::InitializerType::HNSW}, + HistoricalFixtureCase{ + "v0.8.3", "ivf32_hnsw_ip", 96, 130, 4, ivf::InitializerType::HNSW}, + HistoricalFixtureCase{"v0.8.3", "ivf1_flat_matrix_l2", 96, 130, 4}, + HistoricalFixtureCase{"v0.8.3", "hnsw1_l2", 96, 130}, + HistoricalFixtureCase{"v0.8.3", "hnsw5_ip", 96, 130}, + HistoricalFixtureCase{"v0.8.3", "hnsw9_l2", 96, 130}, + HistoricalFixtureCase{"v0.8.3", "qg0_ip", 96, 130}, + HistoricalFixtureCase{"v0.8.3", "qg4_l2", 96, 130}, + HistoricalFixtureCase{"v0.8.3", "qg8_ip", 96, 130}, + HistoricalFixtureCase{"v0.8.3", "qg4_matrix_l2", 96, 130} ), [](const auto& info) { return info.param.name; } ); +// Damaged headers of the earlier layouts that pass the size checks raise FormatError +// rather than the errors of the code that would use them. +TEST(HistoricalPersistenceTest, MalformedEarlierHeadersRaiseFormatError) { + namespace fs = std::filesystem; + const fs::path path = + fs::temp_directory_path() / rabitq_test::process_unique("historical_header"); + struct Case { + const char* fixture; + size_t keep_bytes; // 0 keeps the whole file + std::vector> patches; // little-endian u64 fields + const char* message; + }; + const uint64_t huge = uint64_t{1} << 62; + const Case cases[] = { + // RABQIDX1 v2: the stream fails inside the fixed header. + {"ivf4_flat_l2", 20, {}, "IVF index file is truncated"}, + // An FHT/Kac dimension below 64 at offset 36. + {"ivf4_flat_l2", + 0, + {{36, 32}}, + "Invalid IVF rotator dimensions: Unsupported dimension for FhtKacRotator"}, + // SymphonyQG quantized layout: dimension at 28, padded dimension at 36. + {"qg4_l2", + 0, + {{28, 33}}, + "QuantizedGraph file is invalid: Unsupported dimension for FhtKacRotator"}, + {"qg4_l2", + 0, + {{28, huge}, {36, huge}}, + "QuantizedGraph file describes records too large to address"}, + // The headerless raw layout: dimension at 16, padded dimension at 24. + {"qg0_ip", + 0, + {{16, huge}, {24, huge}}, + "QuantizedGraph file describes records too large to address"}, + }; + for (const auto& c : cases) { + SCOPED_TRACE(c.message); + fs::copy_file( + fs::path(RABITQ_TEST_FIXTURE_ROOT) / "v0.8.3" / + (std::string(c.fixture) + ".index"), + path, + fs::copy_options::overwrite_existing + ); + if (c.keep_bytes != 0) { + fs::resize_file(path, c.keep_bytes); + } + { + std::fstream file(path, std::ios::binary | std::ios::in | std::ios::out); + for (const auto& [offset, value] : c.patches) { + file.seekp(static_cast(offset)); + for (size_t i = 0; i < sizeof(uint64_t); ++i) { + file.put(static_cast(value >> (8 * i))); + } + } + } + const std::string file = path.string(); + try { + if (c.fixture[0] == 'i') { + ivf::IVF index; + index.load(file.c_str()); + } else { + symqg::QuantizedGraph index; + index.load(file.c_str()); + } + ADD_FAILURE() << "accepted"; + } catch (const FormatError& error) { EXPECT_STREQ(error.what(), c.message); } + } + fs::remove(path); +} + +// Truncated files of the earlier layouts raise FormatError like format-2 files. +TEST(HistoricalPersistenceTest, TruncatedEarlierFilesRaiseFormatError) { + namespace fs = std::filesystem; + const fs::path path = + fs::temp_directory_path() / rabitq_test::process_unique("historical_truncated"); + for (const char* name : {"ivf4_flat_l2", "hnsw1_l2", "qg4_l2"}) { + SCOPED_TRACE(name); + fs::copy_file( + fs::path(RABITQ_TEST_FIXTURE_ROOT) / "v0.8.3" / (std::string(name) + ".index"), + path, + fs::copy_options::overwrite_existing + ); + fs::resize_file(path, fs::file_size(path) / 2); + const std::string file = path.string(); + if (name[0] == 'i') { + ivf::IVF index; + EXPECT_THROW(index.load(file.c_str()), FormatError); + } else if (name[0] == 'h') { + hnsw::HierarchicalNSW index; + EXPECT_THROW(index.load(file.c_str()), FormatError); + } else { + symqg::QuantizedGraph index; + EXPECT_THROW(index.load(file.c_str()), FormatError); + } + } + // A metric byte that is neither L2 nor IP, at offset 61 of an RABQIDX1 v2 file. + fs::copy_file( + fs::path(RABITQ_TEST_FIXTURE_ROOT) / "v0.8.3" / "ivf4_flat_l2.index", + path, + fs::copy_options::overwrite_existing + ); + { + std::fstream file(path, std::ios::binary | std::ios::in | std::ios::out); + file.seekp(61); + file.put(static_cast(2)); + } + ivf::IVF bad_metric; + try { + bad_metric.load(path.string().c_str()); + ADD_FAILURE() << "an invalid metric was loaded"; + } catch (const FormatError& error) { + EXPECT_STREQ(error.what(), "Invalid IVF index metadata"); + } + fs::remove(path); + ivf::IVF missing; + try { + missing.load((path.string() + ".missing").c_str()); + ADD_FAILURE() << "a missing file was loaded"; + } catch (const FormatError&) { + ADD_FAILURE() << "a missing file is not a format error"; + } catch (const std::runtime_error&) { // "Cannot open IVF index file" + } +} + +// The checks of the SymphonyQG layout before format 2, on a copy of its v0.8.3 file. +TEST(HistoricalPersistenceTest, RejectsMalformedEarlierGraphFilesWithoutChangingTarget) { + namespace fs = std::filesystem; + const fs::path fixture = fs::path(RABITQ_TEST_FIXTURE_ROOT) / "v0.8.3" / "qg4_l2.index"; + const fs::path path = + fs::temp_directory_path() / rabitq_test::process_unique("historical_qg_malformed"); + symqg::QuantizedGraph target(65, 64, 32, METRIC_IP); + const auto load_patched = [&](std::streamoff offset, size_t value) { + fs::copy_file(fixture, path, fs::copy_options::overwrite_existing); + { + std::fstream file(path, std::ios::binary | std::ios::in | std::ios::out); + ASSERT_TRUE(file.is_open()); + file.seekp(offset); + file.write(reinterpret_cast(&value), sizeof(value)); + ASSERT_TRUE(file.good()); + } + target.load(path.string().c_str()); + }; + // After the magic and version: point count, degree, dimension, padded dimension. + constexpr std::streamoff kPointCountOffset = sizeof(uint64_t) + sizeof(uint32_t); + constexpr std::streamoff kPaddedDimensionOffset = + kPointCountOffset + 3 * sizeof(size_t); + try { + load_patched(kPaddedDimensionOffset, 100); + ADD_FAILURE() << "an invalid padded dimension was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "QuantizedGraph file has an invalid padded dimension"); + } + EXPECT_THROW( + load_patched(kPointCountOffset, std::numeric_limits::max()), FormatError + ); + EXPECT_EQ(target.num_vertices(), 65U); + EXPECT_EQ(target.metric_type(), METRIC_IP); + fs::remove(path); +} + +// The checks of the HNSW-RaBitQ layout before format 2, on a copy of its v0.8.3 file. +TEST(HistoricalPersistenceTest, RejectsMalformedEarlierHnswFilesWithoutChangingTarget) { + namespace fs = std::filesystem; + const fs::path fixture = + fs::path(RABITQ_TEST_FIXTURE_ROOT) / "v0.8.3" / "hnsw5_ip.index"; + const fs::path path = fs::temp_directory_path() / + rabitq_test::process_unique("historical_hnsw_malformed"); + hnsw::HierarchicalNSW target; + target.load(fixture.string().c_str()); + const auto expect_rejected = [&](const auto& damage) { + fs::copy_file(fixture, path, fs::copy_options::overwrite_existing); + damage(); + try { + target.load(path.string().c_str()); + ADD_FAILURE() << "a damaged file was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "HNSW index file is invalid or truncated"); + } + EXPECT_EQ(target.num_stored_points(), 130U); + }; + // The capacity is the first field. + expect_rejected([&] { + std::fstream file(path, std::ios::binary | std::ios::in | std::ios::out); + const size_t capacity = 1; + file.write(reinterpret_cast(&capacity), sizeof(capacity)); + }); + expect_rejected([&] { fs::resize_file(path, fs::file_size(path) - 1); }); + fs::remove(path); +} + +// Files before format 2 keep the HNSW routing graph in .hnsw; loading reads it +// from there, through the same path conversion as the index file. +TEST(HistoricalPersistenceTest, LoadsTheHnswSidecarOfEarlierLayouts) { + namespace fs = std::filesystem; + const fs::path fixture = + fs::path(RABITQ_TEST_FIXTURE_ROOT) / "v0.8.3" / "ivf4_hnsw_l2.index"; + const fs::path directory = + fs::temp_directory_path() / + fs::u8path(rabitq_test::process_unique(u8"historical_sidecar_\u6d4b\u8bd5")); + fs::create_directories(directory); + const fs::path index_path = directory / fs::u8path(u8"\u7d22\u5f15.index"); + const fs::path sidecar = fs::path(index_path.native() + fs::path(".hnsw").native()); + fs::copy_file(fixture, index_path); + const std::string path = index_path.u8string(); + + ivf::IVF index; + try { + index.load(path.c_str()); + ADD_FAILURE() << "an HNSW-routed file without its sidecar was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "Cannot open HNSW centroid sidecar"); + } + + // A graph of another dimension does not fit the index. + { + std::vector centroids(4 * 128, 0.5F); + ivf::HNSWInitializer other(128, 4); + other.add_vectors(centroids.data(), 1); + std::ofstream output(sidecar, std::ios::binary); + other.write(output); + } + try { + index.load(path.c_str()); + ADD_FAILURE() << "a mismatched sidecar was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "HNSW centroid layout does not match IVF dimension"); + } + + fs::copy_file( + fs::path(fixture.native() + fs::path(".hnsw").native()), + sidecar, + fs::copy_options::overwrite_existing + ); + ASSERT_NO_THROW(index.load(path.c_str())); + EXPECT_EQ(index.initializer_type(), ivf::InitializerType::HNSW); + EXPECT_EQ(index.num_points(), 129U); + fs::remove_all(directory); +} + } // namespace } // namespace rabitqlib diff --git a/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp b/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp index 3937f905..1ab304ab 100644 --- a/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp @@ -6,15 +6,14 @@ #include #include #include -#include #include -#include #include #include #include #include #include +#include "format2_file.hpp" #include "rabitqlib/core/metric.hpp" #include "rabitqlib/core/rotator.hpp" #include "rabitqlib/index/ivf/initializer.hpp" @@ -75,8 +74,11 @@ struct Fixture { ); index.construct(data.data(), data.data(), ids.data(), false, 1); index.save(path.c_str()); - std::ifstream file(path + ".hnsw", std::ios::binary); - original.assign(std::istreambuf_iterator(file), {}); + // The routing graph is the HNSW-routing section of the format-2 index file. + const std::string graph = rabitq_test::read_format2_section( + path, rabitq_test::format2::IndexKind::Ivf, rabitq_test::kIvfSections, 260 + ); + original.assign(graph.begin(), graph.end()); level0_record_bytes = read_at(original, 3 * sizeof(size_t)); max_m = read_at(original, kMaxMOffset); max_m0 = read_at(original, kMaxMOffset + sizeof(size_t)); @@ -104,15 +106,20 @@ struct Fixture { ~Fixture() { std::error_code ignored; std::filesystem::remove(path, ignored); - std::filesystem::remove(path + ".hnsw", ignored); } + // Store `bytes` as the index file's routing graph. void save(const std::vector& bytes) const { - std::ofstream file(path + ".hnsw", std::ios::binary | std::ios::trunc); - file.write(bytes.data(), static_cast(bytes.size())); + rabitq_test::replace_format2_section( + path, + rabitq_test::format2::IndexKind::Ivf, + rabitq_test::kIvfSections, + 260, + std::string(bytes.begin(), bytes.end()) + ); } }; -TEST(IvfHnswSidecarGraphTest, PublicLoadRejectsCorruptStructureAndCanRetry) { +TEST(IvfHnswRoutingGraphTest, PublicLoadRejectsCorruptStructureAndCanRetry) { for (const auto metric : {METRIC_L2, METRIC_IP}) { Fixture fixture(metric); const auto reject = [&](const char* name, auto corrupt, const char* message) { @@ -123,7 +130,7 @@ TEST(IvfHnswSidecarGraphTest, PublicLoadRejectsCorruptStructureAndCanRetry) { ivf::IVF loaded; try { loaded.load(fixture.path.c_str()); - FAIL() << "damaged sidecar was accepted"; + FAIL() << "damaged routing graph was accepted"; } catch (const std::runtime_error& error) { EXPECT_STREQ(error.what(), message); } @@ -257,7 +264,7 @@ TEST(IvfHnswSidecarGraphTest, PublicLoadRejectsCorruptStructureAndCanRetry) { reject( "truncated upper length", [&](auto& b) { b.resize(fixture.upper_offsets.back() + 2); }, - "Truncated HNSW centroid sidecar" + "Truncated HNSW centroid graph" ); reject( "truncated upper links", @@ -287,7 +294,7 @@ TEST(IvfHnswSidecarGraphTest, PublicLoadRejectsCorruptStructureAndCanRetry) { reject( "truncated header", [](auto& b) { b.resize(4); }, - "Truncated HNSW centroid sidecar" + "Truncated HNSW centroid graph" ); reject( "trailing byte", @@ -297,7 +304,7 @@ TEST(IvfHnswSidecarGraphTest, PublicLoadRejectsCorruptStructureAndCanRetry) { } } -TEST(IvfHnswSidecarGraphTest, DisconnectedGraphCannotReadEmptyResults) { +TEST(IvfHnswRoutingGraphTest, DisconnectedGraphCannotReadEmptyResults) { for (const auto metric : {METRIC_L2, METRIC_IP}) { Fixture fixture(metric); auto bytes = fixture.original; diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 8835a2c4..815c098f 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -21,6 +21,7 @@ #include #include +#include "format2_file.hpp" #include "rabitqlib/clustering/rabitqkmeans.hpp" #include "rabitqlib/core/encoder.hpp" #include "rabitqlib/core/metric.hpp" @@ -906,7 +907,7 @@ TEST_F(HnswSaveTest, ReportsWriteOrCloseFailure) { index_.save("/dev/full"); FAIL() << "Saving to a full destination must fail"; } catch (const std::runtime_error& error) { - EXPECT_STREQ(error.what(), "HNSW failed to write index file"); + EXPECT_STREQ(error.what(), "HNSW index file cannot be written"); } } @@ -925,25 +926,188 @@ TEST_F(HnswSaveTest, SuccessfulSavePreservesSearchAfterLoading) { TEST_F(HnswSaveTest, RejectsPointCountLargerThanCapacity) { index_.save(path_.c_str()); - { - std::fstream file(path_, std::ios::binary | std::ios::in | std::ios::out); - const size_t invalid_capacity = 1; - file.write( - reinterpret_cast(&invalid_capacity), sizeof(invalid_capacity) - ); - ASSERT_TRUE(file.good()); - } + // The capacity is the first field of the graph parameters. + std::string params = rabitq_test::read_format2_section( + path_, rabitq_test::format2::IndexKind::Hnsw, rabitq_test::kHnswSections, 512 + ); + rabitq_test::format2::put_u64(reinterpret_cast(params.data()), 1); + rabitq_test::replace_format2_section( + path_, + rabitq_test::format2::IndexKind::Hnsw, + rabitq_test::kHnswSections, + 512, + params + ); HierarchicalNSW loaded; try { loaded.load(path_.c_str()); FAIL() << "Invalid HNSW point count must be rejected"; } catch (const std::runtime_error& error) { - EXPECT_NE(std::string(error.what()).find("HNSW"), std::string::npos); + EXPECT_STREQ(error.what(), "HNSW index file has invalid graph parameters"); + } +} + +// A graph with upper layers, saved in format 2 with a known seed. +class HnswFormat2Test : public ::testing::Test { + protected: + static constexpr size_t kDim = 64; + static constexpr size_t kCount = 200; + static constexpr size_t kSeed = 9; + std::vector data_{std::vector(kCount * kDim)}; + std::vector centroids_{std::vector(2 * kDim)}; + std::vector cluster_ids_{std::vector(kCount)}; + HierarchicalNSW index_{kCount, kDim, 5, 4, 16, kSeed, METRIC_IP}; + std::string path_; + + void SetUp() override { + path_ = ::testing::TempDir() + "rabitq_hnsw_format2_" + + ::testing::UnitTest::GetInstance()->current_test_info()->name() + ".index"; + std::mt19937 generator(3); + std::normal_distribution normal; + for (auto& value : data_) { + value = normal(generator); + } + for (size_t i = 0; i < kCount; ++i) { + cluster_ids_[i] = static_cast(i % 2); + } + index_.construct( + 2, centroids_.data(), kCount, data_.data(), cluster_ids_.data(), 1, false + ); + const PID removed = 5; + index_.remove(&removed, 1); + index_.save(path_.c_str()); + } + + void TearDown() override { std::remove(path_.c_str()); } + + [[nodiscard]] rabitq_test::Format2Contents contents() const { + return rabitq_test::read_format2_file( + path_, rabitq_test::format2::IndexKind::Hnsw, rabitq_test::kHnswSections + ); + } + + void expect_rejected( + const rabitq_test::Format2Contents& damaged, const std::string& message + ) { + rabitq_test::write_format2_file(path_, damaged); + HierarchicalNSW loaded; + try { + loaded.load(path_.c_str()); + ADD_FAILURE() << "accepted: " << message; + } catch (const std::runtime_error& error) { + EXPECT_EQ(std::string(error.what()), message); + } } +}; + +TEST_F(HnswFormat2Test, RoundTripsTheGraphAndRecordsTheSeed) { + const auto saved = contents(); + EXPECT_EQ(saved.header.index_kind, rabitq_test::format2::IndexKind::Hnsw); + EXPECT_EQ(saved.header.metric, rabitq_test::format2::Metric::InnerProduct); + EXPECT_EQ(saved.header.code_bits, 5U); + EXPECT_EQ(saved.header.num_points, kCount); + ASSERT_TRUE(saved.header.seed.has_value()); + EXPECT_EQ(*saved.header.seed, kSeed); + std::vector kinds; + for (const auto& entry : saved.sections) { + kinds.push_back(entry.first.kind); + } + EXPECT_EQ(kinds, (std::vector{1, 512, 513, 514, 515, 516})); + EXPECT_FALSE(rabitq_test::Format2Contents(saved).bytes(516).empty()); // upper layers + + HierarchicalNSW loaded; + loaded.load(path_.c_str()); + EXPECT_EQ(loaded.seed(), std::optional(kSeed)); + EXPECT_EQ(loaded.num_points(), kCount - 1); + EXPECT_EQ(loaded.max_elements(), kCount); + EXPECT_EQ( + loaded.search(data_.data(), 4, 5, 32, 1), index_.search(data_.data(), 4, 5, 32, 1) + ); + const std::string first = rabitq_test::Format2Contents(saved).bytes(514); + loaded.save(path_.c_str()); + EXPECT_EQ(contents().bytes(514), first); +} + +TEST_F(HnswFormat2Test, RejectsContentsThatDoNotFitTheGraph) { + const auto saved = contents(); + auto damaged = saved; + damaged.remove(515); + expect_rejected(damaged, "HNSW index file has no section kind 515"); + + damaged = saved; + damaged.bytes(516).pop_back(); + const size_t upper_bytes = rabitq_test::Format2Contents(saved).bytes(516).size(); + expect_rejected( + damaged, + "HNSW index file section kind 516 has " + std::to_string(upper_bytes - 1) + + " bytes; expected " + std::to_string(upper_bytes) + ); + + damaged = saved; + rabitq_test::format2::put_u32( + reinterpret_cast(damaged.bytes(515).data()), 1000 + ); + expect_rejected(damaged, "HNSW index file has a point above the maximum level"); + + damaged = saved; + rabitq_test::format2::put_u32( + reinterpret_cast(damaged.bytes(512).data()) + 60, 1 + ); + expect_rejected(damaged, "HNSW index file has invalid graph parameters"); + + damaged = saved; + damaged.header.metric = rabitq_test::format2::Metric::Cosine; + expect_rejected( + damaged, + "HNSW index file uses the cosine metric, which this release does not support" + ); + + damaged = saved; + damaged.bytes(515).resize(4 * (kCount - 1)); + expect_rejected( + damaged, + "HNSW index file section kind 515 has " + std::to_string(4 * (kCount - 1)) + + " bytes; expected " + std::to_string(4 * kCount) + ); + + // A neighbor beyond the stored points, in the first base-layer record. + damaged = saved; + char* record = damaged.bytes(514).data(); + rabitq_test::format2::put_u32(reinterpret_cast(record), 1); // one neighbor + rabitq_test::format2::put_u32(reinterpret_cast(record) + 4, kCount); + expect_rejected(damaged, "HNSW index file is invalid or truncated"); +} + +TEST_F(HnswFormat2Test, UnknownReadOnlyFeaturesAllowSearchButNoChanges) { + rabitq_test::patch_format2_header(path_, 24, 1, 8); // an unknown ro_compat bit + HierarchicalNSW loaded; + ASSERT_NO_THROW(loaded.load(path_.c_str())); + EXPECT_EQ( + loaded.search(data_.data(), 2, 5, 32, 1), index_.search(data_.data(), 2, 5, 32, 1) + ); + const auto expect_refused = [](const auto& operation, const std::string& name) { + try { + operation(); + ADD_FAILURE() << name << " was allowed"; + } catch (const std::logic_error& error) { + EXPECT_EQ( + std::string(error.what()), + "HNSW index was loaded from a file with features this release cannot " + "update, so it cannot " + + name + ); + } + }; + const PID label = 7; + expect_refused([&] { loaded.add(data_.data(), 1, &label, false, 1); }, "add"); + expect_refused([&] { loaded.resize(kCount + 10); }, "resize"); + expect_refused([&] { loaded.remove(&label, 1); }, "remove"); + expect_refused([&] { loaded.save((path_ + ".copy").c_str()); }, "save"); + EXPECT_EQ(loaded.num_points(), kCount - 1); } -TEST_F(HnswSaveTest, RejectsTruncatedRotatorWithoutLosingExistingIndex) { +TEST_F(HnswSaveTest, RejectsTruncatedFileWithoutLosingExistingIndex) { index_.save(path_.c_str()); HierarchicalNSW loaded; loaded.load(path_.c_str()); @@ -952,7 +1116,7 @@ TEST_F(HnswSaveTest, RejectsTruncatedRotatorWithoutLosingExistingIndex) { std::filesystem::resize_file(path_, std::filesystem::file_size(path_) - 1); try { loaded.load(path_.c_str()); - FAIL() << "Truncated HNSW rotator must be rejected"; + FAIL() << "A truncated HNSW file must be rejected"; } catch (const std::runtime_error& error) { EXPECT_NE(std::string(error.what()).find("HNSW"), std::string::npos); } diff --git a/tests/unit/rabitqlib/index/initializer_test.cpp b/tests/unit/rabitqlib/index/initializer_test.cpp index 93e20555..d25432c2 100644 --- a/tests/unit/rabitqlib/index/initializer_test.cpp +++ b/tests/unit/rabitqlib/index/initializer_test.cpp @@ -10,10 +10,10 @@ #include #include #include -#include #include #include #include +#include #include #include #include @@ -24,58 +24,33 @@ #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" -#include "temp_file.hpp" namespace rabitqlib::ivf { namespace { -TEST(IvfHnswInitializerTest, PersistsUtf8SidecarPaths) { - const auto directory = std::filesystem::temp_directory_path() / - std::filesystem::u8path(u8"rabitq_\u6d4b\u8bd5_\U0001f680"); - std::filesystem::create_directories(directory); - const auto path = (directory / std::filesystem::u8path(u8"\u7d22\u5f15")).u8string(); +TEST(IvfHnswInitializerTest, RoundTripsThroughAStreamAndChecksItsSize) { const std::array centroids{0.0F, 0.0F, 3.0F, 4.0F}; HNSWInitializer initializer(2, 2); initializer.add_vectors(centroids.data(), 1); - std::ofstream ignored_output; - initializer.save(ignored_output, path.c_str()); - EXPECT_TRUE(std::filesystem::is_regular_file(std::filesystem::u8path(path + ".hnsw"))); - HNSWInitializer loaded(2, 2); - std::ifstream ignored_input; - loaded.load(ignored_input, path.c_str()); - std::vector> candidates(1); - loaded.centroids_distances(centroids.data() + 2, 1, candidates); - EXPECT_EQ(candidates[0].id, 1U); - EXPECT_FLOAT_EQ(candidates[0].distance, 0.0F); - std::filesystem::remove(std::filesystem::u8path(path + ".hnsw")); - std::filesystem::remove(directory); -} - -TEST(IvfHnswInitializerTest, ReportsSidecarOpenFailureAndCanRetry) { - const std::string path = ::testing::TempDir() + "rabitq_hnsw_save_failure"; - const std::filesystem::path sidecar(path + ".hnsw"); - std::filesystem::create_directory(sidecar); - const std::array centroids{0.0F, 0.0F, 3.0F, 4.0F}; - HNSWInitializer initializer(2, 2); - initializer.add_vectors(centroids.data(), 1); - std::ofstream ignored_output; - try { - initializer.save(ignored_output, path.c_str()); - FAIL() << "Saving the HNSW sidecar to a directory must fail"; - } catch (const std::runtime_error& error) { - EXPECT_STREQ(error.what(), "Cannot open HNSW index file for writing"); - } - std::filesystem::remove(sidecar); + std::stringstream graph; + initializer.write(graph); + const std::string bytes = graph.str(); - initializer.save(ignored_output, path.c_str()); HNSWInitializer loaded(2, 2); - std::ifstream ignored_input; - loaded.load(ignored_input, path.c_str()); + loaded.read(graph, bytes.size()); std::vector> candidates(1); loaded.centroids_distances(centroids.data() + 2, 1, candidates); EXPECT_EQ(candidates[0].id, 1U); EXPECT_FLOAT_EQ(candidates[0].distance, 0.0F); - std::filesystem::remove(sidecar); + + // The graph is read from the current position and only within its byte count. + std::stringstream shifted("xyz" + bytes); + shifted.seekg(3); + HNSWInitializer offset(2, 2); + EXPECT_NO_THROW(offset.read(shifted, bytes.size())); + std::stringstream short_input(bytes); + HNSWInitializer truncated(2, 2); + EXPECT_THROW(truncated.read(short_input, bytes.size() - 1), std::runtime_error); } TEST(IvfHnswInitializerTest, ReportsBufferedWriteFailure) { @@ -222,16 +197,13 @@ TEST(IvfHnswInitializerTest, ParallelInsertionPreservesCentroidLabelsAfterReload centroids[label * kDim + dim] = static_cast(label * kDim + dim); } } - const rabitq_test::TempFile path_file("rabitq_initializer_labels"); - const std::string& path = path_file.path(); HNSWInitializer initializer(kDim, kCount); initializer.add_vectors(centroids.data(), 4); - std::ofstream ignored_output; - initializer.save(ignored_output, path.c_str()); + std::stringstream graph; + initializer.write(graph); HNSWInitializer loaded(kDim, kCount); - std::ifstream ignored_input; - loaded.load(ignored_input, path.c_str()); + loaded.read(graph, graph.str().size()); for (size_t label = 0; label < kCount; ++label) { for (size_t dim = 0; dim < kDim; ++dim) { EXPECT_FLOAT_EQ( @@ -251,16 +223,13 @@ TEST(IvfHnswInitializerTest, ReloadsInnerProductGraphForRouting) { std::vector centroids(kCount * kDim, 0.0F); centroids[0] = 0.9F; centroids[kDim] = 100.0F; - const rabitq_test::TempFile path_file("rabitq_initializer_ip"); - const std::string& path = path_file.path(); HNSWInitializer initializer(kDim, kCount, METRIC_IP); initializer.add_vectors(centroids.data(), 4); - std::ofstream ignored_output; - initializer.save(ignored_output, path.c_str()); + std::stringstream graph; + initializer.write(graph); HNSWInitializer loaded(kDim, kCount, METRIC_IP); - std::ifstream ignored_input; - loaded.load(ignored_input, path.c_str()); + loaded.read(graph, graph.str().size()); std::array query{}; query[0] = 1.0F; std::vector> candidates(1); @@ -393,31 +362,28 @@ TEST(ParallelForTest, JoinsWorkersBeforeRethrowingFunctionException) { TEST(IvfInitializerTest, FlatInitializerReportsStreamFailures) { constexpr size_t kDim = 8; constexpr size_t kCount = 3; + constexpr size_t kBytes = sizeof(float) * kDim * kCount; std::vector centroids(kCount * kDim, 1.5F); FlatInitializer flat(kDim, kCount); flat.add_vectors(centroids.data(), 1); - const rabitq_test::TempFile path_file("rabitq_flat_initializer.bin"); - const std::string& path = path_file.path(); - { - std::ofstream output(path, std::ios::binary); - ASSERT_NO_THROW(flat.save(output, path.c_str())); - } - { - // A stream that failed to open cannot be written. - std::ofstream unopened(path + ".missing/initializer.bin", std::ios::binary); - EXPECT_THROW(flat.save(unopened, path.c_str()), std::runtime_error); - } - { - std::ifstream input(path, std::ios::binary); - FlatInitializer restored(kDim, kCount); - ASSERT_NO_THROW(restored.load(input, path.c_str())); - } - std::filesystem::resize_file( - std::filesystem::u8path(path), sizeof(float) * kDim * kCount - 1 - ); - std::ifstream truncated(path, std::ios::binary); + std::stringstream stored; + ASSERT_NO_THROW(flat.write(stored)); + ASSERT_EQ(stored.str().size(), kBytes); + + std::ostringstream broken; + broken.setstate(std::ios::badbit); + EXPECT_THROW(flat.write(broken), std::runtime_error); + + FlatInitializer restored(kDim, kCount); + ASSERT_NO_THROW(restored.read(stored, kBytes)); + EXPECT_FLOAT_EQ(restored.centroid(2)[kDim - 1], 1.5F); + + FlatInitializer wrong_size(kDim, kCount); + std::stringstream copy(stored.str()); + EXPECT_THROW(wrong_size.read(copy, kBytes - 1), std::runtime_error); + std::stringstream truncated(stored.str().substr(0, kBytes - 1)); FlatInitializer partial(kDim, kCount); - EXPECT_THROW(partial.load(truncated, path.c_str()), std::runtime_error); + EXPECT_THROW(partial.read(truncated, kBytes), std::runtime_error); } } // namespace diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index cfce20c4..5632641b 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -17,10 +18,13 @@ #include #include #include +#include #include #include #include +#include "blob.hpp" +#include "format2_file.hpp" #include "ivf_search.hpp" #include "rabitqlib/clustering/qgkmeans.hpp" #include "rabitqlib/clustering/rabitqkmeans.hpp" @@ -49,6 +53,8 @@ struct IvfTestAccess { static std::unique_ptr>& rotator(IVF& index) { return index.rotator_; } + static std::optional& seed(IVF& index) { return index.seed_; } + // Rotated centroid as the initializer routes with it. static std::vector centroid(const IVF& index, PID i) { const float* centroid = index.initializer_->centroid(i); @@ -544,7 +550,12 @@ TEST(IvfPersistenceTest, RejectsPointIdsUsingTheSearchBufferMarker) { file.close(); IVF loaded; - EXPECT_THROW(loaded.load(path.c_str()), std::runtime_error); + try { + loaded.load(path.c_str()); + ADD_FAILURE() << "a marked point ID was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "Invalid point ID in IVF index file"); + } } // ── add / remove ──────────────────────────────────────────────────────────────────────── @@ -1687,7 +1698,9 @@ TEST(IvfPersistenceTest, RejectsTrailingBytesAndReportsUnopenableDestination) { loaded.load(path.c_str()); ADD_FAILURE() << "trailing bytes were accepted"; } catch (const std::runtime_error& error) { - EXPECT_STREQ(error.what(), "IVF index file has unexpected trailing data"); + const std::string message = error.what(); + EXPECT_EQ(message.rfind("IVF index file is ", 0), 0U) << message; + EXPECT_NE(message.find(" bytes, but its sections end at "), std::string::npos); } // The failed load leaves the previously loaded index intact. EXPECT_EQ(loaded.num_points(), kCount); @@ -1797,13 +1810,15 @@ TEST(IvfTrainTest, SetCentroidsAndAddMatchConstructByteForByte) { << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") << " bits=" << bits ); IVF constructed(kCount, kDynDim, kDynNumClusters, bits, metric); - // Give the point-count constructor the rotation the seed selects. + // Give the point-count constructor the rotation the seed selects, and the + // seed itself, which the file records. IvfTestAccess::rotator(constructed) .reset( choose_rotator( kDynDim, RotatorType::FhtKacRotator, constructed.padded_dim(), kSeed ) ); + IvfTestAccess::seed(constructed) = kSeed; constructed.construct( points.data(), centroids.data(), assigned.data(), false, 1 ); @@ -2016,10 +2031,393 @@ TEST(IvfTrainTest, DefaultsUseSphericalKMeansForInnerProduct) { } } +uint32_t get_size(const std::string& bytes, size_t offset) { + uint32_t value = 0; + std::memcpy(&value, bytes.data() + offset, sizeof(value)); + return value; +} + +// ── index format 2 ────────────────────────────────────────────────────────────────── + +using rabitq_test::format2::IndexKind; +namespace f2 = rabitq_test::format2; + +// An index of 33 points over four clusters, with point 3 removed, saved at `path`. +std::unique_ptr saved_format2_index( + const std::string& path, + size_t total_bits, + MetricType metric, + InitializerType initializer, + RotatorType rotator = RotatorType::FhtKacRotator +) { + constexpr size_t kCount = 33; + const auto points = random_rows(kCount, 951); + const auto centroids = random_rows(kDynNumClusters, 952); + const auto assigned = nearest_centroids(centroids, points.data(), kCount, metric); + auto index = std::make_unique( + kCount, kDynDim, kDynNumClusters, total_bits, metric, rotator, initializer + ); + index->construct(points.data(), centroids.data(), assigned.data(), false, 1); + const PID removed = 3; + index->remove(&removed, 1); + index->save(path.c_str()); + return index; +} + +TEST(IvfFormat2Test, WritesOneFileWithTheSpecifiedHeaderAndSections) { + const rabitq_test::TempFile path_file("rabitq_ivf_format2_layout.index"); + const std::string& path = path_file.path(); + const struct { + size_t total_bits; + MetricType metric; + InitializerType initializer; + f2::Metric stored_metric; + uint32_t code_bits; + f2::RawType raw_type; + uint32_t routing; + bool rerank; + } cases[] = { + {5, + METRIC_L2, + InitializerType::HNSW, + f2::Metric::L2, + 5, + f2::RawType::None, + 260, + true}, + {1, + METRIC_IP, + InitializerType::Flat, + f2::Metric::InnerProduct, + 1, + f2::RawType::None, + 258, + false}, + {32, + METRIC_L2, + InitializerType::FlatRaBitQ, + f2::Metric::L2, + 1, + f2::RawType::Float32, + 259, + true}, + }; + for (const auto& c : cases) { + const auto index = saved_format2_index(path, c.total_bits, c.metric, c.initializer); + EXPECT_FALSE(std::filesystem::exists(path + ".hnsw")); + const auto contents = + rabitq_test::read_format2_file(path, IndexKind::Ivf, rabitq_test::kIvfSections); + const f2::Header& header = contents.header; + EXPECT_EQ(header.metric, c.stored_metric); + EXPECT_EQ(header.dim, kDynDim); + EXPECT_EQ(header.padded_dim, 96U); + EXPECT_EQ(header.code_bits, c.code_bits); + EXPECT_EQ(header.raw_type, c.raw_type); + EXPECT_EQ(header.id_bits, 32U); + EXPECT_FALSE(header.seed.has_value()); + EXPECT_EQ(header.num_points, 33U); + std::vector kinds; + for (const auto& entry : contents.sections) { + kinds.push_back(entry.first.kind); + } + std::vector expected = {1, 256, 257, c.routing, 261}; + if (c.rerank) { + expected.push_back(262); + } + expected.push_back(263); + EXPECT_EQ(kinds, expected); + + IVF loaded; + loaded.load(path.c_str()); + EXPECT_EQ(loaded.initializer_type(), c.initializer); + EXPECT_EQ(loaded.num_points(), 32U); + std::vector expected_ids(5), ids(5); + std::vector expected_distances(5), distances(5); + const auto query = random_rows(1, 953); + index->search( + query.data(), 5, kDynNumClusters, expected_ids.data(), expected_distances.data() + ); + loaded.search(query.data(), 5, kDynNumClusters, ids.data(), distances.data()); + EXPECT_EQ(ids, expected_ids); + EXPECT_EQ(distances, expected_distances); + } +} + +TEST(IvfFormat2Test, RejectsContentsThatDoNotFitAnIvfIndex) { + const rabitq_test::TempFile path_file("rabitq_ivf_format2_reject.index"); + const std::string& path = path_file.path(); + saved_format2_index(path, 5, METRIC_L2, InitializerType::HNSW); + const auto original = + rabitq_test::read_format2_file(path, IndexKind::Ivf, rabitq_test::kIvfSections); + IVF loaded; + loaded.load(path.c_str()); + + const auto expect_rejected = + [&](const char* name, auto edit, const std::string& message) { + SCOPED_TRACE(name); + auto contents = original; + edit(contents); + rabitq_test::write_format2_file(path, contents); + try { + loaded.load(path.c_str()); + ADD_FAILURE() << "accepted"; + } catch (const std::runtime_error& error) { + EXPECT_EQ(std::string(error.what()), message); + } + EXPECT_EQ(loaded.num_points(), 32U); // the failed load keeps the loaded index + }; + const auto set_u32 = [](std::string& bytes, size_t offset, uint32_t value) { + std::memcpy(bytes.data() + offset, &value, sizeof(value)); + }; + expect_rejected( + "cosine", + [](auto& c) { c.header.metric = f2::Metric::Cosine; }, + "IVF index file uses the cosine metric, which this release does not support" + ); + expect_rejected( + "64-bit IDs", + [](auto& c) { c.header.id_bits = 64; }, + "IVF index file has 64-bit IDs; IVF-RaBitQ stores 32 bits" + ); + expect_rejected( + "uint8 raw vectors", + [](auto& c) { c.header.raw_type = f2::RawType::Uint8; }, + "IVF index file has uint8 raw vectors, which IVF-RaBitQ lacks" + ); + expect_rejected( + "raw vectors with extra bits", + [](auto& c) { c.header.raw_type = f2::RawType::Float32; }, + "IVF index file has 5 code bits, which its raw vectors do not allow" + ); + expect_rejected( + "no points", + [](auto& c) { c.header.num_points = 0; }, + "IVF index file has an invalid point count 0" + ); + expect_rejected( + "missing codes", + [](auto& c) { c.remove(261); }, + "IVF index file has no section kind 261" + ); + expect_rejected( + "short parameters", + [](auto& c) { c.bytes(256).pop_back(); }, + "IVF index file section kind 256 has 15 bytes; expected 16" + ); + expect_rejected( + "unknown routing", + [&](auto& c) { set_u32(c.bytes(256), 8, 9); }, + "IVF index file has invalid routing parameters" + ); + expect_rejected( + "cluster sizes above the point count", + [&](auto& c) { set_u32(c.bytes(257), 0, get_size(c.bytes(257), 0) + 1); }, + "IVF index file cluster sizes exceed its point count" + ); + expect_rejected( + "cluster sizes below the point count", + [&](auto& c) { + std::string& sizes = c.bytes(257); + for (size_t offset = 0; offset < sizes.size(); offset += sizeof(uint64_t)) { + if (get_size(sizes, offset) != 0) { + set_u32(sizes, offset, get_size(sizes, offset) - 1); + break; + } + } + }, + "IVF index file cluster sizes do not add up to its point count" + ); + expect_rejected( + "routing of another initializer", + [](auto& c) { + f2::Section flat; + flat.kind = 258; + flat.flags = f2::kSectionRequired; + c.sections.insert( + c.sections.begin() + 4, {flat, std::string(4 * 4 * 96, '\0')} + ); + }, + "IVF index file has routing data of another initializer" + ); + expect_rejected( + "corrupt rotator", + [](auto& c) { c.bytes(1)[70] ^= 1; }, + "IVF index file has an invalid rotator (data checksum mismatch (corrupt blob))" + ); + expect_rejected( + "rotator of another dimension", + [](auto& c) { + const std::unique_ptr> other( + choose_rotator(64, RotatorType::FhtKacRotator, 96, 1) + ); + std::string blob( + rabitqlib::detail::rotator_blob_bytes(RotatorType::FhtKacRotator, 64, 96), + '\0' + ); + rabitqlib::detail::encode_rotator_blob( + *other, + RotatorType::FhtKacRotator, + 64, + std::nullopt, + reinterpret_cast(blob.data()) + ); + c.bytes(1) = blob; + }, + "IVF index file rotator does not match its dimensions" + ); + expect_rejected( + "point ID out of range", + [&](auto& c) { set_u32(c.bytes(263), 0, 33); }, + "Invalid point ID in IVF index file" + ); + expect_rejected( + "missing rerank data", + [](auto& c) { c.remove(262); }, + "IVF index file has no section kind 262" + ); + expect_rejected( + "short IDs", + [](auto& c) { c.bytes(263).resize(128); }, + "IVF index file section kind 263 has 128 bytes; expected 132" + ); + expect_rejected( + "reserved routing field", + [&](auto& c) { set_u32(c.bytes(256), 12, 1); }, + "IVF index file has invalid routing parameters" + ); + + // One-bit codes have no rerank data, so a rerank section is rejected. + saved_format2_index(path, 1, METRIC_L2, InitializerType::Flat); + auto one_bit = + rabitq_test::read_format2_file(path, IndexKind::Ivf, rabitq_test::kIvfSections); + f2::Section rerank; + rerank.kind = 262; + rerank.flags = f2::kSectionRequired; + one_bit.sections.insert(one_bit.sections.end() - 1, {rerank, std::string(64, '\0')}); + rabitq_test::write_format2_file(path, one_bit); + try { + loaded.load(path.c_str()); + ADD_FAILURE() << "accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), "IVF index file has rerank data that its code bits do not use" + ); + } +} + +TEST(IvfFormat2Test, UnknownReadOnlyFeaturesAllowSearchButNoChanges) { + const rabitq_test::TempFile path_file("rabitq_ivf_format2_read_only.index"); + const std::string& path = path_file.path(); + const auto index = saved_format2_index(path, 5, METRIC_L2, InitializerType::Flat); + rabitq_test::patch_format2_header(path, 24, 1, 8); // an unknown ro_compat bit + IVF loaded; + ASSERT_NO_THROW(loaded.load(path.c_str())); + const auto query = random_rows(1, 954); + std::vector expected_ids(5), ids(5); + std::vector distances(5); + index->search(query.data(), 5, kDynNumClusters, expected_ids.data(), distances.data()); + loaded.search(query.data(), 5, kDynNumClusters, ids.data(), distances.data()); + EXPECT_EQ(ids, expected_ids); + + const auto expect_refused = [](const auto& operation, const std::string& name) { + try { + operation(); + ADD_FAILURE() << name << " was allowed"; + } catch (const std::logic_error& error) { + EXPECT_EQ( + std::string(error.what()), + "IVF index was loaded from a file with features this release cannot " + "update, so it cannot " + + name + ); + } + }; + const auto row = random_rows(1, 955); + const PID cluster = 0; + const PID removed = 0; + expect_refused([&] { loaded.add(row.data(), 1, &cluster, false, 1); }, "add"); + expect_refused([&] { loaded.reserve(0, 64); }, "reserve"); + expect_refused([&] { loaded.remove(&removed, 1); }, "remove"); + expect_refused([&] { loaded.save((path + ".copy").c_str()); }, "save"); + EXPECT_EQ(loaded.num_points(), 32U); + + // Rebuilding replaces the data the unknown features described, so it is writable. + const auto points = random_rows(33, 951); + const auto centroids = random_rows(kDynNumClusters, 952); + const auto assigned = nearest_centroids(centroids, points.data(), 33, METRIC_L2); + loaded.construct(points.data(), centroids.data(), assigned.data(), false, 1); + EXPECT_NO_THROW(loaded.save(path.c_str())); + + // An unknown incompatible feature is rejected outright. + rabitq_test::patch_format2_header(path, 16, 2, 8); + try { + loaded.load(path.c_str()); + ADD_FAILURE() << "an unknown incompatible feature was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_EQ( + std::string(error.what()), + std::string( + "IVF index file uses incompatible features 0x2 that this release " + "does not know; written by rabitqlib " + ) + RABITQ_TEST_LIBRARY_VERSION + ); + } +} + +TEST(IvfFormat2Test, RecordsTheWriterAndSkipsUnknownOptionalSections) { + const rabitq_test::TempFile path_file("rabitq_ivf_format2_optional.index"); + const std::string& path = path_file.path(); + const auto index = saved_format2_index(path, 4, METRIC_IP, InitializerType::FlatRaBitQ); + auto contents = + rabitq_test::read_format2_file(path, IndexKind::Ivf, rabitq_test::kIvfSections); + EXPECT_EQ(contents.writer_version, RABITQ_TEST_LIBRARY_VERSION); + // A later release may add an optional section; this one skips it. + f2::Section future; + future.kind = 900; + contents.sections.emplace_back(future, std::string(100, 'x')); + rabitq_test::write_format2_file(path, contents); + IVF loaded; + ASSERT_NO_THROW(loaded.load(path.c_str())); + const auto query = random_rows(1, 956); + std::vector expected_ids(5), ids(5); + std::vector distances(5); + index->search(query.data(), 5, kDynNumClusters, expected_ids.data(), distances.data()); + loaded.search(query.data(), 5, kDynNumClusters, ids.data(), distances.data()); + EXPECT_EQ(ids, expected_ids); + // save drops the section it does not know. + loaded.save(path.c_str()); + std::vector kinds = rabitq_test::kIvfSections; + kinds.push_back(900); + const auto resaved = rabitq_test::read_format2_file(path, IndexKind::Ivf, kinds); + EXPECT_EQ(resaved.sections.size(), contents.sections.size() - 1); + EXPECT_THROW( + static_cast(rabitq_test::Format2Contents(resaved).bytes(900)), + std::runtime_error + ); +} + +TEST(IvfFormat2Test, FailedWriteReportsAnErrorAndKeepsTheIndex) { + if (!std::filesystem::exists("/dev/full")) { + GTEST_SKIP() << "Requires a device that rejects writes"; + } + const rabitq_test::TempFile path_file("rabitq_ivf_format2_write_failure.index"); + const std::string& path = path_file.path(); + for (const auto initializer : + {InitializerType::Flat, InitializerType::FlatRaBitQ, InitializerType::HNSW}) { + const auto index = saved_format2_index(path, 5, METRIC_L2, initializer); + // /dev/full is not a regular file, so save writes it in place and the write fails. + EXPECT_ANY_THROW(index->save("/dev/full")); + index->save(path.c_str()); + IVF loaded; + ASSERT_NO_THROW(loaded.load(path.c_str())); + EXPECT_EQ(loaded.num_points(), 32U); + } +} + } // namespace } // namespace rabitqlib::ivf -TEST(IvfSidecarValidationTest, RejectsMismatchedVectorStorage) { +TEST(IvfRoutingGraphValidationTest, RejectsMismatchedVectorStorage) { using namespace rabitqlib; for (const auto metric : {METRIC_L2, METRIC_IP}) { for (const size_t sidecar_dim : {64U, 128U}) { @@ -2034,12 +2432,19 @@ TEST(IvfSidecarValidationTest, RejectsMismatchedVectorStorage) { original.save(path.c_str()); ivf::HNSWInitializer other(sidecar_dim, 1, metric); other.add_vectors(center.data(), 1); - std::ofstream ignored; - other.save(ignored, path.c_str()); + std::ostringstream graph; + other.write(graph); + rabitq_test::replace_format2_section( + path, + rabitq_test::format2::IndexKind::Ivf, + rabitq_test::kIvfSections, + 260, + graph.str() + ); ivf::IVF loaded; try { loaded.load(path.c_str()); - FAIL() << "mismatched sidecar must be rejected"; + FAIL() << "mismatched routing graph must be rejected"; } catch (const std::runtime_error& error) { EXPECT_STREQ( error.what(), "HNSW centroid layout does not match IVF dimension" diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index a35edca4..7ba5967a 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -23,6 +23,7 @@ #include #include +#include "format2_file.hpp" #include "rabitqlib/core/data_layout.hpp" #include "rabitqlib/core/metric.hpp" #include "rabitqlib/core/rotator.hpp" @@ -36,6 +37,7 @@ #include "rabitqlib/index/detail/buffer.hpp" #include "rabitqlib/index/detail/exact_nn.hpp" #include "rabitqlib/index/detail/visited_set.hpp" +#include "rabitqlib/index/format_error.hpp" #include "rabitqlib/index/symqg/detail/pipnn.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" #include "rabitqlib/index/types.hpp" @@ -1398,6 +1400,199 @@ TEST(QgConfigurationTest, RejectsOutOfRangeEntryPoint) { EXPECT_EQ(graph.entry_point(), 32U); } +TEST(QgFormat2Test, RecordsTheSeedAndConstructionSettings) { + constexpr size_t kCount = 65, kDim = 64, kDegree = 32; + std::vector data(kCount * kDim); + std::mt19937 generator(5); + std::normal_distribution normal; + for (auto& value : data) { + value = normal(generator); + } + const rabitq_test::TempFile path_file("rabitq_qg_format2.index"); + const std::string& path = path_file.path(); + for (const size_t bits : {0UL, 4UL, 8UL}) { + SCOPED_TRACE(bits); + QuantizedGraph source( + kCount, kDim, kDegree, METRIC_L2, RotatorType::FhtKacRotator, bits, 7 + ); + QGBuilder builder(source, 40, data.data(), 1, QGInitialization::Random, 7); + builder.build(); + source.save(path.c_str()); + + const auto saved = rabitq_test::read_format2_file( + path, rabitq_test::format2::IndexKind::SymphonyQg, rabitq_test::kQgSections + ); + EXPECT_EQ(saved.header.code_bits, bits); + EXPECT_EQ( + saved.header.raw_type, + bits == 0 ? rabitq_test::format2::RawType::Float32 + : rabitq_test::format2::RawType::None + ); + std::vector kinds; + for (const auto& entry : saved.sections) { + kinds.push_back(entry.first.kind); + } + const std::vector expected_kinds = + bits == 0 ? std::vector{1, 768, 770} + : std::vector{1, 768, 769, 770}; + EXPECT_EQ(kinds, expected_kinds); + + QuantizedGraph loaded; + loaded.load(path.c_str()); + EXPECT_EQ(loaded.seed(), std::optional(7)); + EXPECT_EQ(loaded.ef_construction(), std::optional(40)); + EXPECT_EQ(loaded.initialization(), std::optional(QGInitialization::Random)); + EXPECT_EQ(loaded.quantization_bits(), bits); + std::vector expected(5), actual(5); + std::vector expected_distances(5), actual_distances(5); + source.set_ef(64); + loaded.set_ef(64); + source.search(data.data(), 5, expected.data(), expected_distances.data()); + loaded.search(data.data(), 5, actual.data(), actual_distances.data()); + EXPECT_EQ(actual, expected); + EXPECT_EQ(actual_distances, expected_distances); + + // A graph without a builder or a recorded seed stores them as unknown. + QuantizedGraph unknown; + EXPECT_FALSE(unknown.seed().has_value()); + EXPECT_FALSE(unknown.ef_construction().has_value()); + EXPECT_FALSE(unknown.initialization().has_value()); + } +} + +TEST(QgFormat2Test, RejectsContentsThatDoNotFitTheGraph) { + constexpr size_t kCount = 65, kDim = 64, kDegree = 32; + std::vector data(kCount * kDim, 0.5F); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = static_cast((i * 13) % 17) / 8.0F; + } + const rabitq_test::TempFile path_file("rabitq_qg_format2_reject.index"); + const std::string& path = path_file.path(); + QuantizedGraph source( + kCount, kDim, kDegree, METRIC_IP, RotatorType::FhtKacRotator, 4, 3 + ); + QGBuilder builder(source, 40, data.data(), 1, QGInitialization::PiPNN, 3); + builder.build(); + source.save(path.c_str()); + const auto saved = rabitq_test::read_format2_file( + path, rabitq_test::format2::IndexKind::SymphonyQg, rabitq_test::kQgSections + ); + QuantizedGraph target; + target.load(path.c_str()); + const auto expect_rejected = + [&](const char* name, auto edit, const std::string& message) { + SCOPED_TRACE(name); + auto damaged = saved; + edit(damaged); + rabitq_test::write_format2_file(path, damaged); + try { + target.load(path.c_str()); + ADD_FAILURE() << "accepted"; + } catch (const std::runtime_error& error) { + EXPECT_EQ(std::string(error.what()), message); + } + EXPECT_EQ(target.num_vertices(), kCount); + }; + const auto params = [](auto& c) { + return reinterpret_cast(c.bytes(768).data()); + }; + expect_rejected( + "missing centroid", + [](auto& c) { c.remove(769); }, + "QuantizedGraph index file has no section kind 769" + ); + expect_rejected( + "raw vectors with a centroid", + [](auto& c) { + c.header.code_bits = 0; + c.header.raw_type = rabitq_test::format2::RawType::Float32; + }, + "QuantizedGraph index file has a centroid but stores raw vectors" + ); + expect_rejected( + "unknown initialization", + [&](auto& c) { rabitq_test::format2::put_u32(params(c) + 24, 3); }, + "QuantizedGraph index file has invalid graph parameters" + ); + expect_rejected( + "entry point out of range", + [&](auto& c) { rabitq_test::format2::put_u64(params(c) + 8, kCount); }, + "QuantizedGraph index file has invalid graph parameters" + ); + expect_rejected( + "degree bound", + [&](auto& c) { rabitq_test::format2::put_u64(params(c), 48); }, + "QuantizedGraph index file is invalid: QuantizedGraph degree bound must be a " + "positive multiple of 32" + ); + expect_rejected( + "seed above 32 bits", + [](auto& c) { c.header.seed = uint64_t{1} << 40; }, + "QuantizedGraph index file has a seed above 32 bits" + ); + expect_rejected( + "short records", + [](auto& c) { c.bytes(770).pop_back(); }, + "QuantizedGraph index file section kind 770 has " + + std::to_string(rabitq_test::Format2Contents(saved).bytes(770).size() - 1) + + " bytes; expected " + + std::to_string(rabitq_test::Format2Contents(saved).bytes(770).size()) + ); + expect_rejected( + "neighbor out of range", + [&](auto& c) { + std::string& records = c.bytes(770); + // The last 4 bytes are the last vertex's last neighbor ID. + rabitq_test::format2::put_u32( + reinterpret_cast(records.data()) + records.size() - 4, kCount + ); + }, + "QuantizedGraph neighbor ID is invalid" + ); + expect_rejected( + "1-bit codes", + [](auto& c) { c.header.code_bits = 1; }, + "QuantizedGraph index file stores neither raw vectors nor 4- or 8-bit codes with " + "32-bit IDs" + ); +} + +TEST(QgFormat2Test, UnknownReadOnlyFeaturesAllowSearchUntilRebuilt) { + constexpr size_t kCount = 65, kDim = 64, kDegree = 32; + std::vector data(kCount * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = static_cast((i * 7) % 23) / 8.0F; + } + const rabitq_test::TempFile path_file("rabitq_qg_format2_read_only.index"); + const std::string& path = path_file.path(); + QuantizedGraph source( + kCount, kDim, kDegree, METRIC_L2, RotatorType::FhtKacRotator, 8, 2 + ); + QGBuilder(source, 40, data.data(), 1, QGInitialization::PiPNN, 2).build(); + source.save(path.c_str()); + rabitq_test::patch_format2_header(path, 24, 1, 8); // an unknown ro_compat bit + + QuantizedGraph loaded; + ASSERT_NO_THROW(loaded.load(path.c_str())); + std::vector ids(5); + std::vector distances(5); + loaded.set_ef(64); + ASSERT_NO_THROW(loaded.search(data.data(), 5, ids.data(), distances.data())); + try { + loaded.save(path.c_str()); + ADD_FAILURE() << "save was allowed"; + } catch (const std::logic_error& error) { + EXPECT_STREQ( + error.what(), + "QuantizedGraph was loaded from a file with features this release cannot " + "update, so it cannot save" + ); + } + // A builder replaces the data the unknown features described. + QGBuilder(loaded, 40, data.data(), 1, QGInitialization::PiPNN, 2).build(); + EXPECT_NO_THROW(loaded.save(path.c_str())); +} + TEST(QgPersistenceTest, RejectsMalformedPayloadWithoutChangingTarget) { constexpr size_t kCount = 33; constexpr size_t kDim = 64; @@ -1419,38 +1614,57 @@ TEST(QgPersistenceTest, RejectsMalformedPayloadWithoutChangingTarget) { EXPECT_EQ(target.num_vertices(), 65U); EXPECT_EQ(target.metric_type(), METRIC_IP); - source.save(path.c_str()); - { - std::fstream file(path, std::ios::binary | std::ios::in | std::ios::out); - ASSERT_TRUE(file.is_open()); - constexpr std::streamoff kPaddedDimensionOffset = - sizeof(uint64_t) + sizeof(uint32_t) + (3 * sizeof(size_t)); - file.seekp(kPaddedDimensionOffset); - const size_t invalid_padded_dim = kDim * 2; - file.write( - reinterpret_cast(&invalid_padded_dim), sizeof(invalid_padded_dim) - ); - ASSERT_TRUE(file.good()); - } - EXPECT_THROW(target.load(path.c_str()), std::runtime_error); - EXPECT_EQ(target.num_vertices(), 65U); - EXPECT_EQ(target.metric_type(), METRIC_IP); + // Format-2 header fields that do not fit the graph; the earlier layout's are covered + // by the historical persistence tests. + const auto expect_rejected = + [&](size_t offset, uint64_t value, const std::string& message) { + source.save(path.c_str()); + rabitq_test::patch_format2_header(path, offset, value, sizeof(uint64_t)); + try { + target.load(path.c_str()); + ADD_FAILURE() << "accepted: " << message; + } catch (const std::runtime_error& error) { + EXPECT_EQ(std::string(error.what()), message); + } + EXPECT_EQ(target.num_vertices(), 65U); + EXPECT_EQ(target.metric_type(), METRIC_IP); + }; + expect_rejected( + 56, kDim * 2, "QuantizedGraph index file rotator does not match its dimensions" + ); + expect_rejected( + 88, + uint64_t{1} << 32, + "QuantizedGraph index file is invalid: QuantizedGraph point count exceeds the " + "search-buffer ID limit" + ); + // Settings within their limits whose records exceed size_t: 2**31 points with a + // degree bound of 2**30. source.save(path.c_str()); - { - std::fstream file(path, std::ios::binary | std::ios::in | std::ios::out); - ASSERT_TRUE(file.is_open()); - constexpr std::streamoff kPointCountOffset = sizeof(uint64_t) + sizeof(uint32_t); - file.seekp(kPointCountOffset); - const size_t invalid_point_count = std::numeric_limits::max(); - file.write( - reinterpret_cast(&invalid_point_count), sizeof(invalid_point_count) + std::string params = rabitq_test::read_format2_section( + path, rabitq_test::format2::IndexKind::SymphonyQg, rabitq_test::kQgSections, 768 + ); + rabitq_test::format2::put_u64( + reinterpret_cast(params.data()), uint64_t{1} << 30 + ); + rabitq_test::replace_format2_section( + path, + rabitq_test::format2::IndexKind::SymphonyQg, + rabitq_test::kQgSections, + 768, + params + ); + rabitq_test::patch_format2_header(path, 88, uint64_t{1} << 31, sizeof(uint64_t)); + try { + target.load(path.c_str()); + ADD_FAILURE() << "accepted records too large to address"; + } catch (const FormatError& error) { + EXPECT_STREQ( + error.what(), "QuantizedGraph index file describes records too large to address" ); - ASSERT_TRUE(file.good()); } - EXPECT_THROW(target.load(path.c_str()), std::invalid_argument); EXPECT_EQ(target.num_vertices(), 65U); - EXPECT_EQ(target.metric_type(), METRIC_IP); if (std::filesystem::exists("/dev/full")) { EXPECT_THROW(source.save("/dev/full"), std::ios_base::failure); diff --git a/tests/unit/rabitqlib/utils/blob_test.cpp b/tests/unit/rabitqlib/utils/blob_test.cpp new file mode 100644 index 00000000..b7d61100 --- /dev/null +++ b/tests/unit/rabitqlib/utils/blob_test.cpp @@ -0,0 +1,206 @@ +#include "blob.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "crc32.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/core/types.hpp" + +namespace { +using rabitqlib::RotatorType; +namespace detail = rabitqlib::detail; + +std::vector read_fixture(const std::string& name) { + std::ifstream input( + std::string(RABITQ_QUANT_FIXTURE_ROOT) + "/" + name, std::ios::binary + ); + EXPECT_TRUE(input.is_open()) << name; + return {std::istreambuf_iterator(input), std::istreambuf_iterator()}; +} + +std::vector encode( + const rabitqlib::Rotator& rotator, + RotatorType type, + size_t dim, + std::optional seed +) { + std::vector blob(detail::rotator_blob_bytes(type, dim, rotator.size())); + detail::encode_rotator_blob(rotator, type, dim, seed, blob.data()); + return blob; +} + +std::vector rotate(const rabitqlib::Rotator& rotator, size_t dim) { + std::vector input(dim); + for (size_t i = 0; i < dim; ++i) { + input[i] = static_cast(i % 7) - 3.0F; + } + std::vector output(rotator.size()); + rotator.rotate(input.data(), output.data()); + return output; +} + +void put_u32(std::vector& blob, size_t offset, uint32_t value) { + for (size_t i = 0; i < 4; ++i) { + blob[offset + i] = static_cast(value >> (8 * i)); + } +} + +// Recompute the checksum after an edit, so that the check under test fires. +void reseal(std::vector& blob) { + put_u32(blob, 32, 0); + put_u32(blob, 32, detail::crc32(blob.data(), blob.size())); +} + +void expect_rejected(const std::vector& blob, const std::string& message) { + try { + detail::decode_rotator_blob(blob.data(), blob.size()); + ADD_FAILURE() << "expected std::invalid_argument: " << message; + } catch (const std::invalid_argument& error) { + EXPECT_EQ(std::string(error.what()), message); + } +} + +TEST(CoreBlobTest, Crc32MatchesZlib) { + EXPECT_EQ(detail::crc32("", 0), 0U); + EXPECT_EQ(detail::crc32("123456789", 9), 0xCBF43926U); + std::vector data(1024); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = static_cast(i); + } + EXPECT_EQ(detail::crc32(data.data(), data.size()), 0xB70B4C26U); // zlib.crc32 + const uint32_t first = detail::crc32(data.data(), 100); + EXPECT_EQ(detail::crc32(data.data() + 100, data.size() - 100, first), 0xB70B4C26U); +} + +// The golden files rabitqlib.quant writes: the C++ codec reads and rewrites them +// byte for byte, and they hold the same rotation as a freshly seeded rotator. +TEST(CoreBlobTest, ReadsAndRewritesTheQuantGoldenRotators) { + const struct { + const char* name; + size_t dim; + size_t padded_dim; + uint32_t seed; + } cases[] = { + {"rotator_fht_kac_dim64_seed0.bin", 64, 64, 0}, + {"rotator_fht_kac_dim100_padded128_seed42.bin", 100, 128, 42}, + }; + for (const auto& c : cases) { + const std::vector golden = read_fixture(c.name); + const detail::DecodedRotator decoded = + detail::decode_rotator_blob(golden.data(), golden.size()); + EXPECT_EQ(decoded.type, RotatorType::FhtKacRotator); + EXPECT_EQ(decoded.dim, c.dim); + EXPECT_EQ(decoded.padded_dim, c.padded_dim); + ASSERT_TRUE(decoded.seed.has_value()); + EXPECT_EQ(*decoded.seed, c.seed); + EXPECT_EQ( + encode(*decoded.rotator, decoded.type, decoded.dim, decoded.seed), golden + ); + + const std::unique_ptr> fresh( + rabitqlib::choose_rotator( + c.dim, RotatorType::FhtKacRotator, c.padded_dim, c.seed + ) + ); + EXPECT_EQ(rotate(*decoded.rotator, c.dim), rotate(*fresh, c.dim)); + } +} + +TEST(CoreBlobTest, RoundTripsAMatrixRotatorWithoutASeed) { + const std::unique_ptr> rotator( + rabitqlib::choose_rotator(40, RotatorType::MatrixRotator, 64, 7) + ); + const std::vector blob = encode(*rotator, RotatorType::MatrixRotator, 40, {}); + EXPECT_EQ(blob.size(), 64 + (4 * 40 * 64)); + const detail::DecodedRotator decoded = + detail::decode_rotator_blob(blob.data(), blob.size()); + EXPECT_EQ(decoded.type, RotatorType::MatrixRotator); + EXPECT_EQ(decoded.dim, 40U); + EXPECT_EQ(decoded.padded_dim, 64U); + EXPECT_FALSE(decoded.seed.has_value()); + EXPECT_EQ(rotate(*decoded.rotator, 40), rotate(*rotator, 40)); +} + +TEST(CoreBlobTest, EncoderRejectsDimensionsTheFormatDoesNotAllow) { + const std::unique_ptr> rotator( + rabitqlib::choose_rotator(40, RotatorType::MatrixRotator, 64, 7) + ); + std::vector blob( + detail::rotator_blob_bytes(RotatorType::MatrixRotator, 40, 64) + ); + EXPECT_THROW( + detail::encode_rotator_blob( + *rotator, RotatorType::FhtKacRotator, 40, {}, blob.data() + ), + std::invalid_argument + ); + EXPECT_THROW( + detail::encode_rotator_blob( + *rotator, RotatorType::MatrixRotator, 41, {}, blob.data() + ), + std::invalid_argument + ); +} + +TEST(CoreBlobTest, RejectsInvalidBlobsWithTheQuantMessages) { + const std::vector golden = read_fixture("rotator_fht_kac_dim64_seed0.bin"); + const auto edited = [&golden](size_t offset, uint32_t value) { + std::vector blob = golden; + put_u32(blob, offset, value); + reseal(blob); + return blob; + }; + + expect_rejected( + std::vector(golden.begin(), golden.begin() + 39), + "data is too short for a rabitqlib blob" + ); + std::vector blob = golden; + blob[0] = 'X'; + expect_rejected(blob, "data is not a rabitqlib blob (bad magic)"); + expect_rejected(edited(8, 3), "data is not a rotator blob; got kind 3"); + expect_rejected( + edited(12, 2), + "data has unsupported rotator format version 2; this release reads version 1" + ); + expect_rejected( + edited(16, 72), "data has an invalid header size for rotator version 1" + ); + expect_rejected( + std::vector(golden.begin(), golden.end() - 1), + "data length does not match its header (truncated or padded)" + ); + blob = golden; + blob[70] ^= 1; + expect_rejected(blob, "data checksum mismatch (corrupt blob)"); + expect_rejected(edited(20, 2), "data has unknown flags"); + expect_rejected(edited(60, 1), "data has non-zero reserved fields"); + expect_rejected(edited(40, 3), "data has an unknown rotator method 3"); + expect_rejected( + edited(44, 2), "data stores a float64 rotator, which version 1 cannot load" + ); + expect_rejected(edited(44, 9), "data has an unknown scalar type 9"); + expect_rejected( + edited(48, 32), + "data has invalid rotator dimensions: dim must be at least 64 for method " + "'fht_kac'; got 32" + ); + std::vector seeded = + read_fixture("rotator_fht_kac_dim100_padded128_seed42.bin"); + put_u32(seeded, 20, 0); // seed 42 without seed_known + reseal(seeded); + expect_rejected(seeded, "data has a seed without the seed_known flag"); + expect_rejected(edited(52, 96), "data body size does not match the rotator parameters"); +} + +} // namespace