github.com/datastax/jvector @3.0.6

Introduction to approximate nearest neighbor search

Exact nearest neighbor search (k-nearest-neighbor or KNN) is prohibitively expensive at higher dimensions, because approaches to segment the search space that work in 2D or 3D like quadtree or k-d tree devolve to linear scans at higher dimensions. This is one aspect of what is called “the curse of dimensionality.”

With larger datasets, it is almost always more useful to get an approximate answer in logarithmic time, than the exact answer in linear time. This is abbreviated as ANN (approximate nearest neighbor) search.

There are two broad categories of ANN index: * Partition-based indexes, like LSH or IVF or SCANN * Graph indexes, like HNSW or DiskANN

Graph-based indexes tend to be simpler to implement and faster, but more importantly they can be constructed and updated incrementally. This makes them a much better fit for a general-purpose index than partitioning approaches that only work on static datasets that are completely specified up front. That is why all the major commercial vector indexes use graph approaches.

JVector is a graph index in the DiskANN family tree.

JVector Architecture

JVector is a graph-based index that builds on the DiskANN design with composeable extensions.

JVector implements a single-layer graph with nonblocking concurrency control, allowing construction to scale linearly with the number of cores:

The graph is represented by an on-disk adjacency list per node, with additional data stored inline to support two-pass searches, with the first pass powered by lossily compressed representations of the vectors kept in memory, and the second by a more accurate representation read from disk. The first pass can be performed with * Product quantization (PQ), optionally with anisotropic weighting * Binary quantization (BQ) * Fused ADC, where PQ codebooks are transposed and written inline with the graph adjacency list

The second pass can be performed with * Full resolution float32 vectors

This two-pass design reduces memory usage and reduces latency while preserving accuracy.

Additionally, JVector is unique in offering the ability to construct the index itself using two-pass searches, allowing larger-than-memory indexes to be built:

This is important because it allows you to take advantage of logarithmic search within a single index, instead of spilling over to linear-time merging of results from multiple indexes.

JVector step-by-step

All code samples are from SiftSmall 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

First the code:

    public static void siftInMemory(ArrayList<VectorFloat<?>> baseVectors) throws IOException {
        // infer the dimensionality from the first vector
        int originalDimension = baseVectors.get(0).length();
        // wrap the raw vectors in a RandomAccessVectorValues
        RandomAccessVectorValues ravv = new ListRandomAccessVectorValues(baseVectors, originalDimension);

        // score provider using the raw, in-memory vectors
        BuildScoreProvider bsp = BuildScoreProvider.randomAccessScoreProvider(ravv, VectorSimilarityFunction.EUCLIDEAN);
        try (GraphIndexBuilder builder = new GraphIndexBuilder(bsp,
                                                               ravv.dimension(),
                                                               16, // graph degree
                                                               100, // construction search depth
                                                               1.2f, // allow degree overflow during construction by this factor
                                                               1.2f)) // relax neighbor diversity requirement by this factor
        {
            // build the index (in memory)
            OnHeapGraphIndex index = builder.build(ravv);

            // search for a random vector
            VectorFloat<?> q = randomVector(originalDimension);
            SearchResult sr = GraphSearcher.search(q,
                                                   10, // number of results
                                                   ravv, // vectors we're searching, used for scoring
                                                   VectorSimilarityFunction.EUCLIDEAN, // how to score
                                                   index,
                                                   Bits.ALL); // valid ordinals to consider
            for (SearchResult.NodeScore ns : sr.getNodes()) {
                System.out.println(ns);
            }
        }
    }

Commentary: * All indexes assume that you have a source of vectors that has a consistent, fixed dimension (number of float32 components). * Vector sources are usually represented as a subclass of RandomAccessVectorValues, which offers a simple API around getVector / getVectorInto. Do be aware of isValueShared() in a multithreaded context; as a rule of thumb, in-memory RAVV will not use shared values and from-disk RAVV will (as an optimization to avoid allocating a new on-heap vector for every call). * You do NOT have to provide all the vectors to the index at once, but since this is a common scenario when prototyping, a convenience method is provided to do so. We will cover how to build an index incrementally later. * For the overflow Builder parameter, the sweet spot is about 1.2 for in-memory construction and 1.5 for on-disk. (The more overflow is allowed, the fewer recomputations of best edges are required, but the more neighbors will be consulted in every search.) * The alpha parameter controls the tradeoff between edge distance and diversity; usually 1.2 is sufficient for high-dimensional vectors; 2.0 is recommended for 2D or 3D datasets. See the DiskANN paper for more details. * The Bits parameter to GraphSearcher is intended for controlling your resultset based on external predicates and won’t be used in this tutorial.

Step 2: more control over GraphSearcher

Keeping the Builder the same, the updated search code looks like this:

            // search for a random vector using a GraphSearcher and SearchScoreProvider
            VectorFloat<?> q = randomVector(originalDimension);
            try (GraphSearcher searcher = new GraphSearcher(index)) {
                SearchScoreProvider ssp = SearchScoreProvider.exact(q, VectorSimilarityFunction.EUCLIDEAN, ravv);
                SearchResult sr = searcher.search(ssp, 10, Bits.ALL);
                for (SearchResult.NodeScore ns : sr.getNodes()) {
                    System.out.println(ns);
                }
            }

Commentary: * Searcher allocation is modestly expensive since there is a bunch of internal state to initialize at construction time. Therefore, JVector supports pooling searchers (e.g. with ExplicitThreadLocal, covered below). * When managing GraphSearcher instances, you’re also responsible for constructing a SearchScoreProvider, which you can think of as: given a query vector, tell JVector how to compute scores for other nodes in the index. SearchScoreProvider can be exact, as shown here, or a combination of an ApproximateScoreFunction and a Reranker, covered below.

Step 3: Measuring recall

A blisteringly-fast vector index isn’t very useful if it doesn’t return accurate results. As a sanity check, SiftSmall includes a helper method testRecall. Wiring up that to our code mostly involves turning the SearchScoreProvider into a factory lambda:

            Function<VectorFloat<?>, SearchScoreProvider> sspFactory = q -> SearchScoreProvider.exact(q, VectorSimilarityFunction.EUCLIDEAN, ravv);
            testRecall(index, queryVectors, groundTruth, sspFactory);

If you run the code, you will see slight differences in recall every time (printed by testRecall):

(OnHeapGraphIndex) Recall: 0.9898
...
(OnHeapGraphIndex) Recall: 0.9890

This is expected given the approximate nature of the index being created and the perturbations introduced by the multithreaded concurrency of the build call.

Step 4: write and load index to and from disk

The code:

        Path indexPath = Files.createTempFile("siftsmall", ".inline");
        try (GraphIndexBuilder builder = new GraphIndexBuilder(bsp, ravv.dimension(), 16, 100, 1.2f, 1.2f)) {
            // build the index (in memory)
            OnHeapGraphIndex index = builder.build(ravv);
            // write the index to disk with default options
            OnDiskGraphIndex.write(index, ravv, indexPath);
        }

        // on-disk indexes require a ReaderSupplier (not just a Reader) because we will want it to
        // open additional readers for searching
        ReaderSupplier rs = new SimpleMappedReaderSupplier(indexPath);
        OnDiskGraphIndex index = OnDiskGraphIndex.load(rs);
        // measure our recall against the (exactly computed) ground truth
        Function<VectorFloat<?>, SearchScoreProvider> sspFactory = q -> SearchScoreProvider.exact(q, VectorSimilarityFunction.EUCLIDEAN, ravv);
        testRecall(index, queryVectors, groundTruth, sspFactory);

Commentary: * We can write indexes that are constructed in-memory, like this one, to disk with a single method call. * Loading and searching against on-disk indexes require a ReaderSupplier, which supplies RandomAccessReader objects. The RandomAccessReader interface is intended to be extended by the consuming project. For instance, DataStax Astra implements a RandomAccessReader backed by the Cassandra chunk cache. JVector provides two implementations out of the box. * SimpleMappedReader: implemented using FileChannel.map, which means it is compatible with all Java versions that can run JVector, but also means it is limited to 2GB file sizes. SimpleMappedReader is primarily intended for example code. * MemorySegmentReader: implemented using the newer MemorySegment API, with no file size limit, but is limited to Java 22+. (The actual MemorySegmentReader code is compatible with Java 20+, but we left it in the 22+ module for convenience. The motivated reader is welcome to refactor the build to improve this.) If you have no specialized requirements then MemorySegmentReader is recommended for production.

Step 5: use compressed vectors in the search

Compressing the vectors with product quantization is done as follows:

        // compute and write compressed vectors to disk
        Path pqPath = Files.createTempFile("siftsmall", ".pq");
        try (DataOutputStream out = new DataOutputStream(new BufferedOutputStream(Files.newOutputStream(pqPath)))) {
            // Compress the original vectors using PQ. this represents a compression ratio of 128 * 4 / 16 = 32x
            ProductQuantization pq = ProductQuantization.compute(ravv,
                                                                 16, // number of subspaces
                                                                 256, // number of centroids per subspace
                                                                 true); // center the dataset
            // Note: before jvector 3.1.0, encodeAll returned an array of ByteSequence.
            PQVectors pqv = pq.encodeAll(ravv);
            // write the compressed vectors to disk
            pqv.write(out);
        }

• JVector also supports Binary Quantization, but BQ is generally less useful than PQ since it takes such a large toll on search accuracy.

Then we can wire up the compressed vectors to a two-phase search by getting the fast ApproximateScoreFunction from PQVectors, and the Reranker from the index View: ```java ReaderSupplier rs = new MMapReaderSupplier(indexPath); OnDiskGraphIndex index = OnDiskGraphIndex.load(rs); // load the PQVectors that we just wrote to disk try (RandomAccessReader in = new SimpleMappedReader(pqPath)) { PQVectors pqv = PQVectors.load(in); // SearchScoreProvider that does a first pass with the loaded-in-memory PQVectors, // then reranks with the exact vectors that are stored on disk in the index Function<VectorFloat<?>, SearchScoreProvider> sspFactory = q -> { ApproximateScoreFunction asf = pqv.precomputedScoreFunctionFor(q, VectorSimilarityFunction.EUCLIDEAN); Reranker reranker = index.getView().rerankerFor(q, VectorSimilarityFunction.EUCLIDEAN); return new SearchScoreProvider(asf, reranker); }; // measure our recall against the (exactly computed)