Efficient Vector Search

// - ThreadLocalRandom.current().nextDouble() is a double between 0.0 and 1.0
// - Since the logarithm of a number between 0 and 1 is negative, we change the sign
// - mL is an input parameter to the algorithm
//
// As a result, l is a random number between 0 and mL, but values close to 0 
// are much more probable.
int l = (int) (-Math.log(ThreadLocalRandom.current().nextDouble()) * mL);

final class PagedVectorIndex {
    // List of memory segments (pages)
    private final List<MemorySegment> pages = new ArrayList<>();
    private final AtomicLong currentCount = new AtomicLong(0);
    private final int blockSize;
    private final int dimensions;

    /**
     * @param blockSize  how many vectors per page
     * @param dimensions how many elements per vector
     */
    public PagedVectorIndex(int blockSize, int dimensions) {
        this.blockSize = blockSize;
        this.dimensions = dimensions;

        if (blockSize < 1) {
            throw new IllegalArgumentException("Invalid block size, must be > 0");
        }
        if (dimensions < 1) {
            throw new IllegalArgumentException("Invalid dimensions, must be > 0");
        }
        // Start with one page
        addPage();
    }

    private void addPage() {
        // Allocate off-heap memory for 1 block of vectors
        long byteSize = (long) blockSize * dimensions * Float.BYTES;
        MemorySegment newPage = Arena.ofAuto().allocate(byteSize);
        pages.add(newPage);
    }

    //... other class members
}

public long addVector(float[] vector) {
    long internalId = currentCount.getAndIncrement(); // Atomic, gives unique ID

    // 1. Check if we need a new page
    if (internalId % blockSize == 0 && internalId > 0) {
        // Synchronize only the slow, array-modifying operation
        synchronized (this) {
            // Double-check lock: check again in case another thread already added the page
            if (pages.size() <= internalId / blockSize) {
                addPage();
            }
        }
    }

    // 2. Calculate location
    int pageIdx = (int) (internalId / blockSize);
    long offset = (internalId % blockSize) * dimensions * Float.BYTES;

    // 3. Write data (MemorySegment is generally thread-safe for disjoint writes)
    MemorySegment.copy(vector, 0, pages.get(pageIdx), ValueLayout.JAVA_FLOAT, offset, dimensions);

    return internalId;
}

public float[] getVector(long id) {
    if (id < 0 || id >= currentCount.get()) {
        throw new IllegalArgumentException("Illegal vector id");
    }
    int pageIdx = (int) (id / blockSize);
    long pageOffset = (id % blockSize) * dimensions * Float.BYTES;
    var sliceSize = dimensions * Float.BYTES;

    var slice = pages.get(pageIdx).asSlice(pageOffset, sliceSize);
    return slice.toArray(ValueLayout.JAVA_FLOAT);
}

[ Node 0 Data ........................ ] [ Node 1 Data ........................ ]
| Count (int) | N1 | N2 | ... | N_max |  | Count (int) | N1 | N2 | ... | N_max |

final public class PagedGraphIndex {
    // Max neighbors for Layer 0 (usually 2 * M)
    private final int maxConnections;

    // Size of one node's connection data in bytes:
    // 4 bytes (count) + (mMax0 * 8 bytes)
    private final int slotSizeBytes;

    private final int blockSize; // Nodes per page
    private final List<MemorySegment> pages = new ArrayList<>();

    public PagedGraphIndex(int maxConnections, int blockSize) {
        this.maxConnections = maxConnections;
        this.slotSizeBytes = Integer.BYTES + (maxConnections * Long.BYTES);
        this.blockSize = blockSize;
        addPage();
    }

    private void addPage() {
        long totalBytes = (long) blockSize * slotSizeBytes;
        pages.add(Arena.ofAuto().allocate(totalBytes));
    }

    private void ensureCapacity(long vectorId) {
        int pageIdx = (int) (vectorId / blockSize);
        while (pages.size() <= pageIdx) {
            addPage();
        }
    }

    /**
     * Overwrites the connections for a specific node.
     * Used after the Heuristic (Algorithm 4) selects the best neighbors.
     */
    public void setConnections(long vectorId, long[] neighbors) {
        if (neighbors.length > maxConnections) {
            throw new IllegalArgumentException("Too many neighbors");
        }
        ensureCapacity(vectorId);

        // 1. Calculate Offsets
        int pageIdx = (int) (vectorId / blockSize);
        long pageStart = (vectorId % blockSize) * slotSizeBytes;
        MemorySegment page = pages.get(pageIdx);

        // 2. Write Count
        page.set(ValueLayout.JAVA_INT, pageStart, neighbors.length);

        // 3. Write Neighbors
        // We skip 4 bytes (the count) to find the start of the neighbor array
        long neighborArrayOffset = pageStart + Integer.BYTES;

        MemorySegment.copy(
                MemorySegment.ofArray(neighbors), 0,
                page, neighborArrayOffset,
                (long) neighbors.length * Long.BYTES
        );
    }
    // ... other class members ...

}

public class HnswIndex {
    private final HnswConfig config;
    private final PagedVectorIndex layer0;
    private final PagedGraphIndex layer0Connections;

    /*
    Only ~5% of the nodes exist in Layer 1, and ~0.25% in Layer 2. The memory overhead of Java objects here is
    negligible compared to the massive Layer 0.
     */
    private final Map<Integer, Map<Long, Set<Long>>> upperLayerConnections = new HashMap<>();

    private int globalMaxLayer = 0;
    private long globalEnterPoint = 0;

    public HnswIndex(HnswConfig config) {
        this.config = config;
        layer0 = new PagedVectorIndex(config.blockSize(), config.dimensions());
        layer0Connections = new PagedGraphIndex(config.mMax0(), config.blockSize());
    }

    //... other class members ...
}

public long insert(float[] newVector) {
    if (newVector.length != config.dimensions()) {
        throw new IllegalArgumentException(
                "Invalid vector dimension: %d (expected %d)".formatted(newVector.length, config.dimensions()));
    }

    long currentId = globalEnterPoint;
    int maxL = globalMaxLayer;

    //1. Randomly choose a maximum layer
    int vectorMaxLayer = determineMaxLayer();

    //2. Add vector to layer 0
    long newId = layer0.addVector(newVector);

    //3. Zoom down from Top Layer to vectorMaxLayer+1 (Greedy search, ef=1)
    for (int l = globalMaxLayer; l > vectorMaxLayer; l--) {
        currentId = searchLayerGreedy(newVector, currentId, l);
    }

    // 4. From L down to 0, find neighbors and link
    for (int layer = Math.min(vectorMaxLayer, maxL); layer >= 0; layer--) {
        // Search layer with efConstruction
        var neighborCandidates = searchLayer(newVector, currentId, layer, config.efConstruction());

        if (neighborCandidates.isEmpty()) {
            continue;
        }

        // Update entry point for next layer
        currentId = neighborCandidates.peek().vectorId(); // Rough logic, usually closest

        // Select neighbors using heuristic
        var neighbors = selectNeighborsHeuristic(neighborCandidates);

        // Add connections bidirectional
        for (long neighborId : neighbors) {
            // Register connection bidirectionally. Not both connections are guaranteed to be stored,
            // as there will be a pruning process for each node.
            addConnection(newId, neighborId, layer);
            addConnection(neighborId, newId, layer);
        }
    }

    // 5. Update global entry point if new node is in a higher layer
    if (vectorMaxLayer > globalMaxLayer) {
        globalMaxLayer = vectorMaxLayer;
        globalEnterPoint = newId;
    }
    return newId;
}

/**
 * Returns the K-Nearest Neighbors of queryVector
 *
 * @param queryVector the physical vector
 * @param k           the number of vectors to return
 * @return a list of {@link NodeSimilarity} instances representing the search results.
 */
public List<NodeSimilarity> search(float[] queryVector, int k) {
    if (layer0.getVectorCount() == 0) {
        return Collections.emptyList();
    }

    long currObj = globalEnterPoint;

    // 1. Zoom-in Phase: Fast greedy search through upper layers
    // We use ef=1 for speed as we just need a good entry point for the next layer
    for (int l = globalMaxLayer; l >= 1; l--) {
        currObj = searchLayerGreedy(queryVector, currObj, l);
    }

    // 2. Fine Search Phase: Comprehensive search at the ground layer (Layer 0)
    // W is a Max-Heap that keeps track of the 'ef' best candidates
    PriorityQueue<NodeSimilarity> w = searchLayer(queryVector, currObj, 0, config.efSearch());

    // 3. Return top K results from the candidates found in Layer 0
    List<NodeSimilarity> results = new ArrayList<>();
    // W is a max-heap; we want the closest K, so we might need to sort or poll
    while (w.size() > k) {
        var o = w.poll(); // Remove furthest until we have K
    }
    while (!w.isEmpty()) {
        results.add(w.poll());
    }
    Collections.reverse(results);
    return results;
}

    private static VulcanoDb createVulcanoDB() {
        var properties = new Properties();
        properties.setProperty(ConfigProperties.PROPERTY_PATH, path.toString());
        var hnswConfig = HnswConfig.builder()
                .withEfConstruction(500) // max recall
                .withEfSearch(500) // max recall
                .build();
        var axon = AxonDataStore
                .builder()
                .withDocumentWriter(new DefaultDocumentPersister(new TypedProperties(properties)))
                .withVectorIndex("indexedVector", hnswConfig)
                .build();

        return VulcanoDb
                .builder()
                .withDataStore(axon)
                .build();
    }

    private static List<Document> generateSamples() throws URISyntaxException, IOException {
        var loader = Thread.currentThread().getContextClassLoader();
        var url = loader.getResource("rotten-tomatoes-reviews.txt");
        Path path = Paths.get(url.toURI());
        try (var lines = Files.lines(path)) {
            return lines
                    .parallel()
                    .map(txt -> {
                        var embedding = Embedding.of(txt);
                        return Document.builder()
                                .withStringField("originalText", txt)
                                .withVectorField("nonIndexedVector", embedding)
                                .withVectorField("indexedVector", embedding)
                                .build();
                    })
                    .toList();
        }
    }

...
IO.println("Generating embedding for the query...");
var query = "Clever horror movie with great acting";
var queryVector = Embedding.of(query);

IO.println("Non-indexed search:");
long t0 = System.currentTimeMillis();
var nonIndexedQuery = Query.builder().isSimilarTo(queryVector, "nonIndexedVector").build();
vulcanoDB
        .search(nonIndexedQuery, 10)
        .getDocuments()
        .forEach(HnswIndexPerformance::showResult);

long bruteTime = System.currentTimeMillis() - t0;
IO.println("Non-indexed search time: %d ms".formatted(bruteTime));


IO.println("Index search: ");
t0 = System.currentTimeMillis();
var indexedQuery = Query.builder().isSimilarTo(queryVector, "indexedVector").build();
vulcanoDB
        .search(indexedQuery, 10)
        .getDocuments()
        .forEach(HnswIndexPerformance::showResult);
var indexTime = System.currentTimeMillis() - t0;
IO.println("Index search time: %d ms.".formatted(indexTime));

Generating embedding for the query...
Non-indexed search:
[0.83] b0089fc7-ab3c-4b85-8daf-7d422a9f9567 -> a slick , skillful little horror film .
[0.80] 39a5e68b-171d-45f3-a3fa-8a9354512dba -> one of the best silly horror movies of recent memory , with some real shocks in store for unwary viewers .
[0.79] e4bcf721-412a-4392-8919-5e3a551d51f9 -> this is a superior horror flick .
[0.77] 6416021c-c35e-4c90-8f2e-8092ca1c7ba3 -> an exceedingly clever piece of cinema . another great ‘what you don't see' is much more terrifying than what you do see thriller , coupled with some arresting effects , incandescent tones and stupendous performances
[0.77] 0783c1fb-8719-40d5-a211-bead28e9a228 -> clever , brutal and strangely soulful movie .
[0.77] 414227e8-e428-494c-8bbd-9ef9608d9fb7 -> a chilling movie without oppressive gore .
[0.74] a8b28cb8-4884-46b1-9a58-f910286ab829 -> marshall puts a suspenseful spin on standard horror flick formula .
[0.74] c90ecce2-72e4-4321-9436-af5657dd3558 -> the entire movie establishes a wonderfully creepy mood .
[0.73] e95c46b6-4e44-4fa8-894b-1d595bc7ed2e -> a deftly entertaining film , smartly played and smartly directed .
[0.73] 449979ab-56ef-45b0-befc-4062f0ca68e6 -> if nothing else , this movie introduces a promising , unusual kind of psychological horror .
Non-indexed search time: 372 ms
Index search: 
[0.87] 39a5e68b-171d-45f3-a3fa-8a9354512dba -> one of the best silly horror movies of recent memory , with some real shocks in store for unwary viewers .
[0.86] e4bcf721-412a-4392-8919-5e3a551d51f9 -> this is a superior horror flick .
[0.84] 0783c1fb-8719-40d5-a211-bead28e9a228 -> clever , brutal and strangely soulful movie .
[0.83] a8b28cb8-4884-46b1-9a58-f910286ab829 -> marshall puts a suspenseful spin on standard horror flick formula .
[0.83] c90ecce2-72e4-4321-9436-af5657dd3558 -> the entire movie establishes a wonderfully creepy mood .
[0.82] e95c46b6-4e44-4fa8-894b-1d595bc7ed2e -> a deftly entertaining film , smartly played and smartly directed .
[0.82] 65105ae3-32fd-449a-b0ec-414a8d487abf -> my little eye is the best little " horror " movie i've seen in years .
[0.81] e8635913-5f1f-474c-9ad2-9a46ea725483 -> a properly spooky film about the power of spirits to influence us whether we believe in them or not .
[0.81] d0a79c1c-cd0b-41f1-b48b-afdca5121811 -> an elegant film with often surprising twists and an intermingling of naiveté and sophistication .
[0.81] 190ea188-c21e-4def-be0f-822b47b97b18 -> a perceptive , good-natured movie .
Index search time: 22 ms.