diff --git a/README.md b/README.md index 3d7bb2546..a16da866a 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,11 @@ 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. + +For writing completed graphs with FP or NVQ features, see [Deferred batch serialization](./docs/deferred-batch-writer.md). It works with application vector sources and preloaded data; dataset prefetching is optional. + 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/benchmarks-jmh/README.md b/benchmarks-jmh/README.md index 1da3d8d9f..88a0e1c07 100644 --- a/benchmarks-jmh/README.md +++ b/benchmarks-jmh/README.md @@ -133,7 +133,7 @@ The formatter will output a clean table showing: - **Visited_Count**: Average number of nodes visited during search - **Expanded_Count_BaseLayer**: Average number of nodes expanded in base layer -Example output: +Example output:[^machine] ``` k PQ_Subspaces Time_ms Recall ReRanked_Count Visited_Count Expanded_Count_BaseLayer 50 0 19.283 1.000 0.0 3290.8 253.7 @@ -143,3 +143,5 @@ Example output: + +[^machine]: Machine type and core count were not recorded for this formatting example; its timings are illustrative, not a hardware baseline. 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/benchmarking.md b/docs/benchmarking.md index 79ddba82e..ccb97764a 100644 --- a/docs/benchmarking.md +++ b/docs/benchmarking.md @@ -115,3 +115,10 @@ To add a custom fvecs/ivecs dataset: ``` For remote datasets, use `base_url` to specify where files should be downloaded from. The `${VAR}` and `${VAR:-default}` syntax is supported for environment variable expansion. See the example config for details. + +## Final index serialization + +Grid defers inline FP/NVQ serialization until after graph cleanup. The existing +parallel-writer selection enables complete-record batches (default 64); dataset +prefetching is a separate option and is not required. See +[Deferred batch serialization](deferred-batch-writer.md) for activation and caller examples. diff --git a/docs/deferred-batch-writer.md b/docs/deferred-batch-writer.md new file mode 100644 index 000000000..05aa26f8e --- /dev/null +++ b/docs/deferred-batch-writer.md @@ -0,0 +1,127 @@ +# Deferred batch serialization + +The parallel graph writer can serialize complete layer-zero records into bounded +worker buffers. Each record contains its ordinal, inline features, neighbor count, +and adjacency list. By default, a worker writes up to 64 complete records together. +The final partial batch is written without adding records or padding to the format. +This preserves the V6 layout and existing readers. + +## Production sources + +The library writer does not depend on Grid, BenchYAML, dataset preloading, or +#744's buffered-access API. It accepts the existing +`Map>`: callers supply feature states for the +requested original ordinals from their own source. That source can be resident +vectors, a database's flush/compaction source, or another source; it must support +concurrent supplier calls and remain available through final serialization. +The writer does not require the original dataset to be fully resident. + +Complete-record batching also applies when all input vectors are preloaded. Its +benefits come from deferring per-insertion feature writes and grouping final writes, +independently of how input vectors are loaded. Buffered range access is an optional +Grid integration, not a prerequisite for the core writer. + +## Caller migration: early writes versus deferred batches + +Both examples assume the same trained `nvq`, graph `builder`, node `count`, +`writerWorkers`, and output path. The builder's construction scoring is unchanged; +its peer vectors are supplied separately. `vectorForOrdinal` is the application's +lookup of an original input vector, not a harness loader. + +### Old: write each inline feature during insertion + +```java +try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), outputPath) + .withParallelWorkerThreads(writerWorkers) + .withMapper(new OrdinalMapper.IdentityMapper(count - 1)) + .with(new NVQ(nvq)) + .build()) { + for (int ordinal = 0; ordinal < count; ordinal++) { + var vector = vectorForOrdinal.apply(ordinal); + builder.addGraphNode(ordinal, vector); + writer.writeFeaturesInline(ordinal, Feature.singleState( + FeatureId.NVQ_VECTORS, new NVQ.State(nvq.encode(vector)))); + } + builder.cleanup(); + writer.write(Map.of()); // Preserve the earlier feature writes; add adjacency. +} +``` + +### New: supply the features when writing the completed graph + +```java +for (int ordinal = 0; ordinal < count; ordinal++) { + builder.addGraphNode(ordinal, vectorForOrdinal.apply(ordinal)); +} +builder.cleanup(); + +try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), outputPath) + .withParallelWorkerThreads(writerWorkers) + .withMapper(new OrdinalMapper.IdentityMapper(count - 1)) + .with(new NVQ(nvq)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.NVQ_VECTORS, + ordinal -> new NVQ.State(nvq.encode(vectorForOrdinal.apply(ordinal))))); +} +``` + +The second call supplies NVQ payloads during final serialization, together with +adjacency, in batches of up to 64 complete records. No `writeFeaturesInline` pass +is needed. The same pattern works for FP using `InlineVectors` and +`InlineVectors.State`, or another supported inline feature. Supply every configured +inline feature to obtain complete-record batching, including fused PQ if configured. + +Final suppliers run concurrently. `vectorForOrdinal` must remain available through +`write` and provide independent borrowed vectors per worker, or stable vectors. +For an existing `RandomAccessVectorValues`, `threadLocalSupplier()` supplies those +independent views; it does not require buffered hooks. A caller can instead return +already encoded feature states from its own storage, without retaining original FP +vectors for this write. + +Integrations that already call `write` with every inline feature supplier receive +batching when using the parallel writer. Integrations that write features early +must move those calls to final suppliers to use the deferred path. The incremental +API remains compatible: a one-pass source whose values cannot be obtained again +must retain encoded states for final writing or continue using incremental writes. + +## Batch size and memory + +The default is 64 records. To override it, call `withParallelBatchSize(records)` +before `with(...)` on either `OnDiskParallelGraphIndexWriter.Builder` or +`RandomAccessOnDiskGraphIndexWriter.Builder`. The value must be positive. Each active +worker holds one buffer of `recordSize * records` bytes; feature encoders and vector +sources have their own memory costs. No dataset-sized output buffer is allocated. + +The unified `RandomAccessOnDiskGraphIndexWriter.Builder` retains its existing +parallel/sequential selection. Setting the batch size does not enable parallel +writing, and is ignored by the sequential writer. A caller-supplied executor +remains caller-owned and must outlive the write. + +## Dataset prefetching in Grid + +When the parallel writer is selected, Grid defers its inline NVQ or FP feature serialization +until after cleanup. Its final serialization suppliers consume worker-local ranges through +`VectorAccess.openRange`. With the default sequential writer, which calls suppliers on one +thread, Grid keeps encoding these features on the insertion threads and writing them inline. +Sources implementing the buffered-access capability use that implementation; +ordinary sources retain the ordinary-access fallback. Borrowed vectors are consumed +before advancing a cursor, and ranges are released after the workers finish. + +This integrates with #744's optional `useDatasetPrefetching` setting. It does not +change construction scoring, quantization parameters, or the query path. Production +integrations can opt into deferred serialization by supplying their features at the +final write; their existing vector-source and incremental-write APIs remain valid. + +## Running the Grid parallel-write path + +Grid keeps the existing writer-selection setting. To exercise complete-record +parallel batches through BenchYAML, enable it before invoking the benchmark: + +```java +GraphIndexBuilderConfig.getInstance().setParallelBuild(true); +BenchYAML.main(new String[]{"your-index-parameters.yml"}); +``` + +Grid sizes its writer pool from `PhysicalCoreExecutor`; no machine-specific worker +count is embedded in this change. Dataset prefetching is a separate optional setting +in the index YAML. An ordinary preloaded dataset can use the same deferred writer. diff --git a/docs/legacy/jvector-step-by-step.md b/docs/legacy/jvector-step-by-step.md index 176ff38f8..b7424ac2e 100644 --- a/docs/legacy/jvector-step-by-step.md +++ b/docs/legacy/jvector-step-by-step.md @@ -1,5 +1,9 @@ ## JVector step-by-step +> Historical guide: API names and early-write examples below reflect an older +> version. Use the [current tutorials](../tutorials/2-disk-tutorial.md) and +> [deferred serialization guide](../deferred-batch-writer.md) for current construction APIs. + All code samples are from [SiftSmall](https://github.com/jbellis/jvector/blob/main/jvector-examples/src/main/java/io/github/jbellis/jvector/example/SiftSmall.java) in the JVector source repo, which includes the siftsmall dataset as well. Just import the project in your IDE and click Run to try it out! #### Step 1: Build and query an index in memory 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/release notes/4.0.2/747.performance.md b/docs/release notes/4.0.2/747.performance.md new file mode 100644 index 000000000..b1d74dfce --- /dev/null +++ b/docs/release notes/4.0.2/747.performance.md @@ -0,0 +1,29 @@ +### Deferred batch graph serialization + +Parallel graph serialization can now write complete layer-zero records in bounded +batches of 64, with an optional positive batch-size override. When the parallel writer +is selected, Grid supplies inline NVQ or full-precision features during final +serialization, after graph cleanup, and uses buffered vector ranges when the input +source supports them. With the default sequential writer, Grid keeps encoding those +features during insertion. + +The core writer accepts application-supplied feature states independently of the +benchmark harness or dataset prefetching. Complete-record batching applies equally +to fully preloaded inputs; the source does not have to preload the whole dataset. + +The existing V6 layout, readers, ordinary vector sources, and incremental +`writeFeaturesInline` API remain supported. Prewritten features are preserved. +Concurrent feature suppliers must keep their input source available until writing +finishes. The existing parallel-writer selection remains unchanged. + +Recorded construction studies with dataset prefetching reduced ample-memory final +write time from 136.10 to 50.69 seconds on MPNet 10M, 148.50 to 64.02 seconds on +Cohere 10M, and 162.64 to 51.76 seconds on CAP 6M.[^machine] These are separate fresh runs, +not a universal throughput or memory guarantee. See +[Deferred batch serialization](../../deferred-batch-writer.md) for usage. + +Construction tutorials and examples now demonstrate final feature suppliers for +deferred FP/NVQ writing, with per-worker source views and an explicit source-lifetime +requirement. The query algorithms and quantization settings are unchanged. + +[^machine]: GCE VM, 48-core Intel Xeon 6985P-C (96 threads), local SSD RAID10, JDK 23, 128 GiB heap; native Highway, 48 workers. diff --git a/docs/tutorials/2-disk-tutorial.md b/docs/tutorials/2-disk-tutorial.md index 4040ff7de..e9f87bb87 100644 --- a/docs/tutorials/2-disk-tutorial.md +++ b/docs/tutorials/2-disk-tutorial.md @@ -8,7 +8,7 @@ Unlike the previous example, this time we'll use a proper ANN dataset based on O // This is a preconfigured dataset that will be downloaded automatically. DataSet dataset = DataSets.loadDataSet("ada002-100k").orElseThrow(() -> new RuntimeException("Dataset doesn't exist or wasn't configured correctly") -); +).getDataSet(); ``` > [!TIP] @@ -45,7 +45,8 @@ Path graphPath = Files.createTempFile("jvector-example-graph", null); // or whe // Create a writer for the on-heap graph we just built. // Remember to close when done. -GraphIndexWriter writer = GraphIndexWriter.getBuilderFor(GraphIndexWriterTypes.RANDOM_ACCESS_PARALLEL, heapGraph, graphPath) +GraphIndexWriter writer = new OnDiskParallelGraphIndexWriter.Builder(heapGraph, graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) // Let the writer know that we'll also be passing in the actual vector data // to be saved "inline" with the data for each corresponding graph node. .with(new InlineVectors(dim)) @@ -60,13 +61,20 @@ Since we want to include a feature, we need to tell the writer how to obtain the // Supply one map entry for each feature. // The key is a FeatureId enum corresponding to the feature // and the value is a function which generates the feature state for each graph node. +var writerVectors = ravv.threadLocalSupplier(); writer.write(Map.of( FeatureId.INLINE_VECTORS, // we already have a RAVV, so we'll just use that to supply the writer. - nodeId -> new InlineVectors.State(ravv.getVector(nodeId)))); + nodeId -> new InlineVectors.State(writerVectors.get().getVector(nodeId)))); // writer.close() if not using try-with-resources ``` +Supplying every inline feature lets the parallel writer serialize complete records +in batches of up to 64. The source must remain readable until `write` finishes. +`threadLocalSupplier()` provides independent views when a source borrows a scratch +vector; it also works with ordinary preloaded sources. See +[Deferred batch serialization](../deferred-batch-writer.md) for NVQ and migration examples. + At this point the graph index has been written to disk and can be used by creating an instance of `OnDiskGraphIndex`. To do so, we'll need to specify a `RandomAccessReader` implementation which JVector will use to read parts of the file as needed. This interface isn't thread safe, so when creating an `OnDiskGraphIndex` we pass in a `ReaderSupplier` object which can be used by the `OnDiskGraphIndex` to create `RandomAccessReader`s as needed. diff --git a/docs/tutorials/3-larger-than-memory-tutorial.md b/docs/tutorials/3-larger-than-memory-tutorial.md index 1374f3c11..6001c8e0a 100644 --- a/docs/tutorials/3-larger-than-memory-tutorial.md +++ b/docs/tutorials/3-larger-than-memory-tutorial.md @@ -1,7 +1,7 @@ # JVector Tutorial Part 3: Larger-Than-Memory Indexes with Product Quantization In the previous tutorials, we built indexes in-memory then wrote them to disk. However, this requires enough memory to hold all the vectors during construction. In this tutorial we'll build larger-than-memory indexes using two key techniques: -- Write vectors to disk as they are streamed in +- Write complete records after graph cleanup, reading FP payloads from a replayable source - Keep compressed versions of the vectors in the memory to allow JVector to compute similarities during construction. We'll use Product Quantization (PQ). > [!TIP] @@ -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: @@ -24,7 +29,7 @@ We'll change things up with a different dataset: // the base vectors in-memory. DataSet dataset = DataSets.loadDataSet("e5-small-v2-100k").orElseThrow(() -> new RuntimeException("Dataset doesn't exist or wasn't configured correctly") -); +).getDataSet(); // Remember that RAVVs need not be in-memory in the general case. // We will sample from this RAVV to compute the PQ codebooks. @@ -64,44 +69,35 @@ var pqVectors = new MutablePQVectors(pq); BuildScoreProvider bsp = BuildScoreProvider.pqBuildScoreProvider(vsf, pqVectors); ``` -Since we have a BuildScoreProvider, we're now ready to create and use a `GraphIndexBuilder`. We will also create the corresponding `GraphIndexWriter` at the same time. Unlike the last tutorial we created the writer only after building the complete graph, we'll write full-resolution vectors to disk in tandem with adding them to the graph. +The graph builder uses resident compressed vectors for neighbor comparisons. +Original input vectors are consumed for insertion and reread for final FP payloads; +the source must remain available through serialization. It need not preload the +whole dataset. A one-pass source must retain its final feature states or use the +compatible incremental-writing API instead. ```java -// Initialize the graph parameters M, ef etc -// ... - +// Initialize the graph parameters M, ef etc. Path graphPath = Files.createTempFile("jvector-ltm-graph", null); -// remember to close these manually or use try-with-resources +// Close the builder after final writing, or use try-with-resources. GraphIndexBuilder builder = new GraphIndexBuilder(bsp, dim, M, ef, overflow, alpha, addHierarchy); -OnDiskGraphIndexWriter writer = new OnDiskGraphIndexWriter.Builder(builder.getGraph(), graphPath) - .with(new InlineVectors(dim)) - // Since we start with an empty graph, the writer will, by default, - // assume an ordinal mapping of size 0 (which is obviously incorrect). - // This is easy to rectify if you know the number of vectors beforehand, - // if not you may need to implement OrdinalMapper yourself. - .withMapper(new OrdinalMapper.IdentityMapper(ravv.size() - 1)) - .build(); +var vectors = ravv.threadLocalSupplier(); ``` ## Incremental Index Construction We'll proceed to add vectors to the index one at a time. The general procedure is: - Encode the vector using the PQ codebooks created earlier, and add it to the collection of PQ vectors -- Write the full-resolution vector to disk - Add the new vector to the graph ```java for (int ordinal = 0; ordinal < ravv.size(); ordinal++) { - VectorFloat v = ravv.getVector(ordinal); + VectorFloat v = vectors.get().getVector(ordinal); // Encode and add the vector to the working set of PQ vectors, // which allows the graph builder to access it through the BuildScoreProvder. pqVectors.encodeAndSet(ordinal, v); - // Write the feature (full-resolution vector) for a single vector instead of all at once. - writer.writeFeaturesInline(ordinal, Feature.singleState(FeatureId.INLINE_VECTORS, new InlineVectors.State(v))); - builder.addGraphNode(ordinal, v); } ``` @@ -116,12 +112,25 @@ Once we've added all the vectors we can write the graph structure to disk: // Enforces maximum degree constraint among other things. builder.cleanup(); -// Write the graph structure (neighbor lists) to disk. -// No need to pass-in a feature supplier since we wrote the features incrementally -// using writer.writeInLine -writer.write(Map.of()); +// Serialize FP payloads and adjacency together after cleanup. +try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new InlineVectors(dim)) + .withMapper(new OrdinalMapper.IdentityMapper(ravv.size() - 1)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + ordinal -> new InlineVectors.State(vectors.get().getVector(ordinal)))); +} +builder.close(); ``` +The parallel writer batches up to 64 complete records without changing V6. +Callbacks run concurrently, so the per-worker source views above are important for +borrowed vectors. For NVQ reranking, supply `NVQ.State` instead of `InlineVectors.State`; +see [Deferred batch serialization](../deferred-batch-writer.md) and +[`NvqExample.java`](../../jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java). +The construction scorer and two-stage query algorithm are unchanged. + We should also save the PQ vectors we created since we'll need them later on when it's time to run searches: ```java @@ -149,7 +158,7 @@ ReaderSupplier graphSupplier = ReaderSupplierFactory.open(graphPath); OnDiskGraphIndex graph = OnDiskGraphIndex.load(graphSupplier); // except that we also need a reader for the PQ vectors ReaderSupplier pqSupplier = ReaderSupplierFactory.open(pqPath); -RandomAccessReader pqReader = pqSupplier.get() +RandomAccessReader pqReader = pqSupplier.get(); // don't forget to close all of the above! (or just use try-with-resources) // we need to have the PQ vectors in memory @@ -162,7 +171,7 @@ int rerankK = (int) (topK * overqueryFactor); // a friendly reminder that searchers need closing too GraphSearcher searcher = new GraphSearcher(graph); var graphRavv = (RandomAccessVectorValues) searcher.getView(); -VectorFloat query : dataset.getQueryVectors().get(0); +VectorFloat query = dataset.getQueryVectors().get(0); // Two-phase search: // 1. ApproximateScoreFunction (ASF) uses compressed vectors for fast initial search 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/disk/NodeRecordTask.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java index 7e756dab9..5872147da 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java @@ -28,6 +28,7 @@ import java.util.Map; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; +import java.util.function.BooleanSupplier; import java.util.function.IntFunction; /** @@ -38,8 +39,8 @@ * Each task processes a contiguous range of ordinals to reduce task creation overhead. *

* This writes directly to the AsynchronousFileChannel using position-based writes with writeFully - * to ensure all bytes are written before returning. This eliminates race conditions where the OS - * buffer cache hasn't flushed data before subsequent reads occur. + * to ensure each write completes before its buffer is reused. Durability is handled by + * the caller's final channel.force(true). */ class NodeRecordTask implements Callable { private final int startOrdinal; // Inclusive @@ -53,6 +54,7 @@ class NodeRecordTask implements Callable { private final long baseOffset; // Base file offset for L0 (offsets calculated per-ordinal) private final AsynchronousFileChannel channel; private final ByteBuffer buffer; // Thread-local buffer for building record components + private final BooleanSupplier stopRequested; NodeRecordTask(int startOrdinal, int endOrdinal, @@ -65,6 +67,25 @@ class NodeRecordTask implements Callable { long baseOffset, AsynchronousFileChannel channel, ByteBuffer buffer) { + this(startOrdinal, endOrdinal, ordinalMapper, graph, view, inlineFeatures, featureStateSuppliers, + recordSize, baseOffset, channel, buffer, () -> false); + } + + /// Create a task that stops before its next record once `stopRequested` returns true, for + /// example after another task failed or the writing thread was interrupted. A stopped task + /// returns normally; the writer reports the failure that requested the stop. + NodeRecordTask(int startOrdinal, + int endOrdinal, + OrdinalMapper ordinalMapper, + ImmutableGraphIndex graph, + ImmutableGraphIndex.View view, + List inlineFeatures, + Map> featureStateSuppliers, + int recordSize, + long baseOffset, + AsynchronousFileChannel channel, + ByteBuffer buffer, + BooleanSupplier stopRequested) { this.startOrdinal = startOrdinal; this.endOrdinal = endOrdinal; this.ordinalMapper = ordinalMapper; @@ -76,6 +97,7 @@ class NodeRecordTask implements Callable { this.baseOffset = baseOffset; this.channel = channel; this.buffer = buffer; + this.stopRequested = stopRequested; } /** @@ -95,8 +117,8 @@ private static void writeFully(AsynchronousFileChannel channel, ByteBuffer buffe long currentPosition = position; while (buffer.hasRemaining()) { int written = channel.write(buffer, currentPosition).get(); - if (written < 0) { - throw new IOException("Channel closed while writing"); + if (written <= 0) { + throw new IOException("Channel made no progress while writing"); } currentPosition += written; } @@ -107,16 +129,60 @@ public Void call() throws Exception { // Reuse writer and buffer across all ordinals in this range var writer = new ByteBufferIndexWriter(buffer); + int payloadBytes = Integer.BYTES * (2 + graph.getDegree(0)) + + inlineFeatures.stream().mapToInt(Feature::featureSize).sum(); + int padding = recordSize - payloadBytes; + final boolean completeRecord = inlineFeatures.stream().allMatch(f -> featureStateSuppliers.get(f.id()) != null); + final int batch = buffer.capacity() / recordSize; for (int newOrdinal = startOrdinal; newOrdinal < endOrdinal; newOrdinal++) { + if (stopRequested.getAsBoolean()) return null; var originalOrdinal = ordinalMapper.newToOld(newOrdinal); long recordOffset = baseOffset + (long) newOrdinal * recordSize; long currentPosition = recordOffset; + // A complete record can share one bounded buffer and one positioned write. + if (completeRecord) { + if (buffer.position() == buffer.limit()) writer.reset(); + int recordBegin = buffer.position(); + writer.writeInt(originalOrdinal == OrdinalMapper.OMITTED ? -1 : newOrdinal); + if (originalOrdinal == OrdinalMapper.OMITTED) { + for (var feature : inlineFeatures) + for (int i = 0; i < feature.featureSize(); i++) writer.writeByte(0); + writer.writeInt(0); + for (int n = 0; n < graph.getDegree(0); n++) writer.writeInt(-1); + } else { + if (!graph.containsNode(originalOrdinal)) { + throw new IllegalStateException("Missing graph node " + originalOrdinal); + } + for (var feature : inlineFeatures) + feature.writeInline(writer, featureStateSuppliers.get(feature.id()).apply(originalOrdinal)); + var neighbors = view.getNeighborsIterator(0, originalOrdinal); + if (neighbors.size() > graph.getDegree(0)) throw new IllegalStateException("Too many neighbors for " + originalOrdinal); + writer.writeInt(neighbors.size()); + int n = 0; + while (neighbors.hasNext()) { + int mapped = ordinalMapper.oldToNew(neighbors.nextInt()); + if (mapped < 0 || mapped > ordinalMapper.maxOrdinal()) throw new IllegalStateException("Invalid mapped neighbor " + mapped); + writer.writeInt(mapped); n++; + } + for (; n < graph.getDegree(0); n++) writer.writeInt(-1); + } + for (int i = 0; i < padding; i++) writer.writeByte(0); + if (writer.position() != recordBegin + recordSize) throw new IllegalStateException("Wrong serialized record length"); + if (buffer.position() < recordSize * batch && newOrdinal != endOrdinal - 1) continue; + long batchOffset = recordOffset - recordBegin; + buffer.flip(); + // writeFully waits for completion before this worker reuses its buffer. + writeFully(channel, buffer, batchOffset); + buffer.clear(); + continue; + } + // Reset buffer for this ordinal writer.reset(); // Write node ordinal - writer.writeInt(newOrdinal); + writer.writeInt(originalOrdinal == OrdinalMapper.OMITTED ? -1 : newOrdinal); ByteBuffer ordinalData = writer.cloneBuffer(); writeFully(channel, ordinalData, currentPosition); currentPosition += Integer.BYTES; @@ -132,8 +198,10 @@ public Void call() throws Exception { } } ByteBuffer featureData = writer.cloneBuffer(); + // writeFully consumes the buffer; retain its length before writing. + int featureBytes = featureData.remaining(); writeFully(channel, featureData, currentPosition); - currentPosition += featureData.remaining(); + currentPosition += featureBytes; // Write empty neighbor list writer.reset(); @@ -194,6 +262,11 @@ public Void call() throws Exception { ByteBuffer neighborData = writer.cloneBuffer(); writeFully(channel, neighborData, currentPosition); } + if (padding > 0) { + writer.reset(); + for (int i = 0; i < padding; i++) writer.writeByte(0); + writeFully(channel, writer.cloneBuffer(), recordOffset + payloadBytes); + } } return null; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskParallelGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskParallelGraphIndexWriter.java index 4e8e7ac2d..0207225f2 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskParallelGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskParallelGraphIndexWriter.java @@ -70,6 +70,7 @@ public class OnDiskParallelGraphIndexWriter extends RandomAccessOnDiskGraphIndex private final int parallelWorkerThreads; private final boolean parallelUseDirectBuffers; private final ExecutorService parallelExecutor; + private final int parallelBatchSize; /** * Constructs an OnDiskParallelGraphIndexWriter with all parameters including file path @@ -118,6 +119,24 @@ public class OnDiskParallelGraphIndexWriter extends RandomAccessOnDiskGraphIndex int parallelWorkerThreads, boolean parallelUseDirectBuffers, ExecutorService parallelExecutor) + { + this(randomAccessWriter, version, startOffset, graph, oldToNewOrdinals, dimension, features, + filePath, parallelWorkerThreads, parallelUseDirectBuffers, parallelExecutor, + ParallelGraphWriter.DEFAULT_RECORDS_PER_BATCH); + } + + OnDiskParallelGraphIndexWriter(RandomAccessWriter randomAccessWriter, + int version, + long startOffset, + ImmutableGraphIndex graph, + OrdinalMapper oldToNewOrdinals, + int dimension, + EnumMap features, + Path filePath, + int parallelWorkerThreads, + boolean parallelUseDirectBuffers, + ExecutorService parallelExecutor, + int parallelBatchSize) { super(randomAccessWriter, version, startOffset, graph, oldToNewOrdinals, dimension, features); if (filePath == null) { @@ -127,6 +146,7 @@ public class OnDiskParallelGraphIndexWriter extends RandomAccessOnDiskGraphIndex this.parallelWorkerThreads = parallelWorkerThreads; this.parallelUseDirectBuffers = parallelUseDirectBuffers; this.parallelExecutor = parallelExecutor; + this.parallelBatchSize = parallelBatchSize; } /** @@ -153,7 +173,8 @@ protected void writeL0Records(ImmutableGraphIndex.View view, var config = new ParallelGraphWriter.Config( parallelWorkerThreads, parallelUseDirectBuffers, - 4 // Default task multiplier (4x cores) + 4, // Default task multiplier (4x cores) + parallelBatchSize ); try (var parallelWriter = new ParallelGraphWriter( @@ -200,6 +221,7 @@ public static class Builder extends AbstractGraphIndexWriter.Builder */ class ParallelGraphWriter implements AutoCloseable { + // Amortize positioned writes with bounded worker-local buffers. + static final int DEFAULT_RECORDS_PER_BATCH = 64; private final RandomAccessWriter writer; private final ImmutableGraphIndex graph; private final ExecutorService executor; private final boolean ownsExecutor; private final ThreadLocal viewPerThread; - private final ThreadLocal bufferPerThread; + private final ExplicitThreadLocal bufferPerThread; private final CopyOnWriteArrayList allViews = new CopyOnWriteArrayList<>(); private final int recordSize; private final Path filePath; @@ -81,6 +85,7 @@ static class Config { final int workerThreads; final boolean useDirectBuffers; final int taskMultiplier; + final int recordsPerBatch; /** * @param workerThreads number of worker threads for building records (0 = use available processors) @@ -89,6 +94,14 @@ static class Config { * (4x = good balance for most use cases, higher = more fine-grained parallelism) */ public Config(int workerThreads, boolean useDirectBuffers, int taskMultiplier) { + this(workerThreads, useDirectBuffers, taskMultiplier, DEFAULT_RECORDS_PER_BATCH); + } + + public Config(int workerThreads, boolean useDirectBuffers, int taskMultiplier, int recordsPerBatch) { + if (recordsPerBatch <= 0) { + throw new IllegalArgumentException("Record batch size must be positive"); + } + this.recordsPerBatch = recordsPerBatch; this.workerThreads = workerThreads <= 0 ? Runtime.getRuntime().availableProcessors() : workerThreads; this.useDirectBuffers = useDirectBuffers; this.taskMultiplier = taskMultiplier <= 0 ? 4 : taskMultiplier; @@ -143,6 +156,10 @@ public ParallelGraphWriter(RandomAccessWriter writer, this.graph = graph; this.filePath = Objects.requireNonNull(filePath); this.taskMultiplier = config.taskMultiplier; + // Validate buffer size before creating a pool that would otherwise need cleanup. + this.recordSize = Math.addExact(Integer.BYTES * (2 + graph.getDegree(0)), + inlineFeatures.stream().mapToInt(Feature::featureSize).sum()); + final int bufferSize = Math.multiplyExact(recordSize, config.recordsPerBatch); if (externalExecutor != null) { this.executor = externalExecutor; this.ownsExecutor = false; @@ -157,12 +174,6 @@ public ParallelGraphWriter(RandomAccessWriter writer, this.ownsExecutor = true; } - // Compute fixed record size for L0 - this.recordSize = Integer.BYTES // node ordinal - + inlineFeatures.stream().mapToInt(Feature::featureSize).sum() - + Integer.BYTES // neighbor count - + graph.getDegree(0) * Integer.BYTES; // neighbors + padding - // Thread-local views for safe neighbor iteration // CopyOnWriteArrayList handles concurrent additions safely this.viewPerThread = ThreadLocal.withInitial(() -> { @@ -173,9 +184,8 @@ public ParallelGraphWriter(RandomAccessWriter writer, // Thread-local buffers to avoid allocation overhead // Use BIG_ENDIAN to match Java DataOutput specification - final int bufferSize = recordSize; final boolean useDirect = config.useDirectBuffers; - this.bufferPerThread = ThreadLocal.withInitial(() -> { + this.bufferPerThread = ExplicitThreadLocal.withInitial(() -> { ByteBuffer buffer = useDirect ? ByteBuffer.allocateDirect(bufferSize) : ByteBuffer.allocate(bufferSize); buffer.order(java.nio.ByteOrder.BIG_ENDIAN); return buffer; @@ -224,54 +234,59 @@ public void writeL0Records(OrdinalMapper ordinalMapper, try (var channel = AsynchronousFileChannel.open(filePath, opts, null)) { List> taskFutures = new ArrayList<>(numTasks); + // Set on interruption or the first failure, so remaining ranges stop instead of running to completion. + var stop = new AtomicBoolean(); + boolean submitted = false; - // Submit range-based tasks - for (int i = 0; i < numTasks; i++) { - int startOrdinal = i * ordinalsPerTask; - int endOrdinal = Math.min(startOrdinal + ordinalsPerTask, totalOrdinals); - - // Skip if range is empty (can happen with final task) - if (startOrdinal >= totalOrdinals) { - break; - } + // Always join submitted tasks before closing the channel or supplier resources. + try { + // Submit range-based tasks + for (int i = 0; i < numTasks; i++) { + int startOrdinal = i * ordinalsPerTask; + int endOrdinal = Math.min(startOrdinal + ordinalsPerTask, totalOrdinals); - final int start = startOrdinal; - final int end = endOrdinal; + // Skip if range is empty (can happen with final task) + if (startOrdinal >= totalOrdinals) { + break; + } - Future future = executor.submit(() -> { - var view = viewPerThread.get(); - var buffer = bufferPerThread.get(); + final int start = startOrdinal; + final int end = endOrdinal; - var task = new NodeRecordTask( - start, // Start of range (inclusive) - end, // End of range (exclusive) - ordinalMapper, - graph, - view, - inlineFeatures, - featureStateSuppliers, - recordSize, - baseOffset, // Base offset (task calculates per-ordinal offsets) - channel, // Async file channel for position-based writes - buffer // Thread-local buffer - ); + Future future = executor.submit(() -> { + if (stop.get()) return null; + var view = viewPerThread.get(); + var buffer = bufferPerThread.get(); - return task.call(); - }); + var task = new NodeRecordTask( + start, // Start of range (inclusive) + end, // End of range (exclusive) + ordinalMapper, + graph, + view, + inlineFeatures, + featureStateSuppliers, + recordSize, + baseOffset, // Base offset (task calculates per-ordinal offsets) + channel, // Async file channel for position-based writes + buffer, // Thread-local buffer + stop::get + ); - taskFutures.add(future); - } + try { + return task.call(); + } finally { + // External executor threads must not retain this writer's graph views. + viewPerThread.remove(); + } + }); - // Wait for all tasks to complete - each task writes synchronously using writeFully - for (Future taskFuture : taskFutures) { - try { - taskFuture.get(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IOException("Interrupted while writing records", e); - } catch (ExecutionException e) { - throw unwrapExecutionException(e); + taskFutures.add(future); } + submitted = true; + } finally { + if (!submitted) stop.set(true); + awaitTasks(taskFutures, stop); } // Force all writes to disk before closing the channel. @@ -283,23 +298,35 @@ public void writeL0Records(OrdinalMapper ordinalMapper, } } - /** - * Unwraps ExecutionException to throw the underlying cause. - * Handles IOException, RuntimeException, and wraps other exceptions. - * - * @param e the execution exception to unwrap - * @return an IOException wrapping the cause - * @throws RuntimeException if the cause is a RuntimeException - */ - private IOException unwrapExecutionException(ExecutionException e) { - Throwable cause = e.getCause(); - if (cause instanceof IOException) { - return (IOException) cause; - } else if (cause instanceof RuntimeException) { - throw (RuntimeException) cause; - } else { - throw new RuntimeException("Error building node record", cause); + /// Join every task, so no worker still uses the channel or suppliers when this returns. On + /// interruption or the first task failure, set `stop` so the remaining tasks skip or abandon + /// their ranges rather than writing them, then rethrow the first failure. + private void awaitTasks(List> futures, AtomicBoolean stop) throws IOException { + Throwable failure = null; + boolean interrupted = false; + for (var future : futures) { + boolean done = false; + while (!done) { + try { + future.get(); + done = true; + } catch (InterruptedException e) { + interrupted = true; + stop.set(true); + if (failure == null) failure = new IOException("Interrupted while writing records", e); + } catch (ExecutionException e) { + stop.set(true); + if (failure == null) failure = e.getCause(); + else if (failure != e.getCause()) failure.addSuppressed(e.getCause()); + done = true; + } + } } + if (interrupted) Thread.currentThread().interrupt(); + if (failure instanceof IOException) throw (IOException) failure; + if (failure instanceof RuntimeException) throw (RuntimeException) failure; + if (failure instanceof Error) throw (Error) failure; + if (failure != null) throw new IOException("Error building node record", failure); } /** @@ -330,6 +357,7 @@ public void close() throws IOException { view.close(); } allViews.clear(); + bufferPerThread.close(); } catch (IOException e) { throw e; } catch (Exception e) { diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java index e3ce70e21..debb2b5bd 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java @@ -141,7 +141,7 @@ protected abstract void writeL0Records(ImmutableGraphIndex.View view, * Both produce an identical on-disk format. * *

Parallel-specific options ({@link #withParallelWorkerThreads}, {@link #withParallelDirectBuffers}, - * {@link #withExecutor}) are accepted unconditionally but silently ignored when the sequential + * {@link #withExecutor}, {@link #withParallelBatchSize}) are accepted unconditionally but silently ignored when the sequential * writer is selected. */ public static class Builder extends AbstractGraphIndexWriter.Builder { @@ -150,6 +150,7 @@ public static class Builder extends AbstractGraphIndexWriter.Builder 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 24c39ae47..7d41d6c77 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,34 +139,52 @@ 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 — failures are non-fatal and don't block checkpointing - 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 — failures are non-fatal and don't block checkpointing + 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); + } } } 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..50bcf2631 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,23 @@ 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 +425,10 @@ 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; } @@ -427,43 +437,45 @@ private static BuildOnDiskResult buildOnDisk(List> feat throw new IllegalStateException("Bench looks for either NVQ_VECTORS or INLINE_VECTORS feature set for scoring compressed builds."); } - // build the graph incrementally + // The parallel writer calls final suppliers from its workers, so original-vector features + // can be encoded after cleanup. The sequential writer calls them on one thread, so it keeps + // encoding on the insertion threads and writing each node's features inline. + boolean deferFeatures = writers.values().stream() + .allMatch(writer -> writer instanceof OnDiskParallelGraphIndexWriter); 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) -> { + if (!deferFeatures) writeInlineStates(writers, inputSuppliers, node, vector); + builder.addGraphNode(node, vector); }); - }).join(); + } else { + var vv = floatVectors.threadLocalSupplier(); + PhysicalCoreExecutor.pool().submit(() -> { + IntStream.range(0, floatVectors.size()).parallel().forEach(node -> { + var vector = vv.get().getVector(node); + if (!deferFeatures) writeInlineStates(writers, inputSuppliers, node, vector); + builder.addGraphNode(node, vector); + }); + }).join(); + } builder.cleanup(); - // write the edge lists and close the writers - // if our feature set contains Fused PQ, we need a Fused ADC write-time supplier (as we don't have neighbor information during writeInline) + // With deferred features, supply every feature so the parallel writer emits complete records. writers.entrySet().stream().parallel().forEach(entry -> { var writer = entry.getValue(); var features = entry.getKey(); - Map> writeSuppliers; - if (features.contains(FeatureId.FUSED_PQ)) { - writeSuppliers = new EnumMap<>(FeatureId.class); - var view = builder.getGraph().getView(); - writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pq, ordinal)); - } else { - writeSuppliers = Map.of(); - } - try { + try (var vectors = new FinalWriteVectors(floatVectors); + var view = builder.getGraph().getView(); + var closeWriter = writer) { + Map> writeSuppliers = new EnumMap<>(FeatureId.class); + if (deferFeatures) { + inputSuppliers.get(features).forEach((id, supplier) -> + writeSuppliers.put(id, ordinal -> supplier.apply(vectors.get(ordinal)))); + } + if (features.contains(FeatureId.FUSED_PQ)) { + writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pq, ordinal)); + } writer.write(writeSuppliers); - writer.close(); } catch (IOException e) { throw new UncheckedIOException(e); } @@ -495,6 +507,76 @@ private static BuildOnDiskResult buildOnDisk(List> feat return new BuildOnDiskResult(indexes, fileSizes); } + /// Write one node's original-vector features to every writer during insertion. The borrowed + /// vector is encoded and written before the caller advances to another vector. + private static void writeInlineStates( + Map, RandomAccessOnDiskGraphIndexWriter> writers, + Map, Map, Feature.State>>> inputSuppliers, + int node, VectorFloat vector) { + writers.forEach((features, writer) -> { + var states = new EnumMap(FeatureId.class); + inputSuppliers.get(features).forEach((id, supplier) -> states.put(id, supplier.apply(vector))); + try { writer.writeFeaturesInline(node, states); } + catch (IOException e) { throw new UncheckedIOException(e); } + }); + } + + /** Worker-local bounded input ranges; a borrowed vector is consumed before advancing. */ + static final class FinalWriteVectors implements AutoCloseable { + private final RandomAccessVectorValues source; + private final java.util.concurrent.ConcurrentLinkedQueue ranges = + new java.util.concurrent.ConcurrentLinkedQueue<>(); + private final ThreadLocal local = ThreadLocal.withInitial(() -> { + var range = new Range(); + ranges.add(range); + return range; + }); + + FinalWriteVectors(RandomAccessVectorValues source) { this.source = source; } + + VectorFloat get(int ordinal) { + Objects.checkIndex(ordinal, source.size()); + var range = local.get(); + if (range.cursor == null || ordinal != range.ordinal) { + if (range.cursor == null || ordinal != range.ordinal + 1 || ordinal == range.end) { + range.close(); + range.end = (int) Math.min(source.size(), (long) ordinal + 64); + range.cursor = VectorAccess.openRange(source, ordinal, range.end); + } + if (!range.cursor.next() || range.cursor.ordinal() != ordinal) + throw new IllegalStateException("Incomplete final-write input range"); + range.ordinal = ordinal; + } + return range.cursor.vector(); + } + + @Override public void close() { + // writer.write joins its workers before cursors and their read-ahead are released. + RuntimeException failure = null; + try { + for (var range : ranges) { + try { range.close(); } + catch (RuntimeException e) { + if (failure == null) failure = e; + else if (failure != e) failure.addSuppressed(e); + } + } + } finally { ranges.clear(); local.remove(); } + if (failure != null) throw failure; + } + + private static final class Range { + io.github.jbellis.jvector.graph.VectorCursor cursor; + int ordinal = -1, end; + void close() { + if (cursor != null) { + try { cursor.close(); } + finally { cursor = null; ordinal = -1; } + } + } + } + } + private static BuilderWithSuppliers builderWithSuppliers(Set features, ImmutableGraphIndex onHeapGraph, Path outPath, @@ -506,13 +588,18 @@ private static BuilderWithSuppliers builderWithSuppliers(Set features var identityMapper = new OrdinalMapper.IdentityMapper(floatVectors.size() - 1); var builder = new RandomAccessOnDiskGraphIndexWriter.Builder(onHeapGraph, outPath); builder.withMapper(identityMapper); + builder.withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()); 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 +615,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 +647,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 +1215,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 +1226,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 +1238,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/IPCService.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/IPCService.java index 3c125fd2b..92782c316 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/IPCService.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/IPCService.java @@ -23,7 +23,11 @@ import io.github.jbellis.jvector.graph.GraphSearcher; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.SearchResult; -import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; import io.github.jbellis.jvector.graph.similarity.ScoreFunction; import io.github.jbellis.jvector.quantization.CompressedVectors; @@ -194,7 +198,14 @@ private static ImmutableGraphIndex flushGraphIndex(ImmutableGraphIndex index, Ra var testDirectory = Files.createTempDirectory("BenchGraphDir"); var graphPath = testDirectory.resolve("graph.bin"); - OnDiskGraphIndex.write(index, ravv, graphPath); + var vectors = ravv.threadLocalSupplier(); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(index, graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new InlineVectors(ravv.dimension())) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + ordinal -> new InlineVectors.State(vectors.get().getVector(ordinal)))); + } return index; } catch (IOException e) { throw new IOError(e); 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/benchmarks/datasets/SiftSmall.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SiftSmall.java index 320b7f1a7..d029262cb 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SiftSmall.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SiftSmall.java @@ -34,7 +34,8 @@ import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; import io.github.jbellis.jvector.graph.disk.feature.NVQ; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; -import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.graph.disk.OrdinalMapper; import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; @@ -59,7 +60,6 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.List; -import java.util.Map; import java.util.Random; import java.util.concurrent.ThreadLocalRandom; import java.util.function.Function; @@ -155,7 +155,7 @@ public static void siftPersisted(List> baseVectors, List> baseVectors, List> baseVectors, List> baseVectors, List v = baseVectors.get(ordinal); // compress the new vector and add it to the PQVectors pqv.encodeAndSet(ordinal, v); - // write the full vector to disk - writer.writeInline(ordinal, Feature.singleState(FeatureId.INLINE_VECTORS, new InlineVectors.State(v))); - // now add it to the graph -- the previous steps must be completed first since the PQVectors - // and InlineVectorValues are both used during the search that runs as part of addGraphNode construction + // PQ is the resident construction scorer; FP is written after cleanup. builder.addGraphNode(ordinal, v); } @@ -253,8 +245,16 @@ public static void siftDiskAnnLTM(List> baseVectors, List new InlineVectors.State(vectors.get().getVector(ordinal)))); + } pqv.write(pqWriter); } @@ -292,11 +292,6 @@ public static void siftDiskAnnLTMWithNVQ(List> baseVectors, List< Path pqPath = Files.createTempFile("siftsmall", ".pq"); // Builder creation looks mostly the same try (GraphIndexBuilder builder = new GraphIndexBuilder(bsp, ravv.dimension(), 16, 100, 1.2f, 1.2f, false, true); - // explicit Writer for the first time, this is what's behind OnDiskGraphIndex.write - OnDiskGraphIndexWriter writer = new OnDiskGraphIndexWriter.Builder(builder.getGraph(), indexPath) - .with(new NVQ(nvq)) - .withMapper(new OrdinalMapper.IdentityMapper(baseVectors.size() - 1)) - .build(); // output for the compressed vectors IndexWriter pqWriter = new BufferedRandomAccessWriter(pqPath)) { @@ -305,10 +300,7 @@ public static void siftDiskAnnLTMWithNVQ(List> baseVectors, List< VectorFloat v = baseVectors.get(ordinal); // compress the new vector and add it to the PQVectors pqv.encodeAndSet(ordinal, v); - // write the full vector to disk - writer.writeInline(ordinal, Feature.singleState(FeatureId.NVQ_VECTORS, new NVQ.State(nvq.encode(v)))); - // now add it to the graph -- the previous steps must be completed first since the PQVectors - // and InlineVectorValues are both used during the search that runs as part of addGraphNode construction + // PQ is the resident construction scorer; NVQ is written after cleanup. builder.addGraphNode(ordinal, v); } @@ -316,8 +308,16 @@ public static void siftDiskAnnLTMWithNVQ(List> baseVectors, List< // that we don't need to worry about here builder.cleanup(); - // finish writing the index (by filling in the edge lists) and write our completed PQVectors - writer.write(Map.of()); + // Supply all inline features during final complete-record serialization. + var vectors = ravv.threadLocalSupplier(); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), indexPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new NVQ(nvq)) + .withMapper(new OrdinalMapper.IdentityMapper(baseVectors.size() - 1)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.NVQ_VECTORS, + ordinal -> new NVQ.State(nvq.encode(vectors.get().getVector(ordinal))))); + } pqv.write(pqWriter); } @@ -337,6 +337,17 @@ public static void siftDiskAnnLTMWithNVQ(List> baseVectors, List< } } + private static void writeCompleteGraph(ImmutableGraphIndex graph, RandomAccessVectorValues source, Path path) throws IOException { + var vectors = source.threadLocalSupplier(); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, path) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new InlineVectors(source.dimension())) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + ordinal -> new InlineVectors.State(vectors.get().getVector(ordinal)))); + } + } + // // Utilities and main() harness // diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/DiskIntro.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/DiskIntro.java index cfb70da09..6f238d81f 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/DiskIntro.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/DiskIntro.java @@ -34,7 +34,8 @@ import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.SearchResult; import io.github.jbellis.jvector.graph.disk.GraphIndexWriter; -import io.github.jbellis.jvector.graph.disk.GraphIndexWriterTypes; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; @@ -45,12 +46,13 @@ import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.types.VectorFloat; -// Provides the code used in disk-tutorial.md. -// If you edit this file you may also want to edit disk-tutoral.md +// Provides the code used in 2-disk-tutorial.md. +// Keep the tutorial and this example in sync. public class DiskIntro { public static void main(String[] args) throws IOException { // This is a preconfigured dataset that will be downloaded automatically. - DataSet dataset = DataSets.loadDataSet("ada002-100k").orElseThrow(() -> + String datasetName = args.length == 0 ? "ada002-100k" : args[0]; + DataSet dataset = DataSets.loadDataSet(datasetName).orElseThrow(() -> new RuntimeException("Dataset doesn't exist or wasn't configured correctly") ).getDataSet(); @@ -76,22 +78,25 @@ public static void main(String[] args) throws IOException { heapGraph = builder.build(ravv); } + var writerVectors = ravv.threadLocalSupplier(); Path graphPath = Files.createTempFile("jvector-example-graph", null); // or wherever you want to save the graph try ( // Create a writer for the on-heap graph we just built. - GraphIndexWriter writer = GraphIndexWriter.getBuilderFor(GraphIndexWriterTypes.RANDOM_ACCESS_PARALLEL, heapGraph, graphPath) + GraphIndexWriter writer = new OnDiskParallelGraphIndexWriter.Builder(heapGraph, graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) // Let the writer know that we'll also be passing in the actual vector data // to be saved "inline" with the data for each corresponding graph node. .with(new InlineVectors(dim)) .build(); ) { + // Supplying every inline feature enables complete-record batches (default 64). // Supply one map entry for each feature. // The key is a FeatureId enum corresponding to the feature // and the value is a function which generates the feature state for each graph node. writer.write(Map.of( FeatureId.INLINE_VECTORS, // we already have a RAVV, so we'll just use that to supply the writer. - nodeId -> new InlineVectors.State(ravv.getVector(nodeId)))); + nodeId -> new InlineVectors.State(writerVectors.get().getVector(nodeId)))); } // ReaderSupplierFactory automatically picks an available RandomAccessReader implementation diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/LargerThanMemory.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/LargerThanMemory.java index 5f22b1cbc..7afb05da0 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/LargerThanMemory.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/LargerThanMemory.java @@ -17,12 +17,10 @@ package io.github.jbellis.jvector.example.tutorial; import java.io.IOException; -import java.io.UncheckedIOException; import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; import java.util.List; -import java.util.Map; import java.util.stream.IntStream; import io.github.jbellis.jvector.disk.BufferedRandomAccessWriter; @@ -37,7 +35,7 @@ import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.SearchResult; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; -import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; import io.github.jbellis.jvector.graph.disk.OrdinalMapper; import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; @@ -60,11 +58,12 @@ public static void main(String[] args) throws IOException { // The DataSet provided by loadDataSet is in-memory, // but you can apply the same technique even when you don't have // the base vectors in-memory. - DataSet dataset = DataSets.loadDataSet("e5-small-v2-100k").orElseThrow(() -> + String datasetName = args.length == 0 ? "e5-small-v2-100k" : args[0]; + DataSet dataset = DataSets.loadDataSet(datasetName).orElseThrow(() -> new RuntimeException("Dataset doesn't exist or wasn't configured correctly") ).getDataSet(); - // Remember that RAVVs need not be in-memory in the general case. + // RAVVs need not be in-memory, but this source must remain readable through final writing. // We will sample from this RAVV to compute the PQ codebooks. // In general you don't need to have a RAVV over all the vectors to // build PQ codebooks, but you do need a "representative set". @@ -102,36 +101,21 @@ public static void main(String[] args) throws IOException { System.out.println("Building index incrementally..."); try ( - GraphIndexBuilder builder = new GraphIndexBuilder(bsp, dim, M, ef, overflow, alpha, addHierarchy); - // In DiskIntro we created the writer after generating the complete graph and closing the builder, - // but for incremental construction we will build and write in concert. - OnDiskGraphIndexWriter writer = new OnDiskGraphIndexWriter.Builder(builder.getGraph(), graphPath) - .with(new InlineVectors(dim)) - // Since we start with an empty graph, the writer will, by default, - // assume an ordinal mapping of size 0 (which is obviously incorrect). - // This is easy to rectify if you know the number of vectors beforehand, - // if not you may need to implement OrdinalMapper yourself. - .withMapper(new OrdinalMapper.IdentityMapper(ravv.size() - 1)) - .build(); + GraphIndexBuilder builder = new GraphIndexBuilder(bsp, dim, M, ef, overflow, alpha, addHierarchy) ) { + // Independent views also support sources that borrow a scratch vector. + var vectors = ravv.threadLocalSupplier(); // Graph building is best done with threads = number of physical cores // PhysicalCoreExecutor assumes hyperthreading by default, i.e. cores = vCPUs / 2 // If this is not correct, set the system property `jvector.physical_core_count` PhysicalCoreExecutor.pool().submit(() -> { IntStream.range(0, ravv.size()).parallel().forEach(ordinal -> { - VectorFloat v = ravv.getVector(ordinal); + VectorFloat v = vectors.get().getVector(ordinal); // Encode and add the vector to the working set of PQ vectors, // which allows the graph builder to access it through the BuildScoreProvder. pqVectors.encodeAndSet(ordinal, v); - // Write the feature (full-resolution vector) for a single vector instead of all at once. - try { - writer.writeFeaturesInline(ordinal, Feature.singleState(FeatureId.INLINE_VECTORS, new InlineVectors.State(v))); - } catch (IOException e) { - throw new UncheckedIOException(e); - } - builder.addGraphNode(ordinal, v); }); }).join(); @@ -140,9 +124,15 @@ public static void main(String[] args) throws IOException { // Enforces maximum degree constraint among other things. builder.cleanup(); - // No need to pass-in a feature supplier since we wrote the features incrementally - // using writer.writeInLine - writer.write(Map.of()); + // Reread this replayable source for final FP payloads; peers were resident PQ. + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new InlineVectors(dim)) + .withMapper(new OrdinalMapper.IdentityMapper(ravv.size() - 1)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + ordinal -> new InlineVectors.State(vectors.get().getVector(ordinal)))); + } } // PQ codebooks and vectors also need to be saved somewhere! (we'll not concern ourselves with Fused PQ) diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java index 2f0e9b1eb..661d29cb2 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java @@ -17,11 +17,9 @@ package io.github.jbellis.jvector.example.tutorial; import java.io.IOException; -import java.io.UncheckedIOException; import java.nio.file.Files; import java.nio.file.Path; import java.util.List; -import java.util.Map; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -33,7 +31,7 @@ import io.github.jbellis.jvector.graph.ImmutableGraphIndex; import io.github.jbellis.jvector.graph.SearchResult; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; -import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; import io.github.jbellis.jvector.graph.disk.OrdinalMapper; import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; @@ -52,7 +50,8 @@ public class NvqExample { public static void main(String[] args) throws IOException { // Load a preconfigured dataset - var ds = DataSets.loadDataSet("ada002-100k").orElseThrow(() -> + String datasetName = args.length == 0 ? "ada002-100k" : args[0]; + var ds = DataSets.loadDataSet(datasetName).orElseThrow(() -> new RuntimeException("dataset not found")) .getDataSet(); var dim = ds.getDimension(); @@ -62,7 +61,7 @@ public static void main(String[] args) throws IOException { var numSubVectors = 2; // Setup NVQ parameters. - // The base vectors RAVV instance is used only for computing the global mean + // Compute NVQ statistics; the source remains readable for final encoding. var nvq = NVQuantization.compute(base, numSubVectors); // Use this method instead if you don't have all the vectors up-front but can estimate the mean // var nvq = NVQuantization.create(scaledGlobalMean, numSubVectors); @@ -88,16 +87,11 @@ public static void main(String[] args) throws IOException { Path graphPath = Files.createTempFile("jvector-nvq-graph", null); - System.out.println("Building graph in streaming mode..."); + var vectors = base.threadLocalSupplier(); + System.out.println("Building graph incrementally with resident PQ vectors..."); try ( // Create the graph builder using PQ-based scoring var builder = new GraphIndexBuilder(bsp, dim, M, ef, nOv, alpha, addHierarchy); - // Create the on-disk writer configured with NVQ feature - // This allows us to write both the graph structure and NVQ-compressed vectors - var writer = new OnDiskGraphIndexWriter.Builder(builder.getGraph(), graphPath) - .with(new NVQ(nvq)) - .withMapper(new OrdinalMapper.IdentityMapper(base.size() - 1)) - .build(); var pb = new ProgressBar("Build graph", base.size()); ) { @@ -105,22 +99,12 @@ public static void main(String[] args) throws IOException { IntStream.range(0, base.size()) .parallel() .forEach(ordinal -> { - var vec = base.getVector(ordinal); + var vec = vectors.get().getVector(ordinal); // Encode the PQ vector first, then add the graph node pqv.encodeAndSet(ordinal, vec); builder.addGraphNode(ordinal, vec); - // Encode and write NVQ vectors for later re-ranking - var nvqVec = nvq.encode(vec); - Map featureMap = Map.of( - FeatureId.NVQ_VECTORS, new NVQ.State(nvqVec) - ); - try { - writer.writeFeaturesInline(ordinal, featureMap); - } catch (IOException e) { - throw new UncheckedIOException(e); - } pb.step(); }); }).join(); @@ -129,7 +113,16 @@ public static void main(String[] args) throws IOException { // cleanup System.out.println("Cleanup..."); builder.cleanup(); - writer.write(Map.of()); + // NVQ is reranking data; encode it while writing complete records after cleanup. + System.out.println("Writing graph and NVQ vectors..."); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new NVQ(nvq)) + .withMapper(new OrdinalMapper.IdentityMapper(base.size() - 1)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.NVQ_VECTORS, + ordinal -> new NVQ.State(nvq.encode(vectors.get().getVector(ordinal))))); + } } // Search parameters 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 6e56a6e26..dc59696b1 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; private static final String defaultDirectory = "jvector-examples/yaml-configs/index-parameters/"; private static final java.util.regex.Pattern YAML_SCHEMA_VERSION_KEY = diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java index f3728234c..eb3c3cce9 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java @@ -31,6 +31,9 @@ import io.github.jbellis.jvector.quantization.NVQuantization; import io.github.jbellis.jvector.quantization.PQVectors; import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.util.ExplicitThreadLocal; +import io.github.jbellis.jvector.util.ExceptionUtils; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import java.io.IOException; import java.nio.file.Files; @@ -54,9 +57,9 @@ * writer.write(featureSuppliers); * * // Parallel: - * var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath) - * .with(inlineVectors) - * .withParallelWrites(true) // Enable parallel writes + * var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, outputPath) + * .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + * .with(inlineVectors) // Supply every inline feature for complete-record batches * .build(); * writer.write(featureSuppliers); * @@ -193,8 +196,8 @@ private static void verifyIndicesIdentical(OnDiskGraphIndex index1, OnDiskGraphI } // Compare vectors if present (only check layer 0) - if (index1.getFeatureSet().contains(FeatureId.INLINE_VECTORS) || - index1.getFeatureSet().contains(FeatureId.NVQ_VECTORS)) { + // NVQ and other feature payloads are checked by verifySerializedBytes below. + if (index1.getFeatureSet().contains(FeatureId.INLINE_VECTORS)) { int vectorsChecked = 0; int maxToCheck = Math.min(100, index1.size(0)); // Check up to 100 vectors as a sample @@ -240,9 +243,10 @@ public static void benchmarkComparison(ImmutableGraphIndex graph, var nvqFeature = new NVQ(nvq); var fusedPQFeature = new FusedPQ(graph.maxDegree(), pq); + var vectors = floatVectors.threadLocalSupplier(); // Build suppliers for inline features (NVQ only - FUSED_ADC needs neighbors) Map> inlineSuppliers = new EnumMap<>(FeatureId.class); - inlineSuppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(floatVectors.getVector(ordinal)))); + inlineSuppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(vectors.get().getVector(ordinal)))); // FUSED_ADC supplier needs graph view, provided at write time var identityMapper = new OrdinalMapper.IdentityMapper(floatVectors.size() - 1); @@ -254,15 +258,16 @@ public static void benchmarkComparison(ImmutableGraphIndex graph, .with(nvqFeature) .with(fusedPQFeature) .withMapper(identityMapper) - .build()) { + .build(); + var views = ExplicitThreadLocal.withInitial(graph::getView)) { - var view = graph.getView(); Map> writeSuppliers = new EnumMap<>(FeatureId.class); writeSuppliers.put(FeatureId.NVQ_VECTORS, inlineSuppliers.get(FeatureId.NVQ_VECTORS)); - writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pqVectors, ordinal)); + writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(views.get(), pqVectors, ordinal)); writer.write(writeSuppliers); - view.close(); + } catch (Exception e) { + ExceptionUtils.throwIoException(e); } long sequentialTime = System.nanoTime() - sequentialStart; System.out.printf("Sequential write: %.2f ms%n", sequentialTime / 1_000_000.0); @@ -270,18 +275,20 @@ public static void benchmarkComparison(ImmutableGraphIndex graph, // Parallel write long parallelStart = System.nanoTime(); try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, parallelPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) .with(nvqFeature) .with(fusedPQFeature) .withMapper(identityMapper) - .build()) { + .build(); + var views = ExplicitThreadLocal.withInitial(graph::getView)) { - var view = graph.getView(); Map> writeSuppliers = new EnumMap<>(FeatureId.class); writeSuppliers.put(FeatureId.NVQ_VECTORS, inlineSuppliers.get(FeatureId.NVQ_VECTORS)); - writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pqVectors, ordinal)); + writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(views.get(), pqVectors, ordinal)); writer.write(writeSuppliers); - view.close(); + } catch (Exception e) { + ExceptionUtils.throwIoException(e); } long parallelTime = System.nanoTime() - parallelStart; @@ -289,6 +296,23 @@ public static void benchmarkComparison(ImmutableGraphIndex graph, System.out.printf("Speedup: %.2fx%n", (double) sequentialTime / parallelTime); } + private static void verifySerializedBytes(Path sequential, Path parallel) throws IOException { + try (var left = Files.newInputStream(sequential); var right = Files.newInputStream(parallel)) { + byte[] a = new byte[64 * 1024], b = new byte[a.length]; + long offset = 0; + while (true) { + int n = left.readNBytes(a, 0, a.length); + int m = right.readNBytes(b, 0, b.length); + if (n != m) throw new AssertionError("Serialized index lengths differ"); + if (n == 0) break; + int mismatch = Arrays.mismatch(a, 0, n, b, 0, n); + if (mismatch >= 0) throw new AssertionError("Serialized indexes differ at byte " + (offset + mismatch)); + offset += n; + } + } + System.out.println("All serialized features match byte for byte."); + } + /** * Main method to run a benchmark test of sequential vs parallel writes. * @@ -365,6 +389,7 @@ public static void main(String[] args) throws IOException { // Verify that both indices are identical verifyIndicesIdentical(sequentialIndex, parallelIndex); + verifySerializedBytes(sequentialPath, parallelPath); // Close the loaded indices sequentialIndex.close(); 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..d94ae48a9 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetPrefetching.java @@ -0,0 +1,151 @@ +/* + * 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 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/example/TestFinalWriteVectors.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestFinalWriteVectors.java new file mode 100644 index 000000000..5c35708b1 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestFinalWriteVectors.java @@ -0,0 +1,152 @@ +/* + * 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.graph.*; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; +import java.util.*; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicInteger; +import static org.junit.Assert.*; + +public class TestFinalWriteVectors { + private static ListRandomAccessVectorValues values() { + var vectors = new ArrayList>(); + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + for (int i = 0; i < 131; i++) vectors.add(vts.createFloatVector(new float[]{i, -i, i * 2})); + return new ListRandomAccessVectorValues(vectors, 3); + } + + @Test + public void rangesAreBoundedAndReleased() { + var resident = values(); + var opens = new AtomicInteger(); + var closes = new 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) { throw new AssertionError("Point access"); } + public VectorCursor openSelection(int[] ordinals, int offset, int count) { + throw new AssertionError("Selection access"); + } + public VectorCursor openRange(int start, int end) { + assertTrue(end - start <= 64); + opens.incrementAndGet(); + var cursor = VectorAccess.openRange(resident, start, end); + return new VectorCursor() { + boolean closed; + 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() { + if (!closed) { closed = true; closes.incrementAndGet(); cursor.close(); } + } + }; + } + }; + try (var input = new Grid.FinalWriteVectors(source)) { + for (int i = 0; i < source.size(); i++) { + var v = input.get(i); + assertEquals(i, v.get(0), 0); + assertSame("Features for one record reuse the same input", v, input.get(i)); + } + assertEquals(3, opens.get()); + assertEquals(17, input.get(17).get(0), 0); + assertEquals(0, input.get(0).get(0), 0); + } + assertEquals(opens.get(), closes.get()); + } + + @Test(timeout = 60000) + public void ordinaryBorrowedSourcesHaveIndependentWorkerViews() throws Exception { + var resident = values(); + class Borrowed implements RandomAccessVectorValues { + final VectorFloat scratch = VectorizationProvider.getInstance().getVectorTypeSupport().createFloatVector(3); + public int size() { return resident.size(); } + public int dimension() { return 3; } + public boolean isValueShared() { return true; } + public RandomAccessVectorValues copy() { return new Borrowed(); } + public VectorFloat getVector(int ordinal) { + resident.getVectorInto(ordinal, scratch, 0); + return scratch; + } + } + var pool = Executors.newFixedThreadPool(4); + try (var input = new Grid.FinalWriteVectors(new Borrowed())) { + var tasks = new ArrayList>(); + for (int worker = 0; worker < 8; worker++) tasks.add(() -> { + for (int ordinal = 0; ordinal < resident.size(); ordinal++) { + var vector = input.get(ordinal); + Thread.yield(); + assertEquals(ordinal, vector.get(0), 0); + assertEquals(-ordinal, vector.get(1), 0); + assertSame(vector, input.get(ordinal)); + } + return null; + }); + for (var future : pool.invokeAll(tasks)) future.get(); + } finally { pool.shutdownNow(); } + } + + @Test(timeout = 60000) + public void closeReleasesOtherWorkersWhenOneCursorFails() throws Exception { + var resident = values(); + var closes = new 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) { throw new AssertionError("Point access"); } + public VectorCursor openSelection(int[] ordinals, int offset, int count) { throw new AssertionError("Selection 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 false; } + public void close() { + closes.incrementAndGet(); cursor.close(); + if (start == 0) throw new IllegalStateException("cursor close failure"); + } + }; + } + }; + var pool = Executors.newFixedThreadPool(2); + var input = new Grid.FinalWriteVectors(source); + try { + var barrier = new CyclicBarrier(2); + for (var future : pool.invokeAll(List.>of( + () -> { input.get(0); barrier.await(10, TimeUnit.SECONDS); return null; }, + () -> { input.get(64); barrier.await(10, TimeUnit.SECONDS); return null; }))) future.get(); + var failure = assertThrows(IllegalStateException.class, input::close); + assertEquals("cursor close failure", failure.getMessage()); + assertEquals(2, closes.get()); + input.close(); + assertEquals(2, closes.get()); + } finally { + pool.shutdownNow(); pool.awaitTermination(10, TimeUnit.SECONDS); + input.close(); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestGridDeferredFeatureParallelism.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestGridDeferredFeatureParallelism.java new file mode 100644 index 000000000..39522b8ff --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestGridDeferredFeatureParallelism.java @@ -0,0 +1,184 @@ +/* + * 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.TestUtil; +import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; +import io.github.jbellis.jvector.example.reporting.RunArtifacts; +import io.github.jbellis.jvector.example.util.CompressorParameters; +import io.github.jbellis.jvector.example.util.OnDiskGraphIndexCache; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.management.GraphIndexBuilderConfig; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorUtil; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Assume; +import org.junit.Test; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.EnumSet; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.Assert.assertTrue; + +/// Repro for PR #747: with the default writer, Grid serializes all deferred feature work on +/// one thread. +/// +/// ## The bug +/// +/// Before #747, Grid encoded NVQ (and copied full-precision vectors) for each node during +/// insertion, on the parallel insertion threads, and wrote them with `writeInline`. #747 moves +/// that work into the final `writer.write(suppliers)` call for every Grid run. +/// +/// Only the parallel writer calls final suppliers from several threads. It is selected only when +/// `GraphIndexBuilderConfig.isParallelBuild()` is true. The default is `false`, and BenchYAML +/// offers no way to change it, so a plain BenchYAML run (including CI) uses the sequential +/// `OnDiskGraphIndexWriter`. Its `writeL0Records` calls each supplier in turn, on the calling +/// thread. Every vector read and NVQ encode for the whole dataset then runs on one thread after +/// construction. +/// +/// On Cohere-1M (preload, native, 48 workers), build+write went from 74–76 s at `4dcfe917` +/// to 100–104 s at `ec42a94b` with the default writer, about 35% slower. With the parallel writer +/// enabled, #747 is faster than before (78 s against 88–92 s). But only the PR's measurements enable it. +/// +/// ## The test +/// +/// Records the thread of every base-vector read made during final serialization, meaning with a +/// frame from the graph-writer package (`io.github.jbellis.jvector.graph.disk`) on the reading +/// thread's stack. That covers the sequential writer's `writeL0Records` and the parallel writer's +/// `NodeRecordTask`s. No other read in a Grid run goes through that package. The test runs Grid +/// twice on a resident dataset. +/// +/// 1. **Control**, with the parallel writer (`parallelBuild = true`). Grid defers features with +/// this writer, so the final write must read base vectors, from more than one writer thread. +/// This proves the detector sees the final write. If the writers move package or the final +/// write stops reading through the dataset's vectors, the control fails instead of the check +/// below passing vacuously. +/// 2. **Requirement**, with BenchYAML's default writer selection (`parallelBuild = false`). +/// Final-write input must not be confined to one thread. Either no vectors are read during +/// the final write (features were produced during insertion), or they are read from more than +/// one thread. +/// +/// At `ec42a94b` the control passes, and every read in the default run's final write comes from +/// one thread. The control needs at least two writer threads, so it's skipped on a machine with +/// one physical core. +public class TestGridDeferredFeatureParallelism { + private static final int SIZE = 4096, DIMENSIONS = 17; + + /// Package of the graph writers: the sequential writer's `writeL0Records`, and the parallel + /// writer's `NodeRecordTask`s, which run on worker threads whose stacks lack `writeL0Records`. + private static final String WRITER_PACKAGE = "io.github.jbellis.jvector.graph.disk."; + + /// Resident vectors that record which threads read them during final serialization, meaning + /// with a graph-writer frame on the reading thread's stack. + private static final class RecordingVectors implements RandomAccessVectorValues { + final RandomAccessVectorValues resident; + final Set writeThreads = ConcurrentHashMap.newKeySet(); + final AtomicInteger writeReads = new AtomicInteger(); + + RecordingVectors(RandomAccessVectorValues resident) { this.resident = resident; } + + @Override public int size() { return resident.size(); } + @Override public int dimension() { return resident.dimension(); } + @Override public boolean isValueShared() { return false; } + @Override public RandomAccessVectorValues copy() { return this; } + + @Override public VectorFloat getVector(int ordinal) { + boolean inFinalWrite = StackWalker.getInstance().walk(frames -> + frames.anyMatch(f -> f.getClassName().startsWith(WRITER_PACKAGE))); + if (inFinalWrite) { + writeReads.incrementAndGet(); + writeThreads.add(Thread.currentThread().getName()); + } + return resident.getVector(ordinal); + } + } + + /// Run Grid once on `vectors` with the given writer selection, restoring the selection and + /// removing the run's work directory and cached codebook afterwards. + private static void runGrid(List> list, RecordingVectors vectors, boolean parallel) + throws Exception { + String name = "deferred-input-" + UUID.randomUUID(); + DataSet ds = new DataSet() { + public int getDimension() { return DIMENSIONS; } + public String getName() { return name; } + public VectorSimilarityFunction getSimilarityFunction() { return VectorSimilarityFunction.DOT_PRODUCT; } + public RandomAccessVectorValues getBaseRavv() { return vectors; } + public List> getBaseVectors() { return list; } + public List> getQueryVectors() { return List.of(list.get(0)); } + public List> getGroundTruth() { return List.of(List.of(0)); } + }; + var config = GraphIndexBuilderConfig.getInstance(); + boolean previous = config.isParallelBuild(); + var work = Files.createTempDirectory("grid-deferred-input"); + var pqParameters = new CompressorParameters.PQParameters(4, 16, false, -1f); + try { + config.setParallelBuild(parallel); + Grid.runOneGraph(OnDiskGraphIndexCache.initialize(work.resolve("saved")), + List.of(EnumSet.of(FeatureId.NVQ_VECTORS)), 8, 32, 1.2f, true, false, + ignored -> pqParameters, List.of(), Map.of(10, List.of(1.0)), List.of(false), + RunArtifacts.disabled(), ds, work); + } finally { + config.setParallelBuild(previous); + TestUtil.deleteQuietly(work); + // Grid caches codebooks by dataset name in the working directory; remove this run's. + Files.deleteIfExists(Path.of("pq_cache").resolve(pqParameters.idStringFor(ds))); + } + } + + @Test(timeout = 120_000) + public void defaultWriterDoesNotConfineDeferredFeatureInputToOneThread() throws Exception { + Assume.assumeTrue("the control needs at least two writer threads", + PhysicalCoreExecutor.getPhysicalCoreCount() > 1); + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var random = new Random(42); + var list = new ArrayList>(); + for (int n = 0; n < SIZE; n++) { + var v = vts.createFloatVector(DIMENSIONS); + for (int d = 0; d < DIMENSIONS; d++) v.set(d, random.nextFloat() - 0.5f); + VectorUtil.l2normalize(v); + list.add(v); + } + var resident = new ListRandomAccessVectorValues(list, DIMENSIONS); + + var control = new RecordingVectors(resident); + runGrid(list, control, true); + assertTrue("control: the parallel writer's final write read " + control.writeReads.get() + + " base vectors on threads " + control.writeThreads + + "; the detector no longer sees the final write", + control.writeReads.get() > 0 && control.writeThreads.size() > 1); + + var vectors = new RecordingVectors(resident); + runGrid(list, vectors, false); // BenchYAML's default writer selection + int reads = vectors.writeReads.get(); + assertTrue("final serialization read " + reads + " base vectors, all on one thread " + + vectors.writeThreads + ", with the default (sequential) writer", + reads == 0 || vectors.writeThreads.size() > 1); + } +} 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..fe171b872 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBufferedGraphCompatibility.java @@ -0,0 +1,158 @@ +/* + * 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 { + /// Every level's adjacency, one "level node: neighbors" line per node in node order. + private static List adjacency(ImmutableGraphIndex graph) throws Exception { + var lines = new ArrayList(); + try (var view = graph.getView()) { + for (int level = 0; level <= graph.getMaxLevel(); level++) { + var nodes = new ArrayList(); + for (var it = graph.getNodes(level); it.hasNext(); ) nodes.add(it.nextInt()); + Collections.sort(nodes); + for (int node : nodes) { + var neighbors = new StringBuilder(); + for (var it = view.getNeighborsIterator(level, node); it.hasNext(); ) neighbors.append(it.nextInt()).append(' '); + lines.add(level + " " + node + ": " + neighbors); + } + } + } + return lines; + } + + @Test(timeout = 180000) + public void exactPqAndBinaryBuildsWithAllMetricsAndRefinementModes() throws Exception { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var pool = new ForkJoinPool(3); + // One worker makes insertion order, level draws and cleanup order reproducible. + var single = new ForkJoinPool(1); + 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); + var residentBsp = mode == 0 ? BuildScoreProvider.randomAccessScoreProvider(resident, metric) : bsp; + List residentAdjacency; + try (var reference = new GraphIndexBuilder(residentBsp, dimension, 8, 64, 1.2f, 1.2f, + true, refine, single, single)) { + residentAdjacency = adjacency(reference.build(resident)); + } + try (var builder = new GraphIndexBuilder(bsp, dimension, 8, 64, 1.2f, 1.2f, + true, refine, single, single)) { + var graph = builder.build(source); + assertEquals(count, graph.size(0)); + // Deterministic builds: the file-backed source must give the resident graph. + assertEquals("dimension=" + dimension + ", build=" + (mode == 0 ? "exact" : mode == 1 ? "PQ" : "BQ") + + ", metric=" + metric + ", refine=" + refine, + residentAdjacency, adjacency(graph)); + // 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(); single.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/graph/disk/TestDeferredGraphWriter.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestDeferredGraphWriter.java new file mode 100644 index 000000000..3951b4f3e --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestDeferredGraphWriter.java @@ -0,0 +1,192 @@ +/* + * 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.disk; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.LuceneTestCase; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.disk.BufferedRandomAccessWriter; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.feature.*; +import io.github.jbellis.jvector.management.GraphIndexBuilderConfig; +import io.github.jbellis.jvector.quantization.NVQuantization; +import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import org.junit.Test; + +import java.nio.file.Files; +import java.util.*; +import java.util.concurrent.*; +import java.util.function.IntFunction; + +import static org.junit.Assert.*; + +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class TestDeferredGraphWriter extends LuceneTestCase { + private ListRandomAccessVectorValues values() { + var vectors = new ArrayList>(); + var support = VectorizationProvider.getInstance().getVectorTypeSupport(); + var random = new Random(42); + for (int n = 0; n < 259; n++) { + var vector = support.createFloatVector(17); + for (int d = 0; d < 17; d++) vector.set(d, random.nextFloat() - 0.5f); + vectors.add(vector); + } + return new ListRandomAccessVectorValues(vectors, 17); + } + + @Test + public void completeAndPrewrittenFeaturesMatchSequentialV6() throws Exception { + var directory = Files.createTempDirectory("deferred-writer"); + var values = values(); + try (var builder = new GraphIndexBuilder(values, VectorSimilarityFunction.DOT_PRODUCT, + 8, 32, 1.2f, 1.2f, true)) { + var graph = TestUtil.buildSequentially(builder, values); + var pq = ProductQuantization.compute(values, 3, 256, false); + var codes = pq.encodeAll(values, ForkJoinPool.commonPool()); + var nvq = NVQuantization.compute(values, 2); + var frozenView = graph.getView(); + try { + for (boolean holes : new boolean[]{false, true}) { + OrdinalMapper mapper = new OrdinalMapper() { + public int maxOrdinal() { return holes ? values.size() * 2 - 1 : values.size() - 1; } + public int oldToNew(int n) { return (values.size() - 1 - n) * (holes ? 2 : 1); } + public int newToOld(int n) { return holes && n % 2 != 0 ? OMITTED : values.size() - 1 - n / (holes ? 2 : 1); } + }; + for (boolean quantized : new boolean[]{false, true}) { + Feature vectors = quantized ? new NVQ(nvq) : new InlineVectors(values.dimension()); + var fused = new FusedPQ(graph.maxDegree(), pq); + Map> suppliers = new EnumMap<>(FeatureId.class); + suppliers.put(vectors.id(), n -> quantized ? new NVQ.State(nvq.encode(values.getVector(n))) + : new InlineVectors.State(values.getVector(n))); + suppliers.put(fused.id(), n -> new FusedPQ.State(frozenView, codes, n)); + byte[] prefix = new byte[83]; + Arrays.fill(prefix, (byte) 0x5a); + var sequential = directory.resolve("sequential"); + Files.write(sequential, prefix); + try (var writer = new OnDiskGraphIndexWriter.Builder(graph, sequential) + .withStartOffset(prefix.length).withVersion(6).withMapper(mapper) + .with(vectors).with(fused).build()) { + writer.write(suppliers); + } + byte[] expected = Files.readAllBytes(sequential); + for (int batch : new int[]{1, 7, 64, 128}) { + for (boolean direct : new boolean[]{false, true}) { + for (boolean prewritten : new boolean[]{false, true}) { + var parallel = directory.resolve("parallel"); + Files.write(parallel, prefix); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, parallel) + .withParallelWorkerThreads(3).withParallelBatchSize(batch) + .withParallelDirectBuffers(direct).withStartOffset(prefix.length) + .withVersion(6).withMapper(mapper).with(vectors).with(fused).build()) { + if (prewritten) { + for (int ordinal = 0; ordinal < values.size(); ordinal++) { + writer.writeFeaturesInline(mapper.oldToNew(ordinal), + Map.of(vectors.id(), suppliers.get(vectors.id()).apply(ordinal))); + } + writer.write(Map.of(fused.id(), suppliers.get(fused.id()))); + } else writer.write(suppliers); + } + assertArrayEquals("holes=" + holes + " nvq=" + quantized + " batch=" + batch + + " direct=" + direct + " prewritten=" + prewritten, + expected, Files.readAllBytes(parallel)); + } + } + } + } + } + } finally { frozenView.close(); } + } finally { TestUtil.deleteQuietly(directory); } + } + + @Test(timeout = 60000) + public void failedWriteJoinsWorkersAndPreservesExternalExecutor() throws Exception { + var directory = Files.createTempDirectory("failed-deferred-writer"); + var values = values(); + var workers = Executors.newFixedThreadPool(2); + var caller = Executors.newSingleThreadExecutor(); + var release = new CountDownLatch(1); + try (var builder = new GraphIndexBuilder(values, VectorSimilarityFunction.DOT_PRODUCT, + 8, 32, 1.2f, 1.2f, false)) { + var graph = TestUtil.buildSequentially(builder, values); + var file = directory.resolve("records"); + Files.createFile(file); + var blocked = new CountDownLatch(1); + var failed = new CountDownLatch(1); + try (var out = new BufferedRandomAccessWriter(file); + var writer = new ParallelGraphWriter(out, graph, List.of(new InlineVectors(17)), + new ParallelGraphWriter.Config(2, false, 1, 7), file, workers)) { + var writing = caller.submit(() -> { + writer.writeL0Records(new OrdinalMapper.IdentityMapper(values.size() - 1), + List.of(new InlineVectors(17)), Feature.singleStateFactory(FeatureId.INLINE_VECTORS, n -> { + if (n == 0) { failed.countDown(); throw new IllegalStateException("supplier failure"); } + blocked.countDown(); + try { if (!release.await(10, TimeUnit.SECONDS)) throw new AssertionError("Writer cleanup timeout"); } + catch (InterruptedException e) { throw new AssertionError(e); } + return new InlineVectors.State(values.getVector(n)); + }), 0); + return null; + }); + assertTrue(failed.await(10, TimeUnit.SECONDS)); + assertTrue(blocked.await(10, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> writing.get(100, TimeUnit.MILLISECONDS)); + release.countDown(); + var failure = assertThrows(ExecutionException.class, () -> writing.get(10, TimeUnit.SECONDS)); + assertTrue(failure.getCause() instanceof IllegalStateException); + } + assertEquals(42, workers.submit(() -> 42).get().intValue()); + } finally { + release.countDown(); + caller.shutdownNow(); workers.shutdownNow(); + caller.awaitTermination(10, TimeUnit.SECONDS); workers.awaitTermination(10, TimeUnit.SECONDS); + TestUtil.deleteQuietly(directory); + } + } + + @Test + public void batchSizeValidationAndUnifiedBuilderSelection() throws Exception { + var directory = Files.createTempDirectory("writer-batch-option"); + var values = values(); + var settings = GraphIndexBuilderConfig.getInstance(); + boolean wasParallel = settings.isParallelBuild(); + try (var graphBuilder = new GraphIndexBuilder(values, VectorSimilarityFunction.DOT_PRODUCT, + 8, 32, 1.2f, 1.2f, false)) { + var graph = TestUtil.buildSequentially(graphBuilder, values); + assertEquals(64, ParallelGraphWriter.Config.defaultConfig().recordsPerBatch); + assertThrows(IllegalArgumentException.class, () -> new ParallelGraphWriter.Config(1, false, 1, 0)); + for (boolean parallel : new boolean[]{false, true}) { + settings.setParallelBuild(parallel); + var builder = new RandomAccessOnDiskGraphIndexWriter.Builder(graph, directory.resolve("unified")); + assertThrows(IllegalArgumentException.class, () -> builder.withParallelBatchSize(-1)); + try (var writer = builder.withParallelBatchSize(7).withParallelWorkerThreads(2) + .with(new InlineVectors(17)).build()) { + assertEquals(parallel, writer instanceof OnDiskParallelGraphIndexWriter); + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + n -> new InlineVectors.State(values.getVector(n)))); + } + } + var direct = new OnDiskParallelGraphIndexWriter.Builder(graph, directory.resolve("direct")); + assertThrows(IllegalArgumentException.class, () -> direct.withParallelBatchSize(0)); + try (var writer = direct.with(new InlineVectors(17)).build()) { } + } finally { + settings.setParallelBuild(wasParallel); + TestUtil.deleteQuietly(directory); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestNodeRecordTask.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestNodeRecordTask.java new file mode 100644 index 000000000..e21e888de --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestNodeRecordTask.java @@ -0,0 +1,125 @@ +/* + * 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.disk; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.LuceneTestCase; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.AsynchronousFileChannel; +import java.nio.file.Files; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class TestNodeRecordTask extends LuceneTestCase { + @Test + public void testCompleteRecordBatches() throws Exception { + var root = Files.createTempDirectory("node-record-batches"); + var values = new ArrayList>(); + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + for (int n = 0; n < 131; n++) { + var vector = vts.createFloatVector(17); + for (int d = 0; d < 17; d++) { + vector.set(d, (float) Math.sin(n * 17 + d)); + } + values.add(vector); + } + var ravv = new ListRandomAccessVectorValues(values, 17); + try (var builder = new GraphIndexBuilder(ravv, VectorSimilarityFunction.DOT_PRODUCT, + 4, 10, 1.2f, 1.2f, false)) { + var graph = TestUtil.buildSequentially(builder, ravv); + var feature = new InlineVectors(17); + for (boolean holes : new boolean[]{false, true}) { + int count = holes ? 262 : 131; + OrdinalMapper mapper = holes ? new OrdinalMapper() { + public int maxOrdinal() { return 261; } + public int oldToNew(int n) { return 2 * n; } + public int newToOld(int n) { return n % 2 == 0 ? n / 2 : OMITTED; } + } : new OrdinalMapper.IdentityMapper(130); + int payload = Integer.BYTES * (2 + graph.getDegree(0)) + feature.featureSize(); + for (int stride : new int[]{payload, 4096}) { + for (boolean direct : new boolean[]{false, true}) { + var target = root.resolve("records-" + holes + "-" + stride + "-" + direct); + // Preserve bytes before L0 and force full batches plus a partial tail. + final int baseOffset = 83; + byte[] prefix = new byte[baseOffset]; + Arrays.fill(prefix, (byte) 0x5a); + Files.write(target, prefix); + var buffer = direct ? ByteBuffer.allocateDirect(stride * 64) + : ByteBuffer.allocate(stride * 64); + buffer.order(ByteOrder.BIG_ENDIAN); + try (var view = graph.getView(); + var channel = AsynchronousFileChannel.open(target, StandardOpenOption.WRITE)) { + var task = new NodeRecordTask(0, count, mapper, graph, view, List.of(feature), + Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + n -> new InlineVectors.State(values.get(n))), + stride, baseOffset, channel, buffer); + task.call(); + } + byte[] actual = Files.readAllBytes(target); + assertEquals(baseOffset + count * stride, actual.length); + assertArrayEquals(prefix, Arrays.copyOf(actual, baseOffset)); + try (var view = graph.getView()) { + for (int n = 0; n < count; n++) { + var expected = ByteBuffer.allocate(stride).order(ByteOrder.BIG_ENDIAN); + expected.putInt(mapper.newToOld(n) == OrdinalMapper.OMITTED ? -1 : n); + int old = mapper.newToOld(n); + if (old == OrdinalMapper.OMITTED) { + expected.position(expected.position() + feature.featureSize()); + expected.putInt(0); + for (int d = 0; d < graph.getDegree(0); d++) expected.putInt(-1); + } else { + for (int d = 0; d < 17; d++) expected.putFloat(values.get(old).get(d)); + var neighbors = view.getNeighborsIterator(0, old); + expected.putInt(neighbors.size()); + int d = 0; + while (neighbors.hasNext()) { + expected.putInt(mapper.oldToNew(neighbors.nextInt())); + d++; + } + for (; d < graph.getDegree(0); d++) expected.putInt(-1); + } + assertEquals("Record " + n + " holes=" + holes + " direct=" + direct, -1, + Arrays.mismatch(actual, baseOffset + n * stride, + baseOffset + (n + 1) * stride, expected.array(), 0, stride)); + } + } + } + } + } + } finally { + TestUtil.deleteQuietly(root); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestParallelWriterInterrupt.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestParallelWriterInterrupt.java new file mode 100644 index 000000000..b1db13976 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestParallelWriterInterrupt.java @@ -0,0 +1,131 @@ +/* + * 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.disk; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.LuceneTestCase; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Random; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +/// Repro for PR #747: an interrupted parallel write keeps running until every record is written. +/// +/// ## The bug +/// +/// `ParallelGraphWriter.writeL0Records` submits one task per ordinal range, then +/// `awaitTasks` waits on each future in turn. When the waiting thread is interrupted, +/// `awaitTasks` records the interruption and calls `future.get()` again. When a task fails, it +/// keeps the failure and moves on to the next future. Either way it waits for every task. No +/// task is cancelled, so not-yet-started ranges still run and every running supplier runs to +/// the end of its range. +/// +/// Joining before the channel and suppliers are released is right. But nothing stops the work +/// being joined. So an interrupted or already-failed write runs the rest of the final +/// serialization, encoding every remaining record, and only then reports the interruption or +/// failure. For a 10M-node graph that's the whole final write. Before #747, `writeL0Records` +/// threw on interrupt immediately, without joining. +/// +/// ## The test +/// +/// Writes a 6,000-node graph with two writer workers and an `InlineVectors` supplier that takes +/// 1 ms per record, so an uninterrupted write takes about 3 s. After the first 50 records it +/// interrupts the writing thread, then expects the write to fail within 1 s and to have supplied +/// fewer than all records. At `ec42a94b` the write runs to completion, 3.2 s after the +/// interrupt, and supplies every record. +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class TestParallelWriterInterrupt extends LuceneTestCase { + private static final int SIZE = 6000, DIMENSIONS = 17, WORKERS = 2; + private static final long MAX_STOP_MILLIS = 1000; + + @Test(timeout = 120_000) + public void interruptedWriteStopsPromptly() throws Exception { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var random = new Random(42); + var list = new ArrayList>(); + for (int n = 0; n < SIZE; n++) { + var v = vts.createFloatVector(DIMENSIONS); + for (int d = 0; d < DIMENSIONS; d++) v.set(d, random.nextFloat() - 0.5f); + list.add(v); + } + var values = new ListRandomAccessVectorValues(list, DIMENSIONS); + var path = Files.createTempFile("parallel-writer-interrupt", ".idx"); + try (var builder = new GraphIndexBuilder(values, VectorSimilarityFunction.DOT_PRODUCT, + 8, 32, 1.2f, 1.2f, true)) { + var graph = TestUtil.buildSequentially(builder, values); + var supplied = new AtomicInteger(); + var started = new CountDownLatch(50); + var failure = new AtomicReference(); + var interruptedAt = new AtomicReference(); + var finishedAt = new AtomicReference(); + + var writing = new Thread(() -> { + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, path) + .withParallelWorkerThreads(WORKERS) + .with(new InlineVectors(DIMENSIONS)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, ordinal -> { + try { + Thread.sleep(1); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("supplier interrupted", e); + } + supplied.incrementAndGet(); + started.countDown(); + return new InlineVectors.State(values.getVector(ordinal)); + })); + } catch (Throwable t) { + failure.set(t); + } finally { + finishedAt.set(System.nanoTime()); + } + }, "parallel-writer-under-test"); + writing.start(); + assertTrue("writer did not start supplying records", started.await(30, TimeUnit.SECONDS)); + interruptedAt.set(System.nanoTime()); + writing.interrupt(); + writing.join(60_000); + + assertNotNull("write did not finish", finishedAt.get()); + long stopMillis = TimeUnit.NANOSECONDS.toMillis(finishedAt.get() - interruptedAt.get()); + assertTrue("interrupted write took " + stopMillis + " ms to stop and supplied " + + supplied.get() + " of " + SIZE + " records", + stopMillis <= MAX_STOP_MILLIS && supplied.get() < SIZE); + assertNotNull("interrupted write reported success", failure.get()); + } finally { + Files.deleteIfExists(path); + } + } +} 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: - -