jbellis commented on code in PR #12421: URL: https://github.com/apache/lucene/pull/12421#discussion_r1291603910
########## lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentNeighborSet.java: ########## @@ -0,0 +1,292 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.lucene.util.hnsw; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.PrimitiveIterator; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import org.apache.lucene.util.BitSet; +import org.apache.lucene.util.FixedBitSet; + +/** A concurrent set of neighbors. */ +public class ConcurrentNeighborSet { + /** the node id whose neighbors we are storing */ + private final int nodeId; + + /** + * We use a copy-on-write NeighborArray to store the neighbors. Even though updating this is + * expensive, it is still faster than using a concurrent Collection because "iterate through a + * node's neighbors" is a hot loop in adding to the graph, and NeighborArray can do that much + * faster: no boxing/unboxing, all the data is stored sequentially instead of having to follow + * references, and no fancy encoding necessary for node/score. + */ + private final AtomicReference<ConcurrentNeighborArray> neighborsRef; + + private final NeighborSimilarity similarity; + + /** the maximum number of neighbors we can store */ + private final int maxConnections; + + public ConcurrentNeighborSet(int nodeId, int maxConnections, NeighborSimilarity similarity) { + this.nodeId = nodeId; + this.maxConnections = maxConnections; + this.similarity = similarity; + neighborsRef = new AtomicReference<>(new ConcurrentNeighborArray(maxConnections, true)); + } + + public PrimitiveIterator.OfInt nodeIterator() { + // don't use a stream here. stream's implementation of iterator buffers + // very aggressively, which is a big waste for a lot of searches. + return new NeighborIterator(neighborsRef.get()); + } + + public void backlink(Function<Integer, ConcurrentNeighborSet> neighborhoodOf) throws IOException { + NeighborArray neighbors = neighborsRef.get(); + for (int i = 0; i < neighbors.size(); i++) { + int nbr = neighbors.node[i]; + float nbrScore = neighbors.score[i]; + ConcurrentNeighborSet nbrNbr = neighborhoodOf.apply(nbr); + nbrNbr.insert(nodeId, nbrScore); + } + } + + private static class NeighborIterator implements PrimitiveIterator.OfInt { + private final NeighborArray neighbors; + private int i; + + private NeighborIterator(NeighborArray neighbors) { + this.neighbors = neighbors; + i = 0; + } + + @Override + public boolean hasNext() { + return i < neighbors.size(); + } + + @Override + public int nextInt() { + return neighbors.node[i++]; + } + } + + public int size() { + return neighborsRef.get().size(); + } + + public int arrayLength() { + return neighborsRef.get().node.length; + } + + /** + * For each candidate (going from best to worst), select it only if it is closer to target than it + * is to any of the already-selected candidates. This is maintained whether those other neighbors + * were selected by this method, or were added as a "backlink" to a node inserted concurrently + * that chose this one as a neighbor. + */ + public void insertDiverse(NeighborArray candidates) { + BitSet selected = new FixedBitSet(candidates.size()); + for (int i = candidates.size() - 1; i >= 0; i--) { + int cNode = candidates.node[i]; + float cScore = candidates.score[i]; + if (isDiverse(cNode, cScore, candidates, selected)) { + selected.set(i); + } + } + insertMultiple(candidates, selected); + // This leaves the paper's keepPrunedConnection option out; we might want to add that + // as an option in the future. + } + + private void insertMultiple(NeighborArray others, BitSet selected) { + neighborsRef.getAndUpdate( + current -> { + ConcurrentNeighborArray next = current.copy(); Review Comment: Looked at another profile this morning. 99.75% of insertMultiple is score comparisons, for vectors of dimension 256. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: issues-unsubscr...@lucene.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: issues-unsubscr...@lucene.apache.org For additional commands, e-mail: issues-h...@lucene.apache.org
