benwtrent commented on code in PR #14331:
URL: https://github.com/apache/lucene/pull/14331#discussion_r2003974286
##########
lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java:
##########
@@ -51,19 +57,85 @@ protected HnswBuilder createBuilder(KnnVectorValues
mergedVectorValues, int maxO
OnHeapHnswGraph graph;
BitSet initializedNodes = null;
- if (initReader == null) {
+ if (graphReaders.size() == 0) {
graph = new OnHeapHnswGraph(M, maxOrd);
} else {
+
graphReaders.sort(Comparator.comparingInt(GraphReader::graphSize).reversed());
+ GraphReader initGraphReader = graphReaders.get(0);
+ KnnVectorsReader initReader = initGraphReader.reader();
+ MergeState.DocMap initDocMap = initGraphReader.initDocMap();
+ int initGraphSize = initGraphReader.graphSize();
HnswGraph initializerGraph = ((HnswGraphProvider)
initReader).getGraph(fieldInfo.name);
+
if (initializerGraph.size() == 0) {
graph = new OnHeapHnswGraph(M, maxOrd);
} else {
initializedNodes = new FixedBitSet(maxOrd);
- int[] oldToNewOrdinalMap = getNewOrdMapping(mergedVectorValues,
initializedNodes);
+ int[] oldToNewOrdinalMap =
+ getNewOrdMapping(
+ fieldInfo,
+ initReader,
+ initDocMap,
+ initGraphSize,
+ mergedVectorValues,
+ initializedNodes);
graph = InitializedHnswGraphBuilder.initGraph(initializerGraph,
oldToNewOrdinalMap, maxOrd);
}
}
return new HnswConcurrentMergeBuilder(
taskExecutor, numWorker, scorerSupplier, beamWidth, graph,
initializedNodes);
}
+
+ /**
+ * Creates a new mapping from old ordinals to new ordinals and returns the
total number of vectors
+ * in the newly merged segment.
+ *
+ * @param mergedVectorValues vector values in the merged segment
+ * @param initializedNodes track what nodes have been initialized
+ * @return the mapping from old ordinals to new ordinals
+ * @throws IOException If an error occurs while reading from the merge state
+ */
+ private static final int[] getNewOrdMapping(
+ FieldInfo fieldInfo,
+ KnnVectorsReader initReader,
+ MergeState.DocMap initDocMap,
+ int initGraphSize,
+ KnnVectorValues mergedVectorValues,
+ BitSet initializedNodes)
+ throws IOException {
+ KnnVectorValues.DocIndexIterator initializerIterator = null;
+
+ switch (fieldInfo.getVectorEncoding()) {
+ case BYTE -> initializerIterator =
initReader.getByteVectorValues(fieldInfo.name).iterator();
+ case FLOAT32 ->
+ initializerIterator =
initReader.getFloatVectorValues(fieldInfo.name).iterator();
+ }
+
+ IntIntHashMap newIdToOldOrdinal = new IntIntHashMap(initGraphSize);
+ int maxNewDocID = -1;
+ for (int docId = initializerIterator.nextDoc();
+ docId != NO_MORE_DOCS;
+ docId = initializerIterator.nextDoc()) {
+ int newId = initDocMap.get(docId);
+ maxNewDocID = Math.max(newId, maxNewDocID);
+ newIdToOldOrdinal.put(newId, initializerIterator.index());
+ }
+
+ if (maxNewDocID == -1) {
+ return new int[0];
+ }
+ final int[] oldToNewOrdinalMap = new int[initGraphSize];
+ KnnVectorValues.DocIndexIterator mergedVectorIterator =
mergedVectorValues.iterator();
+ for (int newDocId = mergedVectorIterator.nextDoc();
+ newDocId <= maxNewDocID;
+ newDocId = mergedVectorIterator.nextDoc()) {
+ int hashDocIndex = newIdToOldOrdinal.indexOf(newDocId);
+ if (newIdToOldOrdinal.indexExists(hashDocIndex)) {
Review Comment:
Is this stuff around `indexOf` `indexExists`, etc. just performance
improvements over a simple `newIdToOldOrdinal.get(...)`?
Looking at the `IntIntHashMap` its weird that "does not exist" may actually
just be `0`, where `0` is a valid doc id :/.
##########
lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java:
##########
@@ -51,19 +57,85 @@ protected HnswBuilder createBuilder(KnnVectorValues
mergedVectorValues, int maxO
OnHeapHnswGraph graph;
BitSet initializedNodes = null;
- if (initReader == null) {
+ if (graphReaders.size() == 0) {
graph = new OnHeapHnswGraph(M, maxOrd);
} else {
+
graphReaders.sort(Comparator.comparingInt(GraphReader::graphSize).reversed());
+ GraphReader initGraphReader = graphReaders.get(0);
+ KnnVectorsReader initReader = initGraphReader.reader();
+ MergeState.DocMap initDocMap = initGraphReader.initDocMap();
+ int initGraphSize = initGraphReader.graphSize();
HnswGraph initializerGraph = ((HnswGraphProvider)
initReader).getGraph(fieldInfo.name);
+
if (initializerGraph.size() == 0) {
graph = new OnHeapHnswGraph(M, maxOrd);
} else {
initializedNodes = new FixedBitSet(maxOrd);
- int[] oldToNewOrdinalMap = getNewOrdMapping(mergedVectorValues,
initializedNodes);
+ int[] oldToNewOrdinalMap =
+ getNewOrdMapping(
+ fieldInfo,
+ initReader,
+ initDocMap,
+ initGraphSize,
+ mergedVectorValues,
+ initializedNodes);
graph = InitializedHnswGraphBuilder.initGraph(initializerGraph,
oldToNewOrdinalMap, maxOrd);
}
}
return new HnswConcurrentMergeBuilder(
taskExecutor, numWorker, scorerSupplier, beamWidth, graph,
initializedNodes);
}
+
+ /**
+ * Creates a new mapping from old ordinals to new ordinals and returns the
total number of vectors
+ * in the newly merged segment.
+ *
+ * @param mergedVectorValues vector values in the merged segment
+ * @param initializedNodes track what nodes have been initialized
+ * @return the mapping from old ordinals to new ordinals
+ * @throws IOException If an error occurs while reading from the merge state
+ */
+ private static final int[] getNewOrdMapping(
Review Comment:
```suggestion
private static int[] getNewOrdMapping(
```
##########
lucene/core/src/java/org/apache/lucene/util/hnsw/MergingHnswGraphBuilder.java:
##########
@@ -0,0 +1,198 @@
+/*
+ * 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 static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
+
+import java.io.IOException;
+import java.util.HashSet;
+import java.util.Set;
+import org.apache.lucene.util.BitSet;
+
+/**
+ * A graph builder that is used during segments' merging.
+ *
+ * <p>This builder uses a smart algorithm to merge multiple graphs into a
single graph. The
+ * algorithm is based on the idea if we know where we want to insert a node,
we have a good idea of
+ * where we want to insert its neighbours.
+ *
+ * <p>The algorithm is based on the following steps: 1. Get all graphs that
don't have deletions and
+ * sort them by size. 2. Copy the largest graph to the new graph (gL). 3. For
each remaining small
+ * graph (gS) - Find the nodes that best cover gS: join set `j`. These nodes
will be inserted into
+ * gL as usual: by searching gL to find the best candidates: `w` to which
connect the nodes. - For
+ * each remaining node in gS: - we do NOT do search in gL. Instead, we form
`w` by union of the
+ * node's neighbors' in Gs and the node's neighbors' neighbors' in gL.
+ *
+ * <p>We expect the size of join set `j` to be small, 1/5-1/2 of size gS. And
for the rest of the
+ * nodes of gS, we expect savings by not doing extra search in gL.
+ *
+ * @lucene.experimental
+ */
Review Comment:
class comments here twice?
##########
lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java:
##########
@@ -51,19 +57,85 @@ protected HnswBuilder createBuilder(KnnVectorValues
mergedVectorValues, int maxO
OnHeapHnswGraph graph;
BitSet initializedNodes = null;
- if (initReader == null) {
+ if (graphReaders.size() == 0) {
graph = new OnHeapHnswGraph(M, maxOrd);
} else {
+
graphReaders.sort(Comparator.comparingInt(GraphReader::graphSize).reversed());
+ GraphReader initGraphReader = graphReaders.get(0);
+ KnnVectorsReader initReader = initGraphReader.reader();
+ MergeState.DocMap initDocMap = initGraphReader.initDocMap();
+ int initGraphSize = initGraphReader.graphSize();
HnswGraph initializerGraph = ((HnswGraphProvider)
initReader).getGraph(fieldInfo.name);
+
if (initializerGraph.size() == 0) {
graph = new OnHeapHnswGraph(M, maxOrd);
} else {
initializedNodes = new FixedBitSet(maxOrd);
- int[] oldToNewOrdinalMap = getNewOrdMapping(mergedVectorValues,
initializedNodes);
+ int[] oldToNewOrdinalMap =
+ getNewOrdMapping(
+ fieldInfo,
+ initReader,
+ initDocMap,
+ initGraphSize,
+ mergedVectorValues,
+ initializedNodes);
graph = InitializedHnswGraphBuilder.initGraph(initializerGraph,
oldToNewOrdinalMap, maxOrd);
}
}
return new HnswConcurrentMergeBuilder(
taskExecutor, numWorker, scorerSupplier, beamWidth, graph,
initializedNodes);
}
+
+ /**
+ * Creates a new mapping from old ordinals to new ordinals and returns the
total number of vectors
+ * in the newly merged segment.
+ *
+ * @param mergedVectorValues vector values in the merged segment
+ * @param initializedNodes track what nodes have been initialized
+ * @return the mapping from old ordinals to new ordinals
+ * @throws IOException If an error occurs while reading from the merge state
+ */
+ private static final int[] getNewOrdMapping(
+ FieldInfo fieldInfo,
+ KnnVectorsReader initReader,
+ MergeState.DocMap initDocMap,
+ int initGraphSize,
+ KnnVectorValues mergedVectorValues,
+ BitSet initializedNodes)
+ throws IOException {
+ KnnVectorValues.DocIndexIterator initializerIterator = null;
+
+ switch (fieldInfo.getVectorEncoding()) {
+ case BYTE -> initializerIterator =
initReader.getByteVectorValues(fieldInfo.name).iterator();
+ case FLOAT32 ->
+ initializerIterator =
initReader.getFloatVectorValues(fieldInfo.name).iterator();
+ }
+
+ IntIntHashMap newIdToOldOrdinal = new IntIntHashMap(initGraphSize);
+ int maxNewDocID = -1;
+ for (int docId = initializerIterator.nextDoc();
+ docId != NO_MORE_DOCS;
+ docId = initializerIterator.nextDoc()) {
+ int newId = initDocMap.get(docId);
+ maxNewDocID = Math.max(newId, maxNewDocID);
+ newIdToOldOrdinal.put(newId, initializerIterator.index());
Review Comment:
```suggestion
assert newIdToOldOrdinal.containsKey(newId) == false;
newIdToOldOrdinal.put(newId, initializerIterator.index());
```
We should assert if we put `newId` more than once.
##########
lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java:
##########
@@ -208,6 +211,11 @@ private void addVectors(int maxOrd) throws IOException {
}
public void addGraphNode(int node, UpdateableRandomVectorScorer scorer)
throws IOException {
+ addGraphNodeInternal(node, scorer, null);
+ }
+
+ private void addGraphNodeInternal(
+ int node, UpdateableRandomVectorScorer scorer, Set<Integer> eps0) throws
IOException {
Review Comment:
`IntHashSet` instead maybe?
--
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
