diff --git a/README.md b/README.md index 3d7bb2546..f37f0a878 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,9 @@ This is important because it allows you to take advantage of logarithmic search Introductory tutorials for JVector are available in [docs/tutorials](./docs/tutorials/). Start with the [basic tutorial](./docs/tutorials/1-intro-tutorial.md) or review [VectorIntro.java](./jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/VectorIntro.java) for a simple example using JVector. +For file-backed input datasets, see [Batched input vector access](./docs/batched-vector-access.md), +including runnable 100k-vector examples and optional integration with PQ and NVQ. + The older step-by-step guide for JV can be found [here](./docs/legacy/jvector-step-by-step.md). New users should start with the tutorials mentioned earler, but the step-by-step guide contains useful commentary for advanced users. diff --git a/docs/batched-vector-access.md b/docs/batched-vector-access.md new file mode 100644 index 000000000..c31673018 --- /dev/null +++ b/docs/batched-vector-access.md @@ -0,0 +1,281 @@ +# Batched input vector access + +File-backed input vectors can be supplied without loading the entire dataset into +Java heap. `RandomAccessVectorValues` remains the basic contract. Its optional +`BatchedVectorValues` extension lets algorithms request ranges or arbitrary +selections while the source schedules bounded asynchronous I/O. + +This API concerns the **input vector source**. It does not select the stored graph's +reader, writer, layout, or direct-I/O backend. + +## Runnable hello-world examples + +From the repository root, using the JDK required by the examples module (currently +JDK 23), run. The `jdk20` profile enables the Java vector provider and avoids +requiring a native-library build for these examples: + +```bash +mvn -Pjdk20 -pl jvector-examples -am compile exec:exec@hello-batched-vectors +mvn -Pjdk20 -pl jvector-examples -am compile exec:exec@hello-batched-quantization +``` + +By default, both examples use **e5-small-v2-100k** from JVector's existing +[public dataset catalog](https://jvector-datasets-public.s3.us-east-1.amazonaws.com/datasets-clean/catalog_entries.yaml). +This is a real embedding dataset; its cleaned base contains **99,443 vectors with +384 dimensions**. The catalog loader downloads only the base file on first use and +reuses its cache thereafter. It honors catalog cache settings, including +`DATASET_CACHE_DIR`; there is no eager heap preload and no synthetic data generation. +The first invocation needs network access to the catalog and uncached base file. + +- [HelloBatchedVectors](../jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloBatchedVectors.java) + shows a full range scan, arbitrary selection with duplicates, and copying a borrowed vector. +- [HelloBatchedQuantization](../jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloBatchedQuantization.java) + trains a PQ codebook and encodes the full base file with PQ and NVQ through their usual APIs. + +Pass your own fixed-dimension `.fvecs` file using `-DvectorFile`: + +```bash +mvn -Pjdk20 -pl jvector-examples -am compile exec:exec@hello-batched-vectors \ + -DvectorFile=/path/to/ada_002_100k_base.fvecs +mvn -Pjdk20 -pl jvector-examples -am compile exec:exec@hello-batched-quantization \ + -DvectorFile=/path/to/ada_002_100k_base.fvecs +``` + +Alternatively, run either class's `main` in your IDE with the file path as its single +argument. The quantization example needs at least 256 vectors. Caller-supplied files +and downloaded catalog files are never deleted or rewritten by these examples. +Values are read as stored; the source does not +normalize, deduplicate, or scrub them. + +## Open a file + +`FvecFileVectorValues` associates a file with the capability interface: + +```java +try (var source = FvecFileVectorValues.open(Path.of("vectors.fvecs"))) { + System.out.println(source.size()); + System.out.println(source.dimension()); +} +``` + +The file must remain immutable while open. The source uses positional channel reads, +not memory mapping, and uses the OS page cache. Each fvecs record contains a +little-endian dimension integer followed by that many float32 coordinates; every +record must have the same dimension. I/O failures during consumption propagate as +`UncheckedIOException`. + +## Dataset streaming and caches + +Prepare the input through the existing dataset catalog/download cache, then open +its local `.fvecs` path with `FvecFileVectorValues`. The examples use +`loadBaseVectorFile`, which honors catalog cache settings and `DATASET_CACHE_DIR` +without materializing the base vectors. Downloads must finish before the source +is opened; cursors stream a stable local file, not an in-progress remote download. + +## Purpose and benchmark opt-in + +The input file need not be fully preloaded for quantizer training (samples or full +scans for statistics), batch encoding, or streaming FP/NVQ serialization. Training +samples are temporary working memory; returned models do not retain those samples. +Encoded vectors and graph structures keep their existing memory requirements. + +BenchYAML and AutoBenchYAML preserve resident loading by default. To opt in, set this +root field in the index-parameter YAML used by either runner: + +```yaml +useDatasetPrefetching: true +``` + +Omitted or `false` preserves the resident path. The existing VM option +`-Djvector.dataset_loader=buffered` explicitly overrides the YAML flag; `preload` +overrides the flag in the other direction. Neither runner changes a user's YAML files. + +The harness uses file input for bulk consumption in quantized construction and +existing disk storage. Uncompressed construction, or a search configuration that +requires raw input rather than stored vectors/codes, still uses resident input. +The VM option does not bypass this resident-input routing. If an explicit +`buffered` request requires resident input, the harness logs the reason for using preload. +A mixed configuration requiring resident input remains resident; split those cases +into separate configurations to measure dataset prefetching independently. This is +source routing, not a restriction on the vector-access API. + +Grid consumes each input batch through training/encoding and inline serialization; +the original insertion vector is supplied from that batch, while peer scoring uses +resident quantized codes. Query scoring is unchanged from main. Raw-input exact-query configurations retain +resident input; quantized query paths continue to use their existing codes and +stored reranking vectors. Compaction algorithms and query scoring are unchanged. AutoBenchYAML's separate +legacy compaction-query validation still uses resident original vectors; that pass +retains its existing memory requirement. The current compactor processes saved +partition indexes, not the original input file. Any future original-input training, +encoding or serialization step can use the same bulk hooks. + +The source retains bounded reusable read buffers during consumption. It does not +change persistent dataset, index or quantizer caches. File preparation remains +outside construction timing; input consumption is included in the corresponding +library calls. Queries and ground truth remain resident. Files must be complete and +immutable, and buffered loading requires `NO_SCRUB` metadata. Existing direct calls +to `DataSets.loadDataSet(name)` still preload. Application-provided RAVVs need no +changes; future quantizers opt into `VectorAccess` for their bulk input needs. + +Use a distinct AutoBenchYAML output/checkpoint file for each loading mode so earlier +completed results are not reused as a different mode. The harness closes its input +source after the configuration finishes. + +## Random ordinals: fetch an arbitrary selection + +Use a selection when the algorithm knows several requested ordinals ahead of time, +for example a random training sample. The ordinal list can be unsorted and can +contain duplicates. In this small example, every ordinal exists in the default +100k dataset: + +```java +int[] ordinals = {42, 7, 90_000, 42}; +List> retained = new ArrayList<>(); +try (var cursor = VectorAccess.openSelection(source, ordinals, 0, ordinals.length)) { + while (cursor.next()) { + // File cursor storage is reused: copy vectors kept for later computation. + retained.add(cursor.vector().copy()); + System.out.println("Loaded ordinal " + cursor.ordinal()); + } +} +// retained is in order 42, 7, 90000, 42 and remains valid after cursor closure. +``` + +For a training sample, `VectorAccess.copySelected(source, ordinals, executor)` +performs this materialization in parallel, with source-aware scheduling and the +same order/duplicate guarantees. The algorithm selects the ordinals; the source +only retrieves them. For a single unpredictable point request, the existing +`getVector(ordinal)` API remains available; copy a shared result before retaining it. + +## Scan ordinals: process the file in order + +Use a range for a full pass, such as computing statistics or encoding vectors. +Consume each borrowed vector before advancing; no copy is needed in this example: + +```java +int scanned = 0; +double firstCoordinateSum = 0; +try (var cursor = VectorAccess.openRange(source, 0, source.size())) { + while (cursor.next()) { + firstCoordinateSum += cursor.vector().get(0); + scanned++; + } +} +System.out.printf("Scanned %d vectors; sum = %f%n", scanned, firstCoordinateSum); +``` + +For a partition, `openRange(source, 10, 20)` reads ordinals 10 through 19. +`openRange(startInclusive, endExclusive)` uses an exclusive end; equal endpoints +are an empty range, including `openRange(size, size)`. Selection arguments instead +use an array slice: `openSelection(ordinals, offset, count)`. Keep the selection +array unchanged until its cursor closes. + +`VectorAccess` selects the optional batching capability or falls back to ordinary +point access. Algorithms should normally use these utilities rather than inspecting +the concrete file-source type. The runnable `HelloBatchedVectors` example demonstrates +both patterns, including a full scan of the real catalog dataset. + +## Use with quantization and indexes + +The algorithm chooses its vectors and performs computation. `VectorAccess` handles +capability selection and fallback. The source owns file format, scheduling, I/O +workers, buffering, and read-ahead. + +PQ chooses its training sample as before and materializes it with: + +```java +List> training = VectorAccess.copySelected(source, ordinals, executor); +``` + +For capable sources, the utility schedules selected reads in storage order, then +restores the original sample order and duplicates. PQ's sample policy and centroid +learning are unchanged. The point-access fallback retains the existing parallel +path and avoids copying non-shared resident vectors. + +PQ, NVQ, and binary quantization bulk encoding use `VectorAccess.forEach`; NVQ's mean computation uses a +range cursor and preserves its accumulation order. Existing calls to +`ProductQuantization.compute`, `encodeAll`, and `NVQuantization.compute` accept the +file source directly, as shown in the runnable quantization example. + +For a new quantizer, use the full-scan overload without generating all ordinals: + +```java +VectorAccess.forEach(source, executor, (ordinal, vector) -> { + // Encode or consume the borrowed vector here; copy it only if retaining it. + encoded[ordinal] = compressor.encode(vector); +}); +``` + +Use `copySelected` for retained training samples, `openRange` for an ordered +sequential reduction, and the mapped `forEach` overload when output positions +refer to different input ordinals. These utilities work with existing sources +without buffered hooks. Full-precision graph construction retains its existing +resident source in the harness. The optional bulk hooks are not a substitute for +resident raw vectors in comparison-heavy construction. A caller can still use the +ordinary point-access API; no library-wide policy is imposed. +The source does not change distance metrics, quantizer parameters, graph pruning, +or hierarchy. Algorithm restrictions still apply, such as FusedPQ requiring +256 clusters and anisotropic PQ requiring normalized training vectors. + +Other quantizers and indexes can adopt the same utilities independently. Existing +integrations can continue supplying ordinary `RandomAccessVectorValues`; they do +not have to implement batching. No algorithm-specific policy belongs in the file +source, and alternative sources can implement the same interface. + +## Configure bounded read-ahead + +Most callers use `open(path)`. Immutable named options customize individual limits: + +```java +var options = FvecFileVectorValues.Options.defaults() + .withMaxBufferBytes(32L << 20) + .withBatchVectors(64) + .withReadAhead(3); +try (var source = FvecFileVectorValues.open(path, options)) { + // Use VectorAccess or ordinary point access. +} +``` + +The default I/O worker count follows `PhysicalCoreExecutor`, the compute sizing +used by Grid and BenchYAML. It adapts to the machine and honors +`-Djvector.physical_core_count`; there is no fixed core count. Blocking file reads +for range and selection cursors run on a separate source-owned pool. Unpredictable +point reads execute on the requesting thread using exclusively leased channels; +point and asynchronous reads share the I/O concurrency limit. `withIoThreads(count)` remains available for explicit benchmark tuning. The +quantization example uses the shared compute pool and does not shut it down. + +Other defaults are at most 64 vectors per batch and three batches of +read-ahead. Payload capacity is capped at `min(1% of file bytes, 64 MiB)`, with one +record as the minimum. A configured maximum is also capped at 1% of the file. +Copies and cursors share the budget, including queued, in-flight and ready payloads, +point-read buffers. +Read-ahead zero disables speculative batches. Selected reads fetch requested records +without reading the gaps between sparse samples. + +This is an **I/O payload budget**, not a total-process memory limit. Retained training +vectors, encoded outputs, graph data, scratch vectors, JVM/channel internals and the +OS page cache consume additional memory. `statistics()` reports logical source +reads and payload allocation, not physical device I/O. + +## Ownership and closure + +- A cursor has one consumer. Independent cursors can run concurrently. +- A failed advance invalidates the cursor. Its value and ordinal are unavailable; + open a new cursor to retry. Other views remain usable if the source is still open. +- A shared cursor vector is borrowed until the next advance or close. Call `.copy()` + before retaining it. Non-shared resident sources allow retention without copying. +- Point-access views have their own scratch vector and are single-consumer. Use + `copy()` or `threadLocalSupplier()` for concurrent point access. +- Close cursors with try-with-resources, stop consumers, then close the root source. +- Copies share the root's storage and I/O executor. Closing a copy affects that view + only; closing the root invalidates all copies and releases the shared resources. +- If point reads are still running when the root closes, close waits for each read + in progress to finish; later point reads fail with `IllegalStateException`. +- Direct payload buffers become eligible for JVM reclamation after references are + released; closure does not promise immediate native-memory reclamation. + +## Measure input loading + +The [vector prefetching benchmark](vector-prefetch-benchmark.md) compares the exact +buffered source with full preload and demand-only reads, timing sample retrieval, +training and full-set encoding separately. It runs from the command line or IntelliJ. diff --git a/docs/release notes/4.0.2/744.feature.md b/docs/release notes/4.0.2/744.feature.md new file mode 100644 index 000000000..750d2cf28 --- /dev/null +++ b/docs/release notes/4.0.2/744.feature.md @@ -0,0 +1,15 @@ +Added optional batched vector access and a bounded file source for quantizer +training, encoding and existing full-precision/NVQ disk serialization. Capable +algorithms can consume samples or full scans without requiring an upfront preload +of the original base dataset; existing vector sources remain supported. + +BenchYAML and AutoBenchYAML support `useDatasetPrefetching: true` in their index +configuration. The flag defaults to false, preserving resident loading and scoring. +Comparison-heavy uncompressed configurations retain resident input. The existing +`jvector.dataset_loader` VM option overrides the flag, while configurations +requiring resident raw vectors still preload; an explicit buffered request logs +the reason for that fallback. + +See [batched vector access](../../batched-vector-access.md) for ownership, +resource limits, examples and the benchmark opt-in. Encoded vectors, training +working memory and graph structures retain their existing memory requirements. diff --git a/docs/tutorials/3-larger-than-memory-tutorial.md b/docs/tutorials/3-larger-than-memory-tutorial.md index 1374f3c11..3389f6e9a 100644 --- a/docs/tutorials/3-larger-than-memory-tutorial.md +++ b/docs/tutorials/3-larger-than-memory-tutorial.md @@ -14,6 +14,11 @@ In the previous tutorials, we built indexes in-memory then wrote them to disk. H > > Read the [PQ paper](https://ieeexplore.ieee.org/document/5432202) for more details. +For a file-backed vector source with bounded asynchronous reads, see +[Batched input vector access](../batched-vector-access.md). Its hello-world examples +show how to supply a `.fvecs` file directly to PQ and NVQ without preloading the +complete input dataset. + ## Loading the Dataset We'll change things up with a different dataset: diff --git a/docs/vector-prefetch-benchmark.md b/docs/vector-prefetch-benchmark.md new file mode 100644 index 000000000..dfa95ae92 --- /dev/null +++ b/docs/vector-prefetch-benchmark.md @@ -0,0 +1,172 @@ +# Vector prefetching benchmark + +`VectorPrefetchBenchmark` retrieves a reproducible training sample, trains a new +quantizer, calls `encodeAll` on the **entire base candidate set**, then scores the +encoded vectors against one query as a sink. This isolates the input operations +that buffered access can benefit: arbitrary sample retrieval and full-set encoding. +Training and scoring are timed separately; scoring uses encoded resident output, +not the input prefetcher. This is not an index build or recall measurement. + +Each launch compares three arms **one at a time**, in fresh JVMs: + +1. **preload**: the usual `SiftLoader.readFvecs` full resident preload + **outside comparison timing**; the sample is then selected + from resident vectors inside timing. This is the ideal resident-input baseline. +2. **demand**: synchronous positional reads of requested records, without a cache, + I/O executor, batching or speculative read-ahead. Float decoding matches the PR source. +3. **prefetch**: unmodified `FvecFileVectorValues.open(path)` defaults from this PR, + through production `VectorAccess` and quantizer APIs. This comparison includes + the PR's batching/scheduling as well as prefetching; it is not a read-ahead-only ablation. + +Catalog lookup/downloads, selecting ordinal IDs and preparing one query happen +before timing. Each fresh JVM initially has no resident base vectors in Java. The +preload arm loads its complete base set before its comparison timer starts; file +arms start timing before opening their source and include required input reads. By default, +64-bit Linux `posix_fadvise(POSIX_FADV_DONTNEED)` evicts only the base input file; +`mincore` verifies zero resident pages before preload preparation or file-arm timing. No privileged global-cache +drop is used. Verification failure stops the comparison, without retries. Run with +no other readers of this input. Do not use eviction on a live database's files. +After the initial cold start, later passes naturally reuse OS cache pages. + +`--cache uncontrolled` explicitly skips eviction (also usable on other platforms). +It is labeled **not cold**. Neither policy makes JVM startup, code compilation or +storage-controller caches identical across arms: these are fresh-process pipeline +measurements, not warmed JMH kernel measurements. Repeat in different arm orders +using `--mode` for performance conclusions. + +### Command line + +Build with JDK 23 from the repository root: + +```bash +mvn -Pjdk20 -pl jvector-examples -am -DskipTests package +``` + +Use the resulting `jvector-examples/target/jvector-examples-*-jar-with-dependencies.jar` +as `$JV_EXAMPLES`. Select an appropriate caller-controlled heap for the resident +baseline; no machine-specific heap or core count is built into these benchmarks. + +```bash +java --add-modules=jdk.incubator.vector --enable-native-access=ALL-UNNAMED \ + -cp "$JV_EXAMPLES" io.github.jbellis.jvector.example.VectorPrefetchBenchmark \ + --dataset cohere-english-v3-10M --samples 128000 --quantizer pq +``` + +Use any name in your dataset catalog, or `--file /path/to/base.fvecs` instead of +`--dataset`. The default is the public catalog's `e5-small-v2-100k`; sampling defaults +to 128,000 training vectors, or the candidate count if smaller. `DATASET_CACHE_DIR` applies to catalog preparation. Downloaded +files and existing caches are never deleted or rewritten. Files must be immutable. + +### IntelliJ + +Create an Application configuration using the main class +`io.github.jbellis.jvector.example.VectorPrefetchBenchmark`. + +Select the `jvector-examples` module classpath and JDK 23. Use the repository root +as working directory. VM options: + +```text +--add-modules=jdk.incubator.vector --enable-native-access=ALL-UNNAMED +``` + +Program arguments are the same as the command line, e.g. +`--dataset cohere-english-v3-10M --samples 128000 --quantizer pq`. +Heap and provider settings from the launching JVM are inherited by each arm; +debugger/profiler agents are not inherited, avoiding port conflicts. Child JVM +output appears in the same IntelliJ console. Stopping the launcher also stops its +own active child. Start only one microbenchmark configuration at a time. + +### Options and output + +- `--mode all|preload|demand|prefetch`: all is the default; individual modes let you + choose another order or avoid preload when it cannot fit in memory. +- `--quantizer nvq|pq`: NVQ is the default, with one subvector. PQ uses the production + training policy, width approximately eight, 256 clusters, no global centering and + no anisotropic weighting. It needs at least 256 selected candidates. Its internal + random training is independent across arms, so nearest-vector results can differ. +- `--metric DOT_PRODUCT|EUCLIDEAN|COSINE`: dot product is the default. Values are + used as stored; no normalization or scrubbing is performed. +- `--query-file /path/to/queries.fvecs`: use its first query. Otherwise, use the last + base vector as a held-out query and exclude it from all candidate sets. Thus the + default full scan encodes N−1 candidates, and avoids a trivial self-match. +- `--samples N`, `--seed 42`: training sample count and deterministic ordinal seed; all + arms use identical ordinals in identical order. The complete ordinal array is + supplied to `VectorAccess.copySelected` up front; the PR schedules bounded + asynchronous read-ahead across concurrent selection cursors, rather than + enqueueing the entire set at once. Count must fit the candidate set. +- `--progress-seconds 30`: phase transitions and infrequent elapsed-time updates; + `0` disables periodic updates. There is no per-vector logging or monitoring scan. + +Console tables show excluded preload preparation, sample retrieval, training, encoding, +sink and comparison total elapsed seconds, +the winning original ordinal, its score and a sum of all scores. For preload, totals +start **after the entire base set is resident**; its file-load preparation is reported +separately and excluded. For demand and prefetch, totals start before file-source +open and include input reads. All totals finish when the sink completes; JVM launch, +download, ordinal generation, cold-cache preparation, and resource closure are excluded. +This measures the penalty of file-backed input relative to already resident input, +not the time to load and process the dataset from scratch. Sample retrieval includes materialization and the loader's source-open work for +file arms. Training operates on those resident samples. Encoding always operates +on the complete candidate source, so file-arm encoding includes required input I/O. +Source-open time alone is not the time to load the dataset. NVQ's `compute` stage +computes the selected training set's mean; per-vector parameter learning is in +full-set encoding. Use `--samples` equal to the candidate count when evaluating +NVQ's usual full-base mean. Retained training vectors are separate from the +prefetcher's bounded I/O payload budget. +Encoded vectors remain resident for the sink; this is not a bounded-total-memory +encoding pipeline. Preload additionally retains the complete raw base set. + +Source counters report **logical** reads/bytes, not SSD traffic. No graph is built, +no saved quantizer is reused, and the benchmark does not measure index recall or QPS. + +### Quantizer parameters and extension points + +Pass named parameters without changing the harness: + +```text +--quantizer pq --quantizer-options subspaces=128,clusters=256,center=false,anisotropic=-1 +--quantizer nvq --quantizer-options subvectors=2,learn=true +``` + +Full precision is not a built-in benchmark arm yet. Once the required input hooks +are supported by a full-precision vector implementation, an adapter can materialize +or stream those vectors through `VectorAccess` and supply its defined JVector distance +function as the sink. Do not treat current quantized results as measurements of that +future path. + +`--loader-options ioThreads=8,batchVectors=64,readAhead=3,maxBufferBytes=67108864` +configures the existing PR source when using `--mode prefetch`. With no loader +options, the exact PR defaults are used. Run individual modes when passing loader +options because preload/demand reject unused parameters rather than ignore them. + +A future quantizer can supply a public no-argument adapter implementing +`VectorPrefetchBenchmark.Quantizer`: + +```java +public class MyQuantizerBenchmark implements VectorPrefetchBenchmark.Quantizer { + public VectorPrefetchBenchmark.Encoder train(RandomAccessVectorValues source, + Map parameters, + ForkJoinPool executor) { + return VectorPrefetchBenchmark.compressed( + MyQuantizer.compute(source, parameters, executor)); + } +} +``` + +Run with `--quantizer my.package.MyQuantizerBenchmark`. The existing +`VectorCompressor.encodeAll` and `CompressedVectors.scoreFunctionFor` perform +encoding and scoring; there is no algorithm copied into the benchmark. ASH or a +future representation that does not implement those interfaces can instead return +its own `Encoder` and `Encoded`: encoding accepts the source and compute pool; +`Encoded` supplies `count()` and a per-ordinal `scorer(query, metric)`. The adapter +calls that representation's production methods. The score must be finite, with +larger meaning nearer, as for JVector similarities. This also provides the extension +point for full precision once its required hooks are supported. + +Similarly, a public no-argument `VectorPrefetchBenchmark.Loader` adapter opens a +source and returns `VectorPrefetchBenchmark.Input`, exposing `values()` and `close()`. +Use `--mode my.package.MyLoaderBenchmark --loader-options key=value,...`. +The source can implement `BatchedVectorValues` or ordinary `RandomAccessVectorValues`; +the benchmark preserves that capability when restricting a held-out query. A custom +loader must preserve the file's vector count, dimension and ordinal/value mapping. +Adapters are example-harness extension points, not new core-library requirements. diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/FvecFileVectorValues.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/FvecFileVectorValues.java new file mode 100644 index 000000000..8a1815d94 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/FvecFileVectorValues.java @@ -0,0 +1,668 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.disk; + +import io.github.jbellis.jvector.graph.BatchedVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.VectorCursor; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; + +import java.io.EOFException; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.FloatBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.LinkedHashMap; +import java.util.ArrayList; +import java.util.ArrayDeque; +import java.util.Objects; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.LongAdder; +import java.util.function.Supplier; + +/** + * A read-only, fixed-dimension fvecs source with bounded asynchronous range/selection + * reads. The file must remain unchanged while open. Copies share the channel, executor + * and reusable payload budget, including cached point vectors; each copy and cursor + * has its own borrowed vector. + * + *

Close all cursors and stop consumers before closing the source. Closing the root + * releases all resources, including resources used by copies. A point read already in + * progress finishes before its channel is closed; later point reads, including ones + * waiting for a reader, fail with {@link IllegalStateException}. No memory mapping or + * OS cache eviction is performed. Retained training vectors are caller-owned memory. + */ +public final class FvecFileVectorValues implements BatchedVectorValues, AutoCloseable { + private static final VectorTypeSupport VTS = VectorizationProvider.getInstance().getVectorTypeSupport(); + private static final long DEFAULT_MAX_BYTES = 64L << 20; + private static final int MAX_ZERO_READS = 16; + private final Storage storage; + private final boolean ownsStorage; + private final VectorFloat scratch; + private int previousOrdinal = -2; + private Read currentRead; + private boolean closed; + + /** + * Open with at most min(1% of file bytes, 64 MiB) of reusable I/O buffers, + * I/O workers sized by {@link PhysicalCoreExecutor}, and three batches of read-ahead. + * One record is the minimum budget; blocking reads use a separate owned executor. + */ + public static FvecFileVectorValues open(Path path) throws IOException { + return open(path, Options.defaults()); + } + + /** Open a file with named limits. The returned root owns the shared I/O resources. */ + public static FvecFileVectorValues open(Path path, Options options) throws IOException { + Objects.requireNonNull(options, "options"); + return new FvecFileVectorValues(new Storage(path, options.maxBufferBytes, options.ioThreads, + options.batchVectors, options.readAhead), true); + } + + /** + * Immutable source options. Start with {@link #defaults()} and change only the + * limits needed. Buffer capacity is shared across all copies and cursors; + * retained vectors and the OS page cache are outside this payload budget. + */ + public static final class Options { + private final long maxBufferBytes; + private final int ioThreads, batchVectors, readAhead; + + private Options(long maxBufferBytes, int ioThreads, int batchVectors, int readAhead) { + if (maxBufferBytes <= 0 || ioThreads <= 0 || batchVectors <= 0 || readAhead < 0) + throw new IllegalArgumentException("Invalid buffered source limits"); + this.maxBufferBytes = maxBufferBytes; + this.ioThreads = ioThreads; + this.batchVectors = batchVectors; + this.readAhead = readAhead; + } + + /** Defaults: 64 MiB maximum, Grid compute parallelism, 64 vectors/batch, 3 batches ahead. */ + public static Options defaults() { + return new Options(DEFAULT_MAX_BYTES, PhysicalCoreExecutor.getPhysicalCoreCount(), 64, 3); + } + + /** Maximum payload bytes; also capped at 1% of file bytes, with one record minimum. */ + public long maxBufferBytes() { return maxBufferBytes; } + /** Number of asynchronous I/O workers. */ + public int ioThreads() { return ioThreads; } + /** Maximum vectors per batch; reduced when necessary to fit the shared budget. */ + public int batchVectors() { return batchVectors; } + /** Number of speculative batches ahead of demand; zero disables read-ahead. */ + public int readAhead() { return readAhead; } + + /** Return options with a different positive maximum payload budget. */ + public Options withMaxBufferBytes(long bytes) { + return new Options(bytes, ioThreads, batchVectors, readAhead); + } + /** Return options with a different positive I/O worker count. */ + public Options withIoThreads(int threads) { + return new Options(maxBufferBytes, threads, batchVectors, readAhead); + } + /** Return options with a different positive maximum batch size. */ + public Options withBatchVectors(int vectors) { + return new Options(maxBufferBytes, ioThreads, vectors, readAhead); + } + /** Return options with a different nonnegative number of batches ahead. */ + public Options withReadAhead(int batches) { + return new Options(maxBufferBytes, ioThreads, batchVectors, batches); + } + } + + private FvecFileVectorValues(Storage storage, boolean ownsStorage) { + this.storage = storage; + this.ownsStorage = ownsStorage; + scratch = VTS.createFloatVector(storage.dimension); + } + + @Override public int size() { return storage.size; } + @Override public int dimension() { return storage.dimension; } + @Override public boolean isValueShared() { return true; } + + private void ensureOpen() { + if (closed || storage.closed) throw new IllegalStateException("Vector source is closed"); + } + + /** + * Supply one point-access view per thread. Views share the root's I/O resources + * and must stop being used before the root closes. + */ + @Override public Supplier threadLocalSupplier() { + var local = ThreadLocal.withInitial(this::copy); + return local::get; + } + + /** + * Create a view with its own scratch vector. Closing this view leaves other views + * usable; closing the root invalidates every view. Point access is single-consumer. + */ + @Override public FvecFileVectorValues copy() { + ensureOpen(); + return new FvecFileVectorValues(storage, false); + } + + @Override public VectorFloat getVector(int ordinal) { + getVectorInto(ordinal, scratch, 0); + return scratch; + } + + @Override public void getVectorInto(int ordinal, VectorFloat destination, int offset) { + ensureOpen(); + Objects.checkIndex(ordinal, size()); + Objects.checkFromIndexSize(offset, dimension(), destination.length()); + boolean sequential = ordinal == previousOrdinal + 1; + boolean withinCurrent = currentRead != null && ordinal >= currentRead.request.start + && ordinal < currentRead.request.start + currentRead.request.count; + // A point call is not a scan: adjacent graph ordinals must not enter the shared scheduler. + if (storage.pointReaders.length != 0) { + if (currentRead != null) currentRead.consumed = true; + currentRead = null; + storage.readPointInto(ordinal, destination, offset); + previousOrdinal = ordinal; + return; + } + if (!withinCurrent) { + if (currentRead != null) currentRead.consumed = true; + int first = sequential ? ordinal / storage.batchVectors * storage.batchVectors : ordinal; + int count = sequential ? Math.min(storage.batchVectors, size() - first) : 1; + currentRead = storage.decode(new Request(first, count, null), ordinal - first, destination, offset); + storage.prefetchRange(first + count, size(), sequential ? storage.readAhead : 0); + } else if (!storage.tryDecode(currentRead, ordinal - currentRead.request.start, destination, offset)) { + currentRead = storage.decode(currentRead.request, ordinal - currentRead.request.start, destination, offset); + } + previousOrdinal = ordinal; + } + + @Override public VectorCursor openRange(int startInclusive, int endExclusive) { + ensureOpen(); + Objects.checkFromToIndex(startInclusive, endExclusive, size()); + return new Cursor(startInclusive, endExclusive - startInclusive, null); + } + + @Override public VectorCursor openSelection(int[] ordinals, int offset, int count) { + ensureOpen(); + Objects.checkFromIndexSize(offset, count, ordinals.length); + for (int i = offset; i < offset + count; i++) Objects.checkIndex(ordinals[i], size()); + return new Cursor(offset, count, ordinals); + } + + /** Cumulative source I/O counters; bytes are logical input reads, not device traffic. */ + public Statistics statistics() { + synchronized (storage) { + return new Statistics(storage.reads.sum(), storage.bytes.sum(), storage.allocatedBytes, storage.budget); + } + } + + public static final class Statistics { + public final long reads, bytesRead, bufferBytes, bufferLimitBytes; + private Statistics(long reads, long bytes, long allocated, long limit) { + this.reads = reads; bytesRead = bytes; bufferBytes = allocated; bufferLimitBytes = limit; + } + } + + /** + * Close this view. For the root, first close all cursors and stop all consumers; + * root closure shuts down shared I/O and invalidates every copy. + */ + @Override public void close() throws IOException { + if (closed) return; + closed = true; + currentRead = null; + if (ownsStorage) storage.close(); + } + + private final class Cursor implements VectorCursor { + private final int start, count; + private final int[] ordinals; + private final VectorFloat value = VTS.createFloatVector(dimension()); + private int position = -1; + private boolean cursorClosed; + private Read current; + + Cursor(int start, int count, int[] ordinals) { + this.start = start; this.count = count; this.ordinals = ordinals; + prefetch(0); + } + + private Request request(int position) { + int first = position / storage.batchVectors * storage.batchVectors; + return new Request(start + first, Math.min(storage.batchVectors, count - first), ordinals); + } + + private void prefetch(int position) { + for (int i = 0; i <= storage.readAhead; i++) { + long next = (long) position + (long) i * storage.batchVectors; + if (next >= count) break; + storage.schedule(request((int) next)); + } + } + + @Override public boolean next() { + ensureOpen(); + if (cursorClosed) throw new IllegalStateException("Cursor is closed"); + if (position == count) return false; + if (++position == count) { release(); return false; } + try { + int row = position % storage.batchVectors; + if (row == 0) { + if (current != null) current.consumed = true; + prefetch(position); + current = storage.decode(request(position), row, value, 0); + } else if (!storage.tryDecode(current, row, value, 0)) { + current = storage.decode(current.request, row, value, 0); + } + return true; + } catch (RuntimeException | Error failure) { + // A failed advance cannot expose stale data or resume at the next ordinal. + try { close(); } + catch (RuntimeException | Error cleanup) { + if (cleanup != failure) failure.addSuppressed(cleanup); + } + throw failure; + } + } + + private void checkPosition() { + ensureOpen(); + if (cursorClosed || position < 0 || position >= count) throw new IllegalStateException("No current vector"); + } + @Override public int ordinal() { checkPosition(); return ordinals == null ? start + position : ordinals[start + position]; } + @Override public VectorFloat vector() { checkPosition(); return value; } + @Override public boolean isValueShared() { return true; } + private void release() { current = null; storage.discard(ordinals, start, count); } + @Override public void close() { if (!cursorClosed) { cursorClosed = true; release(); } } + } + + private static final class Request { + final int start, count; + final int[] ordinals; + Request(int start, int count, int[] ordinals) { this.start = start; this.count = count; this.ordinals = ordinals; } + int ordinal(int i) { return ordinals == null ? start + i : ordinals[start + i]; } + @Override public int hashCode() { return 31 * (31 * System.identityHashCode(ordinals) + start) + count; } + @Override public boolean equals(Object other) { + if (!(other instanceof Request)) return false; + Request r = (Request) other; + return r.start == start && r.count == count && r.ordinals == ordinals; + } + } + + private static final class Buffer extends ByteBufferReader { + final FloatBuffer floats; + Buffer(int bytes) { + // Read directly into the pooled payload; heap buffers require a temporary native copy. + super(ByteBuffer.allocateDirect(bytes).order(ByteOrder.LITTLE_ENDIAN)); + floats = bb.asFloatBuffer(); + } + @Override public void read(float[] destination, int offset, int count) { + Objects.checkFromIndexSize(offset, count, destination.length); + Objects.checkFromIndexSize(bb.position(), Math.multiplyExact(count, 4), bb.limit()); + floats.position(bb.position() / 4); + floats.get(destination, offset, count); + bb.position(bb.position() + count * 4); + } + @Override public void readFully(float[] destination) { read(destination, 0, destination.length); } + } + + private static final class Read { + final Request request; + final Buffer buffer; + final FutureTask task; + int users; + volatile boolean consumed, invalidated; + Read(Request request, Buffer buffer, Storage storage) { + this.request = request; this.buffer = buffer; + task = new FutureTask<>(() -> { storage.read(request, buffer); return buffer; }) { + @Override protected void done() { synchronized (storage) { storage.notifyAll(); } } + }; + } + } + + /** Exclusively leased point-read resources; channel count is bounded by the I/O limit. */ + private static final class PointReader { + final int index; + final Semaphore permit = new Semaphore(1); + Buffer buffer; + volatile FileChannel channel; + PointReader(int index, int bytes) { this.index = index; buffer = new Buffer(bytes); } + } + + private static final class Storage { + final FileChannel channel; + final Path path; + final PointReader[] pointReaders; + final ThreadLocal pointHint; + final long pointBytes; + final int dimension, size, rowBytes, batchVectors, readAhead, slots; + final long budget; + final ThreadPoolExecutor executor; + final LinkedHashMap cache = new LinkedHashMap<>(16, .75f, true); + final LongAdder reads = new LongAdder(), bytes = new LongAdder(); + final ArrayDeque free = new ArrayDeque<>(); + long allocatedBytes; + int allocated; + volatile boolean closed; + + Storage(Path path, long maxBytes, int threads, int maxBatch, int ahead) throws IOException { + this.path = path; + channel = FileChannel.open(path, StandardOpenOption.READ); + try { + ByteBuffer header = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN); + readFully(channel, header, 0); header.flip(); + dimension = header.getInt(); + if (dimension <= 0 || dimension > Integer.MAX_VALUE / 4 - 1) + throw new IOException("Invalid fvec dimension: " + dimension); + rowBytes = (dimension + 1) * 4; + long length = channel.size(); + if (length % rowBytes != 0 || length / rowBytes > Integer.MAX_VALUE) + throw new IOException("Invalid fvec file length: " + length); + size = (int) (length / rowBytes); + budget = Math.max(rowBytes, Math.min(maxBytes, length / 100)); + // Point reads bypass asynchronous cache scheduling; their buffers share the same budget. + int points = (int) Math.min(threads, Math.min((budget - rowBytes) / rowBytes, + Math.max(1, budget / rowBytes / 16))); + pointBytes = (long) points * rowBytes; + pointReaders = new PointReader[points]; + for (int i = 0; i < points; i++) pointReaders[i] = new PointReader(i, rowBytes); + var nextPointReader = new AtomicInteger(); + pointHint = points == 0 ? null : ThreadLocal.withInitial(() -> + pointReaders[Math.floorMod(nextPointReader.getAndIncrement(), points)]); + long batchBudget = budget - pointBytes; + // Leave room for concurrent consumers and their configured look-ahead. + long concurrentRows = batchBudget / rowBytes / threads / (ahead + 1L); + batchVectors = (int) Math.max(1, Math.min(Math.min(maxBatch, Integer.MAX_VALUE / rowBytes), concurrentRows)); + readAhead = ahead; + slots = (int) Math.min(Integer.MAX_VALUE, batchBudget / (batchVectors * (long) rowBytes)); + allocatedBytes = pointBytes; + executor = new ThreadPoolExecutor(threads, threads, 0, TimeUnit.MILLISECONDS, + new ArrayBlockingQueue<>(slots), runnable -> { + Thread t = new Thread(runnable, "jvector-input-io"); t.setDaemon(true); return t; + }); + } catch (Throwable failure) { channel.close(); throw failure; } + } + + private FileChannel pointChannel(PointReader reader) throws IOException { + // Caller owns the lease; interruption can close only this channel. + if (reader.channel == null || !reader.channel.isOpen()) + reader.channel = FileChannel.open(path, StandardOpenOption.READ); + return reader.channel; + } + + private void read(Request request, Buffer buffer) throws IOException { + PointReader reader = null; + try { + // Bulk and point reads share leases, bounding total I/O without a second global permit. + if (pointReaders.length != 0) reader = borrowPointReader(); + readWithChannel(request, buffer, reader == null ? channel : pointChannel(reader)); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new IOException("Input read interrupted", error); + } finally { + if (reader != null) reader.permit.release(); + } + } + + private void readWithChannel(Request request, Buffer buffer, FileChannel input) throws IOException { + synchronized (buffer) { + buffer.bb.clear(); + boolean contiguous = true; + for (int i = 1; i < request.count && contiguous; i++) + contiguous = request.ordinal(i) == request.ordinal(0) + i; + if (contiguous) { + buffer.bb.limit(request.count * rowBytes); + readFully(input, buffer.bb, (long) request.ordinal(0) * rowBytes); + reads.increment(); bytes.add((long) request.count * rowBytes); + } else { + for (int i = 0; i < request.count; i++) { + buffer.bb.limit((i + 1) * rowBytes); + readFully(input, buffer.bb, (long) request.ordinal(i) * rowBytes); + reads.increment(); bytes.add(rowBytes); + } + } + buffer.bb.flip(); + } + } + + private PointReader borrowPointReader() throws InterruptedException { + var hint = pointHint.get(); + for (int i = 0; i < pointReaders.length; i++) { + var candidate = pointReaders[(hint.index + i) % pointReaders.length]; + if (candidate.permit.tryAcquire()) return candidate; + } + // All readers are busy. Block instead of spinning on cache or queue state. + hint.permit.acquire(); + return hint; + } + + void readPointInto(int ordinal, VectorFloat destination, int offset) { + PointReader reader = null; + try { + reader = borrowPointReader(); + // Close waits for each lease, so a closed source can't reopen this reader's channel. + if (closed) throw new IllegalStateException("Vector source is closed"); + pointHint.set(reader); + var input = pointChannel(reader); + var buffer = reader.buffer; + buffer.bb.clear(); + readFully(input, buffer.bb, (long) ordinal * rowBytes); + buffer.bb.flip(); + if (buffer.readInt() != dimension) + throw new IOException("Unexpected fvec dimension at " + ordinal); + VTS.readFloatVector(buffer, dimension, destination, offset); + reads.increment(); bytes.add(rowBytes); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Input read interrupted", error); + } catch (IOException error) { throw new UncheckedIOException(error); } + finally { + if (reader != null) reader.permit.release(); + } + } + + synchronized Read schedule(Request request) { + if (closed) throw new IllegalStateException("Vector source is closed"); + Read present = cache.get(request); + if (present != null) return present; + Buffer buffer = free.pollFirst(); + if (buffer == null && allocated < slots) { + buffer = new Buffer(batchVectors * rowBytes); + allocated++; allocatedBytes += batchVectors * (long) rowBytes; + } else if (buffer == null) { + Read victim = null; + for (Read r : cache.values()) { + // Never recycle an in-flight or actively decoded payload. + if (r.users == 0 && r.task.isDone() && (victim == null || r.consumed)) { + victim = r; + if (r.consumed) break; + } + } + if (victim != null) { victim.invalidated = true; cache.remove(victim.request); buffer = victim.buffer; } + } + if (buffer == null) return null; + Read read = new Read(request, buffer, this); + cache.put(request, read); + try { executor.execute(read.task); } + catch (RejectedExecutionException e) { + // There can be at most slots reads, including queued and in-flight reads. + cache.remove(request); allocated--; allocatedBytes -= buffer.bb.capacity(); + throw e; + } + return read; + } + + Read decode(Request request, int row, VectorFloat destination, int offset) { + Read read; + synchronized (this) { + while ((read = schedule(request)) == null) { + try { wait(); } + catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new IllegalStateException("Input read interrupted", e); } + } + // Demand owns a lease before waiting, so ready data cannot be recycled + // between I/O completion and decoding. Speculation alone never pins data. + read.users++; + } + try { + read.task.get(); + synchronized (read.buffer) { + read.buffer.seek((long) row * rowBytes); + int dim = read.buffer.readInt(); + if (dim != dimension) throw new IOException("Unexpected fvec dimension " + dim + " at " + request.ordinal(row)); + VTS.readFloatVector(read.buffer, dimension, destination, offset); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Input read interrupted", e); + } catch (ExecutionException e) { + if (e.getCause() instanceof IOException) throw new UncheckedIOException((IOException) e.getCause()); + throw new IllegalStateException("Input read failed", e.getCause()); + } catch (IOException e) { throw new UncheckedIOException(e); } + finally { + synchronized (this) { + read.users--; + if (row == request.count - 1) read.consumed = true; + notifyAll(); + } + } + return read; + } + + boolean tryDecode(Read read, int row, VectorFloat destination, int offset) { + if (read == null || read.invalidated || !read.task.isDone()) return false; + synchronized (read.buffer) { + // Reuse is synchronized on the payload, independently of the scheduler. + // Within a batch, no future or shared scheduler lookup is needed. + if (read.invalidated) return false; + try { + read.buffer.seek((long) row * rowBytes); + int dim = read.buffer.readInt(); + if (dim != dimension) throw new IOException("Unexpected fvec dimension " + dim + " at " + read.request.ordinal(row)); + VTS.readFloatVector(read.buffer, dimension, destination, offset); + return true; + } catch (IOException e) { throw new UncheckedIOException(e); } + } + } + + void prefetchRange(int start, int end, int ahead) { + for (int i = 0; i < ahead; i++) { + long first = (long) start + (long) i * batchVectors; + if (first >= end) break; + schedule(new Request((int) first, Math.min(batchVectors, end - (int) first), null)); + } + } + + void discard(int[] ordinals, int start, int count) { + var pending = new ArrayList(); + synchronized (this) { + for (Read r : cache.values()) { + if (r.request.ordinals == ordinals && r.request.start >= start + && (long) r.request.start < (long) start + count) { + r.consumed = true; + pending.add(r); + } + } + } + // Drain every read before allowing the caller to reuse its ordinal array, + // including when an earlier read failed or this thread was interrupted. + boolean interrupted = false; + RuntimeException failure = null; + for (Read r : pending) { + boolean done = false; + while (!done) { + try { r.task.get(); done = true; } + catch (InterruptedException e) { interrupted = true; } + catch (ExecutionException e) { + done = true; + RuntimeException error = e.getCause() instanceof IOException + ? new UncheckedIOException((IOException) e.getCause()) + : new IllegalStateException("Input read failed", e.getCause()); + if (failure == null) failure = error; else failure.addSuppressed(error); + } + } + } + synchronized (this) { + for (Read r : pending) { + if (cache.get(r.request) == r && r.users == 0) { + r.invalidated = true; + cache.remove(r.request); + free.addLast(r.buffer); + } + } + notifyAll(); + } + if (interrupted) Thread.currentThread().interrupt(); + if (failure != null) throw failure; + } + + void close() throws IOException { + synchronized (this) { if (closed) return; closed = true; notifyAll(); } + executor.shutdown(); + boolean interrupted = false; + while (!executor.isTerminated()) { + try { executor.awaitTermination(1, TimeUnit.SECONDS); } + catch (InterruptedException e) { interrupted = true; } + } + IOException failure = null; + for (PointReader reader : pointReaders) { + // Wait for an in-flight point read to finish; readers that lease afterwards see closed. + reader.permit.acquireUninterruptibly(); + try { if (reader.channel != null) reader.channel.close(); } + catch (IOException error) { + if (failure == null) failure = error; else failure.addSuppressed(error); + } finally { reader.channel = null; reader.buffer = null; reader.permit.release(); } + } + if (pointHint != null) pointHint.remove(); + try { channel.close(); } + catch (IOException error) { + if (failure == null) failure = error; else failure.addSuppressed(error); + } finally { + synchronized (this) { cache.clear(); free.clear(); } + if (interrupted) Thread.currentThread().interrupt(); + } + if (failure != null) throw failure; + } + } + + static void readFully(FileChannel channel, ByteBuffer buffer, long offset) throws IOException { + int start = buffer.position(); + int zeroReads = 0; + while (buffer.hasRemaining()) { + int n = channel.read(buffer, offset + buffer.position() - start); + if (n < 0) throw new EOFException("Truncated fvec record at " + offset); + if (n == 0) { + // A broken channel must not spin forever or prevent the I/O pool from closing. + if (++zeroReads >= MAX_ZERO_READS) + throw new IOException("No progress reading fvec data at " + (offset + buffer.position() - start)); + Thread.onSpinWait(); + } else { + zeroReads = 0; + } + } + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/BatchedVectorValues.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/BatchedVectorValues.java new file mode 100644 index 000000000..5964ef23b --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/BatchedVectorValues.java @@ -0,0 +1,39 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +/** + * Optional range and selection access for a RandomAccessVectorValues source. Algorithms + * describe the vectors they need; the source owns batching, I/O concurrency and read-ahead. + * Implementations must share their I/O memory limit across all copies and cursors. + */ +public interface BatchedVectorValues extends RandomAccessVectorValues { + /** + * Open a cursor for [startInclusive, endExclusive), using an exclusive end. Invalid ranges + * fail before scheduling I/O; equal endpoints produce an empty cursor. Opening + * may start asynchronous reads, but must not materialize the entire range. + * Close the cursor with try-with-resources before closing its source. + */ + VectorCursor openRange(int startInclusive, int endExclusive); + + /** + * Open a cursor for the specified slice, preserving request order and duplicates. + * The caller must not modify the ordinal array until the cursor closes. Invalid + * ordinals fail before scheduling I/O. Memory is bounded independently of count. + */ + VectorCursor openSelection(int[] ordinals, int offset, int count); +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/VectorAccess.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/VectorAccess.java new file mode 100644 index 000000000..7bef30e49 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/VectorAccess.java @@ -0,0 +1,193 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.vector.types.VectorFloat; + +import java.util.Arrays; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.ForkJoinPool; +import java.util.function.IntUnaryOperator; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +/** Source-independent vector access with the existing point-access path as a fallback. */ +public final class VectorAccess { + private VectorAccess() {} + + // Bound mapped-request metadata independently of the dataset size. + private static final int MAX_SELECTION_ORDINALS = 4096; + + @FunctionalInterface + public interface VectorConsumer { + /** The vector is borrowed for this call; copy it before retaining it. */ + void accept(int position, VectorFloat vector); + } + + /** + * Open [startInclusive, endExclusive), including for sources that only implement + * point access. The end ordinal is excluded. + * Algorithms should use this utility to retain the ordinary-source fallback. + */ + public static VectorCursor openRange(RandomAccessVectorValues source, int startInclusive, int endExclusive) { + Objects.checkFromToIndex(startInclusive, endExclusive, source.size()); + if (source instanceof BatchedVectorValues) + return ((BatchedVectorValues) source).openRange(startInclusive, endExclusive); + return pointCursor(source, startInclusive, endExclusive - startInclusive, null); + } + + /** Open a selection in request order, including duplicates. */ + public static VectorCursor openSelection(RandomAccessVectorValues source, int[] ordinals, int offset, int count) { + Objects.checkFromIndexSize(offset, count, ordinals.length); + for (int i = offset; i < offset + count; i++) Objects.checkIndex(ordinals[i], source.size()); + if (source instanceof BatchedVectorValues) + return ((BatchedVectorValues) source).openSelection(ordinals, offset, count); + return pointCursor(source, offset, count, ordinals); + } + + private static VectorCursor pointCursor(RandomAccessVectorValues source, int start, int count, int[] ordinals) { + RandomAccessVectorValues values = source.isValueShared() ? source.copy() : source; + return new VectorCursor() { + int position = -1; + VectorFloat current; + boolean closed; + public boolean next() { + if (closed) throw new IllegalStateException("Cursor is closed"); + if (position == count) return false; + if (++position == count) { current = null; return false; } + try { + current = values.getVector(ordinals == null ? start + position : ordinals[start + position]); + return true; + } catch (RuntimeException | Error failure) { + // Match file cursors: a failed advance invalidates the current value. + close(); + throw failure; + } + } + private void checkPosition() { + if (closed || position < 0 || position >= count) throw new IllegalStateException("No current vector"); + } + public int ordinal() { checkPosition(); return ordinals == null ? start + position : ordinals[start + position]; } + public VectorFloat vector() { checkPosition(); return current; } + public boolean isValueShared() { return values.isValueShared(); } + public void close() { closed = true; current = null; } + }; + } + + /** + * Materialize a selection, retaining its order and duplicates. Only capable sources + * use storage-order scheduling; the sample policy and returned training order stay unchanged. + * Retained vectors are algorithm working memory, separate from the source's I/O budget. + */ + public static List> copySelected(RandomAccessVectorValues source, int[] ordinals, ForkJoinPool executor) { + if (!(source instanceof BatchedVectorValues)) { + var local = source.threadLocalSupplier(); + return executor.submit(() -> IntStream.of(ordinals).parallel().mapToObj(ordinal -> { + var values = local.get(); + var v = values.getVector(ordinal); + return v != null && values.isValueShared() ? v.copy() : v; + }).collect(Collectors.>toList())).join(); + } + // Pack the original position so sorting the reads cannot reorder the training set. + long[] order = new long[ordinals.length]; + for (int i = 0; i < ordinals.length; i++) { + Objects.checkIndex(ordinals[i], source.size()); + order[i] = ((long) ordinals[i] << 32) | (i & 0xffffffffL); + } + Arrays.sort(order); + int[] sorted = new int[order.length]; + for (int i = 0; i < order.length; i++) sorted[i] = (int) (order[i] >>> 32); + VectorFloat[] result = new VectorFloat[ordinals.length]; + int partitions = Math.min(ordinals.length, executor.getParallelism()); + executor.submit(() -> IntStream.range(0, partitions).parallel().forEach(part -> { + int start = (int) ((long) part * sorted.length / partitions); + int end = (int) ((long) (part + 1) * sorted.length / partitions); + try (var cursor = openSelection(source, sorted, start, end - start)) { + int i = start; + while (cursor.next()) { + var v = cursor.vector(); + result[(int) order[i++]] = v != null && cursor.isValueShared() ? v.copy() : v; + } + if (i != end) throw new IllegalStateException("Incomplete vector selection"); + } + })).join(); + return Arrays.asList(result); + } + + /** + * Consume a full scan without materializing an ordinal array. Future compressors can + * use this method for encoding and openSelection/copySelected for training samples. + * Ordinary sources retain their existing thread-local point-access path. + */ + public static void forEach(RandomAccessVectorValues source, ForkJoinPool executor, + VectorConsumer consumer) { + forEach(source, source.size(), executor, consumer); + } + + /** Consume the first count vectors without materializing their ordinals. */ + public static void forEach(RandomAccessVectorValues source, int count, ForkJoinPool executor, + VectorConsumer consumer) { + Objects.checkFromToIndex(0, count, source.size()); + if (!(source instanceof BatchedVectorValues)) { + forEach(source, count, IntUnaryOperator.identity(), executor, consumer); + return; + } + int partitions = Math.min(count, executor.getParallelism()); + executor.submit(() -> IntStream.range(0, partitions).parallel().forEach(part -> { + int start = (int) ((long) part * count / partitions); + int end = (int) ((long) (part + 1) * count / partitions); + try (var cursor = openRange(source, start, end)) { + int i = start; + while (cursor.next()) consumer.accept(i++, cursor.vector()); + if (i != end) throw new IllegalStateException("Incomplete vector range"); + } + })).join(); + } + + /** + * Consume vectors in parallel, with contiguous request partitions for capable sources. + * The callback position is the output position, not the mapped input ordinal. + */ + public static void forEach(RandomAccessVectorValues source, int count, IntUnaryOperator mapping, + ForkJoinPool executor, VectorConsumer consumer) { + if (count < 0) throw new IllegalArgumentException("Negative vector count"); + if (!(source instanceof BatchedVectorValues)) { + var local = source.threadLocalSupplier(); + executor.submit(() -> IntStream.range(0, count).parallel().forEach(i -> + consumer.accept(i, local.get().getVector(mapping.applyAsInt(i))))).join(); + return; + } + int partitions = Math.min(count, executor.getParallelism()); + executor.submit(() -> IntStream.range(0, partitions).parallel().forEach(part -> { + int start = (int) ((long) part * count / partitions); + int end = (int) ((long) (part + 1) * count / partitions); + int[] ordinals = new int[Math.min(MAX_SELECTION_ORDINALS, end - start)]; + for (int first = start; first < end; ) { + int length = Math.min(ordinals.length, end - first); + for (int i = 0; i < length; i++) ordinals[i] = mapping.applyAsInt(first + i); + try (var cursor = openSelection(source, ordinals, 0, length)) { + int i = first; + while (cursor.next()) consumer.accept(i++, cursor.vector()); + if (i != first + length) throw new IllegalStateException("Incomplete vector selection"); + } + // Closing drains reads before reusing the selection array. + first += length; + } + })).join(); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/VectorCursor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/VectorCursor.java new file mode 100644 index 000000000..3f6361867 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/VectorCursor.java @@ -0,0 +1,47 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.vector.types.VectorFloat; + +/** + * A single-consumer cursor over an ordered vector request. Close it to release read-ahead. + * Independent cursors may be consumed concurrently; the source must outlive its cursors. + * Use try-with-resources. A shared vector is borrowed: call {@code vector().copy()} + * before retaining it beyond the next advance or cursor closure. + * I/O failures are reported as {@link java.io.UncheckedIOException}. + */ +public interface VectorCursor extends AutoCloseable { + /** Advance to the next requested vector, or return false when exhausted. */ + boolean next(); + + /** The requested ordinal at the current position; valid after a successful next(). */ + int ordinal(); + + /** + * The current vector (possibly null for a missing value). When isValueShared() is true, + * it is borrowed until the next next() or close(); copy it before retaining it. + */ + VectorFloat vector(); + + /** False permits retaining a vector without copying, as with an in-memory source. */ + boolean isValueShared(); + + /** Release outstanding requests and buffers. Idempotent; does not close the source. */ + @Override + void close(); +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java index 1049069de..8e1f91cf8 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java @@ -110,6 +110,29 @@ static BuildScoreProvider randomAccessScoreProvider(RandomAccessVectorValues rav var vectorsCopy = ravv.threadLocalSupplier(); return new BuildScoreProvider() { + /// Exact scores against `node1`. + /// + /// For a value-shared source the anchor is read on the first comparison and kept as + /// a private copy, so creating a score function reads nothing, later reads through the + /// same views can't overwrite the anchor, and each comparison reads only `node2`. + private ScoreFunction.ExactScoreFunction scoreFunctionFor(int node1) { + if (ravv.isValueShared()) { + return new ScoreFunction.ExactScoreFunction() { + private VectorFloat anchor; + + @Override + public float similarityTo(int node2) { + var a = anchor; + if (a == null) anchor = a = vectors.get().getVector(node1).copy(); + return similarityFunction.compare(a, vectorsCopy.get().getVector(node2)); + } + }; + } + var vector = vectors.get().getVector(node1); + var vc = vectorsCopy.get(); + return node2 -> similarityFunction.compare(vector, vc.getVector(node2)); + } + @Override public boolean isExact() { return true; @@ -119,13 +142,12 @@ public boolean isExact() { public VectorFloat approximateCentroid() { var vv = vectors.get(); var centroid = vts.createFloatVector(vv.dimension()); - for (int i = 0; i < vv.size(); i++) { + int size = vv.size(); + for (int i = 0; i < size; i++) { var v = vv.getVector(i); - if (v != null) { // MapRandomAccessVectorValues is not necessarily dense - VectorUtil.addInPlace(centroid, v); - } + if (v != null) VectorUtil.addInPlace(centroid, v); } - VectorUtil.scale(centroid, 1.0f / vv.size()); + VectorUtil.scale(centroid, 1.0f / size); return centroid; } @@ -137,25 +159,17 @@ public SearchScoreProvider searchProviderFor(VectorFloat vector) { @Override public SearchScoreProvider searchProviderFor(int node1) { - RandomAccessVectorValues randomAccessVectorValues = vectors.get(); - var v = randomAccessVectorValues.getVector(node1); - return searchProviderFor(v); + return new DefaultSearchScoreProvider(scoreFunctionFor(node1)); } @Override public SearchScoreProvider diversityProviderFor(int node1) { - RandomAccessVectorValues randomAccessVectorValues = vectors.get(); - var v = randomAccessVectorValues.getVector(node1); - var vc = vectorsCopy.get(); - return DefaultSearchScoreProvider.exact(v, similarityFunction, vc); + return new DefaultSearchScoreProvider(scoreFunctionFor(node1)); } @Override public ScoreFunction diversityScoreFunctionFor(int node1) { - var v = vectors.get().getVector(node1); - var vc = vectorsCopy.get(); - // don't use ESF.reranker, we need thread safety here - return (ScoreFunction.ExactScoreFunction) node2 -> similarityFunction.compare(v, vc.getVector(node2)); + return scoreFunctionFor(node1); } }; } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java index f0d660301..0ebc4ab51 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java @@ -19,6 +19,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.VectorAccess; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; @@ -26,7 +27,6 @@ import java.io.IOException; import java.util.Objects; import java.util.concurrent.ForkJoinPool; -import java.util.stream.IntStream; /** * Binary Quantization of float vectors: each float is compressed to a single bit, @@ -64,18 +64,11 @@ public CompressedVectors createCompressedVectors(Object[] compressedVectors) { @Override public CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { - var ravvCopy = ravv.threadLocalSupplier(); - var cv = simdExecutor.submit(() -> IntStream.range(0, ravv.size()) - .parallel() - .mapToObj(i -> { - var localRavv = ravvCopy.get(); - VectorFloat v = localRavv.getVector(i); - return v == null - ? new long[compressedVectorSize() / Long.BYTES] - : encode(v); - }) - .toArray(long[][]::new)) - .join(); + var cv = new long[ravv.size()][]; + VectorAccess.forEach(ravv, simdExecutor, (ordinal, vector) -> + cv[ordinal] = vector == null + ? new long[compressedVectorSize() / Long.BYTES] + : encode(vector)); return new ImmutableBQVectors(this, cv); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java index aef0325b9..7b0ab2334 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java @@ -20,6 +20,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.VectorAccess; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.util.Accountable; import io.github.jbellis.jvector.vector.VectorUtil; @@ -33,8 +34,6 @@ import java.util.List; import java.util.Objects; import java.util.concurrent.ForkJoinPool; -import java.util.stream.Collectors; -import java.util.stream.IntStream; import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; import static io.github.jbellis.jvector.vector.VectorUtil.sub; @@ -148,13 +147,12 @@ private NVQuantization(int[][] subvectorSizesAndOffsets, VectorFloat globalMe * @param nSubVectors number of subvectors */ public static NVQuantization compute(RandomAccessVectorValues ravv, int nSubVectors) { - var ravvCopy = ravv.threadLocalSupplier().get(); - var dim = ravvCopy.getVector(0).length(); - var globalMean = vectorTypeSupport.createFloatVector(dim); - for (int i = 0; i < ravvCopy.size(); i++) { - VectorUtil.addInPlace(globalMean, ravvCopy.getVector(i)); + Objects.checkIndex(0, ravv.size()); // Preserve rejection of an empty training source. + var globalMean = vectorTypeSupport.createFloatVector(ravv.dimension()); + try (var cursor = VectorAccess.openRange(ravv, 0, ravv.size())) { + while (cursor.next()) VectorUtil.addInPlace(globalMean, cursor.vector()); } - VectorUtil.scale(globalMean, 1.0f / ravvCopy.size()); + VectorUtil.scale(globalMean, 1.0f / ravv.size()); return create(globalMean, nSubVectors); } @@ -180,17 +178,12 @@ public CompressedVectors createCompressedVectors(Object[] compressedVectors) { */ @Override public NVQVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool parallelExecutor) { - var ravvCopy = ravv.threadLocalSupplier(); - return new NVQVectors(this, - parallelExecutor.submit(() -> IntStream.range(0, ravv.size()) - .parallel() - .mapToObj(i -> { - var localRavv = ravvCopy.get(); - VectorFloat v = localRavv.getVector(i); - return encode(v); - }) - .toArray(QuantizedVector[]::new)) - .join()); + var encoded = new QuantizedVector[ravv.size()]; + // VectorCompressor requires missing ordinals to encode as zero vectors. + var zero = vectorTypeSupport.createFloatVector(originalDimension); + VectorAccess.forEach(ravv, parallelExecutor, + (ordinal, vector) -> encoded[ordinal] = encode(vector == null ? zero : vector)); + return new NVQVectors(this, encoded); } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java index 760eb38ee..5e5f172dc 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java @@ -19,6 +19,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.VectorAccess; import io.github.jbellis.jvector.graph.disk.CompactionContext; import io.github.jbellis.jvector.graph.disk.QuantizationCompactionStrategy; import io.github.jbellis.jvector.graph.similarity.ScoreFunction; @@ -92,7 +93,7 @@ public static PQVectors load(RandomAccessReader in, long offset) throws IOExcept * @return the PQVectors instance */ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { - return encodeAndBuild(pq, vectorCount, IntUnaryOperator.identity(), ravv, simdExecutor); + return encodeAndBuild(pq, vectorCount, IntUnaryOperator.identity(), ravv, simdExecutor, true); } /** @@ -107,16 +108,16 @@ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vect * @return the PQVectors instance */ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, IntUnaryOperator ordinalsMapping, RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { - // Verify that the mapped ordinals are packed, ie. the total range of ordinals does not exceed - // the total vector count as they are mapped from 0 --> vector count - 1. - // This assumes no duplicates are included. - // Note that the total vector count cannot exceed 2B as that is Integer.MAX_VALUE + return encodeAndBuild(pq, vectorCount, ordinalsMapping, ravv, simdExecutor, false); + } + + private static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, + IntUnaryOperator ordinalsMapping, RandomAccessVectorValues ravv, ForkJoinPool simdExecutor, + boolean identityMapping) { + // Validate input ordinals, not output positions; mappings may select or repeat vectors. if (log.isDebugEnabled() || System.getProperties().containsKey("VECTOR_DEBUG")) { - IntStream.range(0, vectorCount).forEach(i -> { - int mappedOrdinal = ordinalsMapping.applyAsInt(i); - if (mappedOrdinal > vectorCount) - throw new IllegalArgumentException("vectorCount=" + mappedOrdinal + " exceeds maximum " + vectorCount); - }); + IntStream.range(0, vectorCount).forEach(i -> + Objects.checkIndex(ordinalsMapping.applyAsInt(i), ravv.size())); log.info("Range check completed for encodeAndBuild with vectorCount={}", vectorCount); } @@ -133,20 +134,17 @@ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vect // Encode the vectors in parallel into the compressed data chunks // The changes are concurrent, but because they are coordinated and do not overlap, we can use parallel streams // and then we are guaranteed safe publication because we join the thread after completion. - var ravvCopy = ravv.threadLocalSupplier(); - simdExecutor.submit(() -> IntStream.range(0, vectorCount) - .parallel() - .forEach(ordinal -> { - // Retrieve the slice and mutate it. - var localRavv = ravvCopy.get(); - var slice = PQVectors.get(chunks, ordinal, layout.fullChunkVectors, pq.getSubspaceCount()); - var vector = localRavv.getVector(ordinalsMapping.applyAsInt(ordinal)); - if (vector != null) - pq.encodeTo(vector, slice); - else - slice.zero(); - })) - .join(); + VectorAccess.VectorConsumer encode = (ordinal, vector) -> { + var slice = PQVectors.get(chunks, ordinal, layout.fullChunkVectors, pq.getSubspaceCount()); + if (vector != null) + pq.encodeTo(vector, slice); + else + slice.zero(); + }; + if (identityMapping) + VectorAccess.forEach(ravv, vectorCount, simdExecutor, encode); + else + VectorAccess.forEach(ravv, vectorCount, ordinalsMapping, simdExecutor, encode); return new ImmutablePQVectors(pq, chunks, vectorCount, layout.fullChunkVectors); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java index 37823db28..1da149475 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java @@ -20,6 +20,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.VectorAccess; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.util.Accountable; import io.github.jbellis.jvector.util.PhysicalCoreExecutor; @@ -167,15 +168,7 @@ static List> extractTrainingVectors(RandomAccessVectorValues ravv ordinalStream = IntStream.of(ordinalArray); } - var ravvCopy = ravv.threadLocalSupplier(); - return parallelExecutor.submit(() -> ordinalStream.parallel() - .mapToObj(targetOrd -> { - var localRavv = ravvCopy.get(); - VectorFloat v = localRavv.getVector(targetOrd); - return localRavv.isValueShared() ? v.copy() : v; - }) - .collect(Collectors.toList())) - .join(); + return VectorAccess.copySelected(ravv, ordinalStream.toArray(), parallelExecutor); } /** diff --git a/jvector-examples/README.md b/jvector-examples/README.md index c782f7db8..46990445b 100644 --- a/jvector-examples/README.md +++ b/jvector-examples/README.md @@ -2,6 +2,21 @@ JVector comes with the following sample programs to try: +### Batched file-backed vectors + +Two hello-world programs use public-catalog `e5-small-v2-100k`, or accept an existing +`.fvecs` file through `-DvectorFile=/path/to/vectors.fvecs`: + +```bash +mvn -Pjdk20 -pl jvector-examples -am compile exec:exec@hello-batched-vectors +mvn -Pjdk20 -pl jvector-examples -am compile exec:exec@hello-batched-quantization +``` + +Run from the repository root with JDK 23. The first demonstrates scans, selections +and vector ownership; the second trains and encodes with PQ and NVQ. See +[Batched input vector access](../docs/batched-vector-access.md) for configuration, +source lifecycle, and integration guidance. + ### SiftSmall A simple benchmark for the sift dataset located in the [siftsmall](./siftsmall) directory in the project root. @@ -87,3 +102,13 @@ RESULT [2,1,0] * `BULKLOAD {localpath}` * Bulk loads a local file in numpy format Rows x Columns + +### Vector prefetching benchmark + +Compare preload, synchronous demand reads, and the PR's buffered source with +`io.github.jbellis.jvector.example.VectorPrefetchBenchmark`. Use `--samples N` for training sample size; +each launch samples, trains, and encodes the entire candidate set, running the +three arms sequentially in fresh JVMs. PQ/NVQ and named quantization parameters are selectable. + +See [command-line and IntelliJ instructions](../docs/vector-prefetch-benchmark.md) +for cold-file-cache verification, timing boundaries and future vector/loader adapters. diff --git a/jvector-examples/pom.xml b/jvector-examples/pom.xml index b86dd37a9..843c8ef25 100644 --- a/jvector-examples/pom.xml +++ b/jvector-examples/pom.xml @@ -13,6 +13,9 @@ JVector Examples 2.21.10 + + + @@ -21,6 +24,7 @@ maven-surefire-plugin false + ${argLine} ${project.parent.basedir} @@ -31,6 +35,34 @@ false ${project.parent.basedir} + + + hello-batched-vectors + exec + + + -classpath + + --add-modules=jdk.incubator.vector + io.github.jbellis.jvector.example.HelloBatchedVectors + ${vectorFile} + + + + + hello-batched-quantization + exec + + + -classpath + + --add-modules=jdk.incubator.vector + io.github.jbellis.jvector.example.HelloBatchedQuantization + ${vectorFile} + + + + org.apache.maven.plugins @@ -245,6 +277,9 @@ jdk20 + + --add-modules jdk.incubator.vector + io.github.jbellis @@ -365,6 +400,9 @@ jdk22 + + --add-modules jdk.incubator.vector --enable-native-access=ALL-UNNAMED + true diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java index 81aec4178..a3d1a80bc 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java @@ -20,7 +20,7 @@ import io.github.jbellis.jvector.example.util.BenchmarkSummarizer; import io.github.jbellis.jvector.example.util.BenchmarkSummarizer.SummaryStats; import io.github.jbellis.jvector.example.util.CheckpointManager; -import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; +import io.github.jbellis.jvector.example.benchmarks.datasets.BaseVectorLoading; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets; import io.github.jbellis.jvector.example.yaml.DatasetCollection; import io.github.jbellis.jvector.example.yaml.MultiConfig; @@ -43,7 +43,7 @@ * Automated benchmark runner for GitHub Actions workflow. * This class is specifically designed to handle the --output argument * for regression testing in the run-bench.yml workflow. - * + * * The benchmark runner supports checkpointing to allow resuming from failures. * It creates a checkpoint file (outputPath + ".checkpoint.json") that records * which datasets have been fully processed. If the benchmark is restarted, @@ -85,7 +85,7 @@ public static void main(String[] args) throws IOException { } String finalConfigPath = configPath; String[] filteredArgs = Arrays.stream(args) - .filter(arg -> !arg.equals("--output") && !arg.equals(finalOutputPath) && + .filter(arg -> !arg.equals("--output") && !arg.equals(finalOutputPath) && !arg.equals("--config") && !arg.equals(finalConfigPath)) .toArray(String[]::new); @@ -121,13 +121,7 @@ public static void main(String[] args) throws IOException { continue; } - logger.info("Loading dataset: {}", datasetName); try { - DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow( - () -> new RuntimeException("Dataset " + datasetName + " not found") - ).getDataSet(); - logger.info("Dataset loaded: {} with {} vectors", datasetName, ds.getBaseVectors().size()); - String normalizedDatasetName = datasetName; if (normalizedDatasetName.endsWith(".hdf5")) { normalizedDatasetName = normalizedDatasetName.substring(0, normalizedDatasetName.length() - ".hdf5".length()); @@ -145,38 +139,56 @@ public static void main(String[] args) throws IOException { config.dataset = normalizedDatasetName; } logger.info("Using configuration: {}", config); + var loading = BaseVectorLoading.forBenchmark(config); - List datasetResults = Grid.runAllAndCollectResults(ds, - config.construction.useSavedIndexIfExists, - config.construction.outDegree, - config.construction.efConstruction, - config.construction.neighborOverflow, - config.construction.addHierarchy, - config.construction.refineFinalGraph, - config.construction.getFeatureSets(), - config.construction.getCompressorParameters(), - config.search.getCompressorParameters(), - config.search.topKOverquery, - config.search.useSearchPruning); - results.addAll(datasetResults); - - logger.info("Benchmark completed for dataset: {}", datasetName); - // Mark dataset as completed and update checkpoint, passing results - checkpointManager.markDatasetCompleted(datasetName, datasetResults); - - // Compaction regression (opt-in via withCompaction: [Yes]) — failures are non-fatal and don't block checkpointing - if (config.isCompactionEnabled()) { - try { - logger.info("Running compaction benchmark for dataset: {}", datasetName); - List datasetCompactionResults = CompactionBench.run(ds); - compactionResults.addAll(datasetCompactionResults); - logger.info("Compaction benchmark completed for dataset: {} ({} configs)", - datasetName, datasetCompactionResults.size()); - } catch (Exception e) { - logger.error("Compaction benchmark failed for dataset {}", datasetName, e); + logger.info("Loading dataset: {}", datasetName); + try (var ds = DataSets.loadDataSet(datasetName, loading).orElseThrow( + () -> new RuntimeException("Dataset " + datasetName + " not found") + ).getDataSet()) { + logger.info("Dataset loaded: {} with {} vectors", datasetName, ds.getBaseVectors().size()); + + List datasetResults = Grid.runAllAndCollectResults(ds, + config.construction.useSavedIndexIfExists, + config.construction.outDegree, + config.construction.efConstruction, + config.construction.neighborOverflow, + config.construction.addHierarchy, + config.construction.refineFinalGraph, + config.construction.getFeatureSets(), + config.construction.getCompressorParameters(), + config.search.getCompressorParameters(), + config.search.topKOverquery, + config.search.useSearchPruning); + results.addAll(datasetResults); + + logger.info("Benchmark completed for dataset: {}", datasetName); + // Mark dataset as completed and update checkpoint, passing results + checkpointManager.markDatasetCompleted(datasetName, datasetResults); + + // Compaction regression is opt-in; failures remain non-fatal. + if (config.isCompactionEnabled()) { + try { + logger.info("Running compaction benchmark for dataset: {}", datasetName); + List datasetCompactionResults; + if (loading == BaseVectorLoading.BUFFERED) { + // This legacy validation queries raw originals; keep its resident source. + try (var resident = DataSets.loadDataSet(datasetName, BaseVectorLoading.PRELOAD) + .orElseThrow(() -> new RuntimeException("Dataset " + datasetName + " not found")) + .getDataSet()) { + datasetCompactionResults = CompactionBench.run(resident); + } + } else { + datasetCompactionResults = CompactionBench.run(ds); + } + compactionResults.addAll(datasetCompactionResults); + logger.info("Compaction benchmark completed for dataset: {} ({} configs)", + datasetName, datasetCompactionResults.size()); + } catch (Exception e) { + logger.error("Compaction benchmark failed for dataset {}", datasetName, e); + } + } else { + logger.info("Skipping compaction benchmark for dataset: {} (withCompaction not enabled)", datasetName); } - } else { - logger.info("Skipping compaction benchmark for dataset: {} (withCompaction not enabled)", datasetName); } } catch (Exception e) { logger.error("Exception while processing dataset {}", datasetName, e); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java index e066a34dc..e5ab3d970 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java @@ -16,7 +16,7 @@ package io.github.jbellis.jvector.example; -import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; +import io.github.jbellis.jvector.example.benchmarks.datasets.BaseVectorLoading; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets; import io.github.jbellis.jvector.example.reporting.RunArtifacts; import io.github.jbellis.jvector.example.reporting.SearchReportingCatalog; @@ -116,27 +116,29 @@ public static void main(String[] args) throws IOException { for (var config : allConfigs) { + var loading = BaseVectorLoading.forBenchmark(config); String datasetName = config.dataset; - DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow( + try (var ds = DataSets.loadDataSet(datasetName, loading).orElseThrow( () -> new RuntimeException("Could not load dataset:" + datasetName) - ).getDataSet(); - // Register dataset info the first time we actually load the dataset for benchmarking - artifacts.registerDataset(datasetName, ds); - - Grid.runAll(ds, - config.construction.useSavedIndexIfExists, - config.construction.outDegree, - config.construction.efConstruction, - config.construction.neighborOverflow, - config.construction.addHierarchy, - config.construction.refineFinalGraph, - config.construction.getFeatureSets(), - config.construction.getCompressorParameters(), - config.search.getCompressorParameters(), - config.search.topKOverquery, - config.search.useSearchPruning, - artifacts); + ).getDataSet()) { + // Register dataset info the first time we actually load the dataset for benchmarking + artifacts.registerDataset(datasetName, ds); + + Grid.runAll(ds, + config.construction.useSavedIndexIfExists, + config.construction.outDegree, + config.construction.efConstruction, + config.construction.neighborOverflow, + config.construction.addHierarchy, + config.construction.refineFinalGraph, + config.construction.getFeatureSets(), + config.construction.getCompressorParameters(), + config.search.getCompressorParameters(), + config.search.topKOverquery, + config.search.useSearchPruning, + artifacts); + } } } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ColdFileCache.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ColdFileCache.java new file mode 100644 index 000000000..36e17df93 --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ColdFileCache.java @@ -0,0 +1,111 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import java.io.IOException; +import java.lang.invoke.MethodHandle; +import java.lang.reflect.Array; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Optional; + +/** Linux file-specific cache eviction and residency verification; never drops global caches. */ +final class ColdFileCache { + private ColdFileCache() {} + + static void evictAndVerify(Path file) throws IOException { + if (!System.getProperty("os.name").equals("Linux") || !System.getProperty("os.arch").matches("amd64|aarch64")) + throw new IOException("Cold file-cache verification requires 64-bit Linux; use --cache uncontrolled explicitly elsewhere"); + try (var nativeCalls = new NativeCalls()) { + var open = nativeCalls.function("open", "JAVA_INT", "ADDRESS", "JAVA_INT"); + var close = nativeCalls.function("close", "JAVA_INT", "JAVA_INT"); + var advise = nativeCalls.function("posix_fadvise", "JAVA_INT", "JAVA_INT", "JAVA_LONG", "JAVA_LONG", "JAVA_INT"); + var mmap = nativeCalls.function("mmap", "ADDRESS", "ADDRESS", "JAVA_LONG", "JAVA_INT", "JAVA_INT", "JAVA_INT", "JAVA_LONG"); + var mincore = nativeCalls.function("mincore", "JAVA_INT", "ADDRESS", "JAVA_LONG", "ADDRESS"); + var munmap = nativeCalls.function("munmap", "JAVA_INT", "ADDRESS", "JAVA_LONG"); + var getpagesize = nativeCalls.function("getpagesize", "JAVA_INT"); + int fd = (int) open.invokeWithArguments(nativeCalls.string(file.toAbsolutePath().toString()), 0); + if (fd < 0) throw new IOException("Cannot open input for cache eviction: " + file); + try { + int error = (int) advise.invokeWithArguments(fd, 0L, 0L, 4); // POSIX_FADV_DONTNEED + if (error != 0) throw new IOException("posix_fadvise failed: " + error); + long length = Files.size(file); + int pageSize = (int) getpagesize.invokeWithArguments(); + long pages = (length + pageSize - 1) / pageSize; + Object resident = nativeCalls.allocate(pages); + // Mapping alone does not fault in pages; mincore only inspects their residency. + Object mapping = mmap.invokeWithArguments(nativeCalls.nullSegment(), length, 1, 1, fd, 0L); + if (nativeCalls.address(mapping) == -1L) throw new IOException("mmap for residency check failed"); + try { + if ((int) mincore.invokeWithArguments(mapping, length, resident) != 0) + throw new IOException("mincore residency check failed"); + long remaining = nativeCalls.residentPages(resident, pages); + if (remaining != 0) throw new IOException("Cold-cache check failed: " + remaining + "/" + pages + + " input pages remain resident; stop other readers or choose --cache uncontrolled"); + System.out.printf("Cold cache verified: 0/%,d input pages resident%n", pages); + } finally { + if ((int) munmap.invokeWithArguments(mapping, length) != 0) throw new IOException("Residency mapping cleanup failed"); + } + } finally { + if ((int) close.invokeWithArguments(fd) != 0) throw new IOException("Cache-check descriptor cleanup failed"); + } + } catch (IOException e) { throw e; } + catch (Throwable e) { throw new IOException("Cold cache requires JDK 23+ with --enable-native-access=ALL-UNNAMED", e); } + } + + // Resolve FFM only for cold-cache setup, so ordinary examples retain their Java 11 target. + // None of this reflection runs inside the measured interval or vector-processing loop. + private static final class NativeCalls implements AutoCloseable { + private final Class arenaType = Class.forName("java.lang.foreign.Arena"); + private final Class segmentType = Class.forName("java.lang.foreign.MemorySegment"); + private final Class layoutType = Class.forName("java.lang.foreign.MemoryLayout"); + private final Class valueType = Class.forName("java.lang.foreign.ValueLayout"); + private final Class descriptorType = Class.forName("java.lang.foreign.FunctionDescriptor"); + private final Class linkerType = Class.forName("java.lang.foreign.Linker"); + private final Class optionType = Class.forName("java.lang.foreign.Linker$Option"); + private final Class lookupType = Class.forName("java.lang.foreign.SymbolLookup"); + private final Object arena, linker, lookup; + NativeCalls() throws Exception { + arena = arenaType.getMethod("ofConfined").invoke(null); + linker = linkerType.getMethod("nativeLinker").invoke(null); + lookup = linkerType.getMethod("defaultLookup").invoke(linker); + } + private Object layout(String name) throws Exception { return valueType.getField(name).get(null); } + MethodHandle function(String name, String result, String... parameters) throws Exception { + var symbol = (Optional) lookupType.getMethod("find", String.class).invoke(lookup, name); + Object layouts = Array.newInstance(layoutType, parameters.length); + for (int i = 0; i < parameters.length; i++) Array.set(layouts, i, layout(parameters[i])); + Object descriptor = descriptorType.getMethod("of", layoutType, layouts.getClass()).invoke(null, layout(result), layouts); + Object options = Array.newInstance(optionType, 0); + return (MethodHandle) linkerType.getMethod("downcallHandle", segmentType, descriptorType, options.getClass()) + .invoke(linker, symbol.orElseThrow(), descriptor, options); + } + Object string(String value) throws Exception { return arenaType.getMethod("allocateFrom", String.class).invoke(arena, value); } + Object allocate(long bytes) throws Exception { return arenaType.getMethod("allocate", long.class).invoke(arena, bytes); } + Object nullSegment() throws Exception { return segmentType.getField("NULL").get(null); } + long address(Object segment) throws Exception { return (long) segmentType.getMethod("address").invoke(segment); } + long residentPages(Object segment, long count) throws Exception { + // Copy only the tiny residency bitmap, not file data. No per-page reflective calls. + byte[] bitmap = (byte[]) segmentType.getMethod("toArray", Class.forName("java.lang.foreign.ValueLayout$OfByte")) + .invoke(segment, layout("JAVA_BYTE")); + long remaining = 0; + for (int i = 0; i < count; i++) if ((bitmap[i] & 1) != 0) remaining++; + return remaining; + } + public void close() throws Exception { arenaType.getMethod("close").invoke(arena); } + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DemandFvecVectorValues.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DemandFvecVectorValues.java new file mode 100644 index 000000000..c71628072 --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DemandFvecVectorValues.java @@ -0,0 +1,135 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.disk.ByteBufferReader; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; + +import java.io.EOFException; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.FloatBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.Objects; +import java.util.concurrent.atomic.LongAdder; + +/** Benchmark control: synchronous point reads, with no executor, cache or read-ahead. */ +final class DemandFvecVectorValues implements RandomAccessVectorValues, AutoCloseable { + private static final VectorTypeSupport VTS = VectorizationProvider.getInstance().getVectorTypeSupport(); + private final FileChannel channel; + private final int dimension, size, recordBytes; + private final boolean owner; + private final Buffer buffer; + private final VectorFloat scratch; + private final LongAdder reads, bytes; + + DemandFvecVectorValues(Path file) throws IOException { + channel = FileChannel.open(file, StandardOpenOption.READ); + try { + var header = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN); + readFully(channel, header, 0); + dimension = header.flip().getInt(); + if (dimension <= 0 || dimension > 100_000) throw new IOException("Invalid fvecs dimension: " + dimension); + recordBytes = Math.multiplyExact(dimension + 1, 4); + long length = channel.size(); + if (length % recordBytes != 0 || length / recordBytes > Integer.MAX_VALUE) + throw new IOException("Invalid fvecs file length: " + length); + size = (int) (length / recordBytes); + owner = true; + buffer = new Buffer(recordBytes); + scratch = VTS.createFloatVector(dimension); + reads = new LongAdder(); + bytes = new LongAdder(); + } catch (Throwable e) { + channel.close(); + throw e; + } + } + + private DemandFvecVectorValues(DemandFvecVectorValues root) { + channel = root.channel; dimension = root.dimension; size = root.size; + recordBytes = root.recordBytes; owner = false; + buffer = new Buffer(recordBytes); scratch = VTS.createFloatVector(dimension); + reads = root.reads; bytes = root.bytes; + } + + @Override public int size() { return size; } + @Override public int dimension() { return dimension; } + @Override public boolean isValueShared() { return true; } + @Override public DemandFvecVectorValues copy() { return new DemandFvecVectorValues(this); } + @Override public java.util.function.Supplier threadLocalSupplier() { + // Views share the owner's channel; closing the owner releases it after all workers finish. + var local = ThreadLocal.withInitial(this::copy); + return local::get; + } + @Override public VectorFloat getVector(int ordinal) { + Objects.checkIndex(ordinal, size); + try { + buffer.payload.clear(); + readFully(channel, buffer.payload, (long) ordinal * recordBytes); + buffer.payload.flip(); + if (buffer.readInt() != dimension) throw new IOException("Dimension mismatch at ordinal " + ordinal); + VTS.readFloatVector(buffer, dimension, scratch, 0); + reads.increment(); bytes.add(recordBytes); + return scratch; + } catch (IOException e) { throw new UncheckedIOException(e); } + } + long reads() { return reads.sum(); } + long bytes() { return bytes.sum(); } + @Override public void close() throws IOException { if (owner) channel.close(); } + + static void readFully(FileChannel channel, ByteBuffer buffer, long offset) throws IOException { + int zeroReads = 0; + while (buffer.hasRemaining()) { + int n = channel.read(buffer, offset); + if (n < 0) throw new EOFException("Incomplete fvecs record"); + if (n == 0) { + // Match the buffered source: a broken channel must not spin forever. + if (++zeroReads >= 16) throw new IOException("No progress reading fvec data at " + offset); + Thread.onSpinWait(); + } else { + zeroReads = 0; + offset += n; + } + } + } + + // Match the PR source's float decoding so the control differs in input scheduling. + private static final class Buffer extends ByteBufferReader { + final ByteBuffer payload; + final FloatBuffer floats; + Buffer(int bytes) { + super(ByteBuffer.allocateDirect(bytes).order(ByteOrder.LITTLE_ENDIAN)); + payload = bb; floats = bb.asFloatBuffer(); + } + @Override public void read(float[] destination, int offset, int count) { + Objects.checkFromIndexSize(offset, count, destination.length); + Objects.checkFromIndexSize(bb.position(), Math.multiplyExact(count, 4), bb.limit()); + floats.position(bb.position() / 4); + floats.get(destination, offset, count); + bb.position(bb.position() + count * 4); + } + @Override public void readFully(float[] destination) { read(destination, 0, destination.length); } + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java index 548fd4b47..d2aa5ceda 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java @@ -37,6 +37,7 @@ import io.github.jbellis.jvector.example.yaml.MetricSelection; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.VectorAccess; import io.github.jbellis.jvector.graph.GraphSearcher; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.*; @@ -50,6 +51,7 @@ import io.github.jbellis.jvector.graph.similarity.ScoreFunction; import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider; import io.github.jbellis.jvector.quantization.CompressedVectors; +import io.github.jbellis.jvector.quantization.BQVectors; import io.github.jbellis.jvector.quantization.NVQuantization; import io.github.jbellis.jvector.quantization.PQVectors; import io.github.jbellis.jvector.quantization.ProductQuantization; @@ -396,15 +398,24 @@ private static BuildOnDiskResult buildOnDisk(List> feat var floatVectors = ds.getBaseRavv(); // Record the encoding time if asked for.... - PQVectors pq = (PQVectors) ((constructionMetrics != null) - ? constructionMetrics.index("PQ").timeEncode(() -> buildCompressor.encodeAll(floatVectors)) - : buildCompressor.encodeAll(floatVectors)); - var bsp = BuildScoreProvider.pqBuildScoreProvider(ds.getSimilarityFunction(), pq); + var encoded = (constructionMetrics != null) + ? constructionMetrics.index(buildCompressor instanceof ProductQuantization ? "PQ" : "BQ") + .timeEncode(() -> buildCompressor.encodeAll(floatVectors)) + : buildCompressor.encodeAll(floatVectors); + if (!(encoded instanceof PQVectors) && !(encoded instanceof BQVectors)) + throw new IllegalArgumentException("Disk construction supports PQ or BQ build compression"); + PQVectors pq = encoded instanceof PQVectors ? (PQVectors) encoded : null; + if (pq == null && featureSets.stream().anyMatch(fs -> fs.contains(FeatureId.FUSED_PQ))) + throw new IllegalArgumentException("FusedPQ requires PQ build compression with 256 clusters"); + var bsp = pq != null + ? BuildScoreProvider.pqBuildScoreProvider(ds.getSimilarityFunction(), pq) + : BuildScoreProvider.bqBuildScoreProvider((BQVectors) encoded); GraphIndexBuilder builder = new GraphIndexBuilder(bsp, floatVectors.dimension(), M, efConstruction, neighborOverflow, 1.2f, addHierarchy, refineFinalGraph); // use the inline vectors index as the score provider for graph construction Map, RandomAccessOnDiskGraphIndexWriter> writers = new HashMap<>(); Map, Map>> suppliers = new HashMap<>(); + Map, Map, Feature.State>>> inputSuppliers = new HashMap<>(); RandomAccessOnDiskGraphIndexWriter scoringWriter = null; int n = 0; for (var features : featureSets) { @@ -415,10 +426,11 @@ private static BuildOnDiskResult buildOnDisk(List> feat } else { graphPath = outputDir.resolve("graph" + n++); } - var bws = builderWithSuppliers(features, builder.getGraph(), graphPath, floatVectors, pq.getCompressor(), constructionMetrics); + var bws = builderWithSuppliers(features, builder.getGraph(), graphPath, floatVectors, pq == null ? null : pq.getCompressor(), constructionMetrics); var writer = bws.builder.build(); writers.put(features, writer); suppliers.put(features, bws.suppliers); + inputSuppliers.put(features, bws.inputSuppliers); if (features.contains(FeatureId.INLINE_VECTORS) || features.contains(FeatureId.NVQ_VECTORS)) { scoringWriter = writer; } @@ -429,23 +441,21 @@ private static BuildOnDiskResult buildOnDisk(List> feat // build the graph incrementally long startTime = System.nanoTime(); - var vv = floatVectors.threadLocalSupplier(); - PhysicalCoreExecutor.pool().submit(() -> { - IntStream.range(0, floatVectors.size()).parallel().forEach(node -> { - writers.forEach((features, writer) -> { - try { - var stateMap = new EnumMap(FeatureId.class); - suppliers.get(features).forEach((featureId, supplier) -> { - stateMap.put(featureId, supplier.apply(node)); - }); - writer.writeInline(node, stateMap); - } catch (IOException e) { - throw new UncheckedIOException(e); - } - }); - builder.addGraphNode(node, vv.get().getVector(node)); + if (ds.useDatasetPrefetching()) { + VectorAccess.forEach(floatVectors, PhysicalCoreExecutor.pool(), (node, vector) -> { + writeInlineStates(writers, suppliers, inputSuppliers, node, vector); + // Keep main's original insertion query, not a decoded quantized substitute. + builder.addGraphNode(node, vector); }); - }).join(); + } else { + var vv = floatVectors.threadLocalSupplier(); + PhysicalCoreExecutor.pool().submit(() -> { + IntStream.range(0, floatVectors.size()).parallel().forEach(node -> { + writeInlineStates(writers, suppliers, null, node, null); + builder.addGraphNode(node, vv.get().getVector(node)); + }); + }).join(); + } builder.cleanup(); // write the edge lists and close the writers @@ -495,6 +505,22 @@ private static BuildOnDiskResult buildOnDisk(List> feat return new BuildOnDiskResult(indexes, fileSizes); } + private static void writeInlineStates( + Map, RandomAccessOnDiskGraphIndexWriter> writers, + Map, Map>> suppliers, + Map, Map, Feature.State>>> inputSuppliers, + int node, VectorFloat vector) { + writers.forEach((features, writer) -> { + var states = new EnumMap(FeatureId.class); + if (inputSuppliers == null) + suppliers.get(features).forEach((id, supplier) -> states.put(id, supplier.apply(node))); + else + inputSuppliers.get(features).forEach((id, supplier) -> states.put(id, supplier.apply(vector))); + try { writer.writeInline(node, states); } + catch (IOException e) { throw new UncheckedIOException(e); } + }); + } + private static BuilderWithSuppliers builderWithSuppliers(Set features, ImmutableGraphIndex onHeapGraph, Path outPath, @@ -508,11 +534,15 @@ private static BuilderWithSuppliers builderWithSuppliers(Set features builder.withMapper(identityMapper); Map> suppliers = new EnumMap<>(FeatureId.class); + Map, Feature.State>> inputSuppliers = new EnumMap<>(FeatureId.class); + // File-backed sources reuse scratch vectors; inline writes need a view per worker. + var writerVectors = floatVectors.threadLocalSupplier(); for (var featureId : features) { switch (featureId) { case INLINE_VECTORS: builder.with(new InlineVectors(floatVectors.dimension())); - suppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(floatVectors.getVector(ordinal))); + inputSuppliers.put(FeatureId.INLINE_VECTORS, InlineVectors.State::new); + suppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(writerVectors.get().getVector(ordinal))); break; case FUSED_PQ: if (pq == null) { @@ -528,12 +558,13 @@ private static BuilderWithSuppliers builderWithSuppliers(Set features ? constructionMetrics.index("NVQ").timeCompute(() -> NVQuantization.compute(floatVectors, nSubVectors)) : NVQuantization.compute(floatVectors, nSubVectors); builder.with(new NVQ(nvq)); - suppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(floatVectors.getVector(ordinal)))); + inputSuppliers.put(FeatureId.NVQ_VECTORS, vector -> new NVQ.State(nvq.encode(vector))); + suppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(writerVectors.get().getVector(ordinal)))); break; } } - return new BuilderWithSuppliers(builder, suppliers); + return new BuilderWithSuppliers(builder, suppliers, inputSuppliers); } public static void setDiagnosticLevel(int diagLevel) { @@ -559,9 +590,17 @@ private static class BuilderWithSuppliers { public final RandomAccessOnDiskGraphIndexWriter.Builder builder; public final Map> suppliers; + public final Map, Feature.State>> inputSuppliers; + public BuilderWithSuppliers(RandomAccessOnDiskGraphIndexWriter.Builder builder, Map> suppliers) { + this(builder, suppliers, Map.of()); + } + + public BuilderWithSuppliers(RandomAccessOnDiskGraphIndexWriter.Builder builder, Map> suppliers, + Map, Feature.State>> inputSuppliers) { this.builder = builder; this.suppliers = suppliers; + this.inputSuppliers = inputSuppliers; } } @@ -1119,6 +1158,7 @@ public static class ConfiguredSystem implements AutoCloseable { ImmutableGraphIndex index; CompressedVectors cv; Set features; + private final Supplier exactVectors; private final ExplicitThreadLocal searchers = ExplicitThreadLocal.withInitial(() -> { return new GraphSearcher(index); @@ -1129,6 +1169,8 @@ public static class ConfiguredSystem implements AutoCloseable { this.index = index; this.cv = cv; this.features = features; + // File-backed sources borrow scratch storage; concurrent scorers need separate views. + this.exactVectors = ds.getBaseRavv().threadLocalSupplier(); } public SearchScoreProvider scoreProviderFor(VectorFloat queryVector, ImmutableGraphIndex.View view) { @@ -1139,7 +1181,7 @@ public SearchScoreProvider scoreProviderFor(VectorFloat queryVector, Immutabl } else { // if we're not compressing then just use the exact score function if (cv == null) { - return DefaultSearchScoreProvider.exact(queryVector, ds.getSimilarityFunction(), ds.getBaseRavv()); + return DefaultSearchScoreProvider.exact(queryVector, ds.getSimilarityFunction(), exactVectors.get()); } asf = cv.precomputedScoreFunctionFor(queryVector, ds.getSimilarityFunction()); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloBatchedQuantization.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloBatchedQuantization.java new file mode 100644 index 000000000..e3010c93a --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloBatchedQuantization.java @@ -0,0 +1,46 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.disk.FvecFileVectorValues; +import io.github.jbellis.jvector.quantization.NVQuantization; +import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; + +/** Train and encode directly from a file through the ordinary PQ and NVQ APIs. */ +public class HelloBatchedQuantization { + public static void main(String[] args) throws Exception { + // Use the same shared compute pool as Grid; the source owns a separate I/O pool. + var executor = PhysicalCoreExecutor.pool(); + try (var source = FvecFileVectorValues.open(HelloVectorFile.path(args))) { + if (source.size() < 256) + throw new IllegalArgumentException("PQ example requires at least 256 vectors"); + System.out.printf("%,d vectors, %d dimensions%n", source.size(), source.dimension()); + + // PQ chooses its training sample; VectorAccess handles efficient selection reads. + int subspaces = Math.max(1, source.dimension() / 8); + var pq = ProductQuantization.compute(source, subspaces, 256, false); + var pqVectors = pq.encodeAll(source, executor); + System.out.printf("PQ encoded %,d vectors (%d subspaces)%n", pqVectors.count(), subspaces); + + // NVQ scans the same file; the source owns batching and read-ahead. + var nvq = NVQuantization.compute(source, 1); + var nvqVectors = nvq.encodeAll(source, executor); + System.out.printf("NVQ encoded %,d vectors%n", nvqVectors.count()); + } + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloBatchedVectors.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloBatchedVectors.java new file mode 100644 index 000000000..221ead423 --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloBatchedVectors.java @@ -0,0 +1,63 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.disk.FvecFileVectorValues; +import io.github.jbellis.jvector.graph.VectorAccess; + +/** Runnable example of range reads, selected reads, and borrowed-vector ownership. */ +public class HelloBatchedVectors { + public static void main(String[] args) throws Exception { + try (var source = FvecFileVectorValues.open(HelloVectorFile.path(args))) { + System.out.printf("%,d vectors, %d dimensions%n", source.size(), source.dimension()); + + // Exclusive end: read ordinals 0 through 4 (or the entire smaller file). + try (var cursor = VectorAccess.openRange(source, 0, Math.min(5, source.size()))) { + while (cursor.next()) { + System.out.printf("Scan %d: first coordinate = %.6f%n", + cursor.ordinal(), cursor.vector().get(0)); + } + } + + // A complete ordered scan consumes each vector without retaining it. + int scanned = 0; + double firstCoordinateSum = 0; + try (var cursor = VectorAccess.openRange(source, 0, source.size())) { + while (cursor.next()) { + firstCoordinateSum += cursor.vector().get(0); + scanned++; + } + } + System.out.printf("Scanned %,d vectors; first-coordinate sum = %.6f%n", + scanned, firstCoordinateSum); + + // Arbitrary ordinals need not be sorted; duplicate requests are preserved. + int[] selected = {source.size() - 1, 0, source.size() / 2, source.size() - 1}; + try (var cursor = VectorAccess.openSelection(source, selected, 0, selected.length)) { + cursor.next(); + // The file cursor reuses its vector; make a copy before retaining it. + var retained = cursor.vector().copy(); + System.out.printf("Selected %d%n", cursor.ordinal()); + while (cursor.next()) System.out.printf("Selected %d%n", cursor.ordinal()); + System.out.printf("Retained first coordinate = %.6f%n", retained.get(0)); + } + + System.out.printf("I/O payload: %,d bytes allocated, %,d bytes limit%n", + source.statistics().bufferBytes, source.statistics().bufferLimitBytes); + } + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloVectorFile.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloVectorFile.java new file mode 100644 index 000000000..a21adf79f --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloVectorFile.java @@ -0,0 +1,36 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.example.benchmarks.datasets.DataSetLoaderSimpleMFD; + +import java.io.IOException; +import java.nio.file.Path; + +/** Resolve a real public-catalog input without materializing its vectors. */ +final class HelloVectorFile { + private HelloVectorFile() {} + + static Path path(String[] args) throws IOException { + if (args.length > 1) throw new IllegalArgumentException("Usage: [vectors.fvecs]"); + if (args.length == 1 && !args[0].isEmpty()) return Path.of(args[0]); + + var loader = new DataSetLoaderSimpleMFD("jvector-examples/yaml-configs/dataset-catalogs"); + return loader.loadBaseVectorFile("e5-small-v2-100k").orElseThrow(() -> + new IllegalArgumentException("Public catalog does not contain e5-small-v2-100k")); + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/VectorPrefetchBenchmark.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/VectorPrefetchBenchmark.java new file mode 100644 index 000000000..ff5e5151f --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/VectorPrefetchBenchmark.java @@ -0,0 +1,452 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.disk.FvecFileVectorValues; +import io.github.jbellis.jvector.example.benchmarks.datasets.DataSetLoaderSimpleMFD; +import io.github.jbellis.jvector.example.util.SiftLoader; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.VectorAccess; +import io.github.jbellis.jvector.quantization.NVQuantization; +import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.quantization.VectorCompressor; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; + +import java.io.IOException; +import java.lang.management.ManagementFactory; +import java.lang.ref.Reference; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Set; +import java.util.SplittableRandom; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; + +/** Measures sample retrieval, quantizer training, full-set encoding and an encoded-score sink. */ +public final class VectorPrefetchBenchmark { + private VectorPrefetchBenchmark() {} + + public static void main(String[] args) throws Exception { + run(args); + } + + private static void run(String[] args) throws Exception { + var options = Options.parse(args); + if (options.contains("help")) { + System.out.println("Usage: VectorPrefetchBenchmark" + + " [--dataset catalog-name | --file vectors.fvecs] [--samples 128000]" + + " [--mode all|preload|demand|prefetch] [--quantizer nvq|pq|adapter-class] [--quantizer-options key=value,...] [--loader-options key=value,...]" + + " [--metric DOT_PRODUCT|EUCLIDEAN|COSINE] [--query-file queries.fvecs]" + + " [--cache cold|uncontrolled] [--seed 42] [--progress-seconds 30]"); + return; + } + if (!options.contains("child")) { + Path file = options.file(); // Finish catalog/download preparation before launching timed arms. + var modes = options.get("mode", "all").equals("all") + ? List.of("preload", "demand", "prefetch") : List.of(options.get("mode", "all")); + var rows = new ArrayList(); + for (var mode : modes) { + options.checkMode(mode); + var command = new ArrayList(); + command.add(Path.of(System.getProperty("java.home"), "bin", "java").toString()); + // Preserve IDE/CLI heap, provider, CPU and native-access settings in each isolated arm. + for (var arg : ManagementFactory.getRuntimeMXBean().getInputArguments()) { + if (!arg.startsWith("-agentlib:") && !arg.startsWith("-javaagent:")) command.add(arg); + } + command.add("-classpath"); command.add(System.getProperty("java.class.path")); + command.add(VectorPrefetchBenchmark.class.getName()); + command.addAll(options.childArguments(file, mode)); + System.out.println("\nStarting " + mode + " in a fresh JVM"); + var child = new ProcessBuilder(command).redirectErrorStream(true).start(); + var stopChild = new Thread(child::destroy, "stop-input-benchmark-arm"); + Runtime.getRuntime().addShutdownHook(stopChild); + try { + try (var output = new java.io.BufferedReader(new java.io.InputStreamReader(child.getInputStream()))) { + String line; + while ((line = output.readLine()) != null) { + System.out.println(line); + if (line.startsWith("| pipeline |")) rows.add(line); + } + } + int exit = child.waitFor(); // One arm at a time, including when launched from IntelliJ. + if (exit != 0) throw new IllegalStateException(mode + " failed (exit " + exit + "); remaining arms not run"); + } catch (InterruptedException e) { + child.destroy(); + child.waitFor(); + Thread.currentThread().interrupt(); + throw e; + } finally { Runtime.getRuntime().removeShutdownHook(stopChild); } + } + System.out.println("\nComparison (seconds; quantizer=" + options.get("quantizer", "nvq") + + ", options=" + options.get("quantizer-options", "defaults") + ", metric=" + options.get("metric", "DOT_PRODUCT") + + ", fresh JVM per arm, cache=" + options.get("cache", "cold") + "):"); + printHeader(); + rows.forEach(System.out::println); + return; + } + measure(options); + } + + private static void measure(Options options) throws Exception { + Path file = Path.of(options.get("file", "")); + String mode = options.get("mode", "all"); + options.checkMode(mode); + var metric = VectorSimilarityFunction.valueOf(options.get("metric", "DOT_PRODUCT").toUpperCase(Locale.ROOT)); + String quantizer = options.get("quantizer", "nvq"); + Quantizer quantizerFactory = quantizer(quantizer); + var quantizerOptions = parameters(options.get("quantizer-options", "")); + var loaderOptions = parameters(options.get("loader-options", "")); + String cache = options.get("cache", "cold"); + if (!cache.equals("cold") && !cache.equals("uncontrolled")) throw new IllegalArgumentException("Unknown cache policy: " + cache); + var executor = PhysicalCoreExecutor.pool(); + int count, dimension; + io.github.jbellis.jvector.vector.types.VectorFloat query; + // Prepare one query outside timing. By default hold out the last base vector in every arm. + boolean heldOut = !options.contains("query-file"); + try (var probe = new DemandFvecVectorValues(file)) { + dimension = probe.dimension(); + count = probe.size() - (heldOut ? 1 : 0); + if (count <= 0) throw new IllegalArgumentException("Need at least one candidate and a query"); + if (heldOut) query = probe.getVector(probe.size() - 1).copy(); + else try (var queries = new DemandFvecVectorValues(Path.of(options.get("query-file", "")))) { + if (queries.size() == 0 || queries.dimension() != dimension) throw new IllegalArgumentException("Invalid query file"); + query = queries.getVector(0).copy(); + } + } + int[] ordinals = sample(count, options.positiveInt("samples", Math.min(count, 128000)), options.longValue("seed", 42)); + int selectedCount = ordinals.length; + if (quantizer.equals("pq") && selectedCount < integer(quantizerOptions, "clusters", 256)) throw new IllegalArgumentException("PQ requires at least the configured cluster count of selected vectors"); + System.out.printf("Pipeline: %s, encoding candidates=%,d, training samples=%,d, dimensions=%d, compute workers=%d%n", + mode, count, selectedCount, dimension, executor.getParallelism()); + System.out.printf("Quantizer=%s, metric=%s, sample seed=%d, query=%s, provider=%s%n", + quantizer, metric, options.longValue("seed", 42), heldOut ? "held-out final base vector" : "first query-file vector", + VectorizationProvider.getInstance().getClass().getSimpleName()); + System.out.println("Quantizer options: " + quantizerOptions + "; loader options: " + loaderOptions); + System.out.println("Input: " + file.toAbsolutePath()); + if (cache.equals("cold")) ColdFileCache.evictAndVerify(file); + else System.out.println("OS file cache: uncontrolled (NOT a cold-cache measurement)"); + + try (var progress = new Progress(options.nonnegativeInt("progress-seconds", 30))) { + long begin = System.nanoTime(); + double preloadSeconds = 0; + progress.phase(mode.equals("preload") ? "preload preparation (excluded)" : "sample retrieval"); + try (var input = loader(mode).open(file, loaderOptions)) { + var original = input.values(); + if (original.size() != count + (heldOut ? 1 : 0) || original.dimension() != dimension) + throw new IllegalStateException("Loader changed source dimensions or vector count"); + RandomAccessVectorValues vectors = prefix(original, count); + if (mode.equals("preload")) { + preloadSeconds = seconds(System.nanoTime() - begin); + System.out.printf(Locale.ROOT, "Resident preload preparation: %.3f s (excluded from comparison)%n", preloadSeconds); + // The resident baseline measures work after all input vectors are available. + begin = System.nanoTime(); + progress.phase("sample retrieval"); + } + RandomAccessVectorValues trainingVectors = new ListRandomAccessVectorValues( + VectorAccess.copySelected(vectors, ordinals, executor), dimension); + long loaded = System.nanoTime(); + progress.phase("train"); + Encoder encoder = quantizerFactory.train(trainingVectors, quantizerOptions, executor); + long trained = System.nanoTime(); + progress.phase("encode"); + var encoded = encoder.encode(vectors, executor); + if (encoded.count() != vectors.size()) throw new IllegalStateException("Encoding changed candidate count"); + long encodedAt = System.nanoTime(); + progress.phase("nearest-vector sink"); + var scorer = encoded.scorer(query, metric); + int nearest = -1; + double best = Double.NEGATIVE_INFINITY; + double checksum = 0; + for (int i = 0; i < vectors.size(); i++) { + double score = scorer.applyAsDouble(i); + if (!Double.isFinite(score)) throw new IllegalStateException("Non-finite score at selected position " + i); + checksum += score; + if (score > best) { best = score; nearest = i; } + } + long finished = System.nanoTime(); + progress.phase("complete"); + int ordinal = nearest; // Encoding and scoring always retain full base ordinals. + printHeader(); + System.out.printf(Locale.ROOT, "| %s | %s | %s | %.3f | %.3f | %.3f | %.3f | %.3f | %d | %.8f | %.9f |%n", + "pipeline", mode, mode.equals("preload") ? String.format(Locale.ROOT, "%.3f", preloadSeconds) : "—", seconds(loaded - begin), seconds(trained - loaded), + seconds(encodedAt - trained), seconds(finished - encodedAt), seconds(finished - begin), ordinal, best, checksum); + String statistics = input.statistics(); + if (!statistics.isEmpty()) System.out.println(statistics); + Reference.reachabilityFence(input); // Keep all resident preload vectors live through the sink. + } + + } + } + + + /** Factory for any representation (PQ, NVQ, ASH or future types), without changing core APIs. */ + public interface Quantizer { + Encoder train(RandomAccessVectorValues vectors, java.util.Map parameters, + java.util.concurrent.ForkJoinPool executor); + } + + /** Prepared representation: encode the candidate set through its normal production API. */ + @FunctionalInterface + public interface Encoder { + Encoded encode(RandomAccessVectorValues vectors, java.util.concurrent.ForkJoinPool executor); + } + + /** A minimal nearest-vector sink contract, independent of a particular vector class. */ + public interface Encoded { + int count(); + java.util.function.IntToDoubleFunction scorer( + io.github.jbellis.jvector.vector.types.VectorFloat query, VectorSimilarityFunction metric); + } + + /** Adapter for existing production compressors; other vector types can supply Encoded directly. */ + public static Encoder compressed(VectorCompressor compressor) { + return (vectors, executor) -> { + var encoded = compressor.encodeAll(vectors, executor); + return new Encoded() { + public int count() { return encoded.count(); } + public java.util.function.IntToDoubleFunction scorer( + io.github.jbellis.jvector.vector.types.VectorFloat query, VectorSimilarityFunction metric) { + var function = encoded.scoreFunctionFor(query, metric); + return function::similarityTo; + } + }; + }; + } + + /** Future vector sources can expose BatchedVectorValues; ordinary RAVV also works. */ + public interface Loader { + Input open(Path file, java.util.Map parameters) throws Exception; + } + + /** Owns source resources until both quantization and the sink have finished. */ + public interface Input extends AutoCloseable { + RandomAccessVectorValues values(); + default String statistics() { return ""; } + @Override void close() throws Exception; + } + + private static Quantizer quantizer(String name) throws Exception { + if (name.equals("nvq")) return (vectors, parameters, executor) -> { + requireKeys(parameters, Set.of("subvectors", "learn")); + var nvq = NVQuantization.compute(vectors, integer(parameters, "subvectors", 1)); + nvq.learn = bool(parameters, "learn", true); + return compressed(nvq); + }; + if (name.equals("pq")) return (vectors, parameters, executor) -> { + requireKeys(parameters, Set.of("subspaces", "clusters", "center", "anisotropic")); + return compressed(ProductQuantization.compute(vectors, integer(parameters, "subspaces", Math.max(1, vectors.dimension() / 8)), + integer(parameters, "clusters", 256), bool(parameters, "center", false), + Float.parseFloat(parameters.getOrDefault("anisotropic", "-1")), executor, executor)); + }; + return (Quantizer) Class.forName(name).getConstructor().newInstance(); + } + + private static Loader loader(String name) throws Exception { + if (name.equals("preload")) return (file, parameters) -> { + requireKeys(parameters, Set.of()); + var all = SiftLoader.readFvecs(file.toString()); + var values = new ListRandomAccessVectorValues(all, all.get(0).length()); + return new Input() { + public RandomAccessVectorValues values() { return values; } + public void close() { } + }; + }; + if (name.equals("demand")) return (file, parameters) -> { + requireKeys(parameters, Set.of()); + var source = new DemandFvecVectorValues(file); + return new Input() { + public RandomAccessVectorValues values() { return source; } + public void close() throws IOException { source.close(); } + public String statistics() { return String.format("Source logical reads=%,d, bytes=%,d (demand-only)", source.reads(), source.bytes()); } + }; + }; + if (name.equals("prefetch")) return (file, parameters) -> { + requireKeys(parameters, Set.of("ioThreads", "batchVectors", "readAhead", "maxBufferBytes")); + var options = FvecFileVectorValues.Options.defaults(); + if (parameters.containsKey("ioThreads")) options = options.withIoThreads(integer(parameters, "ioThreads", 1)); + if (parameters.containsKey("batchVectors")) options = options.withBatchVectors(integer(parameters, "batchVectors", 64)); + if (parameters.containsKey("readAhead")) options = options.withReadAhead(integer(parameters, "readAhead", 3)); + if (parameters.containsKey("maxBufferBytes")) options = options.withMaxBufferBytes(Long.parseLong(parameters.get("maxBufferBytes"))); + var source = FvecFileVectorValues.open(file, options); + return new Input() { + public RandomAccessVectorValues values() { return source; } + public void close() throws IOException { source.close(); } + public String statistics() { + var stats = source.statistics(); + return String.format("Source logical reads=%,d, bytes=%,d, payload bytes=%,d, payload limit=%,d", + stats.reads, stats.bytesRead, stats.bufferBytes, stats.bufferLimitBytes); + } + }; + }; + return (Loader) Class.forName(name).getConstructor().newInstance(); + } + + private static java.util.Map parameters(String input) { + var values = new LinkedHashMap(); + if (!input.isEmpty()) for (var entry : input.split(",", -1)) { + var pair = entry.split("=", 2); + if (pair.length != 2 || pair[0].isEmpty() || pair[1].isEmpty() || values.putIfAbsent(pair[0], pair[1]) != null) + throw new IllegalArgumentException("Expected unique key=value parameters: " + input); + } + return java.util.Collections.unmodifiableMap(values); + } + private static void requireKeys(java.util.Map parameters, Set keys) { + for (var key : parameters.keySet()) if (!keys.contains(key)) throw new IllegalArgumentException("Unknown parameter: " + key); + } + private static int integer(java.util.Map parameters, String key, int fallback) { + return Integer.parseInt(parameters.getOrDefault(key, Integer.toString(fallback))); + } + private static boolean bool(java.util.Map parameters, String key, boolean fallback) { + String value = parameters.getOrDefault(key, Boolean.toString(fallback)); + if (!value.equals("true") && !value.equals("false")) throw new IllegalArgumentException("Invalid boolean for " + key); + return Boolean.parseBoolean(value); + } + private static void printHeader() { + System.out.println("| Workload | Arm | Preload s (excluded) | Sample s | Train s | Encode s | Sink s | Comparison total s | Nearest ordinal | Score | Score sum |"); + System.out.println("|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|"); + } + + private static double seconds(long nanos) { return nanos / 1e9; } + + /** Uniform selection without replacement; preserve the same shuffled request order in every arm. */ + static int[] sample(int population, int count, long seed) { + if (count <= 0 || count > population) throw new IllegalArgumentException("Sample count must be in [1, " + population + "]"); + var random = new SplittableRandom(seed); + var used = new HashSet(); + var selected = new int[count]; + for (int j = population - count, i = 0; j < population; j++, i++) { + int candidate = random.nextInt(j + 1); + int ordinal = used.add(candidate) ? candidate : j; + used.add(ordinal); selected[i] = ordinal; + } + for (int i = count - 1; i > 0; i--) { + int j = random.nextInt(i + 1), swap = selected[i]; + selected[i] = selected[j]; selected[j] = swap; + } + return selected; + } + + private static RandomAccessVectorValues prefix(RandomAccessVectorValues source, int count) { + if (source.size() == count) return source; + if (source instanceof io.github.jbellis.jvector.graph.BatchedVectorValues) + return new Prefix((io.github.jbellis.jvector.graph.BatchedVectorValues) source, count); + return new RandomAccessVectorValues() { + public int size() { return count; } + public int dimension() { return source.dimension(); } + public boolean isValueShared() { return source.isValueShared(); } + public RandomAccessVectorValues copy() { return prefix(source.copy(), count); } + public io.github.jbellis.jvector.vector.types.VectorFloat getVector(int ordinal) { + java.util.Objects.checkIndex(ordinal, count); return source.getVector(ordinal); + } + }; + } + + /** Restrict a file source without losing its batched range/selection capability. */ + private static final class Prefix implements io.github.jbellis.jvector.graph.BatchedVectorValues { + private final io.github.jbellis.jvector.graph.BatchedVectorValues source; + private final int size; + Prefix(io.github.jbellis.jvector.graph.BatchedVectorValues source, int size) { this.source = source; this.size = size; } + public int size() { return size; } + public int dimension() { return source.dimension(); } + public boolean isValueShared() { return true; } + public RandomAccessVectorValues copy() { return prefix(source.copy(), size); } + public io.github.jbellis.jvector.vector.types.VectorFloat getVector(int ordinal) { + java.util.Objects.checkIndex(ordinal, size); return source.getVector(ordinal); + } + public io.github.jbellis.jvector.graph.VectorCursor openRange(int start, int end) { + java.util.Objects.checkFromToIndex(start, end, size); return source.openRange(start, end); + } + public io.github.jbellis.jvector.graph.VectorCursor openSelection(int[] ordinals, int offset, int count) { + java.util.Objects.checkFromIndexSize(offset, count, ordinals.length); + for (int i = offset; i < offset + count; i++) java.util.Objects.checkIndex(ordinals[i], size); + return source.openSelection(ordinals, offset, count); + } + } + + private static final class Progress implements AutoCloseable { + private volatile String phase = "starting"; + private volatile long since = System.nanoTime(); + private final java.util.concurrent.ScheduledExecutorService timer; + Progress(int seconds) { + timer = seconds == 0 ? null : Executors.newSingleThreadScheduledExecutor(r -> { + var thread = new Thread(r, "input-benchmark-progress"); thread.setDaemon(true); return thread; + }); + if (timer != null) timer.scheduleAtFixedRate(() -> System.out.printf(Locale.ROOT, "Still %s (%.0f s)%n", + phase, seconds(System.nanoTime() - since)), seconds, seconds, TimeUnit.SECONDS); + } + void phase(String name) { + phase = name; since = System.nanoTime(); System.out.println("Phase: " + name); + } + public void close() { if (timer != null) timer.shutdownNow(); } + } + + private static final class Options { + private final LinkedHashMap values; + Options(LinkedHashMap values) { this.values = values; } + private static final Set KEYS = Set.of("dataset", "file", "samples", "mode", "quantizer", "metric", + "query-file", "cache", "seed", "progress-seconds", "help", "child", "quantizer-options", "loader-options"); + static Options parse(String[] args) { + var values = new LinkedHashMap(); + for (int i = 0; i < args.length; i++) { + if (!args[i].startsWith("--") || !KEYS.contains(args[i].substring(2))) + throw new IllegalArgumentException("Unknown option: " + args[i]); + String key = args[i].substring(2); + String value = key.equals("help") || key.equals("child") ? "true" + : ++i < args.length ? args[i] : null; + if (value == null || value.startsWith("--")) throw new IllegalArgumentException("Missing value for --" + key); + if (values.putIfAbsent(key, value) != null) throw new IllegalArgumentException("Duplicate --" + key); + } + if (values.containsKey("file") && values.containsKey("dataset")) throw new IllegalArgumentException("Choose --file or --dataset"); + return new Options(values); + } + boolean contains(String key) { return values.containsKey(key); } + String get(String key, String fallback) { return values.getOrDefault(key, fallback); } + long longValue(String key, long fallback) { return Long.parseLong(get(key, Long.toString(fallback))); } + int nonnegativeInt(String key, int fallback) { + int n = Integer.parseInt(get(key, Integer.toString(fallback))); + if (n < 0) throw new IllegalArgumentException("--" + key + " must be nonnegative"); return n; + } + int positiveInt(String key, int fallback) { + int n = nonnegativeInt(key, fallback); + if (n == 0) throw new IllegalArgumentException("--" + key + " must be positive"); return n; + } + void checkMode(String mode) { + if (mode.equals("all")) throw new IllegalArgumentException("A child must select one loader"); + try { loader(mode); } catch (Exception e) { throw new IllegalArgumentException("Unknown loader: " + mode, e); } + } + Path file() throws IOException { + if (contains("file")) return Path.of(get("file", "")).toAbsolutePath(); + var loader = new DataSetLoaderSimpleMFD("jvector-examples/yaml-configs/dataset-catalogs"); + String dataset = get("dataset", "e5-small-v2-100k"); + return loader.loadBaseVectorFile(dataset).orElseThrow(() -> new IllegalArgumentException("Unknown catalog dataset: " + dataset)).toAbsolutePath(); + } + List childArguments(Path file, String mode) { + var args = new ArrayList(); + values.forEach((key, value) -> { + if (!Set.of("dataset", "file", "mode", "child").contains(key)) { args.add("--" + key); args.add(value); } + }); + args.addAll(List.of("--file", file.toString(), "--mode", mode, "--child")); return args; + } + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/BaseVectorLoading.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/BaseVectorLoading.java new file mode 100644 index 000000000..79965a0ef --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/BaseVectorLoading.java @@ -0,0 +1,64 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example.benchmarks.datasets; + +import java.util.Locale; +import java.util.logging.Logger; +import io.github.jbellis.jvector.example.yaml.MultiConfig; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; + +/** Base-vector storage policy for benchmark harnesses, independent of quantization. */ +public enum BaseVectorLoading { + BUFFERED, + PRELOAD; + + private static final Logger LOG = Logger.getLogger(BaseVectorLoading.class.getName()); + + /** Existing VM option is an explicit override of the YAML flag. */ + public static BaseVectorLoading forBenchmark(boolean useDatasetPrefetching) { + String value = System.getProperty("jvector.dataset_loader", + useDatasetPrefetching ? "buffered" : "preload"); + try { return valueOf(value.toUpperCase(Locale.ROOT)); } + catch (IllegalArgumentException e) { + throw new IllegalArgumentException("jvector.dataset_loader must be buffered or preload, got: " + value, e); + } + } + + public static BaseVectorLoading forBenchmark() { return forBenchmark(false); } + + /** + * The VM option overrides the YAML flag. Configurations requiring resident raw + * vectors retain preload so construction and query scoring stay unchanged. + */ + public static BaseVectorLoading forBenchmark(MultiConfig config) { + var loading = forBenchmark(config.useDatasetPrefetching); + if (loading == PRELOAD) return loading; + if (config.construction.compression.stream().anyMatch(c -> "None".equals(c.type))) + return residentInput("uncompressed construction"); + boolean exactSearch = config.search.compression.stream().anyMatch(c -> "None".equals(c.type)); + boolean needsRawInput = exactSearch && config.construction.getFeatureSets().stream().anyMatch(features -> + !features.contains(FeatureId.FUSED_PQ)); + return needsRawInput ? residentInput("raw-input exact search") : BUFFERED; + } + + private static BaseVectorLoading residentInput(String reason) { + if ("buffered".equalsIgnoreCase(System.getProperty("jvector.dataset_loader"))) + LOG.warning("Using preload instead of jvector.dataset_loader=buffered for " + reason + + "; retaining resident input to preserve scoring behavior."); + return PRELOAD; + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/BufferedDataSet.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/BufferedDataSet.java new file mode 100644 index 000000000..260f542d4 --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/BufferedDataSet.java @@ -0,0 +1,74 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example.benchmarks.datasets; + +import io.github.jbellis.jvector.disk.FvecFileVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.VectorFloat; + +import java.io.IOException; +import java.nio.file.Path; +import java.util.AbstractList; +import java.util.List; + +/** Resident queries/ground truth with bounded, file-backed base vectors. */ +final class BufferedDataSet implements DataSet { + private final String name; + private final VectorSimilarityFunction similarity; + private final FvecFileVectorValues base; + private final List> queries; + private final List> groundTruth; + private final List> baseList; + + BufferedDataSet(String name, VectorSimilarityFunction similarity, Path file, + List> queries, List> groundTruth) throws IOException { + this.name = name; + this.similarity = similarity; + this.queries = queries; + this.groundTruth = groundTruth; + this.base = FvecFileVectorValues.open(file); + try { + if (base.size() == 0 || queries.isEmpty() || queries.size() != groundTruth.size()) + throw new IllegalArgumentException("Empty or mismatched dataset facets for " + name); + for (var query : queries) + if (query.length() != base.dimension()) + throw new IllegalArgumentException("Base/query dimensions differ for " + name); + var views = base.threadLocalSupplier(); + // Compatibility list is lazy; retained entries must not alias a view's scratch vector. + baseList = new AbstractList<>() { + @Override public VectorFloat get(int ordinal) { return views.get().getVector(ordinal).copy(); } + @Override public int size() { return base.size(); } + }; + } catch (RuntimeException | Error failure) { + try { base.close(); } catch (IOException closeFailure) { failure.addSuppressed(closeFailure); } + throw failure; + } + System.out.printf("%n%s: %d base and %d query vectors loaded, dimensions %d (buffered)%n", + name, base.size(), queries.size(), base.dimension()); + } + + @Override public boolean useDatasetPrefetching() { return true; } + @Override public int getDimension() { return base.dimension(); } + @Override public RandomAccessVectorValues getBaseRavv() { return base; } + @Override public String getName() { return name; } + @Override public VectorSimilarityFunction getSimilarityFunction() { return similarity; } + @Override public List> getBaseVectors() { return baseList; } + @Override public List> getQueryVectors() { return queries; } + @Override public List> getGroundTruth() { return groundTruth; } + @Override public void close() throws IOException { base.close(); } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java index a33a40d31..daa91d1d0 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java @@ -20,12 +20,17 @@ import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.types.VectorFloat; +import java.io.IOException; import java.util.*; /** * This provides a uniform way to access vector test data, regardless of where it comes from or how it is implemented. */ -public interface DataSet { +public interface DataSet extends AutoCloseable { + + /** Release file-backed input resources after all benchmark consumers stop. */ + @Override + default void close() throws IOException {} /** * Get dimensions of the vectors in this dataset. @@ -33,6 +38,9 @@ public interface DataSet { */ int getDimension(); + /** True only for the harness's explicit bounded-input mode. */ + default boolean useDatasetPrefetching() { return false; } + /** * Get a random-access view of base vectors. * @return base vectors diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoader.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoader.java index 932ea2dc7..e6b8a75cb 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoader.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoader.java @@ -47,4 +47,12 @@ public interface DataSetLoader { * @return a {@link DataSetInfo} handle for the dataset, if found */ Optional loadDataSet(String dataSetName); + + /** Choose base-vector loading for a harness run; existing calls retain preload behavior. */ + default Optional loadDataSet(String dataSetName, BaseVectorLoading loading) { + var result = loadDataSet(dataSetName); + if (result.isPresent() && loading != BaseVectorLoading.PRELOAD) + throw new UnsupportedOperationException("Loader does not support " + loading); + return result; + } } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFD.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFD.java index bef54863b..d32918635 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFD.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFD.java @@ -385,8 +385,34 @@ public DataSetLoaderSimpleMFD(String catalogUrl, String localPath, boolean check } } + /** + * Prepare only a catalog entry's base-vector file for a file-backed source. + * Uses the same cache and download rules as loadDataSet, without downloading + * queries/ground truth or materializing, normalizing, or scrubbing vectors. + * + * @param dataSetName the catalog dataset name + * @return the local base-file path, or empty if the name is not in this catalog + * @throws IOException if the base file cannot be prepared + */ + public Optional loadBaseVectorFile(String dataSetName) throws IOException { + var entry = catalog.get(dataSetName); + if (entry == null) return Optional.empty(); + String filename = entry.fields.get("base"); + if (filename == null) + throw new IOException("Dataset '" + dataSetName + "' has no base file in catalog"); + ensureFileAvailable(filename, entry.cacheDir, + entry.baseUrl != null ? entry.baseUrl : remoteBasePath); + return Optional.of(entry.cacheDir.resolve(filename)); + } + @Override public Optional loadDataSet(String dataSetName) { + return loadDataSet(dataSetName, BaseVectorLoading.PRELOAD); + } + + @Override + public Optional loadDataSet(String dataSetName, BaseVectorLoading loading) { + java.util.Objects.requireNonNull(loading, "loading"); var entry = catalog.get(dataSetName); if (entry == null) return Optional.empty(); @@ -424,9 +450,19 @@ public Optional loadDataSet(String dataSetName) { "Dataset '%s' was found in dataset catalog, but no metadata entry was found in dataset-metadata.yml. ", dataSetName))); return Optional.of(new DataSetInfo(props, () -> { - var baseVectors = SiftLoader.readFvecs(effectiveCacheDir.resolve(baseFile).toString()); + if (loading == BaseVectorLoading.BUFFERED && props.loadBehavior() != DataSetProperties.LoadBehavior.NO_SCRUB) + throw new IllegalArgumentException("Buffered loading requires NO_SCRUB metadata for " + dataSetName); var queryVectors = SiftLoader.readFvecs(effectiveCacheDir.resolve(queryFile).toString()); var gtVectors = SiftLoader.readIvecs(effectiveCacheDir.resolve(gtFile).toString()); + if (loading == BaseVectorLoading.BUFFERED) { + try { + return new BufferedDataSet(dataSetName, props.similarityFunction().orElseThrow(), + effectiveCacheDir.resolve(baseFile), queryVectors, gtVectors); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + var baseVectors = SiftLoader.readFvecs(effectiveCacheDir.resolve(baseFile).toString()); return DataSetUtils.processDataSet(dataSetName, props, baseVectors, queryVectors, gtVectors); })); } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSets.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSets.java index 94b2e1c77..0c327238d 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSets.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSets.java @@ -62,14 +62,26 @@ public static Optional loadDataSet(String dataSetName) { /// @param loaders the loaders to try, in priority order /// @return a lazy {@link DataSetInfo} handle, or empty if no loader recognises the name public static Optional loadDataSet(String dataSetName, Collection loaders) { - logger.info("loading dataset [{}]", dataSetName); + return loadDataSet(dataSetName, loaders, BaseVectorLoading.PRELOAD); + } + + /// Select base-vector loading explicitly for a benchmark harness. + public static Optional loadDataSet(String dataSetName, BaseVectorLoading loading) { + return loadDataSet(dataSetName, defaultLoaders, loading); + } + + public static Optional loadDataSet(String dataSetName, + Collection loaders, + BaseVectorLoading loading) { + java.util.Objects.requireNonNull(loading, "loading"); + logger.info("loading dataset [{}] (base vector loading: {})", dataSetName, loading); if (dataSetName.endsWith(".hdf5")) { throw new InvalidParameterException("DataSet names are not meant to be file names. Did you mean " + dataSetName.replace(".hdf5", "") + "? "); } for (DataSetLoader loader : loaders) { logger.trace("trying loader [{}]", loader.getClass().getSimpleName()); - Optional dataSetLoaded = loader.loadDataSet(dataSetName); + Optional dataSetLoaded = loader.loadDataSet(dataSetName, loading); if (dataSetLoaded.isPresent()) { logger.info("dataset [{}] found with loader [{}]", dataSetName, loader.getClass().getSimpleName()); return dataSetLoaded; diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/MultiConfig.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/MultiConfig.java index 1cfcc3c25..851d2b2ab 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/MultiConfig.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/MultiConfig.java @@ -36,6 +36,8 @@ public class MultiConfig { public ConstructionParameters construction; public SearchParameters search; public String dataset; + /** Opt into bounded bulk input loading; omitted settings preserve resident loading. */ + public boolean useDatasetPrefetching = false; /** Optional; when [Yes], AutoBenchYAML also runs CompactionBench. Absent defaults to [No]. */ public List withCompaction; diff --git a/jvector-examples/src/main/resources/log4j2.xml b/jvector-examples/src/main/resources/log4j2.xml index 83c77bced..0f9260830 100644 --- a/jvector-examples/src/main/resources/log4j2.xml +++ b/jvector-examples/src/main/resources/log4j2.xml @@ -1,12 +1,12 @@ - + - + diff --git a/jvector-examples/src/main/resources/logback.xml b/jvector-examples/src/main/resources/logback.xml index 0a7d8846a..57e7fb10b 100644 --- a/jvector-examples/src/main/resources/logback.xml +++ b/jvector-examples/src/main/resources/logback.xml @@ -15,5 +15,5 @@ - + diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/AutoBenchYAMLTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/AutoBenchYAMLTest.java new file mode 100644 index 000000000..ccdb5da8c --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/AutoBenchYAMLTest.java @@ -0,0 +1,58 @@ +/* + * Copyright DataStax, Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy at http://www.apache.org/licenses/LICENSE-2.0 + */ +package io.github.jbellis.jvector.example; + +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import static org.junit.Assert.*; + +public class AutoBenchYAMLTest { + @Rule public TemporaryFolder temporary = new TemporaryFolder(); + + @Test(timeout = 60000) + public void invalidConfigurationDoesNotAbortRemainingDatasets() throws Exception { + Path work = temporary.newFolder("work").toPath(); + Path yaml = work.resolve("jvector-examples/yaml-configs"); + Files.createDirectories(yaml); + Files.writeString(yaml.resolve("datasets.yml"), "regression-tests:\n - first\n - second\n"); + Path config = work.resolve("invalid.yml"); + Files.writeString(config, "unknownConfigurationProperty: true\n"); + String classpath = Arrays.stream(System.getProperty("surefire.test.class.path", + System.getProperty("java.class.path")).split(java.io.File.pathSeparator)) + .map(p -> Path.of(p).toAbsolutePath().toString()) + .collect(Collectors.joining(java.io.File.pathSeparator)); + Path log = work.resolve("console.log"); + var command = new java.util.ArrayList(); + command.add(Path.of(System.getProperty("java.home"), "bin", "java").toString()); + // The log is decoded as UTF-8, including on Windows with a different default charset. + command.add("-Dfile.encoding=UTF-8"); + // Keep the Java 11 test runnable without optional vector or FFM modules. + if (ModuleLayer.boot().findModule("jdk.incubator.vector").isPresent()) + command.add("--add-modules=jdk.incubator.vector"); + if (Runtime.version().feature() >= 22) + command.add("--enable-native-access=ALL-UNNAMED"); + command.addAll(Arrays.asList("-cp", classpath, AutoBenchYAML.class.getName(), + "--output", work.resolve("results.json").toString(), "--config", config.toString())); + Process process = new ProcessBuilder(command).directory(work.toFile()).redirectErrorStream(true) + .redirectOutput(log.toFile()).start(); + try { + assertTrue("AutoBenchYAML did not exit", process.waitFor(30, TimeUnit.SECONDS)); + String output = Files.readString(log); + assertEquals(output, 0, process.exitValue()); + assertTrue(output, output.contains("Exception while processing dataset first")); + assertTrue(output, output.contains("Exception while processing dataset second")); + } finally { + if (process.isAlive()) { process.destroyForcibly(); process.waitFor(); } + } + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/DemandFvecReadProgressTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/DemandFvecReadProgressTest.java new file mode 100644 index 000000000..70f00a282 --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/DemandFvecReadProgressTest.java @@ -0,0 +1,82 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.jbellis.jvector.example; + +import org.junit.Test; +import java.io.EOFException; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.MappedByteBuffer; +import java.nio.channels.*; +import static org.junit.Assert.*; + +public class DemandFvecReadProgressTest { + @Test(timeout = 1000) public void repeatedZeroReadsFailInsteadOfSpinning() { + var channel = new TestChannel(-1); + var failure = assertThrows(IOException.class, () -> + DemandFvecVectorValues.readFully(channel, ByteBuffer.allocate(1), 123)); + assertTrue(failure.getMessage().contains("123")); + assertEquals(16, channel.calls); + } + + @Test(timeout = 1000) public void progressResetsTheZeroReadLimitAndPreservesOffsets() throws Exception { + var channel = new TestChannel(15); + var buffer = ByteBuffer.allocate(4); + buffer.position(2); + DemandFvecVectorValues.readFully(channel, buffer, 123); + assertEquals(32, channel.calls); + assertEquals(4, buffer.position()); + assertEquals((byte) 123, buffer.get(2)); + assertEquals((byte) 124, buffer.get(3)); + } + + @Test(timeout = 1000) public void eofStillFailsImmediately() { + var channel = new TestChannel(0) { + @Override public int read(ByteBuffer buffer, long offset) { return -1; } + }; + assertThrows(EOFException.class, () -> + DemandFvecVectorValues.readFully(channel, ByteBuffer.allocate(1), 0)); + } + + private static class TestChannel extends FileChannel { + final int zerosPerByte; + int calls, zeros; + TestChannel(int zerosPerByte) { this.zerosPerByte = zerosPerByte; } + @Override public int read(ByteBuffer buffer, long offset) { + calls++; + if (zerosPerByte < 0 || zeros++ < zerosPerByte) return 0; + zeros = 0; + buffer.put((byte) offset); + return 1; + } + @Override public int read(ByteBuffer buffer) { throw new UnsupportedOperationException(); } + @Override public long read(ByteBuffer[] buffers, int offset, int count) { throw new UnsupportedOperationException(); } + @Override public int write(ByteBuffer buffer) { throw new UnsupportedOperationException(); } + @Override public int write(ByteBuffer buffer, long offset) { throw new UnsupportedOperationException(); } + @Override public long write(ByteBuffer[] buffers, int offset, int count) { throw new UnsupportedOperationException(); } + @Override public long position() { throw new UnsupportedOperationException(); } + @Override public FileChannel position(long position) { throw new UnsupportedOperationException(); } + @Override public long size() { throw new UnsupportedOperationException(); } + @Override public FileChannel truncate(long size) { throw new UnsupportedOperationException(); } + @Override public void force(boolean metadata) { throw new UnsupportedOperationException(); } + @Override public long transferTo(long position, long count, WritableByteChannel target) { throw new UnsupportedOperationException(); } + @Override public long transferFrom(ReadableByteChannel source, long position, long count) { throw new UnsupportedOperationException(); } + @Override public MappedByteBuffer map(MapMode mode, long position, long size) { throw new UnsupportedOperationException(); } + @Override public FileLock lock(long position, long size, boolean shared) { throw new UnsupportedOperationException(); } + @Override public FileLock tryLock(long position, long size, boolean shared) { throw new UnsupportedOperationException(); } + @Override protected void implCloseChannel() {} + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/VectorPrefetchBenchmarkTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/VectorPrefetchBenchmarkTest.java new file mode 100644 index 000000000..d352bc9f1 --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/VectorPrefetchBenchmarkTest.java @@ -0,0 +1,174 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.disk.FvecFileVectorValues; +import io.github.jbellis.jvector.example.util.SiftLoader; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.VectorAccess; +import io.github.jbellis.jvector.quantization.NVQuantization; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.Arrays; +import java.util.HashSet; +import java.util.concurrent.ForkJoinPool; + +import static org.junit.Assert.*; +import static org.junit.Assume.assumeTrue; + +public class VectorPrefetchBenchmarkTest { + @Rule public TemporaryFolder temporary = new TemporaryFolder(); + + @Test public void sampleIsRepeatableUniqueAndOrdered() { + int[] selected = VectorPrefetchBenchmark.sample(1000, 128, 42); + assertArrayEquals(selected, VectorPrefetchBenchmark.sample(1000, 128, 42)); + assertFalse(Arrays.equals(selected, VectorPrefetchBenchmark.sample(1000, 128, 43))); + var seen = new HashSet(); + for (int ordinal : selected) { assertTrue(ordinal >= 0 && ordinal < 1000); assertTrue(seen.add(ordinal)); } + assertEquals(128, seen.size()); + assertEquals(1000, Arrays.stream(VectorPrefetchBenchmark.sample(1000, 1000, 42)).distinct().count()); + } + + @Test public void sampleRejectsImpossibleCounts() { + assertThrows(IllegalArgumentException.class, () -> VectorPrefetchBenchmark.sample(100, 101, 42)); + assertThrows(IllegalArgumentException.class, () -> VectorPrefetchBenchmark.sample(100, 0, 42)); + } + + @Test public void demandCopiesDoNotOverwriteEachOther() throws Exception { + Path file = input(); + var preload = SiftLoader.readFvecs(file.toString()); + try (var source = new DemandFvecVectorValues(file); var copy = source.copy()) { + var first = source.getVector(5); + var other = copy.getVector(7); + assertEquals(preload.get(5), first); + assertEquals(preload.get(7), other); + assertThrows(IndexOutOfBoundsException.class, () -> source.getVector(source.size())); + } + } + + @Test public void fullScanAndSampleProduceSameNVQAcrossAllSources() throws Exception { + Path file = input(); + var preload = new ListRandomAccessVectorValues(SiftLoader.readFvecs(file.toString()), 16); + var pool = new ForkJoinPool(2); + try (var demand = new DemandFvecVectorValues(file); var prefetch = FvecFileVectorValues.open(file)) { + var sources = java.util.List.of(preload, demand, prefetch); + for (boolean sampled : new boolean[]{false, true}) { + var ordinals = VectorPrefetchBenchmark.sample(preload.size(), 32, 42); + io.github.jbellis.jvector.quantization.NVQVectors expected = null; + for (var source : sources) { + var vectors = sampled ? new ListRandomAccessVectorValues(VectorAccess.copySelected(source, ordinals, pool), 16) : source; + var encoded = NVQuantization.compute(vectors, 1).encodeAll(source, pool); + assertEquals(preload.size(), encoded.count()); + if (expected == null) expected = encoded; + else assertEquals(expected, encoded); + var scorer = encoded.scoreFunctionFor(preload.getVector(100), VectorSimilarityFunction.DOT_PRODUCT); + for (int i = 0; i < encoded.count(); i++) assertTrue(Float.isFinite(scorer.similarityTo(i))); + } + } + } finally { pool.shutdown(); } + } + + @Test public void coldCacheEvictsAndVerifiesOnlyInputFile() throws Exception { + assumeTrue(Runtime.version().feature() >= 23); + assumeTrue(System.getProperty("os.name").equals("Linux")); + assumeTrue(System.getProperty("os.arch").matches("amd64|aarch64")); + ColdFileCache.evictAndVerify(input()); + } + + @Test public void customLoaderAndNonCompressorRepresentationAreSelectable() throws Exception { + Path file = input(); + var args = new String[]{"--file", file.toString(), "--samples", "32", + "--mode", CustomLoader.class.getName(), "--quantizer", CustomRepresentation.class.getName(), + "--quantizer-options", "example=accepted", "--loader-options", "example=accepted", + "--cache", "uncontrolled", "--progress-seconds", "0", "--child"}; + CustomLoader.closed = false; + CustomRepresentation.trainedCount = CustomRepresentation.encodedCount = 0; + VectorPrefetchBenchmark.main(args); + assertTrue(CustomLoader.closed); + assertEquals(32, CustomRepresentation.trainedCount); + assertEquals(255, CustomRepresentation.encodedCount); + } + + @Test public void builtInPQAndNVQAcceptNamedParameters() throws Exception { + Path file = input(); + for (var quantizer : new String[]{"pq", "nvq"}) { + VectorPrefetchBenchmark.main(new String[]{"--file", file.toString(), + "--samples", "32", "--mode", "prefetch", "--quantizer", quantizer, + "--quantizer-options", quantizer.equals("pq") ? "subspaces=2,clusters=16,center=false,anisotropic=-1" + : "subvectors=2,learn=false", + "--loader-options", "ioThreads=2,batchVectors=8,readAhead=1,maxBufferBytes=1024", + "--cache", "uncontrolled", "--progress-seconds", "0", "--child"}); + } + } + + public static final class CustomLoader implements VectorPrefetchBenchmark.Loader { + static boolean closed; + public VectorPrefetchBenchmark.Input open(Path file, java.util.Map parameters) throws Exception { + assertEquals("accepted", parameters.get("example")); + var source = new DemandFvecVectorValues(file); + return new VectorPrefetchBenchmark.Input() { + public io.github.jbellis.jvector.graph.RandomAccessVectorValues values() { return source; } + public void close() throws Exception { source.close(); closed = true; } + }; + } + } + + // Exercise a representation that has a score function but does not implement VectorCompressor. + public static final class CustomRepresentation implements VectorPrefetchBenchmark.Quantizer { + static int trainedCount, encodedCount; + public VectorPrefetchBenchmark.Encoder train(io.github.jbellis.jvector.graph.RandomAccessVectorValues source, + java.util.Map parameters, ForkJoinPool executor) { + assertEquals("accepted", parameters.get("example")); + trainedCount = source.size(); + return (values, pool) -> { + encodedCount = values.size(); + var ordinals = java.util.stream.IntStream.range(0, values.size()).toArray(); + var vectors = VectorAccess.copySelected(values, ordinals, pool); + return new VectorPrefetchBenchmark.Encoded() { + public int count() { return vectors.size(); } + public java.util.function.IntToDoubleFunction scorer( + io.github.jbellis.jvector.vector.types.VectorFloat query, VectorSimilarityFunction metric) { + return ordinal -> metric.compare(query, vectors.get(ordinal)); + } + }; + }; + } + } + + private Path input() throws Exception { + Path file = temporary.newFile().toPath(); + var data = ByteBuffer.allocate(256 * 17 * 4).order(ByteOrder.LITTLE_ENDIAN); + for (int i = 0; i < 256; i++) { + data.putInt(16); + for (int j = 0; j < 16; j++) data.putFloat((float) Math.sin(i * 0.1 + j * 0.7)); + } + data.flip(); + try (var out = FileChannel.open(file, StandardOpenOption.WRITE)) { + while (data.hasRemaining()) out.write(data); + out.force(true); // Make test pages clean before asking the kernel to evict them. + } + return file; + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java index 33379dd57..bcc9032ac 100644 --- a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java @@ -73,6 +73,33 @@ public void setUp() throws IOException { // Basic loading // ======================================================================== + @Test + public void preparesBaseFileWithoutLoadingOtherFacets() throws IOException { + writeTestCatalog(cacheDir); + Path catalog = cacheDir.resolve("catalog_entries.yaml"); + Files.writeString(catalog, "_defaults:\n cache_dir: '" + + cacheDir.toString().replace("'", "''") + "'\n" + Files.readString(catalog)); + Path base = cacheDir.resolve("test_base.fvecs"); + writeTestFvecs(base, 4, new float[][] {{1, 2, 3, 4}}); + byte[] before = Files.readAllBytes(base); + var loader = new DataSetLoaderSimpleMFD(null, cacheDir.toString(), false, testMetadata); + assertEquals(base, loader.loadBaseVectorFile("test-ds").orElseThrow()); + assertArrayEquals(before, Files.readAllBytes(base)); + assertFalse(Files.exists(cacheDir.resolve("test_query.fvecs"))); + assertFalse(Files.exists(cacheDir.resolve("test_gt.ivecs"))); + assertFalse(loader.loadBaseVectorFile("unknown").isPresent()); + } + + @Test + public void missingBaseFileIsReportedWithoutMaterializingDataset() throws IOException { + writeTestCatalog(cacheDir); + Path catalog = cacheDir.resolve("catalog_entries.yaml"); + Files.writeString(catalog, "_defaults:\n cache_dir: '" + + cacheDir.toString().replace("'", "''") + "'\n" + Files.readString(catalog)); + var loader = new DataSetLoaderSimpleMFD(null, cacheDir.toString(), false, testMetadata); + assertThrows(IOException.class, () -> loader.loadBaseVectorFile("test-ds")); + } + @Test public void loadsDatasetFromLocalCatalogAndFiles() throws IOException { writeTestCatalog(cacheDir); @@ -93,6 +120,79 @@ public void loadsDatasetFromLocalCatalogAndFiles() throws IOException { assertEquals(4, ds.getDimension()); } + @Test + public void bufferedModePreservesValuesAndClosesItsSource() throws Exception { + writeTestCatalog(cacheDir); + writeTestDataFiles(cacheDir); + var loader = new DataSetLoaderSimpleMFD(null, cacheDir.toString(), false, testMetadata); + var info = DataSets.loadDataSet("test-ds", java.util.List.of(loader), BaseVectorLoading.BUFFERED).orElseThrow(); + var ds = info.getDataSet(); + assertSame(ds, info.getDataSet()); + var values = ds.getBaseRavv(); + var file = assertInstanceOf(io.github.jbellis.jvector.disk.FvecFileVectorValues.class, values); + assertEquals(0, file.statistics().bufferBytes, "Preparing a dataset must not scan its base file"); + var retained = ds.getBaseVectors().get(0); + ds.getBaseVectors().get(1); + assertEquals(1f, retained.get(0)); + assertEquals(0f, retained.get(1)); + try (ds; var preload = loader.loadDataSet("test-ds").orElseThrow().getDataSet()) { + assertEquals(preload.getGroundTruth(), ds.getGroundTruth()); + assertEquals(preload.getSimilarityFunction(), ds.getSimilarityFunction()); + var consumers = java.util.concurrent.Executors.newFixedThreadPool(4); + try { + var tasks = new java.util.ArrayList>(); + for (int worker = 0; worker < 16; worker++) { + tasks.add(() -> { + try (var view = file.copy()) { + for (int pass = 0; pass < 20; pass++) { + for (int ordinal = 0; ordinal < ds.getBaseVectors().size(); ordinal++) { + var expected = preload.getBaseVectors().get(ordinal); + var actual = view.getVector(ordinal); + for (int col = 0; col < expected.length(); col++) + assertEquals(Float.floatToRawIntBits(expected.get(col)), Float.floatToRawIntBits(actual.get(col))); + } + } + } + return null; + }); + } + for (var result : consumers.invokeAll(tasks)) result.get(); + } finally { consumers.shutdown(); } + } + assertThrows(IllegalStateException.class, () -> values.getVector(0)); + assertEquals(1f, retained.get(0), "Compatibility list returns an independent vector"); + } + + @Test + public void bufferedModeRejectsLegacyScrubbing() throws IOException { + writeTestCatalog(cacheDir); + writeTestDataFiles(cacheDir); + Path metadataFile = tempFolder.newFile("legacy_metadata.yml").toPath(); + Files.writeString(metadataFile, "test-ds:\n similarity_function: COSINE\n load_behavior: LEGACY_SCRUB\n"); + var loader = new DataSetLoaderSimpleMFD(null, cacheDir.toString(), false, + DataSetMetadataReader.load(metadataFile.toString())); + var info = loader.loadDataSet("test-ds", BaseVectorLoading.BUFFERED).orElseThrow(); + var failure = assertThrows(IllegalArgumentException.class, info::getDataSet); + assertTrue(failure.getMessage().contains("NO_SCRUB")); + } + + @Test + public void benchmarkDefaultsToPreloadWithExplicitBufferedOptIn() { + String previous = System.getProperty("jvector.dataset_loader"); + try { + System.clearProperty("jvector.dataset_loader"); + assertEquals(BaseVectorLoading.PRELOAD, BaseVectorLoading.forBenchmark()); + assertEquals(BaseVectorLoading.BUFFERED, BaseVectorLoading.forBenchmark(true)); + System.setProperty("jvector.dataset_loader", "preload"); + assertEquals(BaseVectorLoading.PRELOAD, BaseVectorLoading.forBenchmark()); + System.setProperty("jvector.dataset_loader", "typo"); + assertThrows(IllegalArgumentException.class, BaseVectorLoading::forBenchmark); + } finally { + if (previous == null) System.clearProperty("jvector.dataset_loader"); + else System.setProperty("jvector.dataset_loader", previous); + } + } + @Test public void returnsEmptyForUnknownDataset() throws IOException { writeTestCatalog(cacheDir); diff --git a/jvector-examples/yaml-configs/index-parameters/cap-10M.yml b/jvector-examples/yaml-configs/index-parameters/cap-10M.yml index c4c285d18..c9c6c09d8 100644 --- a/jvector-examples/yaml-configs/index-parameters/cap-10M.yml +++ b/jvector-examples/yaml-configs/index-parameters/cap-10M.yml @@ -2,6 +2,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: cap-10M +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/cap-1M.yml b/jvector-examples/yaml-configs/index-parameters/cap-1M.yml index 0bd9d5de1..6aa8454b5 100644 --- a/jvector-examples/yaml-configs/index-parameters/cap-1M.yml +++ b/jvector-examples/yaml-configs/index-parameters/cap-1M.yml @@ -2,6 +2,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: cap-1M +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-10M.yml b/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-10M.yml index 7973ce0f8..17b3dd54a 100644 --- a/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-10M.yml +++ b/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-10M.yml @@ -2,6 +2,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: cohere-english-v3-10M +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-1M.yml b/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-1M.yml index e053edfc4..56fb2849f 100644 --- a/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-1M.yml +++ b/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-1M.yml @@ -2,6 +2,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: cohere-english-v3-1M +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/colbert-1M.yml b/jvector-examples/yaml-configs/index-parameters/colbert-1M.yml index d0de5e0f4..66b47a056 100644 --- a/jvector-examples/yaml-configs/index-parameters/colbert-1M.yml +++ b/jvector-examples/yaml-configs/index-parameters/colbert-1M.yml @@ -2,6 +2,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: colbert-1M +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/default.yml b/jvector-examples/yaml-configs/index-parameters/default.yml index b56e27ed0..e0a4ed3a5 100644 --- a/jvector-examples/yaml-configs/index-parameters/default.yml +++ b/jvector-examples/yaml-configs/index-parameters/default.yml @@ -2,6 +2,8 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: default +# Prefetch bounded batches instead of preloading the entire dataset into memory. +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml index f2370a063..90eac34e8 100644 --- a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml +++ b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml @@ -2,6 +2,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: dpr-gemma-10M +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml index 96e92556e..354c06436 100644 --- a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml +++ b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml @@ -2,6 +2,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: dpr-gemma-1M +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/glove-100-angular.yml b/jvector-examples/yaml-configs/index-parameters/glove-100-angular.yml index 5f80f4ae2..1b03b1edd 100644 --- a/jvector-examples/yaml-configs/index-parameters/glove-100-angular.yml +++ b/jvector-examples/yaml-configs/index-parameters/glove-100-angular.yml @@ -3,6 +3,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: glove-100-angular +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/glove-200-angular.yml b/jvector-examples/yaml-configs/index-parameters/glove-200-angular.yml index faf5bf3fa..1630baaba 100644 --- a/jvector-examples/yaml-configs/index-parameters/glove-200-angular.yml +++ b/jvector-examples/yaml-configs/index-parameters/glove-200-angular.yml @@ -3,6 +3,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: glove-200-angular +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/glove-25-angular.yml b/jvector-examples/yaml-configs/index-parameters/glove-25-angular.yml index 953f454e2..bc2293a9d 100644 --- a/jvector-examples/yaml-configs/index-parameters/glove-25-angular.yml +++ b/jvector-examples/yaml-configs/index-parameters/glove-25-angular.yml @@ -3,6 +3,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: glove-25-angular +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/glove-50-angular.yml b/jvector-examples/yaml-configs/index-parameters/glove-50-angular.yml index a62eede12..309ef9769 100644 --- a/jvector-examples/yaml-configs/index-parameters/glove-50-angular.yml +++ b/jvector-examples/yaml-configs/index-parameters/glove-50-angular.yml @@ -3,6 +3,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: glove-50-angular +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/lastfm-64-dot.yml b/jvector-examples/yaml-configs/index-parameters/lastfm-64-dot.yml index 045963849..c2e2fb98a 100644 --- a/jvector-examples/yaml-configs/index-parameters/lastfm-64-dot.yml +++ b/jvector-examples/yaml-configs/index-parameters/lastfm-64-dot.yml @@ -3,6 +3,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: lastfm-64-dot +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/nytimes-256-angular.yml b/jvector-examples/yaml-configs/index-parameters/nytimes-256-angular.yml index 755e99017..4e2db21a7 100644 --- a/jvector-examples/yaml-configs/index-parameters/nytimes-256-angular.yml +++ b/jvector-examples/yaml-configs/index-parameters/nytimes-256-angular.yml @@ -3,6 +3,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: nytimes-256-angular +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean.yml b/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean.yml index 04d4b6fd8..94362ec70 100644 --- a/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean.yml +++ b/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean.yml @@ -3,6 +3,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 dataset: sift-128-euclidean +useDatasetPrefetching: false construction: outDegree: [32] diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecCloseDuringPointReads.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecCloseDuringPointReads.java new file mode 100644 index 000000000..459ea391f --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecCloseDuringPointReads.java @@ -0,0 +1,119 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.disk; + +import org.junit.Assume; +import org.junit.Test; + +import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.ClosedChannelException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.SplittableRandom; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CountDownLatch; + + +/// Closing the root while point reads are in flight must fail those reads cleanly, as the rest +/// of the source does ("Vector source is closed" or a closed-channel I/O error), and must leave +/// no file descriptors open on the input file. +/// +/// The race is timing-dependent, so the test repeats close-under-load until it sees a +/// violation or the round limit is reached. A pass is not proof of absence. +public class TestFvecCloseDuringPointReads { + private static final int ROWS = 10_000, DIMENSIONS = 16, READERS = 12, ROUNDS = 300; + + private static Path writeFvecs() throws Exception { + var file = Files.createTempFile("close-during-point-reads", ".fvecs"); + var bytes = ByteBuffer.allocate(ROWS * (DIMENSIONS + 1) * 4).order(ByteOrder.LITTLE_ENDIAN); + for (int row = 0; row < ROWS; row++) { + bytes.putInt(DIMENSIONS); + for (int d = 0; d < DIMENSIONS; d++) bytes.putFloat(row); + } + Files.write(file, bytes.array()); + return file; + } + + /// Expected failures after close: the source reports closure, or the leased channel was + /// closed underneath the read. + private static boolean isCleanClosureFailure(Throwable t) { + if (t instanceof IllegalStateException) return true; + return t instanceof UncheckedIOException && t.getCause() instanceof ClosedChannelException; + } + + private static long openDescriptorsOn(Path file) throws Exception { + var real = file.toRealPath(); + long count = 0; + try (var fds = Files.list(Path.of("/proc/self/fd"))) { + for (var fd : (Iterable) fds::iterator) { + try { if (Files.readSymbolicLink(fd).equals(real)) count++; } + catch (java.io.IOException ignored) { } // descriptor closed while listing + } + } + return count; + } + + @Test(timeout = 120_000) + public void closeUnderPointLoadFailsCleanlyAndLeaksNoDescriptors() throws Exception { + Assume.assumeTrue("needs /proc/self/fd", Files.isDirectory(Path.of("/proc/self/fd"))); + var file = writeFvecs(); + try { + for (int round = 0; round < ROUNDS; round++) { + var unexpected = new ConcurrentLinkedQueue(); + var root = FvecFileVectorValues.open(file); + var started = new CountDownLatch(READERS); + List readers = new ArrayList<>(); + for (int r = 0; r < READERS; r++) { + var view = root.copy(); + long seed = round * 1000L + r; + var thread = new Thread(() -> { + var random = new SplittableRandom(seed); + started.countDown(); + try { + while (true) view.getVector(random.nextInt(ROWS)); + } catch (Throwable t) { + if (!isCleanClosureFailure(t)) unexpected.add(t); + } + }, "point-reader-" + r); + readers.add(thread); + thread.start(); + } + started.await(); + Thread.sleep(1); + root.close(); + for (var thread : readers) thread.join(); + + long leaked = openDescriptorsOn(file); + if (!unexpected.isEmpty() || leaked != 0) { + var first = unexpected.peek(); + var error = new AssertionError("round " + round + ": " + unexpected.size() + + " point read(s) failed with an unexpected exception after close" + + (first == null ? "" : " (first: " + first + ")") + + "; " + leaked + " file descriptor(s) still open on the input"); + if (first != null) error.initCause(first); + throw error; + } + } + } finally { + Files.delete(file); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecFileVectorValues.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecFileVectorValues.java new file mode 100644 index 000000000..b06a8986c --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecFileVectorValues.java @@ -0,0 +1,304 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.disk; + +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.VectorAccess; +import io.github.jbellis.jvector.quantization.NVQuantization; +import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.Executors; +import java.util.concurrent.ForkJoinPool; +import static org.junit.Assert.*; + +public class TestFvecFileVectorValues { + private static float value(int row, int column) { return column == 0 ? -0.0f : (row * 7 + column * 13) % 201 / 100f - 1; } + private static Path file(int rows, int dimension) throws IOException { + Path path = Files.createTempFile("vector-source-test", ".fvecs"); + var bytes = ByteBuffer.allocate(rows * (dimension + 1) * 4).order(ByteOrder.LITTLE_ENDIAN); + for (int row = 0; row < rows; row++) { + bytes.putInt(dimension); + for (int col = 0; col < dimension; col++) bytes.putFloat(value(row, col)); + } + Files.write(path, bytes.array()); return path; + } + private static void equal(int row, VectorFloat v) { + for (int col = 0; col < v.length(); col++) assertEquals(Float.floatToRawIntBits(value(row, col)), Float.floatToRawIntBits(v.get(col))); + } + + @Test public void optionsAreImmutableAndValidateLimits() throws Exception { + var defaults = FvecFileVectorValues.Options.defaults(); + var options = defaults.withMaxBufferBytes(1024).withIoThreads(2) + .withBatchVectors(7).withReadAhead(0); + assertEquals(64L << 20, defaults.maxBufferBytes()); + assertEquals(PhysicalCoreExecutor.pool().getParallelism(), defaults.ioThreads()); + assertEquals(64, defaults.batchVectors()); + assertEquals(3, defaults.readAhead()); + assertEquals(1024, options.maxBufferBytes()); + assertEquals(2, options.ioThreads()); + assertEquals(7, options.batchVectors()); + assertEquals(0, options.readAhead()); + assertThrows(IllegalArgumentException.class, () -> defaults.withMaxBufferBytes(0)); + assertThrows(IllegalArgumentException.class, () -> defaults.withIoThreads(0)); + assertThrows(IllegalArgumentException.class, () -> defaults.withBatchVectors(0)); + assertThrows(IllegalArgumentException.class, () -> defaults.withReadAhead(-1)); + var path = file(257, 17); + try (var source = FvecFileVectorValues.open(path, options)) { + assertEquals(257 * 18L * 4 / 100, source.statistics().bufferLimitBytes); + try (var cursor = VectorAccess.openRange(source, 7, 10)) { + int row = 7; + while (cursor.next()) { assertEquals(row, cursor.ordinal()); equal(row++, cursor.vector()); } + assertEquals(10, row); + } + try (var cursor = source.openRange(source.size(), source.size())) { + assertFalse(cursor.next()); + } + assertThrows(IndexOutOfBoundsException.class, () -> source.openRange(10, 7)); + assertThrows(NullPointerException.class, () -> FvecFileVectorValues.open(path, null)); + } finally { Files.delete(path); } + } + + @Test public void copiedCursorVectorSurvivesAdvanceAndSourceClose() throws Exception { + var path = file(5, 17); + VectorFloat retained; + try (var source = FvecFileVectorValues.open(path)) { + try (var cursor = VectorAccess.openRange(source, 1, 3)) { + assertTrue(cursor.isValueShared()); + assertTrue(cursor.next()); + var borrowed = cursor.vector(); + retained = borrowed.copy(); + assertTrue(cursor.next()); + assertSame(borrowed, cursor.vector()); + equal(2, borrowed); + equal(1, retained); + assertFalse(cursor.next()); + } + } finally { Files.delete(path); } + equal(1, retained); + } + + @Test(timeout=60000) public void concurrentRangeSelectionAndPointReadsWithOneRecordBudget() throws Exception { + for (int dimension : new int[] {17, 768, 3072}) { + var path = file(257, dimension); + try (var source = FvecFileVectorValues.open(path, FvecFileVectorValues.Options.defaults() + .withMaxBufferBytes((dimension + 1L) * 4).withIoThreads(12))) { + var executor = Executors.newFixedThreadPool(48); + try { + var tasks = new ArrayList>(); + for (int worker = 0; worker < 48; worker++) { + final int seed = worker; + tasks.add(() -> { + try (var copy = source.copy()) { + for (int i = 0; i < 32; i++) equal((i + seed) % source.size(), copy.getVector((i + seed) % source.size())); + } + try (var cursor = source.openRange(seed, seed + 17)) { + int ordinal = seed; + while (cursor.next()) { assertEquals(ordinal, cursor.ordinal()); equal(ordinal++, cursor.vector()); } + } + int[] ordinals = {256, seed, 0, 256, seed}; + try (var cursor = source.openSelection(ordinals, 0, ordinals.length)) { + int i = 0; + while (cursor.next()) { assertEquals(ordinals[i], cursor.ordinal()); equal(ordinals[i++], cursor.vector()); } + } + return null; + }); + } + for (var f : executor.invokeAll(tasks)) f.get(); + } finally { executor.shutdown(); } + assertTrue(source.statistics().bufferBytes <= source.statistics().bufferLimitBytes); + assertEquals((dimension + 1L) * 4, source.statistics().bufferBytes); + assertThrows(IndexOutOfBoundsException.class, () -> source.openRange(250, 258)); + assertThrows(IndexOutOfBoundsException.class, () -> source.openSelection(new int[] {257}, 0, 1)); + } finally { Files.delete(path); } + } + } + + @Test(timeout = 60000) public void concurrentRandomPointsAndScansRespectTheSharedBudget() throws Exception { + var path = file(65537, 17); + try { + for (int ahead : new int[] {0, 3}) { + try (var source = FvecFileVectorValues.open(path, FvecFileVectorValues.Options.defaults() + .withMaxBufferBytes(32768).withIoThreads(16).withBatchVectors(7).withReadAhead(ahead))) { + var executor = Executors.newFixedThreadPool(48); + try { + var tasks = new ArrayList>(); + for (int worker = 0; worker < 48; worker++) { + final int seed = worker; + tasks.add(() -> { + try (var copy = source.copy()) { + for (int i = 0; i < 256; i++) { + int ordinal = Math.floorMod(seed * 8191 + i * 7919, source.size()); + equal(ordinal, copy.getVector(ordinal)); + } + } + int first = seed * 127; + try (var cursor = source.openRange(first, first + 64)) { + int ordinal = first; + while (cursor.next()) equal(ordinal++, cursor.vector()); + assertEquals(first + 64, ordinal); + } + int[] ordinals = {source.size() - 1, seed, 0, seed, 32768}; + try (var cursor = source.openSelection(ordinals, 0, ordinals.length)) { + int i = 0; + while (cursor.next()) equal(ordinals[i++], cursor.vector()); + assertEquals(ordinals.length, i); + } + assertTrue(source.statistics().bufferBytes <= source.statistics().bufferLimitBytes); + return null; + }); + } + for (var future : executor.invokeAll(tasks)) future.get(); + } finally { executor.shutdown(); } + assertTrue(source.statistics().bufferBytes <= source.statistics().bufferLimitBytes); + } + } + } finally { Files.delete(path); } + } + + @Test(timeout = 10000) public void interruptedPointReadDoesNotCloseSharedSource() throws Exception { + var path = file(4096, 17); + try (var source = FvecFileVectorValues.open(path); + var copy = source.copy()) { + Thread.currentThread().interrupt(); + try { + copy.getVector(127); + fail("Interrupted point read must fail"); + } catch (RuntimeException expected) { + assertTrue(Thread.currentThread().isInterrupted()); + } finally { Thread.interrupted(); } + // Clearing the caller interruption permits reuse; other views and cursors remain usable. + equal(127, copy.getVector(127)); + equal(127, source.getVector(127)); + try (var cursor = source.openRange(0, 128)) { + int ordinal = 0; + while (cursor.next()) equal(ordinal++, cursor.vector()); + assertEquals(128, ordinal); + } + assertTrue(source.statistics().bufferBytes <= source.statistics().bufferLimitBytes); + } finally { Thread.interrupted(); Files.delete(path); } + } + + @Test(timeout=60000) public void materializationOrderAndEncodingMatchResidentSource() throws Exception { + var path = file(2049, 17); + var pool = new ForkJoinPool(8); + try (var source = FvecFileVectorValues.open(path)) { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + List> values = new ArrayList<>(); + for (int row = 0; row < source.size(); row++) { + float[] v = new float[17]; for (int col = 0; col < 17; col++) v[col] = value(row, col); + values.add(vts.createFloatVector(v)); + } + var resident = new ListRandomAccessVectorValues(values, 17); + int[] ordinals = {2048, 0, 99, 2048, 10, 9, 0}; + var materialized = VectorAccess.copySelected(source, ordinals, pool); + for (int i = 0; i < ordinals.length; i++) equal(ordinals[i], materialized.get(i)); + // Closing releases selection identity: the caller can reuse and modify its array. + try (var cursor = source.openSelection(ordinals, 0, ordinals.length)) { + assertTrue(cursor.next()); equal(ordinals[0], cursor.vector()); + } + ordinals[0] = 42; + try (var cursor = source.openSelection(ordinals, 0, ordinals.length)) { + assertTrue(cursor.next()); equal(42, cursor.vector()); + } + var pq = ProductQuantization.compute(resident, 2, 16, false); + assertEquals(pq.encodeAll(resident, pool), pq.encodeAll(source, pool)); + var nvq = NVQuantization.compute(resident, 1); + var nvqFile = NVQuantization.compute(source, 1); + assertEquals(nvq, nvqFile); + var residentNVQ = nvq.encodeAll(resident, pool); + var fileNVQ = nvq.encodeAll(source, pool); + assertEquals(residentNVQ, fileNVQ); + try (var cursor = source.openRange(0, 0)) { assertFalse(cursor.next()); } + var copy = source.copy(); copy.close(); equal(0, source.getVector(0)); + assertThrows(IllegalStateException.class, () -> copy.getVector(0)); + } finally { pool.shutdown(); Files.delete(path); } + } + + @Test public void errorsAndCloseInvalidateCopies() throws Exception { + var path = file(5, 17); + var source = FvecFileVectorValues.open(path); + var copy = source.copy(); + var cursor = source.openRange(0, 3); cursor.close(); cursor.close(); + source.close(); source.close(); + assertThrows(IllegalStateException.class, () -> copy.getVector(0)); + assertThrows(IllegalStateException.class, () -> source.openRange(0, 1)); + Files.write(path, new byte[] {1, 2}); + assertThrows(IOException.class, () -> FvecFileVectorValues.open(path)); + Files.write(path, ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN).putInt(-1).array()); + assertThrows(IOException.class, () -> FvecFileVectorValues.open(path)); + Files.write(path, ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN).putInt(Integer.MAX_VALUE).array()); + assertThrows(IOException.class, () -> FvecFileVectorValues.open(path)); + assertThrows(IndexOutOfBoundsException.class, () -> + NVQuantization.compute(new ListRandomAccessVectorValues(new ArrayList<>(), 17), 1)); + Files.delete(path); + final Path malformed = file(5, 17); + try { + byte[] bytes = Files.readAllBytes(malformed); + ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).putInt(18 * 4, 18); + Files.write(malformed, bytes); + try (var bad = FvecFileVectorValues.open(malformed)) { + assertThrows(UncheckedIOException.class, () -> bad.getVector(1)); + } + } finally { Files.delete(malformed); } + } + + @Test(timeout = 10000) public void failedAdvanceInvalidatesCursorAndReleasesItsReads() throws Exception { + // Cover first selection, middle of a batch, and transition to the next batch. + for (int invalidOrdinal : new int[] {1, 2}) { + var path = file(257, 17); + var bytes = ByteBuffer.wrap(Files.readAllBytes(path)).order(ByteOrder.LITTLE_ENDIAN); + bytes.putInt(invalidOrdinal * 18 * 4, 99); + Files.write(path, bytes.array()); + try (var source = FvecFileVectorValues.open(path, FvecFileVectorValues.Options.defaults() + .withBatchVectors(2).withReadAhead(3))) { + try (var cursor = source.openRange(0, 5)) { + for (int ordinal = 0; ordinal < invalidOrdinal; ordinal++) { + assertTrue(cursor.next()); equal(ordinal, cursor.vector()); + } + assertThrows(UncheckedIOException.class, cursor::next); + assertThrows(IllegalStateException.class, cursor::vector); + assertThrows(IllegalStateException.class, cursor::ordinal); + assertThrows(IllegalStateException.class, cursor::next); + } + // The source and other cursors remain usable after this cursor fails. + try (var cursor = source.openSelection(new int[] {invalidOrdinal}, 0, 1)) { + assertThrows(UncheckedIOException.class, cursor::next); + assertThrows(IllegalStateException.class, cursor::vector); + assertThrows(IllegalStateException.class, cursor::next); + } + try (var cursor = source.openRange(5, 7)) { + assertTrue(cursor.next()); equal(5, cursor.vector()); + assertTrue(cursor.next()); equal(6, cursor.vector()); + } + equal(8, source.getVector(8)); + assertTrue(source.statistics().bufferBytes <= source.statistics().bufferLimitBytes); + } finally { Files.delete(path); } + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecPointCache.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecPointCache.java new file mode 100644 index 000000000..c464220b1 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecPointCache.java @@ -0,0 +1,109 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.disk; + +import org.junit.Test; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.file.Files; +import static org.junit.Assert.*; + +public class TestFvecPointCache { + @Test public void openingNarrowLargeFileKeepsMetadataBounded() throws Exception { + // Only metadata and the first record are used; the remainder is a sparse extent. + var file = Files.createTempFile("point-cache-metadata", ".fvecs"); + try { + try (var channel = java.nio.channels.FileChannel.open(file, + java.nio.file.StandardOpenOption.WRITE)) { + var header = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN); + header.putInt(1).putFloat(123f).flip(); + while (header.hasRemaining()) channel.write(header); + channel.position(8_000_000_000L - 1); + channel.write(ByteBuffer.wrap(new byte[1])); + } + var bean = java.lang.management.ManagementFactory.getThreadMXBean(); + org.junit.Assume.assumeTrue(bean instanceof com.sun.management.ThreadMXBean); + var allocations = (com.sun.management.ThreadMXBean) bean; + allocations.setThreadAllocatedMemoryEnabled(true); + long thread = Thread.currentThread().getId(); + long before = allocations.getThreadAllocatedBytes(thread); + try (var source = FvecFileVectorValues.open(file, + FvecFileVectorValues.Options.defaults().withIoThreads(2))) { + long bytes = allocations.getThreadAllocatedBytes(thread) - before; + assertTrue("Constructor allocated " + bytes + " bytes", bytes < 8L << 20); + assertEquals(1_000_000_000, source.size()); + assertEquals(123f, source.getVector(0).get(0), 0f); + assertTrue(source.statistics().bufferBytes <= source.statistics().bufferLimitBytes); + } + } finally { Files.delete(file); } + } + + @Test public void adjacentPointRequestsDoNotTriggerScanReadAhead() throws Exception { + int rows = 10000, dimensions = 17; + var file = Files.createTempFile("point-scan-separation", ".fvecs"); + var bytes = ByteBuffer.allocate(rows * (dimensions + 1) * 4).order(ByteOrder.LITTLE_ENDIAN); + for (int row = 0; row < rows; row++) { + bytes.putInt(dimensions); + for (int d = 0; d < dimensions; d++) bytes.putFloat(row + d / 32f); + } + Files.write(file, bytes.array()); + var source = FvecFileVectorValues.open(file, FvecFileVectorValues.Options.defaults() + .withIoThreads(4).withBatchVectors(8).withReadAhead(3)); + try { + for (int ordinal = 100; ordinal < 103; ordinal++) + assertEquals(ordinal, source.getVector(ordinal).get(0), 0f); + } finally { + source.close(); + Files.delete(file); + } + // Closing drains speculative reads; exactly the three requested records should have been read. + assertEquals(3, source.statistics().reads); + } + + @Test public void pointFallbackKeepsScratchAndResourcesBounded() throws Exception { + int rows = 10000, dimensions = 17; + var file = Files.createTempFile("point-cache-check", ".fvecs"); + var bytes = ByteBuffer.allocate(rows * (dimensions + 1) * 4).order(ByteOrder.LITTLE_ENDIAN); + for (int row = 0; row < rows; row++) { + bytes.putInt(dimensions); + for (int d = 0; d < dimensions; d++) bytes.putFloat(row + d / 32f); + } + Files.write(file, bytes.array()); + try (var source = FvecFileVectorValues.open(file, FvecFileVectorValues.Options.defaults() + .withIoThreads(4).withBatchVectors(8).withReadAhead(0).withMaxBufferBytes(65536))) { + try (var other = source.copy()) { + var first = source.getVector(100); + long before = source.statistics().reads; + var second = other.getVector(100); + assertNotSame(first, second); + assertEquals(before + 1, source.statistics().reads); + source.getVector(101); + assertEquals(100f, second.get(0), 0f); + before = source.statistics().reads; + assertEquals(100f, other.getVector(100).get(0), 0f); + assertEquals(before + 1, source.statistics().reads); + for (int i = 0; i < 20000; i++) { + int ordinal = (i * 71) % rows; + var vector = other.getVector(ordinal); + for (int d = 0; d < dimensions; d++) + assertEquals(ordinal + d / 32f, vector.get(d), 0f); + } + assertTrue(source.statistics().bufferBytes <= source.statistics().bufferLimitBytes); + } + } finally { Files.delete(file); } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecReadProgress.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecReadProgress.java new file mode 100644 index 000000000..06b1f0880 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestFvecReadProgress.java @@ -0,0 +1,82 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.jbellis.jvector.disk; + +import org.junit.Test; +import java.io.EOFException; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.MappedByteBuffer; +import java.nio.channels.*; +import static org.junit.Assert.*; + +public class TestFvecReadProgress { + @Test(timeout = 1000) public void repeatedZeroReadsFailInsteadOfSpinning() { + var channel = new TestChannel(-1); + var failure = assertThrows(IOException.class, () -> + FvecFileVectorValues.readFully(channel, ByteBuffer.allocate(1), 123)); + assertTrue(failure.getMessage().contains("123")); + assertEquals(16, channel.calls); + } + + @Test(timeout = 1000) public void progressResetsTheZeroReadLimitAndPreservesOffsets() throws Exception { + var channel = new TestChannel(15); + var buffer = ByteBuffer.allocate(4); + buffer.position(2); + FvecFileVectorValues.readFully(channel, buffer, 123); + assertEquals(32, channel.calls); + assertEquals(4, buffer.position()); + assertEquals((byte) 123, buffer.get(2)); + assertEquals((byte) 124, buffer.get(3)); + } + + @Test(timeout = 1000) public void eofStillFailsImmediately() { + var channel = new TestChannel(0) { + @Override public int read(ByteBuffer buffer, long offset) { return -1; } + }; + assertThrows(EOFException.class, () -> + FvecFileVectorValues.readFully(channel, ByteBuffer.allocate(1), 0)); + } + + private static class TestChannel extends FileChannel { + final int zerosPerByte; + int calls, zeros; + TestChannel(int zerosPerByte) { this.zerosPerByte = zerosPerByte; } + @Override public int read(ByteBuffer buffer, long offset) { + calls++; + if (zerosPerByte < 0 || zeros++ < zerosPerByte) return 0; + zeros = 0; + buffer.put((byte) offset); + return 1; + } + @Override public int read(ByteBuffer buffer) { throw new UnsupportedOperationException(); } + @Override public long read(ByteBuffer[] buffers, int offset, int count) { throw new UnsupportedOperationException(); } + @Override public int write(ByteBuffer buffer) { throw new UnsupportedOperationException(); } + @Override public int write(ByteBuffer buffer, long offset) { throw new UnsupportedOperationException(); } + @Override public long write(ByteBuffer[] buffers, int offset, int count) { throw new UnsupportedOperationException(); } + @Override public long position() { throw new UnsupportedOperationException(); } + @Override public FileChannel position(long position) { throw new UnsupportedOperationException(); } + @Override public long size() { throw new UnsupportedOperationException(); } + @Override public FileChannel truncate(long size) { throw new UnsupportedOperationException(); } + @Override public void force(boolean metadata) { throw new UnsupportedOperationException(); } + @Override public long transferTo(long position, long count, WritableByteChannel target) { throw new UnsupportedOperationException(); } + @Override public long transferFrom(ReadableByteChannel source, long position, long count) { throw new UnsupportedOperationException(); } + @Override public MappedByteBuffer map(MapMode mode, long position, long size) { throw new UnsupportedOperationException(); } + @Override public FileLock lock(long position, long size, boolean shared) { throw new UnsupportedOperationException(); } + @Override public FileLock tryLock(long position, long size, boolean shared) { throw new UnsupportedOperationException(); } + @Override protected void implCloseChannel() {} + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestBufferedExactQueries.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestBufferedExactQueries.java new file mode 100644 index 000000000..f98c5ec01 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestBufferedExactQueries.java @@ -0,0 +1,115 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; +import io.github.jbellis.jvector.graph.*; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.vector.*; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; +import java.nio.file.Files; +import java.util.*; +import java.util.concurrent.*; +import static org.junit.Assert.*; + +public class TestBufferedExactQueries { + private static final java.util.concurrent.atomic.AtomicInteger pointReads = new java.util.concurrent.atomic.AtomicInteger(); + private static class Borrowed implements RandomAccessVectorValues { + private final RandomAccessVectorValues source; + private final VectorFloat scratch; + Borrowed(RandomAccessVectorValues source) { + this.source = source; + this.scratch = VectorizationProvider.getInstance().getVectorTypeSupport().createFloatVector(source.dimension()); + } + public int size() { return source.size(); } + public int dimension() { return source.dimension(); } + public boolean isValueShared() { return true; } + public RandomAccessVectorValues copy() { return new Borrowed(source); } + public VectorFloat getVector(int ordinal) { + pointReads.incrementAndGet(); + source.getVectorInto(ordinal, scratch, 0); + return scratch; + } + } + + /** Every worker must score through a copy, never through the source's shared root view. */ + @Test(timeout = 60000) + public void concurrentFullPrecisionScorersUseIndependentViews() throws Exception { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var vectors = new ArrayList>(); + var random = new Random(42); + for (int i = 0; i < 128; i++) { + var v = vts.createFloatVector(17); + for (int j = 0; j < 17; j++) v.set(j, random.nextFloat() * 2 - 1); + VectorUtil.l2normalize(v); + vectors.add(v); + } + var resident = new ListRandomAccessVectorValues(vectors, 17); + var root = new Borrowed(resident) { + @Override public VectorFloat getVector(int ordinal) { + throw new AssertionError("Shared root used by scorer"); + } + }; + var pool = new ForkJoinPool(3); + var workers = Executors.newFixedThreadPool(8); + var path = Files.createTempFile("exact-scorer", ".jvector"); + try (var builder = new GraphIndexBuilder(BuildScoreProvider.randomAccessScoreProvider(resident, + VectorSimilarityFunction.DOT_PRODUCT), 17, 8, 32, 1.2f, 1.2f, true, true, pool, pool)) { + OnDiskGraphIndex.write(builder.build(resident), resident, path); + try (var index = OnDiskGraphIndex.load(ReaderSupplierFactory.open(path))) { + for (boolean stored : new boolean[] {false, true}) { + for (var metric : VectorSimilarityFunction.values()) { + pointReads.set(0); + DataSet dataset = new DataSet() { + public boolean useDatasetPrefetching() { return stored; } + public String getName() { return "point-only"; } + public int getDimension() { return 17; } + public VectorSimilarityFunction getSimilarityFunction() { return metric; } + public List> getBaseVectors() { return vectors; } + public List> getQueryVectors() { return List.of(vectors.get(0)); } + public List> getGroundTruth() { return List.of(List.of(0)); } + public RandomAccessVectorValues getBaseRavv() { return root; } + }; + try (var system = new Grid.ConfiguredSystem(dataset, index, null, Set.of(FeatureId.INLINE_VECTORS))) { + var tasks = new ArrayList>(); + for (int worker = 0; worker < 8; worker++) { + final int id = worker; + tasks.add(() -> { + try (var view = index.getView()) { + for (int iteration = 0; iteration < 100; iteration++) { + var query = vectors.get((iteration + id) % vectors.size()); + var scores = system.scoreProviderFor(query, view).scoreFunction(); + int ordinal = (iteration * 7 + id) % vectors.size(); + assertEquals(metric.compare(query, vectors.get(ordinal)), scores.similarityTo(ordinal), 0f); + } + } + return null; + }); + } + for (var task : workers.invokeAll(tasks)) task.get(); + assertTrue("Exact scorer bypassed original input", pointReads.get() > 0); + } + } + } + } + } finally { workers.shutdown(); pool.shutdown(); Files.delete(path); } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetLoaderOverride.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetLoaderOverride.java new file mode 100644 index 000000000..a04c01959 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetLoaderOverride.java @@ -0,0 +1,98 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.example.benchmarks.datasets.BaseVectorLoading; +import io.github.jbellis.jvector.example.yaml.ConstructionParameters; +import io.github.jbellis.jvector.example.yaml.MultiConfig; +import io.github.jbellis.jvector.example.yaml.SearchParameters; +import io.github.jbellis.jvector.example.yaml.Compression; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.logging.Handler; +import java.util.logging.LogRecord; +import java.util.logging.Logger; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +/** The VM override selects loading without changing resident construction/query routing. */ +public class TestDatasetLoaderOverride { + private static final String PROPERTY = "jvector.dataset_loader"; + private String previous; + private final List warnings = new ArrayList<>(); + private final Logger logger = Logger.getLogger(BaseVectorLoading.class.getName()); + private final Handler handler = new Handler() { + @Override public void publish(LogRecord record) { warnings.add(record.getMessage()); } + @Override public void flush() { } + @Override public void close() { } + }; + + @Before public void saveProperty() { + previous = System.getProperty(PROPERTY); + logger.addHandler(handler); + } + + @After public void restoreProperty() { + logger.removeHandler(handler); + if (previous == null) System.clearProperty(PROPERTY); + else System.setProperty(PROPERTY, previous); + } + + private static MultiConfig config(String build, String search) { + var config = new MultiConfig(); + config.construction = new ConstructionParameters(); + config.search = new SearchParameters(); + var bc = new Compression(); bc.type = build; + var sc = new Compression(); sc.type = search; + config.construction.compression = List.of(bc); + config.construction.fusedGraph = List.of(false); + config.construction.reranking = List.of("NVQ"); + config.search.compression = List.of(sc); + return config; + } + + @Test public void explicitBufferedPropertyRetainsExactSearchRoutingAndWarns() { + System.setProperty(PROPERTY, "buffered"); + assertEquals(BaseVectorLoading.PRELOAD, BaseVectorLoading.forBenchmark(config("PQ", "None"))); + assertTrue(warnings.stream().anyMatch(message -> message.contains("raw-input exact search"))); + } + + @Test public void explicitBufferedPropertyRetainsUncompressedConstructionAndWarns() { + System.setProperty(PROPERTY, "buffered"); + assertEquals(BaseVectorLoading.PRELOAD, BaseVectorLoading.forBenchmark(config("None", "PQ"))); + assertTrue(warnings.stream().anyMatch(message -> message.contains("uncompressed construction"))); + } + + @Test public void explicitBufferedPropertyOverridesDisabledFlagForQuantizedWorkloads() { + System.setProperty(PROPERTY, "buffered"); + assertEquals(BaseVectorLoading.BUFFERED, BaseVectorLoading.forBenchmark(config("PQ", "PQ"))); + assertTrue(warnings.isEmpty()); + } + + @Test public void explicitPreloadPropertyOverridesEnabledFlag() { + System.setProperty(PROPERTY, "preload"); + var config = config("PQ", "PQ"); + config.useDatasetPrefetching = true; + assertEquals(BaseVectorLoading.PRELOAD, BaseVectorLoading.forBenchmark(config)); + assertTrue(warnings.isEmpty()); + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetPrefetching.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetPrefetching.java new file mode 100644 index 000000000..5379cb0ef --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetPrefetching.java @@ -0,0 +1,169 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.example.benchmarks.datasets.*; +import io.github.jbellis.jvector.example.yaml.*; +import io.github.jbellis.jvector.example.util.CompressorParameters; +import io.github.jbellis.jvector.example.util.OnDiskGraphIndexCache; +import io.github.jbellis.jvector.example.reporting.RunArtifacts; +import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.graph.*; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.vector.*; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; +import java.nio.file.*; +import java.util.*; +import java.util.concurrent.atomic.AtomicInteger; +import static org.junit.Assert.*; + +public class TestDatasetPrefetching { + private MultiConfig config(String build, String search, boolean fused, String rerank) { + var config = new MultiConfig(); + config.construction = new ConstructionParameters(); + config.search = new SearchParameters(); + var bc = new Compression(); bc.type = build; + var sc = new Compression(); sc.type = search; + config.construction.compression = List.of(bc); + config.construction.fusedGraph = List.of(fused); + config.construction.reranking = List.of(rerank); + config.search.compression = List.of(sc); + return config; + } + @Test public void yamlDefaultsToResidentAndOptInIsParsed() throws Exception { + var file = Files.createTempFile("prefetch-flag", ".yml"); + try { + Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: 6\n"); + assertFalse(MultiConfig.getConfig(file.toString()).useDatasetPrefetching); + Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: 6\nuseDatasetPrefetching: false\n"); + assertFalse(MultiConfig.getConfig(file.toString()).useDatasetPrefetching); + Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: 6\nuseDatasetPrefetching: true\n"); + assertTrue(MultiConfig.getConfig(file.toString()).useDatasetPrefetching); + } finally { Files.delete(file); } + } + @Test public void datasetPrefetchingAndCompactionOptionsAreIndependent() throws Exception { + var file = Files.createTempFile("prefetch-compaction-flags", ".yml"); + try { + String common = "yamlSchemaVersion: 1\nonDiskIndexVersion: 6\nuseDatasetPrefetching: true\n"; + Files.writeString(file, common); + var config = MultiConfig.getConfig(file.toString()); + assertTrue(config.useDatasetPrefetching); + assertFalse(config.isCompactionEnabled()); + Files.writeString(file, common + "withCompaction: [Yes]\n"); + config = MultiConfig.getConfig(file.toString()); + assertTrue(config.useDatasetPrefetching); + assertTrue(config.isCompactionEnabled()); + Files.writeString(file, common + "withCompaction: [No]\n"); + config = MultiConfig.getConfig(file.toString()); + assertTrue(config.useDatasetPrefetching); + assertFalse(config.isCompactionEnabled()); + } finally { Files.delete(file); } + } + @Test public void versionedYamlRejectsUnknownSettings() throws Exception { + var file = Files.createTempFile("prefetch-unknown-setting", ".yml"); + try { + Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: 6\nuseDatasetPrefetching: false\nunknownSetting: true\n"); + assertThrows(org.yaml.snakeyaml.error.YAMLException.class, + () -> MultiConfig.getConfig(file.toString())); + } finally { Files.delete(file); } + } + @Test public void sourceRoutingPreservesRawComparisonWorkloadsWithoutRejectingThem() { + String previous = System.getProperty("jvector.dataset_loader"); + try { + System.clearProperty("jvector.dataset_loader"); + var c = config("PQ", "PQ", true, "NVQ"); + assertEquals(BaseVectorLoading.PRELOAD, BaseVectorLoading.forBenchmark(c)); + c.useDatasetPrefetching = true; + assertEquals(BaseVectorLoading.BUFFERED, BaseVectorLoading.forBenchmark(c)); + c.construction.compression.get(0).type = "None"; + assertEquals(BaseVectorLoading.PRELOAD, BaseVectorLoading.forBenchmark(c)); + c = config("PQ", "None", false, "NVQ"); c.useDatasetPrefetching = true; + assertEquals(BaseVectorLoading.PRELOAD, BaseVectorLoading.forBenchmark(c)); + c.construction.reranking = List.of("FP"); + assertEquals(BaseVectorLoading.PRELOAD, BaseVectorLoading.forBenchmark(c)); + System.setProperty("jvector.dataset_loader", "preload"); + assertEquals(BaseVectorLoading.PRELOAD, BaseVectorLoading.forBenchmark(c)); + } finally { + if (previous == null) System.clearProperty("jvector.dataset_loader"); + else System.setProperty("jvector.dataset_loader", previous); + } + } + @Test(timeout=60000) public void gridFeedsOriginalBulkVectorsToInsertionAndBothStorageFormats() throws Exception { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var vectors = new ArrayList>(); var random = new Random(42); + for (int n=0;n<128;n++) { + var vector = vts.createFloatVector(17); + for (int d=0;d<17;d++) vector.set(d,random.nextFloat()-0.5f); + VectorUtil.l2normalize(vector); vectors.add(vector); + } + var resident = new ListRandomAccessVectorValues(vectors,17); + var ranges = new AtomicInteger(); var selections = new AtomicInteger(); + BatchedVectorValues input = new BatchedVectorValues() { + public int size(){return 128;} + public int dimension(){return 17;} + public boolean isValueShared(){return false;} + public RandomAccessVectorValues copy(){throw new AssertionError("Unshared source copied");} + public VectorFloat getVector(int ordinal){throw new AssertionError("Grid reread a point instead of consuming its input batch");} + public VectorCursor openRange(int a,int b){ranges.incrementAndGet();return VectorAccess.openRange(resident,a,b);} + public VectorCursor openSelection(int[] ordinals,int a,int n){selections.incrementAndGet();return VectorAccess.openSelection(resident,ordinals,a,n);} + }; + // A previous test run must not supply this run's training codebook. + var datasetName = "bulk-wiring-" + UUID.randomUUID(); + DataSet ds = new DataSet() { + public boolean useDatasetPrefetching(){return true;} + public int getDimension(){return 17;} + public String getName(){return datasetName;} + public VectorSimilarityFunction getSimilarityFunction(){return VectorSimilarityFunction.DOT_PRODUCT;} + public RandomAccessVectorValues getBaseRavv(){return input;} + public List> getBaseVectors(){return vectors;} + public List> getQueryVectors(){return List.of(vectors.get(0));} + public List> getGroundTruth(){return List.of(List.of(0));} + }; + var cacheFile = Path.of("pq_cache").resolve(new CompressorParameters.PQParameters(4,16,false,-1f).idStringFor(ds)); + var work = Files.createTempDirectory("bulk-grid-wiring"); + try { + Grid.runOneGraph(OnDiskGraphIndexCache.initialize(work.resolve("saved")), List.of(EnumSet.of(FeatureId.INLINE_VECTORS),EnumSet.of(FeatureId.NVQ_VECTORS)), + 8,32,1.2f,true,false, ignored -> new CompressorParameters.PQParameters(4,16,false,-1f), + List.of(),Map.of(10,List.of(1.0)),List.of(false),RunArtifacts.disabled(),ds,work); + assertTrue(selections.get()>0); assertTrue(ranges.get()>0); + int seen=0; + try(var files=Files.list(work.resolve("saved"))) { + for(var file:files.collect(java.util.stream.Collectors.toList())) { + try(var graph=OnDiskGraphIndex.load(ReaderSupplierFactory.open(file));var view=graph.getView()) { + assertEquals(128,graph.size(0)); + if(graph.getFeatures().containsKey(FeatureId.INLINE_VECTORS)) { + var stored=(RandomAccessVectorValues)view; + for(int n=0;n<128;n++) { + var actual=stored.getVector(n); + for(int d=0;d<17;d++)assertEquals(vectors.get(n).get(d),actual.get(d),0f); + } + } else assertTrue(graph.getFeatures().containsKey(FeatureId.NVQ_VECTORS)); + seen++; + } + } + } + assertEquals(2,seen); + } finally { + try { Files.deleteIfExists(cacheFile); } + finally { + try(var paths=Files.walk(work)){for(var path:paths.sorted(Comparator.reverseOrder()).collect(java.util.stream.Collectors.toList()))Files.delete(path);} + } + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBufferedGraphCompatibility.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBufferedGraphCompatibility.java new file mode 100644 index 000000000..73b761d79 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBufferedGraphCompatibility.java @@ -0,0 +1,128 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.disk.*; +import io.github.jbellis.jvector.graph.disk.*; +import io.github.jbellis.jvector.graph.disk.feature.*; +import io.github.jbellis.jvector.graph.similarity.*; +import io.github.jbellis.jvector.quantization.*; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.*; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.file.*; +import java.util.*; +import java.util.concurrent.ForkJoinPool; +import java.util.function.IntFunction; +import static org.junit.Assert.*; + +public class TestBufferedGraphCompatibility { + @Test(timeout = 180000) + public void exactPqAndBinaryBuildsWithAllMetricsAndRefinementModes() throws Exception { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var pool = new ForkJoinPool(3); + try { + for (int dimension : new int[] {17, 768}) { + int count = 257; + var path = Files.createTempFile("graph-input", ".fvecs"); + var bytes = ByteBuffer.allocate(count * (dimension + 1) * 4).order(ByteOrder.LITTLE_ENDIAN); + var vectors = new ArrayList>(); + var random = new Random(42); + for (int i = 0; i < count; i++) { + var v = vts.createFloatVector(dimension); + for (int j = 0; j < dimension; j++) v.set(j, random.nextFloat() * 2 - 1); + VectorUtil.l2normalize(v); + bytes.putInt(dimension); + for (int j = 0; j < dimension; j++) bytes.putFloat(v.get(j)); + vectors.add(v); + } + Files.write(path, bytes.array()); + var resident = new ListRandomAccessVectorValues(vectors, dimension); + try (var source = FvecFileVectorValues.open(path)) { + var pq = ProductQuantization.compute(resident, Math.max(1, dimension / 8), 256, false, + io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED, pool, pool); + var pqv = pq.encodeAll(source, pool); + var bqv = (BQVectors) new BinaryQuantization(dimension).encodeAll(source, pool); + var nvq = NVQuantization.compute(source, 3); + for (var metric : VectorSimilarityFunction.values()) { + for (boolean refine : new boolean[] {false, true}) { + for (int mode = 0; mode < 3; mode++) { + var bsp = mode == 0 ? BuildScoreProvider.randomAccessScoreProvider(source, metric) + : mode == 1 ? BuildScoreProvider.pqBuildScoreProvider(metric, pqv) + : BuildScoreProvider.bqBuildScoreProvider(bqv); + try (var builder = new GraphIndexBuilder(bsp, dimension, 8, 64, 1.2f, 1.2f, + true, refine, pool, pool)) { + var graph = builder.build(source); + assertEquals(count, graph.size(0)); + // A file-backed shared source must not overwrite insertion/diversity queries. + var result = GraphSearcher.search(vectors.get(5), 10, resident, metric, graph, Bits.ALL); + assertTrue("dimension=" + dimension + ", build=" + (mode == 0 ? "exact" : mode == 1 ? "PQ" : "BQ") + + ", metric=" + metric + ", refine=" + refine + ", provider=" + + VectorizationProvider.getInstance().getClass().getSimpleName() + + ", returned=" + Arrays.toString(Arrays.stream(result.getNodes()).mapToInt(n -> n.node).toArray()), + Arrays.stream(result.getNodes()).anyMatch(n -> n.node == 5)); + // Reuse the SAME graph/books to isolate source and serialization behavior. + for (boolean fused : new boolean[] {false, true}) { + for (boolean quantized : new boolean[] {false, true}) { + byte[] expected = null; + SearchResult expectedAnswer = null; + for (var values : List.of(resident, source)) { + var graphFile = Files.createTempFile("graph-source", ".jvector"); + try { + var writerBuilder = new RandomAccessOnDiskGraphIndexWriter.Builder(graph, graphFile); + if (fused) writerBuilder.with(new FusedPQ(graph.maxDegree(), pq)); + if (quantized) writerBuilder.with(new NVQ(nvq)); + else writerBuilder.with(new InlineVectors(dimension)); + var local = values.threadLocalSupplier(); + Map> suppliers = new EnumMap<>(FeatureId.class); + if (fused) suppliers.put(FeatureId.FUSED_PQ, + i -> new FusedPQ.State(graph.getView(), pqv, i)); + if (quantized) suppliers.put(FeatureId.NVQ_VECTORS, + i -> new NVQ.State(nvq.encode(local.get().getVector(i)))); + else suppliers.put(FeatureId.INLINE_VECTORS, + i -> new InlineVectors.State(local.get().getVector(i))); + try (var writer = writerBuilder.build()) { writer.write(suppliers); } + var actual = Files.readAllBytes(graphFile); + if (expected == null) expected = actual; + else assertArrayEquals(expected, actual); + try (var index = OnDiskGraphIndex.load(ReaderSupplierFactory.open(graphFile)); + var view = index.getView(); var searcher = new GraphSearcher(index)) { + var reranker = view.rerankerFor(vectors.get(5), metric); + var scores = fused + ? new DefaultSearchScoreProvider(view.approximateScoreFunctionFor(vectors.get(5), metric), reranker) + : new DefaultSearchScoreProvider(reranker); + var answer = searcher.search(scores, 10, Bits.ALL); + if (expectedAnswer == null) expectedAnswer = answer; + else assertEquals(expectedAnswer, answer); + } + } finally { Files.delete(graphFile); } + } + } + } + } + } + } + } + } finally { Files.delete(path); } + } + } finally { pool.shutdown(); } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBuildScoreProviderAnchorReads.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBuildScoreProviderAnchorReads.java new file mode 100644 index 000000000..da55e3a26 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBuildScoreProviderAnchorReads.java @@ -0,0 +1,79 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.disk.FvecFileVectorValues; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.file.Files; +import java.nio.file.Path; + +import static org.junit.Assert.assertTrue; + +/// A full-precision build score function compares one fixed anchor against many nodes. For a +/// value-shared file source it should read the anchor once, not once per comparison. +public class TestBuildScoreProviderAnchorReads { + private static final int ROWS = 2000, DIMENSIONS = 16, COMPARISONS = 1000; + + private static Path writeFvecs() throws Exception { + var file = Files.createTempFile("anchor-reads", ".fvecs"); + var bytes = ByteBuffer.allocate(ROWS * (DIMENSIONS + 1) * 4).order(ByteOrder.LITTLE_ENDIAN); + for (int row = 0; row < ROWS; row++) { + bytes.putInt(DIMENSIONS); + for (int d = 0; d < DIMENSIONS; d++) bytes.putFloat(row + d / 32f); + } + Files.write(file, bytes.array()); + return file; + } + + private static void assertAnchorReadOnce(String kind, long reads) { + // One read per compared node, plus one for the anchor and a little slack. + long limit = COMPARISONS + 8; + assertTrue(kind + " made " + reads + " source reads for " + COMPARISONS + + " comparisons against one anchor; expected at most " + limit, reads <= limit); + } + + @Test public void diversityScoreFunctionReadsAnchorOnce() throws Exception { + var file = writeFvecs(); + try (var source = FvecFileVectorValues.open(file)) { + var bsp = BuildScoreProvider.randomAccessScoreProvider(source, VectorSimilarityFunction.EUCLIDEAN); + var sf = bsp.diversityScoreFunctionFor(0); + long before = source.statistics().reads; + for (int node = 1; node <= COMPARISONS; node++) sf.similarityTo(node); + assertAnchorReadOnce("diversityScoreFunctionFor", source.statistics().reads - before); + } finally { + Files.delete(file); + } + } + + @Test public void searchProviderReadsAnchorOnce() throws Exception { + var file = writeFvecs(); + try (var source = FvecFileVectorValues.open(file)) { + var bsp = BuildScoreProvider.randomAccessScoreProvider(source, VectorSimilarityFunction.EUCLIDEAN); + var sf = bsp.searchProviderFor(0).scoreFunction(); + long before = source.statistics().reads; + for (int node = 1; node <= COMPARISONS; node++) sf.similarityTo(node); + assertAnchorReadOnce("searchProviderFor", source.statistics().reads - before); + } finally { + Files.delete(file); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBuildScoreProviderBorrowedVectors.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBuildScoreProviderBorrowedVectors.java new file mode 100644 index 000000000..89f929ddb --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBuildScoreProviderBorrowedVectors.java @@ -0,0 +1,137 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.graph.similarity.ScoreFunction; +import io.github.jbellis.jvector.vector.*; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import java.util.*; +import org.junit.Test; +import static org.junit.Assert.*; + +public class TestBuildScoreProviderBorrowedVectors { + private static class Borrowed implements RandomAccessVectorValues { + final RandomAccessVectorValues source; + final VectorFloat scratch; + Borrowed(RandomAccessVectorValues source) { + this.source = source; + scratch = VectorizationProvider.getInstance().getVectorTypeSupport().createFloatVector(source.dimension()); + } + public int size() { return source.size(); } + public int dimension() { return source.dimension(); } + public boolean isValueShared() { return true; } + public RandomAccessVectorValues copy() { return new Borrowed(source); } + public VectorFloat getVector(int ordinal) { source.getVectorInto(ordinal, scratch, 0); return scratch; } + } + @Test public void borrowedNodeScorersDoNotMaterializeAnchorsAtCreation() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var resident = new ListRandomAccessVectorValues(List.of( + vts.createFloatVector(new float[] {1f, 0f}), + vts.createFloatVector(new float[] {0f, 1f})), 2); + var reads = new java.util.concurrent.atomic.AtomicInteger(); + class Counting extends Borrowed { + Counting() { super(resident); } + @Override public RandomAccessVectorValues copy() { return new Counting(); } + @Override public VectorFloat getVector(int ordinal) { + reads.incrementAndGet(); + return super.getVector(ordinal); + } + } + var bsp = BuildScoreProvider.randomAccessScoreProvider(new Counting(), VectorSimilarityFunction.DOT_PRODUCT); + var search = bsp.searchProviderFor(0).scoreFunction(); + var diversity = bsp.diversityProviderFor(0).scoreFunction(); + var direct = bsp.diversityScoreFunctionFor(0); + assertEquals(0, reads.get()); + float expected = VectorSimilarityFunction.DOT_PRODUCT.compare(resident.getVector(0), resident.getVector(1)); + assertEquals(expected, search.similarityTo(1), 0f); + assertEquals(expected, diversity.similarityTo(1), 0f); + assertEquals(expected, direct.similarityTo(1), 0f); + assertEquals(6, reads.get()); + } + + @Test public void centroidNormalizesItsSizeSnapshot() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var resident = new ListRandomAccessVectorValues(Arrays.asList( + vts.createFloatVector(new float[] {1f, 0f}), null, + vts.createFloatVector(new float[] {0f, 1f})), 2); + var sizes = new java.util.concurrent.atomic.AtomicInteger(); + var ranges = new java.util.concurrent.atomic.AtomicInteger(); + var source = new BatchedVectorValues() { + public int size() { return sizes.getAndIncrement() == 0 ? 3 : 4; } + public int dimension() { return 2; } + public boolean isValueShared() { return false; } + public RandomAccessVectorValues copy() { throw new UnsupportedOperationException(); } + public VectorFloat getVector(int ordinal) { return resident.getVector(ordinal); } + public VectorCursor openRange(int start, int end) { + ranges.incrementAndGet(); + throw new AssertionError("Unexpected bulk centroid optimization"); + } + public VectorCursor openSelection(int[] ordinals, int offset, int count) { + throw new AssertionError("Selection"); + } + }; + var centroid = BuildScoreProvider.randomAccessScoreProvider(source, + VectorSimilarityFunction.DOT_PRODUCT).approximateCentroid(); + assertEquals(0, ranges.get()); + assertEquals(1f / 3, centroid.get(0), 0f); + assertEquals(1f / 3, centroid.get(1), 0f); + } + + @Test public void residentAnchorsAreNotCopied() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var zero = vts.createFloatVector(new float[] {1f, 0f, 0f}); + var one = vts.createFloatVector(new float[] {0f, 1f, 0f}); + var source = new ListRandomAccessVectorValues(Arrays.asList(zero, one), 3); + var bsp = BuildScoreProvider.randomAccessScoreProvider(source, VectorSimilarityFunction.DOT_PRODUCT); + var search = bsp.searchProviderFor(0).scoreFunction(); + var diversity = bsp.diversityProviderFor(0).scoreFunction(); + var direct = bsp.diversityScoreFunctionFor(0); + zero.set(0, 0f); + zero.set(1, 1f); + float expected = VectorSimilarityFunction.DOT_PRODUCT.compare(zero, one); + assertEquals(expected, search.similarityTo(1), 0f); + assertEquals(expected, diversity.similarityTo(1), 0f); + assertEquals(expected, direct.similarityTo(1), 0f); + } + + @Test public void nodeScorersRetainTheirAnchorAcrossOtherFactoryCalls() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + for (int dimensions : new int[] {3, 17, 768}) { + var rows = new ArrayList>(); + for (int i = 0; i < 3; i++) { + var v = vts.createFloatVector(dimensions); + v.set(i == 1 ? 1 : 0, 1f); + if (i == 2) v.set(2, 1f); + VectorUtil.l2normalize(v); + rows.add(v); + } + var resident = new ListRandomAccessVectorValues(rows, dimensions); + for (var metric : VectorSimilarityFunction.values()) { + var bsp = BuildScoreProvider.randomAccessScoreProvider(new Borrowed(resident), metric); + var search = bsp.searchProviderFor(0).scoreFunction(); + var diversity = bsp.diversityProviderFor(0).scoreFunction(); + var direct = bsp.diversityScoreFunctionFor(0); + bsp.searchProviderFor(1); // Reuses the same thread's borrowed source scratch. + float expected = metric.compare(rows.get(0), rows.get(2)); + assertEquals("search " + metric, expected, search.similarityTo(2), 0f); + assertEquals("diversity " + metric, expected, diversity.similarityTo(2), 0f); + assertEquals("direct " + metric, expected, direct.similarityTo(2), 0f); + } + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorAccess.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorAccess.java new file mode 100644 index 000000000..dd95a13c7 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorAccess.java @@ -0,0 +1,248 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; +import java.util.Arrays; +import java.util.concurrent.ForkJoinPool; +import static org.junit.Assert.*; + +public class TestVectorAccess { + @Test public void fallbackPreservesSharedValuesAndDuplicates() { + RandomAccessVectorValues source = new RandomAccessVectorValues() { + final VectorFloat value = VectorizationProvider.getInstance().getVectorTypeSupport().createFloatVector(1); + public int size() { return 8; } + public int dimension() { return 1; } + public boolean isValueShared() { return true; } + public VectorFloat getVector(int ordinal) { value.set(0, ordinal); return value; } + public RandomAccessVectorValues copy() { return new RandomAccessVectorValues() { + final VectorFloat v = value.copy(); + public int size() { return 8; } + public int dimension() { return 1; } + public boolean isValueShared() { return true; } + public VectorFloat getVector(int ordinal) { v.set(0, ordinal); return v; } + public RandomAccessVectorValues copy() { return this; } + }; } + }; + var pool = new ForkJoinPool(4); + try { + var result = VectorAccess.copySelected(source, new int[] {7, 1, 7, 0}, pool); + source.getVector(2); + for (int i = 0; i < 4; i++) assertEquals(new float[] {7, 1, 7, 0}[i], result.get(i).get(0), 0); + var cursor = VectorAccess.openSelection(source, new int[] {5, 2, 5}, 0, 3); + try (cursor) { + assertThrows(IllegalStateException.class, cursor::vector); + for (int expected : new int[] {5, 2, 5}) { assertTrue(cursor.next()); assertEquals(expected, cursor.ordinal()); assertEquals(expected, cursor.vector().get(0), 0); } + assertFalse(cursor.next()); assertFalse(cursor.next()); + assertThrows(IllegalStateException.class, cursor::vector); + } + assertThrows(IllegalStateException.class, cursor::next); + } finally { pool.shutdown(); } + } + + @Test public void fallbackRangeUsesExclusiveEndAndRetainedCopiesStayStable() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var source = new ListRandomAccessVectorValues(Arrays.asList( + vts.createFloatVector(new float[] {0}), vts.createFloatVector(new float[] {1}), + vts.createFloatVector(new float[] {2}), vts.createFloatVector(new float[] {3})), 1); + VectorFloat retained; + try (var cursor = VectorAccess.openRange(source, 1, 3)) { + assertTrue(cursor.next()); assertEquals(1, cursor.ordinal()); + retained = cursor.vector().copy(); + assertTrue(cursor.next()); assertEquals(2, cursor.ordinal()); + assertFalse(cursor.next()); + } + assertEquals(1f, retained.get(0), 0); + try (var cursor = VectorAccess.openRange(source, 4, 4)) { assertFalse(cursor.next()); } + assertThrows(IndexOutOfBoundsException.class, () -> VectorAccess.openRange(source, 3, 1)); + } + + @Test public void residentFallbackDoesNotCopyVectors() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var vector = vts.createFloatVector(new float[] {3}); + var source = new ListRandomAccessVectorValues(Arrays.asList(vector), 1); + var pool = new ForkJoinPool(2); + try { assertSame(vector, VectorAccess.copySelected(source, new int[] {0}, pool).get(0)); } + finally { pool.shutdown(); } + assertThrows(IndexOutOfBoundsException.class, () -> VectorAccess.openRange(source, 1, 2)); + assertThrows(IndexOutOfBoundsException.class, () -> VectorAccess.openSelection(source, new int[] {1}, 0, 1)); + try (var cursor = VectorAccess.openRange(source, 0, 0)) { assertFalse(cursor.next()); } + } + + @Test public void fullScanFallbackVisitsEachPositionOnceIncludingEmptySources() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var vectors = java.util.stream.IntStream.range(0, 19) + .>mapToObj(i -> vts.createFloatVector(new float[] {i})).collect(java.util.stream.Collectors.toList()); + var source = new ListRandomAccessVectorValues(vectors, 1); + var visits = new java.util.concurrent.atomic.AtomicIntegerArray(source.size()); + var pool = new ForkJoinPool(3); + try { + VectorAccess.forEach(source, pool, (ordinal, vector) -> { + assertEquals(ordinal, vector.get(0), 0f); + visits.incrementAndGet(ordinal); + }); + for (int i = 0; i < source.size(); i++) assertEquals(1, visits.get(i)); + var empty = new ListRandomAccessVectorValues(java.util.List.of(), 1); + VectorAccess.forEach(empty, pool, (ordinal, vector) -> fail("Empty source invoked consumer")); + } finally { pool.shutdown(); } + } + + @Test public void fullScanClosesItsCursorWhenConsumerFails() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var resident = new ListRandomAccessVectorValues(java.util.List.of( + vts.createFloatVector(new float[] {1}), vts.createFloatVector(new float[] {2})), 1); + var closed = new java.util.concurrent.atomic.AtomicInteger(); + BatchedVectorValues source = new BatchedVectorValues() { + public int size() { return resident.size(); } + public int dimension() { return resident.dimension(); } + public boolean isValueShared() { return false; } + public RandomAccessVectorValues copy() { return this; } + public VectorFloat getVector(int ordinal) { return resident.getVector(ordinal); } + public VectorCursor openSelection(int[] ordinals, int offset, int count) { + throw new AssertionError("Full scan should use range access"); + } + public VectorCursor openRange(int start, int end) { + var cursor = VectorAccess.openRange(resident, start, end); + return new VectorCursor() { + public boolean next() { return cursor.next(); } + public int ordinal() { return cursor.ordinal(); } + public VectorFloat vector() { return cursor.vector(); } + public boolean isValueShared() { return cursor.isValueShared(); } + public void close() { cursor.close(); closed.incrementAndGet(); } + }; + } + }; + var pool = new ForkJoinPool(1); + try { + assertThrows(IllegalStateException.class, () -> VectorAccess.forEach(source, pool, + (ordinal, vector) -> { throw new IllegalStateException("sink failed"); })); + assertEquals(1, closed.get()); + VectorAccess.forEach(source, pool, (ordinal, vector) -> assertEquals(ordinal + 1, vector.get(0), 0f)); + assertEquals(2, closed.get()); + } finally { pool.shutdown(); } + } + + @Test public void unsharedPointSourceDoesNotNeedCopy() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var values = java.util.List.of(vts.createFloatVector(new float[] {1}), vts.createFloatVector(new float[] {2})); + RandomAccessVectorValues source = new RandomAccessVectorValues() { + public int size() { return values.size(); } + public int dimension() { return 1; } + public boolean isValueShared() { return false; } + public VectorFloat getVector(int ordinal) { return values.get(ordinal); } + public RandomAccessVectorValues copy() { throw new UnsupportedOperationException("No copy needed"); } + }; + try (var cursor = VectorAccess.openRange(source, 0, 2)) { + assertTrue(cursor.next()); assertSame(values.get(0), cursor.vector()); + assertTrue(cursor.next()); assertSame(values.get(1), cursor.vector()); + assertFalse(cursor.next()); + } + try (var cursor = VectorAccess.openSelection(source, new int[] {1, 0, 1}, 0, 3)) { + for (int ordinal : new int[] {1, 0, 1}) { + assertTrue(cursor.next()); assertEquals(ordinal, cursor.ordinal()); assertSame(values.get(ordinal), cursor.vector()); + } + assertFalse(cursor.next()); + } + var pool = new ForkJoinPool(2); + try { + assertSame(values.get(1), VectorAccess.copySelected(source, new int[] {1}, pool).get(0)); + VectorAccess.forEach(source, pool, (ordinal, vector) -> assertSame(values.get(ordinal), vector)); + } finally { pool.shutdown(); } + } + + @Test public void failedPointAdvanceInvalidatesCursor() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var value = vts.createFloatVector(new float[] {1}); + var failure = new IllegalStateException("Source read failed"); + RandomAccessVectorValues source = new RandomAccessVectorValues() { + public int size() { return 3; } + public int dimension() { return 1; } + public boolean isValueShared() { return false; } + public RandomAccessVectorValues copy() { throw new UnsupportedOperationException(); } + public VectorFloat getVector(int ordinal) { + if (ordinal == 1) throw failure; + return value; + } + }; + for (boolean selection : new boolean[] {false, true}) { + for (int start : new int[] {0, 1}) { + try (var cursor = selection + ? VectorAccess.openSelection(source, new int[] {0, 1, 2}, start, 3 - start) + : VectorAccess.openRange(source, start, 3)) { + if (start == 0) { assertTrue(cursor.next()); assertSame(value, cursor.vector()); } + assertSame(failure, assertThrows(IllegalStateException.class, cursor::next)); + assertThrows(IllegalStateException.class, cursor::vector); + assertThrows(IllegalStateException.class, cursor::ordinal); + assertThrows(IllegalStateException.class, cursor::next); + } + } + } + try (var cursor = VectorAccess.openRange(source, 2, 3)) { + assertTrue(cursor.next()); assertSame(value, cursor.vector()); assertFalse(cursor.next()); + } + } + @Test public void mappedScanBoundsSelectionMetadataAndPreservesOutputPositions() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + int count = 1_000_003; + var largest = new java.util.concurrent.atomic.AtomicInteger(); + var reads = new java.util.concurrent.atomic.AtomicInteger(); + BatchedVectorValues source = new BatchedVectorValues() { + public int size() { return count; } + public int dimension() { return 1; } + public boolean isValueShared() { return true; } + public RandomAccessVectorValues copy() { throw new AssertionError("Point copy used"); } + public VectorFloat getVector(int ordinal) { throw new AssertionError("Point read used"); } + public VectorCursor openRange(int start, int end) { throw new AssertionError("Mapped scan used range"); } + public VectorCursor openSelection(int[] ordinals, int offset, int length) { + largest.accumulateAndGet(length, Math::max); + assertTrue("Dataset-sized request", length <= 4096); + final int[] snapshot = Arrays.copyOfRange(ordinals, offset, offset + length); + final VectorFloat value = vts.createFloatVector(1); + return new VectorCursor() { + int position = -1; + public boolean next() { + if (++position == length) return false; + assertEquals(snapshot[position], ordinals[offset + position]); + value.set(0, snapshot[position]); reads.incrementAndGet(); return true; + } + public int ordinal() { return snapshot[position]; } + public VectorFloat vector() { return value; } + public boolean isValueShared() { return true; } + public void close() { + // The caller cannot overwrite this request before the cursor closes. + assertArrayEquals(snapshot, Arrays.copyOfRange(ordinals, offset, offset + length)); + } + }; + } + }; + var visits = new java.util.concurrent.atomic.AtomicIntegerArray(count); + var pool = new ForkJoinPool(2); + try { + // A permutation followed by duplicates exercises both mapping and output positions. + VectorAccess.forEach(source, count, i -> (count - 1 - i) / 2, pool, (position, vector) -> { + assertEquals((count - 1 - position) / 2, vector.get(0), 0f); + visits.incrementAndGet(position); + }); + assertEquals(count, reads.get()); + assertEquals(4096, largest.get()); + for (int i = 0; i < count; i++) assertEquals(1, visits.get(i)); + } finally { pool.shutdown(); } + } + +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/TestVectorSourceCompatibility.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/TestVectorSourceCompatibility.java new file mode 100644 index 000000000..8df90b90d --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/TestVectorSourceCompatibility.java @@ -0,0 +1,269 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.quantization; + +import io.github.jbellis.jvector.disk.FvecFileVectorValues; +import io.github.jbellis.jvector.graph.*; +import io.github.jbellis.jvector.vector.*; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.*; +import java.util.concurrent.ForkJoinPool; + +import static org.junit.Assert.*; + +/** Compare the same quantizer on file, resident and borrowed point-only sources. */ +public class TestVectorSourceCompatibility { + private static final int ROWS = 257; + private static final int[] DIMENSIONS = {1, 2, 3, 7, 8, 9, 15, 16, 17, 31, 32, 33, 64, 65, 768, 1024, 1536, 3072}; + + private static final class Input implements AutoCloseable { + final Path path; + final List> vectors = new ArrayList<>(); + final ListRandomAccessVectorValues resident; + final FvecFileVectorValues file; + Input(int dimension) throws Exception { + path = Files.createTempFile("source-compatibility", ".fvecs"); + var bytes = ByteBuffer.allocate(ROWS * (dimension + 1) * 4).order(ByteOrder.LITTLE_ENDIAN); + var random = new Random(42); + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + for (int row = 0; row < ROWS; row++) { + var vector = vts.createFloatVector(dimension); + for (int col = 0; col < dimension; col++) vector.set(col, random.nextFloat() * 2 - 1); + VectorUtil.l2normalize(vector); // anisotropic training requires unit vectors + bytes.putInt(dimension); + for (int col = 0; col < dimension; col++) bytes.putFloat(vector.get(col)); + vectors.add(vector); + } + Files.write(path, bytes.array()); + resident = new ListRandomAccessVectorValues(vectors, dimension); + file = FvecFileVectorValues.open(path, FvecFileVectorValues.Options.defaults() + .withIoThreads(2).withBatchVectors(7).withReadAhead(3)); + } + RandomAccessVectorValues pointOnly() { return new Borrowed(resident); } + public void close() throws Exception { try { file.close(); } finally { Files.delete(path); } } + } + + /** Models existing integrations that supply one borrowed vector per view, without hooks. */ + private static final class Borrowed implements RandomAccessVectorValues { + final RandomAccessVectorValues source; + final VectorFloat scratch; + Borrowed(RandomAccessVectorValues source) { + this.source = source; + scratch = VectorizationProvider.getInstance().getVectorTypeSupport().createFloatVector(source.dimension()); + } + public int size() { return source.size(); } + public int dimension() { return source.dimension(); } + public boolean isValueShared() { return true; } + public RandomAccessVectorValues copy() { return new Borrowed(source); } + public VectorFloat getVector(int ordinal) { + source.getVectorInto(ordinal, scratch, 0); + return scratch; + } + } + + private static void encodingMatches(ProductQuantization pq, Input input, ForkJoinPool pool) { + var expected = pq.encodeAll(input.resident, pool); + assertEquals(expected, pq.encodeAll(input.file, pool)); + assertEquals(expected, pq.encodeAll(input.pointOnly(), pool)); + for (int i = 0; i < ROWS; i++) assertEquals(pq.encode(input.vectors.get(i)), expected.get(i)); + assertScoresEqual(expected, pq.encodeAll(input.file, pool), input.vectors.get(0)); + } + + private static void assertScoresEqual(CompressedVectors a, CompressedVectors b, VectorFloat query) { + for (var metric : VectorSimilarityFunction.values()) { + var scoreA = a.scoreFunctionFor(query.copy(), metric); + var scoreB = b.scoreFunctionFor(query.copy(), metric); + var tableA = a.precomputedScoreFunctionFor(query.copy(), metric); + var tableB = b.precomputedScoreFunctionFor(query.copy(), metric); + for (int i = 0; i < a.count(); i++) { + assertEquals(scoreA.similarityTo(i), scoreB.similarityTo(i), 0f); + assertEquals(tableA.similarityTo(i), tableB.similarityTo(i), 0f); + } + } + } + + @Test(timeout = 180000) + public void pqWidthsCenteringAndClusterCounts() throws Exception { + var pool = new ForkJoinPool(3); + try { + // Widths on both sides of SIMD boundaries and uneven final subspaces. + for (int dimension : DIMENSIONS) { + try (var input = new Input(dimension)) { + for (boolean centered : new boolean[] {false, true}) { + int m = Math.max(1, dimension / 8); + int k = dimension == 8 ? 256 : 16; + var pq = ProductQuantization.compute(input.resident, m, k, centered, + io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED, pool, pool); + encodingMatches(pq, input, pool); + } + } + } + } finally { pool.shutdown(); } + } + + @Test(timeout = 180000) + public void anisotropicTrainingAndRefinementUseTheSameSamples() throws Exception { + var pool = new ForkJoinPool(3); + try (var input = new Input(17)) { + assertEquals(ProductQuantization.extractTrainingVectors(input.resident, pool), + ProductQuantization.extractTrainingVectors(input.file, pool)); + assertEquals(ProductQuantization.extractTrainingVectors(input.resident, pool), + ProductQuantization.extractTrainingVectors(input.pointOnly(), pool)); + encodingMatches(ProductQuantization.compute(input.file, 3, 31, true, -1f, pool, pool), input, pool); + var pq = ProductQuantization.compute(input.file, 3, 16, false, 0.2f, pool, pool); + encodingMatches(pq, input, pool); + var refined = pq.refine(input.file, 1, 0.2f, pool, pool); + encodingMatches(refined, input, pool); + } finally { pool.shutdown(); } + } + + @Test(timeout = 180000) + public void nvqUnevenSubvectorsAndLearningModes() throws Exception { + var pool = new ForkJoinPool(3); + try { + for (int dimension : DIMENSIONS) { + try (var input = new Input(dimension)) { + for (int parts : new int[] {1, Math.min(3, dimension)}) { + var nvq = NVQuantization.compute(input.resident, parts); + assertEquals(nvq, NVQuantization.compute(input.file, parts)); + assertEquals(nvq, NVQuantization.compute(input.pointOnly(), parts)); + for (boolean learn : new boolean[] {false, true}) { + nvq.learn = learn; + var expected = nvq.encodeAll(input.resident, pool); + assertEquals(expected, nvq.encodeAll(input.file, pool)); + assertEquals(expected, nvq.encodeAll(input.pointOnly(), pool)); + assertScoresEqual(expected, nvq.encodeAll(input.file, pool), input.vectors.get(0)); + } + } + } + } + } finally { pool.shutdown(); } + } + + @Test(timeout = 60000) + public void binaryEncodingAndFutureScanConsumerWorkWithBothSourceContracts() throws Exception { + var pool = new ForkJoinPool(3); + try (var input = new Input(65)) { + var bq = new BinaryQuantization(65); + assertEquals(bq.encodeAll(input.resident, pool), bq.encodeAll(input.file, pool)); + assertEquals(bq.encodeAll(input.resident, pool), bq.encodeAll(input.pointOnly(), pool)); + for (var source : List.of(input.resident, input.file, input.pointOnly())) { + // A future algorithm needs only this source-independent scan, not file type checks. + var copied = new VectorFloat[ROWS]; + VectorAccess.forEach(source, pool, (ordinal, vector) -> copied[ordinal] = vector.copy()); + assertEquals(input.vectors, Arrays.asList(copied)); + } + } finally { pool.shutdown(); } + } + + @Test(timeout = 60000) + public void mappedPqEncodingPreservesOutputPositionsAndDuplicates() throws Exception { + String previousDebug = System.getProperty("VECTOR_DEBUG"); + System.setProperty("VECTOR_DEBUG", "true"); + var pool = new ForkJoinPool(3); + try (var input = new Input(17)) { + var pq = ProductQuantization.compute(input.resident, 3, 16, false, + io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED, pool, pool); + int[] ordinals = {256, 5, 256, 0, 3, 1}; + for (var source : List.of(input.file, input.pointOnly())) { + var encoded = PQVectors.encodeAndBuild(pq, ordinals.length, i -> ordinals[i], source, pool); + for (int i = 0; i < ordinals.length; i++) + assertEquals(pq.encode(input.vectors.get(ordinals[i])), encoded.get(i)); + } + var empty = new ListRandomAccessVectorValues(Collections.emptyList(), 17); + assertThrows(IllegalArgumentException.class, () -> pq.encodeAll(empty, pool)); + assertEquals(0, new BinaryQuantization(17).encodeAll(empty, pool).count()); + assertEquals(0, NVQuantization.compute(input.resident, 1).encodeAll(empty, pool).count()); + } finally { + pool.shutdown(); + if (previousDebug == null) System.clearProperty("VECTOR_DEBUG"); + else System.setProperty("VECTOR_DEBUG", previousDebug); + } + } + + @Test public void nvqMissingOrdinalsEncodeAsZeroVectors() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var first = vts.createFloatVector(new float[] {1, 2, 3}); + var last = vts.createFloatVector(new float[] {3, 2, 1}); + var sparse = new ListRandomAccessVectorValues(Arrays.asList(first, null, last), 3); + var dense = new ListRandomAccessVectorValues(List.of(first, last), 3); + var pool = new ForkJoinPool(2); + BatchedVectorValues batched = new BatchedVectorValues() { + public int size() { return sparse.size(); } + public int dimension() { return sparse.dimension(); } + public boolean isValueShared() { return false; } + public RandomAccessVectorValues copy() { return this; } + public VectorFloat getVector(int ordinal) { return sparse.getVector(ordinal); } + public VectorCursor openRange(int start, int end) { return VectorAccess.openRange(sparse, start, end); } + public VectorCursor openSelection(int[] ordinals, int offset, int count) { + return VectorAccess.openSelection(sparse, ordinals, offset, count); + } + }; + try { + var nvq = NVQuantization.compute(dense, 2); + var zero = vts.createFloatVector(3); + for (boolean learn : new boolean[] {false, true}) { + nvq.learn = learn; + for (var source : List.of(sparse, batched)) { + var encoded = nvq.encodeAll(source, pool); + assertEquals(3, encoded.count()); + assertEquals(nvq.encode(first), encoded.get(0)); + assertEquals(nvq.encode(zero), encoded.get(1)); + assertEquals(nvq.encode(last), encoded.get(2)); + } + } + } finally { pool.shutdown(); } + } + @Test public void pqIdentityEncodingUsesRangesAndHonorsPrefixCount() { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var vectors = new ArrayList>(); + for (int i = 0; i < 5003; i++) + vectors.add(vts.createFloatVector(new float[] {i % 2, i % 2})); + var resident = new ListRandomAccessVectorValues(vectors, 2); + var ranges = new java.util.concurrent.atomic.AtomicInteger(); + BatchedVectorValues source = new BatchedVectorValues() { + public int size() { return resident.size(); } + public int dimension() { return 2; } + public boolean isValueShared() { return false; } + public RandomAccessVectorValues copy() { throw new AssertionError("Point copy used"); } + public VectorFloat getVector(int ordinal) { throw new AssertionError("Point read used"); } + public VectorCursor openSelection(int[] ordinals, int offset, int count) { + throw new AssertionError("Identity encoding materialized a selection"); + } + public VectorCursor openRange(int start, int end) { + ranges.incrementAndGet(); return VectorAccess.openRange(resident, start, end); + } + }; + var pq = new ProductQuantization(new VectorFloat[] { + vts.createFloatVector(new float[] {0, 0, 1, 1})}, 2, new int[][] {{2, 0}}, null, -1f); + var pool = new ForkJoinPool(3); + try { + for (int count : new int[] {5000, 5003}) { + assertEquals(PQVectors.encodeAndBuild(pq, count, resident, pool), + PQVectors.encodeAndBuild(pq, count, source, pool)); + } + assertEquals(6, ranges.get()); + } finally { pool.shutdown(); } + } + +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/TestVectorSourceSampling.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/TestVectorSourceSampling.java new file mode 100644 index 000000000..741bc0465 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/TestVectorSourceSampling.java @@ -0,0 +1,52 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.quantization; + +import io.github.jbellis.jvector.disk.FvecFileVectorValues; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.concurrent.ForkJoinPool; +import static org.junit.Assert.*; + +public class TestVectorSourceSampling { + @Test(timeout=60000) public void floydSampleAndTrainingOrderAreUnchanged() throws Exception { + int count = 129003; + var path = Files.createTempFile("pq-sample-source", ".fvecs"); + var buffer = ByteBuffer.allocate(count * 12).order(ByteOrder.LITTLE_ENDIAN); + var values = new ArrayList>(); + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + for (int i = 0; i < count; i++) { + buffer.putInt(2).putFloat(i).putFloat(-i); + values.add(vts.createFloatVector(new float[] {i, -i})); + } + Files.write(path, buffer.array()); + var pool = new ForkJoinPool(48); + try (var source = FvecFileVectorValues.open(path)) { + var expected = ProductQuantization.extractTrainingVectors(new ListRandomAccessVectorValues(values, 2), pool); + var actual = ProductQuantization.extractTrainingVectors(source, pool); + assertEquals(128000, actual.size()); + for (int i = 0; i < actual.size(); i++) assertEquals(expected.get(i), actual.get(i)); + assertTrue(source.statistics().bufferBytes <= source.statistics().bufferLimitBytes); + } finally { pool.shutdown(); Files.delete(path); } + } +} diff --git a/pom.xml b/pom.xml index 28d6528d9..12cbe3aea 100644 --- a/pom.xml +++ b/pom.xml @@ -110,6 +110,16 @@ org.apache.maven.plugins maven-javadoc-plugin 3.12.0 + + + + + apiNote + a + API Note: + + + attach-javadocs @@ -123,14 +133,6 @@ --add-modules=jdk.incubator.vector 22 - - - - apiNote - a - API Note: - -