You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by to...@apache.org on 2020/01/12 06:11:35 UTC
[lucene-solr] branch jira/lucene-9004-aknn-2 created (now f7c2ae2)
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a change to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git.
at f7c2ae2 Use 'Lucene90' Codec for TestKnnGraph
This branch includes the following new commits:
new 6b1f096 Add the codec, field type, values writer/reader, and query for approximate knn vector search.
new f22b4b5 Support segment merging
new 444c1d3 Implemented TestKnnGraph unit tests
new b005dcb Merge pull request #1 from msokolov/mocobeta/LUCENE-9004-2
new 624632c Add MoreLikeThis like factory method
new 7cdafc5 Added KnnGraphTester; fixed some small bugs and measured recall
new 3639b22 Merge pull request #2 from msokolov/LUCENE-9004-aknn-2
new b4a8979 Mocobeta/lucene 9004 2 (#3)
new e4d348b fix single node scenario (#4)
new f7c2ae2 Use 'Lucene90' Codec for TestKnnGraph
The 10 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
[lucene-solr] 07/10: Merge pull request #2 from
msokolov/LUCENE-9004-aknn-2
Posted by to...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a commit to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git
commit 3639b2247fd0fc4b10858222c5f969bb844268a3
Merge: 624632c 7cdafc5
Author: Tomoko Uchida <to...@gmail.com>
AuthorDate: Sat Jan 11 00:30:45 2020 +0900
Merge pull request #2 from msokolov/LUCENE-9004-aknn-2
Added KnnGraphTester; fixed some small bugs and measured recall
.../codecs/lucene90/Lucene90KnnGraphWriter.java | 2 -
.../apache/lucene/util/hnsw/DeferredNeighbor.java | 22 +-
.../org/apache/lucene/util/hnsw/HNSWGraph.java | 47 +--
.../apache/lucene/util/hnsw/HNSWGraphReader.java | 11 +-
.../apache/lucene/util/hnsw/HNSWGraphWriter.java | 21 +-
.../apache/lucene/util/hnsw/KnnGraphTester.java | 404 +++++++++++++++++++++
6 files changed, 461 insertions(+), 46 deletions(-)
[lucene-solr] 06/10: Added KnnGraphTester;
fixed some small bugs and measured recall
Posted by to...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a commit to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git
commit 7cdafc57a3b58e1c9b1995ebcec489e53398a55e
Author: Michael Sokolov <so...@amazon.com>
AuthorDate: Wed Dec 18 12:46:30 2019 -0500
Added KnnGraphTester; fixed some small bugs and measured recall
---
.../codecs/lucene90/Lucene90KnnGraphWriter.java | 2 -
.../apache/lucene/util/hnsw/DeferredNeighbor.java | 22 +-
.../org/apache/lucene/util/hnsw/HNSWGraph.java | 47 +--
.../apache/lucene/util/hnsw/HNSWGraphReader.java | 11 +-
.../apache/lucene/util/hnsw/HNSWGraphWriter.java | 21 +-
.../apache/lucene/util/hnsw/KnnGraphTester.java | 404 +++++++++++++++++++++
6 files changed, 461 insertions(+), 46 deletions(-)
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
index 8644c50..2417dad 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
@@ -165,7 +165,6 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
}
}
- /*
@Override
public void merge(MergeState mergeState) throws IOException {
for (KnnGraphReader reader : mergeState.knnGraphReaders) {
@@ -180,7 +179,6 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
}
finish();
}
- */
/**
* Merges the segment HNSW graphs by constructing a new merged graph using HNSWGraph
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/DeferredNeighbor.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/DeferredNeighbor.java
index 17e9295..961f00c 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/DeferredNeighbor.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/DeferredNeighbor.java
@@ -69,7 +69,27 @@ final class DeferredNeighbor extends Neighbor {
@Override
public int compareTo(Neighbor o) {
- return this.docId() < o.docId() ? -1 : 1;
+ return this.docId() - o.docId();
}
+
+ @Override
+ public boolean equals(Object o) {
+ return o instanceof Neighbor && this.docId() == ((Neighbor) o).docId();
+ }
+
+ @Override
+ public int hashCode() {
+ return 39 + 61 * docId;
+ }
+
+ @Override
+ public String toString() {
+ if (query != null) {
+ return "(" + docId() + ", " + distance() + ")";
+ } else {
+ return "(" + docId() + ", DEFERRED)";
+ }
+ }
+
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
index e2acd0b..c4d8332 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
@@ -23,6 +23,7 @@ import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
+import java.util.TreeSet;
import org.apache.lucene.index.VectorValues;
import org.apache.lucene.util.Accountable;
@@ -50,39 +51,32 @@ public final class HNSWGraph implements Accountable {
/**
* Searches the nearest neighbors for a specified query at a level.
* @param query search query vector
- * @param ep enter point to fhe level
+ * @param results on entry, has enter points to this level. On exit, the nearest neighbors in this level
* @param ef the number of nodes to be searched
* @param level graph level
* @param vectorValues vector values
* @return nearest neighbors
*/
- public NearestNeighbors searchLayer(float[] query, Neighbor ep, int ef, int level, VectorValues vectorValues) throws IOException {
+ void searchLayer(float[] query, FurthestNeighbors results, int ef, int level, VectorValues vectorValues) throws IOException {
if (level >= layers.size()) {
throw new IllegalArgumentException("layer does not exist for the level: " + level);
}
Layer layer = layers.get(level);
- if (!layer.getNodes().contains(ep.docId())) {
- throw new IllegalArgumentException("enter point " + ep.docId() + "does not exist at layer " + level);
+ TreeSet<Neighbor> candidates = new TreeSet<>();
+ for (Neighbor n : results) {
+ // TODO: candidates should get neighbors of neighbors
+ candidates.add(n);
}
-
- Set<Integer> visited = new HashSet<>(ep.docId());
- NearestNeighbors candidates = new NearestNeighbors(ef, ep);
- FurthestNeighbors results = new FurthestNeighbors(ef, ep);
-
- if (ep.isDeferred()) {
- ep.prepareQuery(query, vectorValues, distFunc);
- }
-
+ // set of docids that have been visited by search on this layer, used to avoid backtracking
+ Set<Integer> visited = new HashSet<>();
+ // We want to efficiently pop the best (nearest, least distance) candidate, so use NearestNeighbors,
+ // but we don't want to overflow the heap and lose the best candidate!
while (candidates.size() > 0) {
- Neighbor c = candidates.pop();
+ Neighbor c = candidates.pollFirst();
Neighbor f = results.top();
- if (c.isDeferred()) {
- c.prepareQuery(query, vectorValues, distFunc);
- }
- if (f.isDeferred()) {
- f.prepareQuery(query, vectorValues, distFunc);
- }
+ assert c.isDeferred() == false;
+ assert f.isDeferred() == false;
if (c.distance() > f.distance()) {
break;
}
@@ -91,15 +85,11 @@ public final class HNSWGraph implements Accountable {
continue;
}
visited.add(e.docId());
-
- f = results.top();
- if (f.isDeferred()) {
- f.prepareQuery(query, vectorValues, distFunc);
- }
+ assert f.isDeferred() == false;
float dist = distance(query, e.docId(), vectorValues);
if (dist < f.distance() || results.size() < ef) {
Neighbor n = new ImmutableNeighbor(e.docId(), dist);
- candidates.insertWithOverflow(n);
+ candidates.add(n);
Neighbor popped = results.insertWithOverflow(n);
if (popped != null && popped != n) {
f = results.top();
@@ -109,7 +99,7 @@ public final class HNSWGraph implements Accountable {
}
//System.out.println("level=" + level + ", visited nodes=" + visited.size());
- return pickNearestNeighbor(results);
+ //return pickNearestNeighbor(results);
}
private float distance(float[] query, int docId, VectorValues vectorValues) throws IOException {
@@ -120,7 +110,7 @@ public final class HNSWGraph implements Accountable {
return VectorValues.distance(query, other, distFunc);
}
- private NearestNeighbors pickNearestNeighbor(FurthestNeighbors queue) {
+ static NearestNeighbors pickNearestNeighbor(Neighbors queue) {
NearestNeighbors nearests = new NearestNeighbors(queue.size());
Set<Integer> addedDocs = new HashSet<>();
int ef = queue.size();
@@ -226,7 +216,6 @@ public final class HNSWGraph implements Accountable {
throw new IllegalArgumentException("layer does not exist for level: " + level);
}
layer.connectNodes(node1, node2, dist);
- layer.connectNodes(node2, node1, dist);
/*
// ensure friends size <= maxConnections
//
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
index b92e96b..ef11386 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
@@ -18,8 +18,10 @@
package org.apache.lucene.util.hnsw;
import java.io.IOException;
+import java.util.HashSet;
import java.util.Map;
import java.util.Objects;
+import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import org.apache.lucene.index.FieldInfo;
@@ -47,18 +49,19 @@ public final class HNSWGraphReader {
public Neighbors searchNeighbors(float[] query, int ef, VectorValues vectorValues) throws IOException {
HNSWGraph hnsw = get(field, context, false);
+ // TODO: refactor into HNSWGraph
int enterPoint = hnsw.getFirstEnterPoint();
if (!vectorValues.seek(enterPoint)) {
throw new IllegalStateException("enterPoint=" + enterPoint + " has no vector value");
}
- Neighbors neighbors;
Neighbor ep = new ImmutableNeighbor(enterPoint, VectorValues.distance(query, vectorValues.vectorValue(), distFunc));
+ FurthestNeighbors neighbors = new FurthestNeighbors(ef, ep);
for (int l = hnsw.topLevel(); l > 0; l--) {
- neighbors = hnsw.searchLayer(query, ep, 1, l, vectorValues);
- ep = neighbors.top();
+ hnsw.searchLayer(query, neighbors, 1, l, vectorValues);
}
- return hnsw.searchLayer(query, ep, ef, 0, vectorValues);
+ hnsw.searchLayer(query, neighbors, ef, 0, vectorValues);
+ return neighbors;
}
public static long loadGraphs(String field, IndexReader reader, boolean forceReload) throws IOException {
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
index dafbb05..2760481 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
@@ -20,7 +20,9 @@ package org.apache.lucene.util.hnsw;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
+import java.util.HashSet;
import java.util.List;
+import java.util.Set;
import org.apache.lucene.index.VectorValues;
import org.apache.lucene.util.Accountable;
@@ -110,11 +112,10 @@ public final class HNSWGraphWriter implements Accountable {
Neighbor ep = new ImmutableNeighbor(enterPoint, distFromEp);
int level = levelGenerator.randomLevel();
-
+ FurthestNeighbors results = new FurthestNeighbors(efConst, ep);
// down to the level from the hnsw's top level
for (int l = hnsw.topLevel(); l > level; l--) {
- NearestNeighbors neighbors = hnsw.searchLayer(value, ep, 1, l, vectorValues);
- ep = neighbors.top();
+ hnsw.searchLayer(value, results, efConst, l, vectorValues);
}
// down to level 0 with placing the doc to each layer
@@ -125,14 +126,14 @@ public final class HNSWGraphWriter implements Accountable {
continue;
}
- NearestNeighbors neighbors = hnsw.searchLayer(value, ep, efConst, l, vectorValues);
- ep = neighbors.top();
+ hnsw.searchLayer(value, results, efConst, l, vectorValues);
int maxConnections = l == 0 ? maxConn0 : maxConn;
- for (int i = 0; i < maxConnections; i++) {
- if (neighbors.size() == 0) {
- break;
- }
- Neighbor n = neighbors.pop();
+ NearestNeighbors neighbors = new NearestNeighbors(maxConnections, results.top());
+ for (Neighbor n : results) {
+ neighbors.insertWithOverflow(n);
+ }
+ for (Neighbor n : neighbors) {
+ // TODO: limit *total* num connections by pruning (shrinking)
hnsw.connectNodes(l, docId, n.docId(), n.distance(), maxConnections);
}
}
diff --git a/lucene/core/src/test/org/apache/lucene/util/hnsw/KnnGraphTester.java b/lucene/core/src/test/org/apache/lucene/util/hnsw/KnnGraphTester.java
new file mode 100644
index 0000000..ad21eca
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/util/hnsw/KnnGraphTester.java
@@ -0,0 +1,404 @@
+/*
+ * 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.BufferedOutputStream;
+import java.io.BufferedInputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Random;
+import java.util.Set;
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.VectorField;
+import org.apache.lucene.document.StringField;
+import org.apache.lucene.document.StoredField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.DocValues;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.KnnGraphValues;
+import org.apache.lucene.index.VectorValues;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.KnnGraphQuery;
+import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FSDirectory;
+import org.apache.lucene.util.PriorityQueue;
+
+/** For testing indexing and search performance of a knn-graph using random vectors */
+public class KnnGraphTester {
+
+ private final static String KNN_FIELD = "knn";
+ private final static String ID_FIELD = "id";
+
+ private Random random;
+ private int numDocs;
+ private int dim;
+ private int topK;
+ private int numProbe;
+ private float[] vectors;
+ private int[] nabors;
+
+ public KnnGraphTester() {
+ // set defaults
+ numDocs = 10_000;
+ dim = 256;
+ topK = 10;
+ numProbe = 500;
+ random = new Random();
+ }
+
+ public static void main(String... args) throws Exception {
+ if (args.length != 2) {
+ usage();
+ }
+ switch (args[0]) {
+ case "-generate":
+ new KnnGraphTester().create(args[1]);
+ break;
+ case "-search":
+ new KnnGraphTester().search(args[1]);
+ break;
+ case "-stats":
+ new KnnGraphTester().stats(args[1]);
+ break;
+ default:
+ usage();
+ }
+ }
+
+ private void search(String dataFile) throws IOException {
+ readDataFile(dataFile);
+ Path indexPath = Paths.get("knn_test_index");
+ createIndex(indexPath);
+ // topK = 25;
+ testSearch(indexPath, 1000);
+ //GraphSearch.VERBOSE = true;
+ //testSearch(indexPath, 1);
+ }
+
+ private void stats(String dataFile) throws IOException {
+ readDataFile(dataFile);
+ Path indexPath = Paths.get("knn_test_index");
+ createIndex(indexPath);
+ printFanoutHist(indexPath);
+ }
+
+ private void printFanoutHist(Path indexPath) throws IOException {
+ try (Directory dir = FSDirectory.open(indexPath);
+ DirectoryReader reader = DirectoryReader.open(dir)) {
+ // int[] globalHist = new int[reader.maxDoc()];
+ for (LeafReaderContext context : reader.leaves()) {
+ LeafReader leafReader = context.reader();
+ KnnGraphValues knnValues = leafReader.getKnnGraphValues(KNN_FIELD);
+ System.out.printf("Leaf %d has %d documents and %d graph layers\n",
+ context.ord, leafReader.maxDoc(), knnValues.getTopLevel());
+ for (int i = 0; i < knnValues.getTopLevel(); i++) {
+ printLayerFanout(knnValues, i, leafReader.maxDoc());
+ // reset for next layer:
+ knnValues = leafReader.getKnnGraphValues(KNN_FIELD);
+ }
+ }
+ /*
+ System.out.println("Whole index fanout");
+ printHist(globalHist, maxFanout);
+ */
+ }
+ }
+
+ private void printLayerFanout(KnnGraphValues knnValues, int layer, int numDocs) throws IOException {
+ int min = Integer.MAX_VALUE, max = 0, total = 0;
+ int count = 0;
+ int[] leafHist = new int[numDocs];
+ while(knnValues.nextDoc() != KnnGraphValues.NO_MORE_DOCS) {
+ // ok to call when the doc is not in a layer? we should get zero?
+ if (knnValues.getMaxLevel() < layer) {
+ //++leafHist[0];
+ continue;
+ }
+ int n = knnValues.getFriends(layer).length;
+ ++leafHist[n];
+ max = Math.max(max, n);
+ min = Math.min(min, n);
+ if (n > 0) {
+ ++count;
+ total += n;
+ }
+ }
+ System.out.printf("Layer %d size=%d, Fanout min=%d, mean=%.2f, max=%d\n", layer, count, min, total / (float) count, max);
+ printHist(leafHist, max, count, 10);
+ }
+
+ private void printHist(int[] hist, int max, int count, int nbuckets) {
+ System.out.print("%");
+ for (int i=0; i <= nbuckets; i ++) {
+ System.out.printf("%4d", i * 100 / nbuckets);
+ }
+ System.out.printf("\n %4d", hist[0]);
+ int total = 0, ibucket = 1;
+ for (int i = 1; i <= max && ibucket <= nbuckets; i++) {
+ total += hist[i];
+ if (total >= count * ibucket / nbuckets) {
+ System.out.printf("%4d", i);
+ ++ibucket;
+ }
+ }
+ System.out.println("");
+ }
+
+ private void testSearch(Path indexPath, int numIters) throws IOException {
+ float[][] targets = new float[numIters][];
+ TopDocs[] results = new TopDocs[numIters];
+ for (int i = 0; i < numIters; i++) {
+ targets[i] = new float[dim];
+ randomVector(targets[i]);
+ }
+ System.out.println("running " + numIters + " targets; topK=" + topK + ", numProbe=" + numProbe);
+ long start = System.nanoTime();
+ try (Directory dir = FSDirectory.open(indexPath);
+ DirectoryReader reader = DirectoryReader.open(dir)) {
+ IndexSearcher searcher = new IndexSearcher(reader);
+ int result = 0;
+ for (int i = 0; i < numIters; i++) {
+ KnnGraphQuery query = new KnnGraphQuery(KNN_FIELD, targets[i], numProbe);
+ results[i] = searcher.search(query, topK);
+ for (ScoreDoc scoreDoc : results[i].scoreDocs) {
+ int id = searcher.doc(scoreDoc.doc).getFields().get(0).numericValue().intValue();
+ scoreDoc.doc = id;
+ }
+ }
+ }
+ long elapsed = (System.nanoTime() - start) / 1_000_000; // ns -> ms
+ System.out.println("completed " + numIters + " searches in " + elapsed + " ms: " + (1000 * numIters / elapsed) + " QPS");
+ System.out.println("checking results");
+ checkResults(targets, results);
+ }
+
+ private void checkResults(float[][] targets, TopDocs[] results) {
+ int[] expected = new int[topK];
+ int totalMatches = 0;
+ for (int i = 0; i < results.length; i++) {
+ if (results[i].scoreDocs.length != topK) {
+ System.err.println("search " + i + " got " + results[i].scoreDocs.length + " results, expecting " + topK);
+ }
+ getActualNN(targets[i], 0, expected, 0);
+ int matched = compareNN(expected, results[i]);
+ totalMatches += matched;
+ }
+ System.out.println("total matches = " + totalMatches);
+ System.out.println("Average overlap = " + (100.0 * totalMatches / (results.length * topK)) + "%");
+ }
+
+ int compareNN(int[] expected, TopDocs results) {
+ int matched = 0;
+ int i = 0;
+ /*
+ System.out.print("expected=");
+ for (int j = 0; j < expected.length; j++) {
+ System.out.print(expected[j]);
+ System.out.print(", ");
+ }
+ System.out.print('\n');
+ System.out.println("results=");
+ for (int j = 0; j < results.scoreDocs.length; j++) {
+ System.out.print("" + results.scoreDocs[j].doc + ":" + results.scoreDocs[j].score + ", ");
+ }
+ System.out.print('\n');
+ */
+ Set<Integer> expectedSet = new HashSet<>();
+ for (int doc : expected) {
+ expectedSet.add(doc);
+ }
+ for (ScoreDoc scoreDoc : results.scoreDocs) {
+ if (expectedSet.contains(scoreDoc.doc)) {
+ ++matched;
+ }
+ }
+ return matched;
+ }
+
+ /** Find the closest this.topK vectors in this.vectors to a target vector by exhaustive
+ * comparison to all vectors.
+ * @param targetArray an array containing the target vector
+ * @param targetOffset offset of the target vector in the target parameter
+ * @param nn an array in which to write the resulting nearest neighbor vector indexes
+ * @param nnOfset the offset in the nearest neighbor array at which to start writing
+ */
+ void getActualNN(float[] targetArray, int targetOffset, int[] nn, int nnOffset) {
+ final ScoreDocQueue queue = new ScoreDocQueue(topK);
+ assert queue.size() == topK : " queue.size()=" + queue.size();
+ int vectorOffset = 0;
+ ScoreDoc bottom = queue.top();
+ for (int j = 0; j < numDocs; j++) {
+ if (targetArray == vectors && targetOffset == vectorOffset) {
+ continue;
+ }
+ float d = distance(targetArray, targetOffset, vectorOffset, bottom.score);
+ if (d < bottom.score) {
+ bottom.doc = j;
+ bottom.score = d;
+ bottom = queue.updateTop();
+ bottom = queue.top();
+ }
+ vectorOffset += dim;
+ }
+ assert queue.size() == topK;
+ nnOffset += topK;
+ for (int k = 1; k <= topK; k++) {
+ ScoreDoc scoreDoc = queue.pop();
+ //System.out.println("" + scoreDoc.doc + ":" + scoreDoc.score);
+ nn[nnOffset - k] = scoreDoc.doc;
+ }
+ }
+
+ private void randomVector(float[] vector) {
+ for(int i =0; i < vector.length; i++) {
+ vector[i] = random.nextFloat();
+ }
+ }
+
+ private void createIndex(Path indexPath) throws IOException {
+ IndexWriterConfig iwc = new IndexWriterConfig()
+ .setOpenMode(IndexWriterConfig.OpenMode.CREATE);
+ System.out.println("creating index in " + indexPath);
+ long start = System.nanoTime();
+ try (FSDirectory dir = FSDirectory.open(indexPath);
+ IndexWriter iw = new IndexWriter(dir, iwc)) {
+ for (int i = 0; i < numDocs; i++) {
+ float[] vector = new float[dim];
+ System.arraycopy(vectors, i * dim, vector, 0, dim);
+ Document doc = new Document();
+ doc.add(new VectorField(KNN_FIELD, vector, VectorValues.DistanceFunction.EUCLIDEAN));
+ doc.add(new StoredField(ID_FIELD, i));
+ iw.addDocument(doc);
+ }
+ }
+ long elapsed = System.nanoTime() - start;
+ System.out.println("Indexed " + numDocs + " documents in " + elapsed / 1_000_000 + "ms");
+ }
+
+ private void create(String dataFile) throws IOException {
+ generateRandomVectors(dim * numDocs);
+ System.out.println("Generated " + numDocs + " random vectors");
+ computeNearest();
+ writeDataFile(dataFile);
+ }
+
+ private void readDataFile(String dataFile) throws IOException {
+ try (InputStream in = Files.newInputStream(Paths.get(dataFile));
+ BufferedInputStream bin = new BufferedInputStream(in);
+ DataInputStream din = new DataInputStream(bin)) {
+ numDocs = din.readInt();
+ dim = din.readInt();
+ topK = din.readInt();
+ vectors = new float[numDocs * dim];
+ for (int i = 0; i < vectors.length; i++) {
+ vectors[i] = din.readFloat();
+ }
+ nabors = new int[numDocs * topK];
+ for (int i = 0; i < nabors.length; i++) {
+ nabors[i] = din.readInt();
+ }
+ }
+ }
+
+ private void writeDataFile(String dataFile) throws IOException {
+ try (OutputStream out = Files.newOutputStream(Paths.get(dataFile));
+ BufferedOutputStream bout = new BufferedOutputStream(out);
+ DataOutputStream dout = new DataOutputStream(bout)) {
+ dout.writeInt(numDocs);
+ dout.writeInt(dim);
+ dout.writeInt(topK);
+ for (int i = 0; i < vectors.length; i++) {
+ dout.writeFloat(vectors[i]);
+ }
+ for (int i = 0; i < nabors.length; i++) {
+ dout.writeInt(nabors[i]);
+ }
+ }
+ }
+
+ private void generateRandomVectors(int size) {
+ System.out.println("Allocating " + size * 4 / 1024 / 1024 + "MB");
+ vectors = new float[size];
+ randomVector(vectors);
+ }
+
+ private void computeNearest() {
+ nabors = new int[topK * numDocs];
+ System.out.println("finding nearest...");
+ for (int i = 0; i < numDocs; i++) {
+ if (i % 1000 == 1) {
+ System.out.println(" " + (i - 1));
+ }
+ getActualNN(vectors, i * dim, nabors, i * topK);
+ }
+ }
+
+ private float distance(float[] target, int targetOffset, int vectorOffset, float scoreToBeat) {
+ float total = 0;
+ for (int i = 0; i < dim; i++) {
+ float d = target[targetOffset++] - vectors[vectorOffset++];
+ total += d * d;
+ if (total > scoreToBeat) {
+ // return early since every dimension of the score is positive; it can only increase
+ return Float.MAX_VALUE;
+ }
+ }
+ return total;
+ }
+
+ private static void usage() {
+ String error = "Usage: TestKnnGraph -generate|-search {datafile}";
+ System.err.println(error);
+ System.exit(1);
+ }
+
+ private static class ScoreDocQueue extends PriorityQueue<ScoreDoc> {
+ ScoreDocQueue(int size) {
+ super(size, () -> new ScoreDoc(-1, Float.MAX_VALUE));
+ }
+
+ @Override
+ protected boolean lessThan(ScoreDoc a, ScoreDoc b) {
+ if (a.score > b.score) {
+ return true;
+ } else if (a.score < b.score) {
+ return false;
+ } else {
+ return a.doc > b.doc;
+ }
+ }
+ }
+
+}
[lucene-solr] 05/10: Add MoreLikeThis like factory method
Posted by to...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a commit to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git
commit 624632cc308c42af4d166446841b110f1bffddad
Author: Tomoko Uchida <to...@apache.org>
AuthorDate: Sat Dec 21 19:45:47 2019 +0900
Add MoreLikeThis like factory method
---
.../codecs/lucene90/Lucene90FieldInfosFormat.java | 2 +-
.../org/apache/lucene/index/MultiVectorValues.java | 145 +++++++++++++++++++++
.../java/org/apache/lucene/index/VectorValues.java | 16 ++-
.../org/apache/lucene/search/KnnGraphQuery.java | 27 +++-
.../org/apache/lucene/search/KnnScoreWeight.java | 12 +-
.../org/apache/lucene/util/hnsw/HNSWGraph.java | 6 +-
.../apache/lucene/util/hnsw/HNSWGraphReader.java | 8 +-
7 files changed, 205 insertions(+), 11 deletions(-)
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java
index e91ab88..1ddcf97 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java
@@ -241,8 +241,8 @@ public final class Lucene90FieldInfosFormat extends FieldInfosFormat {
case NONE:
case MANHATTAN:
case EUCLIDEAN:
- return (byte)distFunc.getId();
case COSINE:
+ return (byte)distFunc.getId();
default:
// BUG
throw new AssertionError("unhandled DistanceFunction: " + distFunc);
diff --git a/lucene/core/src/java/org/apache/lucene/index/MultiVectorValues.java b/lucene/core/src/java/org/apache/lucene/index/MultiVectorValues.java
new file mode 100644
index 0000000..b8abaaa
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/index/MultiVectorValues.java
@@ -0,0 +1,145 @@
+/*
+ * 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.index;
+
+import java.io.IOException;
+import java.util.List;
+
+public class MultiVectorValues {
+
+ /** Returns a vector values for a reader */
+ public static VectorValues getVectorValues(final IndexReader r, final String field) throws IOException {
+ final List<LeafReaderContext> leaves = r.leaves();
+ final int size = leaves.size();
+ if (size == 0) {
+ return null;
+ } else if (size == 1) {
+ return leaves.get(0).reader().getVectorValues(field);
+ }
+
+ boolean anyReal = false;
+ for(LeafReaderContext leaf : leaves) {
+ FieldInfo fieldInfo = leaf.reader().getFieldInfos().fieldInfo(field);
+ if (fieldInfo != null) {
+ if (fieldInfo.getVectorNumDimensions() != 0) {
+ anyReal = true;
+ break;
+ }
+ }
+ }
+
+ if (anyReal == false) {
+ return null;
+ }
+
+
+ return new VectorValues() {
+ private int nextLeaf;
+ private VectorValues currentValues;
+ private LeafReaderContext currentLeaf;
+ private int docID = -1;
+
+ @Override
+ public int docID() {
+ return docID;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ while (true) {
+ while (currentValues == null) {
+ if (nextLeaf == leaves.size()) {
+ docID = NO_MORE_DOCS;
+ return docID;
+ }
+ currentLeaf = leaves.get(nextLeaf);
+ currentValues = currentLeaf.reader().getVectorValues(field);
+ nextLeaf++;
+ }
+
+ int newDocID = currentValues.nextDoc();
+
+ if (newDocID == NO_MORE_DOCS) {
+ currentValues = null;
+ continue;
+ } else {
+ docID = currentLeaf.docBase + newDocID;
+ return docID;
+ }
+ }
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ if (target <= docID) {
+ throw new IllegalArgumentException("can only advance beyond current document: on docID=" + docID + " but targetDocID=" + target);
+ }
+ int readerIndex = ReaderUtil.subIndex(target, leaves);
+ if (readerIndex >= nextLeaf) {
+ if (readerIndex == leaves.size()) {
+ currentValues = null;
+ docID = NO_MORE_DOCS;
+ return docID;
+ }
+ currentLeaf = leaves.get(readerIndex);
+ currentValues = currentLeaf.reader().getVectorValues(field);
+ nextLeaf = readerIndex+1;
+ if (currentValues == null) {
+ return nextDoc();
+ }
+ }
+ int newDocID = currentValues.advance(target - currentLeaf.docBase);
+ if (newDocID == NO_MORE_DOCS) {
+ currentValues = null;
+ return nextDoc();
+ } else {
+ docID = currentLeaf.docBase + newDocID;
+ return docID;
+ }
+ }
+
+
+ @Override
+ public float[] vectorValue() throws IOException {
+ return currentValues.vectorValue();
+ }
+
+ @Override
+ public boolean seek(int target) throws IOException {
+ nextLeaf = 0;
+ for (int i = 0; i < leaves.size(); i++) {
+ currentLeaf = leaves.get(i);
+ currentValues = currentLeaf.reader().getVectorValues(field);
+ if (currentValues.seek(target - currentLeaf.docBase)) {
+ return true;
+ }
+ nextLeaf++;
+ }
+ return false;
+ }
+
+ @Override
+ public long cost() {
+ // TODO
+ return 0;
+ }
+ };
+ }
+
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/index/VectorValues.java b/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
index b0aed22..e83d174 100644
--- a/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
+++ b/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
@@ -153,8 +153,20 @@ public abstract class VectorValues extends DocIdSetIterator {
COSINE(3) {
@Override
float distance(float[] v1, float[] v2) {
- // TODO
- return 0.0f;
+ assert v1.length == v2.length;
+ if (Arrays.equals(v1, v2)) {
+ return 0.0f;
+ }
+ float sum = 0.0f;
+ float squareSum1 = 0.0f;
+ float squareSum2 = 0.0f;
+ int dim = v1.length;
+ for (int i = 0; i < dim; i++) {
+ sum += v1[i] * v2[i];
+ squareSum1 += v1[i] * v1[i];
+ squareSum2 += v2[i] * v2[i];
+ }
+ return 1.0f - sum / ((float)Math.sqrt(squareSum1) * (float)Math.sqrt(squareSum2));
}
};
diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java b/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
index 0ee5ab2..b1c14f6 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
@@ -20,9 +20,17 @@ package org.apache.lucene.search;
import java.io.IOException;
import java.util.Arrays;
import java.util.Objects;
+import java.util.Set;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.FieldInfos;
import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.MultiFields;
+import org.apache.lucene.index.MultiVectorValues;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.RamUsageEstimator;
import org.apache.lucene.util.hnsw.HNSWGraphReader;
@@ -71,16 +79,31 @@ public class KnnGraphQuery extends Query implements Accountable {
* @param ef number of per-segment candidates to be scored/collected. the collector does not return results exceeding {@code ef}.
* increasing this value leads higher recall at the expense of the search speed.
* @param reader index reader
+ * @param forceReload if true forcibly reloads kNN graph
*/
- public KnnGraphQuery(String field, float[] queryVector, int ef, IndexReader reader) throws IOException {
+ public KnnGraphQuery(String field, float[] queryVector, int ef, IndexReader reader, boolean forceReload) throws IOException {
this.field = field;
this.queryVector = queryVector;
this.ef = ef;
if (reader != null) {
- this.bytesUsed = HNSWGraphReader.reloadGraph(field, reader);
+ this.bytesUsed = HNSWGraphReader.loadGraphs(field, reader, forceReload);
}
}
+ public static KnnGraphQuery like(String field, int docId, int ef, IndexReader reader, boolean forceReload) throws IOException {
+ FieldInfo fi = FieldInfos.getMergedFieldInfos(reader).fieldInfo(field);
+ int numDimensions = fi.getVectorNumDimensions();
+ if (numDimensions == 0) {
+ throw new IllegalArgumentException("Doc " + docId + " has no vector values.");
+ }
+ VectorValues vectorValues = MultiVectorValues.getVectorValues(reader, field);
+ boolean found = vectorValues.seek(docId);
+ if (!found) {
+ throw new IllegalArgumentException("Doc " + docId + " has no vector values.");
+ }
+ return new KnnGraphQuery(field, vectorValues.vectorValue(), ef, reader, forceReload);
+ }
+
@Override
public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
return new KnnScoreWeight(this, boost, scoreMode, field, queryVector, ef);
diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java b/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java
index 73c8071..02b6a5c 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java
@@ -107,7 +107,17 @@ class KnnScoreWeight extends ConstantScoreWeight {
score = 0.0f;
} else {
doc = next.docId();
- score = 1.0f / (next.distance() / numDimensions + 0.01f);
+ switch (fi.getVectorDistFunc()) {
+ case MANHATTAN:
+ case EUCLIDEAN:
+ score = 1.0f / (next.distance() / numDimensions + 0.01f);
+ break;
+ case COSINE:
+ score = 1.0f - next.distance();
+ break;
+ default:
+ break;
+ }
}
}
return doc;
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
index 6722e76..e2acd0b 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
@@ -40,6 +40,7 @@ public final class HNSWGraph implements Accountable {
private final List<Layer> layers;
private boolean frozen = false;
+ private long bytesUsed;
public HNSWGraph(VectorValues.DistanceFunction distFunc) {
this.distFunc = distFunc;
@@ -261,7 +262,10 @@ public final class HNSWGraph implements Accountable {
@Override
public long ramBytesUsed() {
- return RamUsageEstimator.sizeOfCollection(layers);
+ if (bytesUsed == 0) {
+ bytesUsed = RamUsageEstimator.sizeOfCollection(layers);
+ }
+ return bytesUsed;
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
index 8a9d0dd..b92e96b 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
@@ -61,19 +61,19 @@ public final class HNSWGraphReader {
return hnsw.searchLayer(query, ep, ef, 0, vectorValues);
}
- public static long reloadGraph(String field, IndexReader reader) throws IOException {
+ public static long loadGraphs(String field, IndexReader reader, boolean forceReload) throws IOException {
long bytesUsed = 0L;
for (LeafReaderContext ctx : reader.leaves()) {
- HNSWGraph hnsw = get(field, ctx, true);
+ HNSWGraph hnsw = get(field, ctx, forceReload);
bytesUsed += hnsw.ramBytesUsed();
}
return bytesUsed;
}
- private static HNSWGraph get(String field, LeafReaderContext context, boolean reload) throws IOException {
+ private static HNSWGraph get(String field, LeafReaderContext context, boolean forceReload) throws IOException {
GraphKey key = new GraphKey(field, context.id());
IOException[] exc = new IOException[]{null};
- if (reload) {
+ if (forceReload) {
cache.put(key, load(field, context));
} else {
cache.computeIfAbsent(key, (k -> {
[lucene-solr] 10/10: Use 'Lucene90' Codec for TestKnnGraph
Posted by to...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a commit to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git
commit f7c2ae249ebcf9da12e66ed87704e41fd9768faa
Author: Tomoko Uchida <to...@apache.org>
AuthorDate: Sun Jan 12 14:59:51 2020 +0900
Use 'Lucene90' Codec for TestKnnGraph
---
.../core/src/test/org/apache/lucene/index/TestKnnGraph.java | 11 +++++------
1 file changed, 5 insertions(+), 6 deletions(-)
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java b/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java
index fe84c22..9579ead 100644
--- a/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java
+++ b/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java
@@ -54,7 +54,7 @@ public class TestKnnGraph extends LuceneTestCase {
*/
public void testBasic() throws Exception {
try (Directory dir = newDirectory();
- IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null).setCodec(Codec.forName("Lucene90")))) {
int numDoc = atLeast(10);
int dimension = atLeast(3);
float[][] values = new float[numDoc][];
@@ -73,7 +73,7 @@ public class TestKnnGraph extends LuceneTestCase {
public void testSingleDocument() throws Exception {
try (Directory dir = newDirectory();
- IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null).setCodec(Codec.forName("Lucene90")))) {
float[][] values = new float[][]{new float[]{0, 1, 2}};
add(iw, 0, values[0]);
assertConsistentGraph(iw, values);
@@ -84,8 +84,7 @@ public class TestKnnGraph extends LuceneTestCase {
public void testSingleDocRecall() throws Exception {
try (Directory dir = newDirectory();
- IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null)
- .setCodec(Codec.forName("Lucene90")))) {
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null).setCodec(Codec.forName("Lucene90")))) {
float[][] values = new float[][]{new float[]{0, 1, 2}};
add(iw, 0, values[0]);
assertConsistentGraph(iw, values);
@@ -102,7 +101,7 @@ public class TestKnnGraph extends LuceneTestCase {
*/
public void testMerge() throws Exception {
try (Directory dir = newDirectory();
- IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null).setCodec(Codec.forName("Lucene90")))) {
int numDoc = atLeast(100);
int dimension = atLeast(10);
float[][] values = new float[numDoc][];
@@ -134,7 +133,7 @@ public class TestKnnGraph extends LuceneTestCase {
*/
public void testSearch() throws Exception {
try (Directory dir = newDirectory();
- IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null).setCodec(Codec.forName("Lucene90")))) {
// Add a document for every cartesian point in an NxN square so we can
// easily know which are the nearest neighbors to every point. Insert by iterating
// using a prime number that is not a divisor of N*N so that we will hit each point once,
[lucene-solr] 01/10: Add the codec, field type, values writer/reader,
and query for approximate knn vector search.
Posted by to...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a commit to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git
commit 6b1f096dcd733e3e00da7103a67e42c8c3c7ccc8
Author: Tomoko Uchida <to...@apache.org>
AuthorDate: Sun Dec 1 21:21:51 2019 +0900
Add the codec, field type, values writer/reader, and query for approximate knn vector search.
---
.../lucene/codecs/lucene80/Lucene80Codec.java | 6 +
.../lucene/codecs/simpletext/SimpleTextCodec.java | 7 +
.../simpletext/SimpleTextFieldInfosFormat.java | 27 +-
.../codecs/uniformsplit/TestBlockWriter.java | 3 +
.../sharedterms/STBlockReaderTest.java | 3 +
.../src/java/org/apache/lucene/codecs/Codec.java | 5 +-
.../java/org/apache/lucene/codecs/FilterCodec.java | 5 +
.../org/apache/lucene/codecs/KnnGraphFormat.java | 79 ++++
.../org/apache/lucene/codecs/KnnGraphReader.java | 50 +++
.../org/apache/lucene/codecs/KnnGraphWriter.java | 39 ++
.../codecs/lucene50/Lucene50FieldInfosFormat.java | 3 +-
.../codecs/lucene60/Lucene60FieldInfosFormat.java | 3 +-
.../lucene/codecs/lucene84/Lucene84Codec.java | 6 +
.../Lucene90Codec.java} | 43 ++-
.../Lucene90FieldInfosFormat.java} | 56 ++-
.../codecs/lucene90/Lucene90KnnGraphFormat.java | 61 +++
.../codecs/lucene90/Lucene90KnnGraphReader.java | 426 +++++++++++++++++++++
.../codecs/lucene90/Lucene90KnnGraphWriter.java | 175 +++++++++
.../java/org/apache/lucene/document/FieldType.java | 30 ++
.../org/apache/lucene/document/VectorField.java | 49 +++
.../java/org/apache/lucene/index/CodecReader.java | 31 ++
.../apache/lucene/index/DefaultIndexingChain.java | 80 ++++
.../java/org/apache/lucene/index/FieldInfo.java | 50 ++-
.../java/org/apache/lucene/index/FieldInfos.java | 79 +++-
.../org/apache/lucene/index/FilterCodecReader.java | 6 +
.../org/apache/lucene/index/FilterLeafReader.java | 10 +
.../java/org/apache/lucene/index/IndexWriter.java | 6 +-
.../apache/lucene/index/IndexableFieldType.java | 10 +
.../org/apache/lucene/index/KnnGraphValues.java | 116 ++++++
.../apache/lucene/index/KnnGraphValuesWriter.java | 333 ++++++++++++++++
.../java/org/apache/lucene/index/LeafReader.java | 8 +
.../apache/lucene/index/MergeReaderWrapper.java | 10 +
.../apache/lucene/index/ParallelLeafReader.java | 14 +
.../apache/lucene/index/SegmentCoreReaders.java | 9 +
.../org/apache/lucene/index/SegmentReader.java | 6 +
.../lucene/index/SlowCodecReaderWrapper.java | 34 ++
.../java/org/apache/lucene/index/VectorValues.java | 230 +++++++++++
.../org/apache/lucene/search/KnnGraphQuery.java | 124 ++++++
.../org/apache/lucene/search/KnnScoreWeight.java | 151 ++++++++
.../apache/lucene/util/hnsw/DeferredNeighbor.java | 75 ++++
.../apache/lucene/util/hnsw/FurthestNeighbors.java | 30 ++
.../org/apache/lucene/util/hnsw/HNSWGraph.java | 253 ++++++++++++
.../apache/lucene/util/hnsw/HNSWGraphBuilder.java | 111 ++++++
.../apache/lucene/util/hnsw/HNSWGraphReader.java | 147 +++++++
.../apache/lucene/util/hnsw/ImmutableNeighbor.java | 55 +++
.../java/org/apache/lucene/util/hnsw/Layer.java | 101 +++++
.../apache/lucene/util/hnsw/NearestNeighbors.java | 37 ++
.../java/org/apache/lucene/util/hnsw/Neighbor.java | 45 +++
.../org/apache/lucene/util/hnsw/Neighbors.java | 45 +++
.../lucene/util/hnsw/RandomLevelGenerator.java | 42 ++
.../services/org.apache.lucene.codecs.Codec | 1 +
.../apache/lucene/index/TestIndexableField.java | 10 +
.../lucene/index/TestPendingSoftDeletes.java | 12 +-
.../lucene/index/TestSegmentToThreadMapping.java | 10 +
.../search/highlight/TermVectorLeafReader.java | 14 +-
.../apache/lucene/index/memory/MemoryIndex.java | 16 +-
.../lucene/index/BaseIndexFileFormatTestCase.java | 3 +-
.../apache/lucene/index/MismatchedLeafReader.java | 2 +
.../apache/lucene/index/RandomPostingsTester.java | 4 +-
.../java/org/apache/lucene/search/QueryUtils.java | 12 +
.../solr/index/SlowCompositeReaderWrapper.java | 10 +
.../java/org/apache/solr/schema/SchemaField.java | 11 +
.../solr/search/CollapsingQParserPlugin.java | 2 +
.../src/java/org/apache/solr/search/Insanity.java | 3 +-
.../apache/solr/uninverting/UninvertingReader.java | 3 +-
.../test/org/apache/solr/search/TestDocSet.java | 12 +
66 files changed, 3395 insertions(+), 54 deletions(-)
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/codecs/lucene80/Lucene80Codec.java b/lucene/backward-codecs/src/java/org/apache/lucene/codecs/lucene80/Lucene80Codec.java
index bc4e5f3..596d8c9 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/codecs/lucene80/Lucene80Codec.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/codecs/lucene80/Lucene80Codec.java
@@ -20,6 +20,7 @@ import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.CompoundFormat;
import org.apache.lucene.codecs.DocValuesFormat;
import org.apache.lucene.codecs.FieldInfosFormat;
+import org.apache.lucene.codecs.KnnGraphFormat;
import org.apache.lucene.codecs.LiveDocsFormat;
import org.apache.lucene.codecs.NormsFormat;
import org.apache.lucene.codecs.PointsFormat;
@@ -128,4 +129,9 @@ public class Lucene80Codec extends Codec {
public final NormsFormat normsFormat() {
return normsFormat;
}
+
+ @Override
+ public final KnnGraphFormat knnGraphFormat() {
+ return KnnGraphFormat.EMPTY;
+ }
}
diff --git a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextCodec.java b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextCodec.java
index 109fec9..6e94f37 100644
--- a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextCodec.java
+++ b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextCodec.java
@@ -21,6 +21,7 @@ import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.CompoundFormat;
import org.apache.lucene.codecs.DocValuesFormat;
import org.apache.lucene.codecs.FieldInfosFormat;
+import org.apache.lucene.codecs.KnnGraphFormat;
import org.apache.lucene.codecs.LiveDocsFormat;
import org.apache.lucene.codecs.NormsFormat;
import org.apache.lucene.codecs.PointsFormat;
@@ -100,4 +101,10 @@ public final class SimpleTextCodec extends Codec {
public PointsFormat pointsFormat() {
return pointsFormat;
}
+
+ @Override
+ public KnnGraphFormat knnGraphFormat() {
+ // TODO
+ return KnnGraphFormat.EMPTY;
+ }
}
diff --git a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextFieldInfosFormat.java b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextFieldInfosFormat.java
index f52a84a..ca17ef4 100644
--- a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextFieldInfosFormat.java
+++ b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextFieldInfosFormat.java
@@ -30,6 +30,7 @@ import org.apache.lucene.index.FieldInfos;
import org.apache.lucene.index.IndexFileNames;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.SegmentInfo;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.store.ChecksumIndexInput;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.IOContext;
@@ -67,6 +68,8 @@ public class SimpleTextFieldInfosFormat extends FieldInfosFormat {
static final BytesRef DATA_DIM_COUNT = new BytesRef(" data dimensional count ");
static final BytesRef INDEX_DIM_COUNT = new BytesRef(" index dimensional count ");
static final BytesRef DIM_NUM_BYTES = new BytesRef(" dimensional num bytes ");
+ static final BytesRef VECTOR_NUM_DIMS = new BytesRef(" vector number of dimensions ");
+ static final BytesRef VECTOR_DIST_FUNC = new BytesRef(" vector distance function ");
static final BytesRef SOFT_DELETES = new BytesRef(" soft-deletes ");
@Override
@@ -147,12 +150,22 @@ public class SimpleTextFieldInfosFormat extends FieldInfosFormat {
int dimensionalNumBytes = Integer.parseInt(readString(DIM_NUM_BYTES.length, scratch));
SimpleTextUtil.readLine(input, scratch);
+ assert StringHelper.startsWith(scratch.get(), VECTOR_NUM_DIMS);
+ int vectorNumDimensions = Integer.parseInt(readString(VECTOR_NUM_DIMS.length, scratch));
+
+ SimpleTextUtil.readLine(input, scratch);
+ assert StringHelper.startsWith(scratch.get(), VECTOR_DIST_FUNC);
+ String distFunc = readString(VECTOR_DIST_FUNC.length, scratch);
+ VectorValues.DistanceFunction vectorDistFunc = distanceFunction(distFunc);
+
+ SimpleTextUtil.readLine(input, scratch);
assert StringHelper.startsWith(scratch.get(), SOFT_DELETES);
boolean isSoftDeletesField = Boolean.parseBoolean(readString(SOFT_DELETES.length, scratch));
infos[i] = new FieldInfo(name, fieldNumber, storeTermVector,
omitNorms, storePayloads, indexOptions, docValuesType, dvGen, Collections.unmodifiableMap(atts),
- dataDimensionalCount, indexDimensionalCount, dimensionalNumBytes, isSoftDeletesField);
+ dataDimensionalCount, indexDimensionalCount, dimensionalNumBytes,
+ vectorNumDimensions, vectorDistFunc, isSoftDeletesField);
}
SimpleTextUtil.checkFooter(input);
@@ -172,6 +185,10 @@ public class SimpleTextFieldInfosFormat extends FieldInfosFormat {
public DocValuesType docValuesType(String dvType) {
return DocValuesType.valueOf(dvType);
}
+
+ public VectorValues.DistanceFunction distanceFunction(String distFunc) {
+ return VectorValues.DistanceFunction.valueOf(distFunc);
+ }
private String readString(int offset, BytesRefBuilder scratch) {
return new String(scratch.bytes(), offset, scratch.length()-offset, StandardCharsets.UTF_8);
@@ -253,6 +270,14 @@ public class SimpleTextFieldInfosFormat extends FieldInfosFormat {
SimpleTextUtil.write(out, Integer.toString(fi.getPointNumBytes()), scratch);
SimpleTextUtil.writeNewline(out);
+ SimpleTextUtil.write(out, VECTOR_NUM_DIMS);
+ SimpleTextUtil.write(out, Integer.toString(fi.getVectorNumDimensions()), scratch);
+ SimpleTextUtil.writeNewline(out);
+
+ SimpleTextUtil.write(out, VECTOR_DIST_FUNC);
+ SimpleTextUtil.write(out, fi.getVectorDistFunc().name(), scratch);
+ SimpleTextUtil.writeNewline(out);
+
SimpleTextUtil.write(out, SOFT_DELETES);
SimpleTextUtil.write(out, Boolean.toString(fi.isSoftDeletesField()), scratch);
SimpleTextUtil.writeNewline(out);
diff --git a/lucene/codecs/src/test/org/apache/lucene/codecs/uniformsplit/TestBlockWriter.java b/lucene/codecs/src/test/org/apache/lucene/codecs/uniformsplit/TestBlockWriter.java
index e3e1d06..52ef9f1 100644
--- a/lucene/codecs/src/test/org/apache/lucene/codecs/uniformsplit/TestBlockWriter.java
+++ b/lucene/codecs/src/test/org/apache/lucene/codecs/uniformsplit/TestBlockWriter.java
@@ -24,6 +24,7 @@ import org.apache.lucene.codecs.lucene84.MockTermStateFactory;
import org.apache.lucene.index.DocValuesType;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.store.ByteBuffersDataOutput;
import org.apache.lucene.store.ByteBuffersIndexOutput;
import org.apache.lucene.util.BytesRef;
@@ -120,6 +121,8 @@ public class TestBlockWriter extends LuceneTestCase {
0,
0,
0,
+ 0,
+ VectorValues.DistanceFunction.NONE,
true
);
}
diff --git a/lucene/codecs/src/test/org/apache/lucene/codecs/uniformsplit/sharedterms/STBlockReaderTest.java b/lucene/codecs/src/test/org/apache/lucene/codecs/uniformsplit/sharedterms/STBlockReaderTest.java
index f4111a4..753152399 100644
--- a/lucene/codecs/src/test/org/apache/lucene/codecs/uniformsplit/sharedterms/STBlockReaderTest.java
+++ b/lucene/codecs/src/test/org/apache/lucene/codecs/uniformsplit/sharedterms/STBlockReaderTest.java
@@ -43,6 +43,7 @@ import org.apache.lucene.index.ImpactsEnum;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.store.ByteBuffersDirectory;
import org.apache.lucene.store.DataInput;
import org.apache.lucene.store.Directory;
@@ -195,6 +196,8 @@ public class STBlockReaderTest extends LuceneTestCase {
0,
0,
0,
+ 0,
+ VectorValues.DistanceFunction.NONE,
false
);
}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/Codec.java b/lucene/core/src/java/org/apache/lucene/codecs/Codec.java
index 07797c6..cbc1984 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/Codec.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/Codec.java
@@ -57,7 +57,7 @@ public abstract class Codec implements NamedSPILoader.NamedSPI {
}
// TODO: should we use this, or maybe a system property is better?
- static Codec defaultCodec = LOADER.lookup("Lucene84");
+ static Codec defaultCodec = LOADER.lookup("Lucene90");
}
private final String name;
@@ -110,6 +110,9 @@ public abstract class Codec implements NamedSPILoader.NamedSPI {
/** Encodes/decodes points index */
public abstract PointsFormat pointsFormat();
+
+ /** Encodes/decodes knn graph */
+ public abstract KnnGraphFormat knnGraphFormat();
/** looks up a codec by name */
public static Codec forName(String name) {
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/FilterCodec.java b/lucene/core/src/java/org/apache/lucene/codecs/FilterCodec.java
index 9abd8d4..0553fd5 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/FilterCodec.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/FilterCodec.java
@@ -108,4 +108,9 @@ public abstract class FilterCodec extends Codec {
public PointsFormat pointsFormat() {
return delegate.pointsFormat();
}
+
+ @Override
+ public KnnGraphFormat knnGraphFormat() {
+ return delegate.knnGraphFormat();
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphFormat.java
new file mode 100644
index 0000000..d053301
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphFormat.java
@@ -0,0 +1,79 @@
+/*
+ * 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.codecs;
+
+import java.io.IOException;
+
+import org.apache.lucene.index.KnnGraphValues;
+import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.index.SegmentWriteState;
+import org.apache.lucene.index.VectorValues;
+
+/**
+ * Encodes/decodes per-document vector and indexed knn-graph for approximate nearest neighbor search.
+ */
+public abstract class KnnGraphFormat {
+
+ /** Sole constructor */
+ protected KnnGraphFormat() {}
+
+ /**
+ * Returns a {@link KnnGraphWriter} to write the vectors and knn-graph to the index.
+ */
+ public abstract KnnGraphWriter fieldsWriter(SegmentWriteState state) throws IOException;
+
+ /**
+ * Returns a {@link KnnGraphReader} to read the vectors and knn-graph from the index.
+ */
+ public abstract KnnGraphReader fieldsReader(SegmentReadState state) throws IOException;
+
+ public static KnnGraphFormat EMPTY = new KnnGraphFormat() {
+ @Override
+ public KnnGraphWriter fieldsWriter(SegmentWriteState state) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public KnnGraphReader fieldsReader(SegmentReadState state) throws IOException {
+ return new KnnGraphReader() {
+ @Override
+ public void checkIntegrity() throws IOException {
+ }
+
+ @Override
+ public VectorValues getVectorValues(String field) throws IOException {
+ return VectorValues.EMPTY;
+ }
+
+ @Override
+ public KnnGraphValues getGraphValues(String field) throws IOException {
+ return KnnGraphValues.EMPTY;
+ }
+
+ @Override
+ public void close() throws IOException {
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return 0;
+ }
+ };
+ }
+ };
+}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphReader.java b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphReader.java
new file mode 100644
index 0000000..fbf5bfd
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphReader.java
@@ -0,0 +1,50 @@
+/*
+ * 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.codecs;
+
+import java.io.Closeable;
+import java.io.IOException;
+
+import org.apache.lucene.index.KnnGraphValues;
+import org.apache.lucene.index.VectorValues;
+import org.apache.lucene.util.Accountable;
+
+/**
+ * Reads vectors and knn-graph from an index.
+ */
+public abstract class KnnGraphReader implements Closeable, Accountable {
+
+ /** Sole constructor */
+ protected KnnGraphReader() {}
+
+ /**
+ * Checks consistency of this reader.
+ * <p>
+ * Note that this may be costly in terms of I/O, e.g.
+ * may involve computing a checksum value against large data files.
+ * @lucene.internal
+ */
+ public abstract void checkIntegrity() throws IOException;
+
+ /** Returns the {@link VectorValues} for the given {@code field} */
+ public abstract VectorValues getVectorValues(String field) throws IOException;
+
+ /** Returns the {@link KnnGraphValues} for the given {@code field} */
+ public abstract KnnGraphValues getGraphValues(String field) throws IOException;
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java
new file mode 100644
index 0000000..c0b5638
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java
@@ -0,0 +1,39 @@
+/*
+ * 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.codecs;
+
+import java.io.Closeable;
+import java.io.IOException;
+
+import org.apache.lucene.index.FieldInfo;
+
+/**
+ * Writes vectors and knn-graph to an index.
+ */
+public abstract class KnnGraphWriter implements Closeable {
+
+ /** Sole constructor */
+ protected KnnGraphWriter() {}
+
+ /** Write all values contained in the provided reader */
+ public abstract void writeField(FieldInfo fieldInfo, KnnGraphReader values) throws IOException;
+
+ /** Called once at the end before close */
+ public abstract void finish() throws IOException;
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene50/Lucene50FieldInfosFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene50/Lucene50FieldInfosFormat.java
index 0ad0cad..8256076 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene50/Lucene50FieldInfosFormat.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene50/Lucene50FieldInfosFormat.java
@@ -31,6 +31,7 @@ import org.apache.lucene.index.FieldInfos;
import org.apache.lucene.index.IndexFileNames;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.SegmentInfo;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.store.ChecksumIndexInput;
import org.apache.lucene.store.DataOutput;
import org.apache.lucene.store.Directory;
@@ -148,7 +149,7 @@ public final class Lucene50FieldInfosFormat extends FieldInfosFormat {
lastAttributes = attributes;
try {
infos[i] = new FieldInfo(name, fieldNumber, storeTermVector, omitNorms, storePayloads,
- indexOptions, docValuesType, dvGen, attributes, 0, 0, 0, false);
+ indexOptions, docValuesType, dvGen, attributes, 0, 0, 0, 0, VectorValues.DistanceFunction.NONE, false);
infos[i].checkConsistency();
} catch (IllegalStateException e) {
throw new CorruptIndexException("invalid fieldinfo for field: " + name + ", fieldNumber=" + fieldNumber, input, e);
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60FieldInfosFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60FieldInfosFormat.java
index e2ca9ee..1d44887 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60FieldInfosFormat.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60FieldInfosFormat.java
@@ -31,6 +31,7 @@ import org.apache.lucene.index.FieldInfos;
import org.apache.lucene.index.IndexFileNames;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.SegmentInfo;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.store.ChecksumIndexInput;
import org.apache.lucene.store.DataOutput;
import org.apache.lucene.store.Directory;
@@ -164,7 +165,7 @@ public final class Lucene60FieldInfosFormat extends FieldInfosFormat {
try {
infos[i] = new FieldInfo(name, fieldNumber, storeTermVector, omitNorms, storePayloads,
indexOptions, docValuesType, dvGen, attributes,
- pointDataDimensionCount, pointIndexDimensionCount, pointNumBytes, isSoftDeletesField);
+ pointDataDimensionCount, pointIndexDimensionCount, pointNumBytes, 0, VectorValues.DistanceFunction.NONE, isSoftDeletesField);
infos[i].checkConsistency();
} catch (IllegalStateException e) {
throw new CorruptIndexException("invalid fieldinfo for field: " + name + ", fieldNumber=" + fieldNumber, input, e);
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene84/Lucene84Codec.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene84/Lucene84Codec.java
index e3f061a..b965a32 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene84/Lucene84Codec.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene84/Lucene84Codec.java
@@ -23,6 +23,7 @@ import org.apache.lucene.codecs.CompoundFormat;
import org.apache.lucene.codecs.DocValuesFormat;
import org.apache.lucene.codecs.FieldInfosFormat;
import org.apache.lucene.codecs.FilterCodec;
+import org.apache.lucene.codecs.KnnGraphFormat;
import org.apache.lucene.codecs.LiveDocsFormat;
import org.apache.lucene.codecs.NormsFormat;
import org.apache.lucene.codecs.PointsFormat;
@@ -136,6 +137,11 @@ public class Lucene84Codec extends Codec {
return new Lucene60PointsFormat();
}
+ @Override
+ public final KnnGraphFormat knnGraphFormat() {
+ return KnnGraphFormat.EMPTY;
+ }
+
/** Returns the postings format that should be used for writing
* new segments of <code>field</code>.
*
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene84/Lucene84Codec.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90Codec.java
similarity index 87%
copy from lucene/core/src/java/org/apache/lucene/codecs/lucene84/Lucene84Codec.java
copy to lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90Codec.java
index e3f061a..6a911ab 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene84/Lucene84Codec.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90Codec.java
@@ -14,7 +14,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.lucene.codecs.lucene84;
+package org.apache.lucene.codecs.lucene90;
import java.util.Objects;
@@ -23,6 +23,7 @@ import org.apache.lucene.codecs.CompoundFormat;
import org.apache.lucene.codecs.DocValuesFormat;
import org.apache.lucene.codecs.FieldInfosFormat;
import org.apache.lucene.codecs.FilterCodec;
+import org.apache.lucene.codecs.KnnGraphFormat;
import org.apache.lucene.codecs.LiveDocsFormat;
import org.apache.lucene.codecs.NormsFormat;
import org.apache.lucene.codecs.PointsFormat;
@@ -35,63 +36,64 @@ import org.apache.lucene.codecs.lucene50.Lucene50LiveDocsFormat;
import org.apache.lucene.codecs.lucene50.Lucene50StoredFieldsFormat;
import org.apache.lucene.codecs.lucene50.Lucene50StoredFieldsFormat.Mode;
import org.apache.lucene.codecs.lucene50.Lucene50TermVectorsFormat;
-import org.apache.lucene.codecs.lucene60.Lucene60FieldInfosFormat;
import org.apache.lucene.codecs.lucene60.Lucene60PointsFormat;
import org.apache.lucene.codecs.lucene70.Lucene70SegmentInfoFormat;
import org.apache.lucene.codecs.lucene80.Lucene80NormsFormat;
+import org.apache.lucene.codecs.lucene84.Lucene84PostingsFormat;
import org.apache.lucene.codecs.perfield.PerFieldDocValuesFormat;
import org.apache.lucene.codecs.perfield.PerFieldPostingsFormat;
/**
- * Implements the Lucene 8.4 index format, with configurable per-field postings
- * and docvalues formats.
+ * Implements the Lucene 9.0 index format
* <p>
* If you want to reuse functionality of this codec in another codec, extend
* {@link FilterCodec}.
*
- * @see org.apache.lucene.codecs.lucene84 package documentation for file format details.
+ * @see org.apache.lucene.codecs.lucene90 package documentation for file format details.
*
* @lucene.experimental
*/
-public class Lucene84Codec extends Codec {
+public class Lucene90Codec extends Codec {
private final TermVectorsFormat vectorsFormat = new Lucene50TermVectorsFormat();
- private final FieldInfosFormat fieldInfosFormat = new Lucene60FieldInfosFormat();
+ private final FieldInfosFormat fieldInfosFormat = new Lucene90FieldInfosFormat();
private final SegmentInfoFormat segmentInfosFormat = new Lucene70SegmentInfoFormat();
private final LiveDocsFormat liveDocsFormat = new Lucene50LiveDocsFormat();
private final CompoundFormat compoundFormat = new Lucene50CompoundFormat();
private final PostingsFormat defaultFormat;
-
+
private final PostingsFormat postingsFormat = new PerFieldPostingsFormat() {
@Override
public PostingsFormat getPostingsFormatForField(String field) {
- return Lucene84Codec.this.getPostingsFormatForField(field);
+ return Lucene90Codec.this.getPostingsFormatForField(field);
}
};
-
+
private final DocValuesFormat docValuesFormat = new PerFieldDocValuesFormat() {
@Override
public DocValuesFormat getDocValuesFormatForField(String field) {
- return Lucene84Codec.this.getDocValuesFormatForField(field);
+ return Lucene90Codec.this.getDocValuesFormatForField(field);
}
};
-
+
+ private final KnnGraphFormat knnGraphFormat = new Lucene90KnnGraphFormat();
+
private final StoredFieldsFormat storedFieldsFormat;
- /**
+ /**
* Instantiates a new codec.
*/
- public Lucene84Codec() {
+ public Lucene90Codec() {
this(Mode.BEST_SPEED);
}
-
- /**
+
+ /**
* Instantiates a new codec, specifying the stored fields compression
* mode to use.
* @param mode stored fields compression mode to use for newly
* flushed/merged segments.
*/
- public Lucene84Codec(Mode mode) {
- super("Lucene84");
+ public Lucene90Codec(Mode mode) {
+ super("Lucene90");
this.storedFieldsFormat = new Lucene50StoredFieldsFormat(Objects.requireNonNull(mode));
this.defaultFormat = new Lucene84PostingsFormat();
}
@@ -136,6 +138,11 @@ public class Lucene84Codec extends Codec {
return new Lucene60PointsFormat();
}
+ @Override
+ public final KnnGraphFormat knnGraphFormat() {
+ return knnGraphFormat;
+ }
+
/** Returns the postings format that should be used for writing
* new segments of <code>field</code>.
*
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60FieldInfosFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java
similarity index 84%
copy from lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60FieldInfosFormat.java
copy to lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java
index e2ca9ee..b846bc0 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60FieldInfosFormat.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java
@@ -14,7 +14,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.lucene.codecs.lucene60;
+package org.apache.lucene.codecs.lucene90;
import java.io.IOException;
@@ -31,6 +31,7 @@ import org.apache.lucene.index.FieldInfos;
import org.apache.lucene.index.IndexFileNames;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.SegmentInfo;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.store.ChecksumIndexInput;
import org.apache.lucene.store.DataOutput;
import org.apache.lucene.store.Directory;
@@ -39,7 +40,7 @@ import org.apache.lucene.store.IndexInput;
import org.apache.lucene.store.IndexOutput;
/**
- * Lucene 6.0 Field Infos format.
+ * Lucene 9.0 Field Infos format.
* <p>Field names are stored in the field info file, with suffix <tt>.fnm</tt>.
* <p>FieldInfos (.fnm) --> Header,FieldsCount, <FieldName,FieldNumber,
* FieldBits,DocValuesBits,DocValuesGen,Attributes,DimensionCount,DimensionNumBytes> <sup>FieldsCount</sup>,Footer
@@ -97,14 +98,23 @@ import org.apache.lucene.store.IndexOutput;
* <li>Attributes: a key-value map of codec-private attributes.</li>
* <li>PointDimensionCount, PointNumBytes: these are non-zero only if the field is
* indexed as points, e.g. using {@link org.apache.lucene.document.LongPoint}</li>
+ * <li>VectorNumDimensions: it is non-zero if the field is indexed as vectors.</li>
+ * <li>VectorDistFunction: a byte containing distance function used for similarity calculation.
+ * <ul>
+ * <li>0: no distance function is defined for this field.</li>
+ * <li>1: MANHATTAN distance. ({@link org.apache.lucene.index.VectorValues.DistanceFunction#MANHATTAN})</li>
+ * <li>2: EUCLIDEAN distance. ({@link org.apache.lucene.index.VectorValues.DistanceFunction#EUCLIDEAN})</li>
+ * <li>3: COSINE distance. ({@link org.apache.lucene.index.VectorValues.DistanceFunction#COSINE})</li>
+ * </ul>
+ * </li>
* </ul>
*
* @lucene.experimental
*/
-public final class Lucene60FieldInfosFormat extends FieldInfosFormat {
+public final class Lucene90FieldInfosFormat extends FieldInfosFormat {
/** Sole constructor. */
- public Lucene60FieldInfosFormat() {
+ public Lucene90FieldInfosFormat() {
}
@Override
@@ -115,9 +125,9 @@ public final class Lucene60FieldInfosFormat extends FieldInfosFormat {
FieldInfo infos[] = null;
try {
int version = CodecUtil.checkIndexHeader(input,
- Lucene60FieldInfosFormat.CODEC_NAME,
- Lucene60FieldInfosFormat.FORMAT_START,
- Lucene60FieldInfosFormat.FORMAT_CURRENT,
+ Lucene90FieldInfosFormat.CODEC_NAME,
+ Lucene90FieldInfosFormat.FORMAT_START,
+ Lucene90FieldInfosFormat.FORMAT_CURRENT,
segmentInfo.getId(), segmentSuffix);
final int size = input.readVInt(); //read in the size
@@ -153,18 +163,20 @@ public final class Lucene60FieldInfosFormat extends FieldInfosFormat {
int pointNumBytes;
int pointIndexDimensionCount = pointDataDimensionCount;
if (pointDataDimensionCount != 0) {
- if (version >= Lucene60FieldInfosFormat.FORMAT_SELECTIVE_INDEXING) {
+ if (version >= Lucene90FieldInfosFormat.FORMAT_SELECTIVE_INDEXING) {
pointIndexDimensionCount = input.readVInt();
}
pointNumBytes = input.readVInt();
} else {
pointNumBytes = 0;
}
+ final int vectorNumDimensions = input.readVInt();
+ final VectorValues.DistanceFunction vectorDistFunc = getDistFunc(input, input.readByte());
try {
infos[i] = new FieldInfo(name, fieldNumber, storeTermVector, omitNorms, storePayloads,
indexOptions, docValuesType, dvGen, attributes,
- pointDataDimensionCount, pointIndexDimensionCount, pointNumBytes, isSoftDeletesField);
+ pointDataDimensionCount, pointIndexDimensionCount, pointNumBytes, vectorNumDimensions, vectorDistFunc, isSoftDeletesField);
infos[i].checkConsistency();
} catch (IllegalStateException e) {
throw new CorruptIndexException("invalid fieldinfo for field: " + name + ", fieldNumber=" + fieldNumber, input, e);
@@ -224,6 +236,26 @@ public final class Lucene60FieldInfosFormat extends FieldInfosFormat {
}
}
+ private static byte distFuncByte(VectorValues.DistanceFunction distFunc) {
+ switch (distFunc) {
+ case NONE:
+ case MANHATTAN:
+ case COSINE:
+ return (byte)distFunc.getId();
+ default:
+ // BUG
+ throw new AssertionError("unhandled DistanceFunction: " + distFunc);
+ }
+ }
+
+ private static VectorValues.DistanceFunction getDistFunc(IndexInput input, byte b) throws IOException {
+ try {
+ return VectorValues.DistanceFunction.fromId(b);
+ } catch (IllegalArgumentException e) {
+ throw new CorruptIndexException("invalid distance function: " + b, input);
+ }
+ }
+
static {
// We "mirror" IndexOptions enum values with the constants below; let's try to ensure if we add a new IndexOption while this format is
// still used for writing, we remember to fix this encoding:
@@ -270,7 +302,7 @@ public final class Lucene60FieldInfosFormat extends FieldInfosFormat {
public void write(Directory directory, SegmentInfo segmentInfo, String segmentSuffix, FieldInfos infos, IOContext context) throws IOException {
final String fileName = IndexFileNames.segmentFileName(segmentInfo.name, segmentSuffix, EXTENSION);
try (IndexOutput output = directory.createOutput(fileName, context)) {
- CodecUtil.writeIndexHeader(output, Lucene60FieldInfosFormat.CODEC_NAME, Lucene60FieldInfosFormat.FORMAT_CURRENT, segmentInfo.getId(), segmentSuffix);
+ CodecUtil.writeIndexHeader(output, Lucene90FieldInfosFormat.CODEC_NAME, Lucene90FieldInfosFormat.FORMAT_CURRENT, segmentInfo.getId(), segmentSuffix);
output.writeVInt(infos.size());
for (FieldInfo fi : infos) {
fi.checkConsistency();
@@ -296,6 +328,8 @@ public final class Lucene60FieldInfosFormat extends FieldInfosFormat {
output.writeVInt(fi.getPointIndexDimensionCount());
output.writeVInt(fi.getPointNumBytes());
}
+ output.writeVInt(fi.getVectorNumDimensions());
+ output.writeByte(distFuncByte(fi.getVectorDistFunc()));
}
CodecUtil.writeFooter(output);
}
@@ -305,7 +339,7 @@ public final class Lucene60FieldInfosFormat extends FieldInfosFormat {
static final String EXTENSION = "fnm";
// Codec header
- static final String CODEC_NAME = "Lucene60FieldInfos";
+ static final String CODEC_NAME = "Lucene90FieldInfos";
static final int FORMAT_START = 0;
static final int FORMAT_SOFT_DELETES = 1;
static final int FORMAT_SELECTIVE_INDEXING = 2;
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphFormat.java
new file mode 100644
index 0000000..0b60250
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphFormat.java
@@ -0,0 +1,61 @@
+/*
+ * 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.codecs.lucene90;
+
+import java.io.IOException;
+
+import org.apache.lucene.codecs.KnnGraphFormat;
+import org.apache.lucene.codecs.KnnGraphReader;
+import org.apache.lucene.codecs.KnnGraphWriter;
+import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.index.SegmentWriteState;
+
+/**
+ * Lucene 9.0 vector and knn graph format, which encodes the vector values and the knn graphs
+ * in a Hierarchical NSW graph for fast approximate nearest neighbor searching.
+ * See <a href="https://arxiv.org/abs/1603.09320">this paper</a> for details.
+ *
+ */
+public final class Lucene90KnnGraphFormat extends KnnGraphFormat {
+
+ static final String META_CODEC_NAME = "Lucene90KnnGraphFormatMeta";
+ static final String VECTOR_DATA_CODEC_NAME = "Lucene90VectorFormatData";
+ static final String GRAPH_DATA_CODEC_NAME = "Lucene90KnnGraphFormatData";
+
+ static final String META_EXTENSION = "gri";
+ static final String VECTOR_DATA_EXTENSION = "vec";
+ static final String GRAPH_DATA_EXTENSION = "grh";
+
+ static final int VERSION_START = 0;
+ static final int VERSION_CURRENT = VERSION_START;
+
+ /** Sole constructor */
+ public Lucene90KnnGraphFormat() {
+ }
+
+ @Override
+ public KnnGraphWriter fieldsWriter(SegmentWriteState state) throws IOException {
+ return new Lucene90KnnGraphWriter(state);
+ }
+
+ @Override
+ public KnnGraphReader fieldsReader(SegmentReadState state) throws IOException {
+ return new Lucene90KnnGraphReader(state);
+ }
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
new file mode 100644
index 0000000..e53e260
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
@@ -0,0 +1,426 @@
+/*
+ * 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.codecs.lucene90;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.lucene.codecs.CodecUtil;
+import org.apache.lucene.codecs.KnnGraphReader;
+import org.apache.lucene.index.CorruptIndexException;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.FieldInfos;
+import org.apache.lucene.index.IndexFileNames;
+import org.apache.lucene.index.KnnGraphValues;
+import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.index.VectorValues;
+import org.apache.lucene.store.ChecksumIndexInput;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.IntsRef;
+import org.apache.lucene.util.RamUsageEstimator;
+
+/**
+ * Reads vector values and knn graphs from the index segments.
+ */
+public final class Lucene90KnnGraphReader extends KnnGraphReader {
+
+ private final FieldInfos fieldInfos;
+ private final Map<String, KnnGraphEntry> graphs = new HashMap<>();
+ private final IndexInput vectorData, graphData;
+ private final int maxDoc;
+ private long ramBytesUsed;
+
+ public Lucene90KnnGraphReader(SegmentReadState state) throws IOException {
+ this.fieldInfos = state.fieldInfos;
+ this.maxDoc = state.segmentInfo.maxDoc();
+ this.ramBytesUsed = RamUsageEstimator.shallowSizeOfInstance(getClass());
+
+ String metaFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, Lucene90KnnGraphFormat.META_EXTENSION);
+ int versionMeta = -1;
+ try (ChecksumIndexInput meta = state.directory.openChecksumInput(metaFileName, state.context)) {
+ Throwable priorE = null;
+ try {
+ versionMeta = CodecUtil.checkIndexHeader(meta,
+ Lucene90KnnGraphFormat.META_CODEC_NAME,
+ Lucene90KnnGraphFormat.VERSION_START,
+ Lucene90KnnGraphFormat.VERSION_CURRENT,
+ state.segmentInfo.getId(),
+ state.segmentSuffix);
+ readFields(meta, state.fieldInfos);
+ } catch (Throwable exception) {
+ priorE = exception;
+ } finally {
+ CodecUtil.checkFooter(meta, priorE);
+ }
+ }
+
+ boolean success = false;
+
+ String vectorDataFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, Lucene90KnnGraphFormat.VECTOR_DATA_EXTENSION);
+ this.vectorData = state.directory.openInput(vectorDataFileName, state.context);
+ try {
+ int versionVectorData = CodecUtil.checkIndexHeader(vectorData,
+ Lucene90KnnGraphFormat.VECTOR_DATA_CODEC_NAME,
+ Lucene90KnnGraphFormat.VERSION_START,
+ Lucene90KnnGraphFormat.VERSION_CURRENT,
+ state.segmentInfo.getId(),
+ state.segmentSuffix);
+ if (versionMeta != versionVectorData) {
+ throw new CorruptIndexException("Format versions mismatch: meta=" + versionMeta + ", vector data=" + versionVectorData, vectorData);
+ }
+ CodecUtil.retrieveChecksum(vectorData);
+
+ success = true;
+ } finally {
+ if (!success) {
+ IOUtils.closeWhileHandlingException(this.vectorData);
+ }
+ }
+
+ success = false;
+ String graphDataFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, Lucene90KnnGraphFormat.GRAPH_DATA_EXTENSION);
+ this.graphData = state.directory.openInput(graphDataFileName, state.context);
+ try {
+ int versionGraphData = CodecUtil.checkIndexHeader(graphData,
+ Lucene90KnnGraphFormat.GRAPH_DATA_CODEC_NAME,
+ Lucene90KnnGraphFormat.VERSION_START,
+ Lucene90KnnGraphFormat.VERSION_CURRENT,
+ state.segmentInfo.getId(),
+ state.segmentSuffix);
+ if (versionMeta != versionGraphData) {
+ throw new CorruptIndexException("Format versions mismatch: meta=" + versionMeta + ", graph data=" + versionGraphData, vectorData);
+ }
+
+ CodecUtil.retrieveChecksum(graphData);
+
+ success = true;
+ } finally {
+ if (!success) {
+ IOUtils.closeWhileHandlingException(this.graphData);
+ }
+ }
+ }
+
+ private void readFields(ChecksumIndexInput meta, FieldInfos infos) throws IOException {
+ for (int fieldNumber = meta.readInt(); fieldNumber != -1; fieldNumber = meta.readInt()) {
+ FieldInfo info = infos.fieldInfo(fieldNumber);
+ if (info == null) {
+ throw new CorruptIndexException("Invalid field number: " + fieldNumber, meta);
+ }
+
+ long vectorDataOffset = meta.readVLong();
+ long vectorDataLength = meta.readVLong();
+ long graphDataOffset = meta.readVLong();
+ long graphDataLength = meta.readVLong();
+ int topLevel = meta.readInt();
+ int epCount = meta.readInt();
+ int[] entryPoints = new int[epCount];
+ for (int i = 0; i < epCount; i++) {
+ entryPoints[i] = meta.readVInt();
+ }
+ int offsetCount = meta.readInt();
+ final Map<Integer, Integer> docToOrd = new HashMap<>();
+ final Map<Integer, Long> docToOffsets = new HashMap<>();
+ for (int i = 0; i < offsetCount; i++) {
+ int doc = meta.readVInt();
+ docToOrd.put(doc, i);
+ docToOffsets.put(doc, meta.readVLong());
+ }
+
+ KnnGraphEntry entry = new KnnGraphEntry(vectorDataOffset, vectorDataLength, graphDataOffset, graphDataLength, topLevel, entryPoints, Map.copyOf(docToOrd), Map.copyOf(docToOffsets));
+ ramBytesUsed += RamUsageEstimator.shallowSizeOfInstance(KnnGraphEntry.class);
+ ramBytesUsed += RamUsageEstimator.shallowSizeOf(entry.docToVectorOrd);
+ ramBytesUsed += RamUsageEstimator.shallowSizeOf(entry.docToGraphOffset);
+ graphs.put(info.name, entry);
+ }
+ }
+
+
+ @Override
+ public void checkIntegrity() throws IOException {
+ // TODO
+ }
+
+ @Override
+ public VectorValues getVectorValues(String field) throws IOException {
+ FieldInfo info = fieldInfos.fieldInfo(field);
+ if (info == null) {
+ return VectorValues.EMPTY;
+ }
+ int numDims = info.getVectorNumDimensions();
+ if (numDims <= 0) {
+ return VectorValues.EMPTY;
+ }
+
+ final KnnGraphEntry entry = graphs.get(field);
+ if (entry == null) {
+ return VectorValues.EMPTY;
+ }
+
+ final IndexInput bytesSlice = vectorData.slice("vector-data", entry.vectorDataOffset, entry.vectorDataLength);
+ return new RandomAccessVectorValuesReader(maxDoc, numDims, entry, bytesSlice);
+ }
+
+ @Override
+ public KnnGraphValues getGraphValues(String field) throws IOException {
+ final KnnGraphEntry entry = graphs.get(field);
+ if (entry == null) {
+ return KnnGraphValues.EMPTY;
+ }
+ final IndexInput bytesSlice = graphData.slice("knn-graph-data", entry.graphDataOffset, entry.graphDataLength);
+ return new IndexedKnnGraphReader(maxDoc, graphs.get(field), bytesSlice);
+ }
+
+ @Override
+ public void close() throws IOException {
+ vectorData.close();
+ graphData.close();
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return 0;
+ }
+
+ private static class KnnGraphEntry {
+ final long vectorDataOffset;
+ final long vectorDataLength;
+ final long graphDataOffset;
+ final long graphDataLength;
+ final int topLevel;
+ final int[] enterPoints;
+ final Map<Integer, Integer> docToVectorOrd;
+ final Map<Integer, Long> docToGraphOffset;
+
+ KnnGraphEntry(long vectorDataOffset, long vectorDataLength, long graphDataOffset, long graphDataLength, int topLevel, int[] enterPoints,
+ Map<Integer, Integer> docToVectorOrd, Map<Integer, Long> docToGraphOffset) {
+ this.vectorDataOffset = vectorDataOffset;
+ this.vectorDataLength = vectorDataLength;
+ this.graphDataOffset = graphDataOffset;
+ this.graphDataLength = graphDataLength;
+ this.topLevel = topLevel;
+ this.enterPoints = enterPoints;
+ this.docToVectorOrd = docToVectorOrd;
+ this.docToGraphOffset = docToGraphOffset;
+ }
+ }
+
+ private static class FriendsRef implements Accountable {
+ final int maxLevel;
+ final int[][] friendsAtLevel;
+
+ long ramBytesUsed;
+
+ FriendsRef(int maxLevel) {
+ if (maxLevel < 0) {
+ this.maxLevel = -1;
+ friendsAtLevel = new int[0][0];
+ } else {
+ this.maxLevel = maxLevel;
+ this.friendsAtLevel = new int[maxLevel+1][];
+ }
+ this.ramBytesUsed = RamUsageEstimator.shallowSizeOfInstance(getClass());
+ }
+
+ void setFriendsAtLevel(int level, int[] friends) {
+ assert level >= 0 && level <= this.maxLevel;
+ friendsAtLevel[level] = friends;
+ ramBytesUsed += RamUsageEstimator.sizeOf(friends);
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return ramBytesUsed;
+ }
+ }
+
+ /** Read the vector values from the index input stream. This allows random access to the underlying vectors data. */
+ private final static class RandomAccessVectorValuesReader extends VectorValues {
+
+ final int maxDoc;
+ final int numDims;
+ final KnnGraphEntry entry;
+ final IndexInput dataIn;
+
+ int doc = -1;
+ byte[] binaryValue;
+
+ RandomAccessVectorValuesReader(int maxDoc, int numDims, KnnGraphEntry entry, IndexInput dataIn) {
+ this.maxDoc = maxDoc;
+ this.numDims = numDims;
+ this.entry = entry;
+ this.dataIn = dataIn;
+ this.binaryValue = new byte[Float.BYTES * numDims];
+ }
+
+ @Override
+ public float[] vectorValue() throws IOException {
+ if (doc == NO_MORE_DOCS) {
+ return null;
+ }
+ return VectorValues.decode(binaryValue, numDims);
+ }
+
+ @Override
+ public boolean seek(int target) throws IOException {
+ if (target > maxDoc) {
+ doc = NO_MORE_DOCS;
+ return false;
+ }
+ doc = target;
+ Integer ord = entry.docToVectorOrd.get(doc);
+ if (ord == null) {
+ return false;
+ }
+ int offset = Float.BYTES * numDims * ord;
+ assert offset >= 0;
+ dataIn.seek(offset);
+ dataIn.readBytes(binaryValue, 0, binaryValue.length);
+ return true;
+ }
+
+ @Override
+ public int docID() {
+ return doc;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ return advance(doc + 1);
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ int _target = target;
+ boolean found;
+ do {
+ found = seek(_target++);
+ } while (!found && doc != NO_MORE_DOCS);
+ return doc;
+ }
+
+ @Override
+ public long cost() {
+ return maxDoc;
+ }
+ }
+
+ /** Read the knn graph from the index input stream */
+ private final static class IndexedKnnGraphReader extends KnnGraphValues {
+
+ final int maxDoc;
+ final KnnGraphEntry entry;
+ final IndexInput dataIn;
+
+ int doc = -1;
+ int maxLevel = -1;
+ FriendsRef friendsRef = null;
+
+ IndexedKnnGraphReader(int maxDoc, KnnGraphEntry entry, IndexInput dataIn) {
+ this.maxDoc = maxDoc;
+ this.entry = entry;
+ this.dataIn = dataIn;
+ }
+
+ @Override
+ public int getTopLevel() {
+ return entry.topLevel;
+ }
+
+ @Override
+ public int getMaxLevel() {
+ return maxLevel;
+ }
+
+ @Override
+ public IntsRef getFriends(int level) {
+ if (level < 0 || level > maxLevel) {
+ throw new IllegalArgumentException("level must be >= 0 or <= maxLevel (=" + maxLevel + "); got " + level);
+ }
+ int[] friends = friendsRef.friendsAtLevel[level];
+ return new IntsRef(friends, 0, friends.length);
+ }
+
+ @Override
+ public boolean advanceExact(int target) throws IOException {
+ advance(target);
+ return doc == target;
+ }
+
+ @Override
+ public int docID() {
+ return doc;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ return advance(doc + 1);
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ if (doc > target || target > maxDoc) {
+ return doc = NO_MORE_DOCS;
+ }
+ doc = target - 1;
+ Long offset;
+ boolean found;
+ do {
+ offset = entry.docToGraphOffset.get(++doc);
+ found = offset != null;
+ } while (!found && doc < maxDoc);
+ if (doc == maxDoc) {
+ return doc = NO_MORE_DOCS;
+ }
+ dataIn.seek(offset);
+ int maxLevel = dataIn.readInt();
+ assert maxLevel > 0;
+ this.maxLevel = maxLevel;
+ this.friendsRef = readFriends(maxLevel);
+ return doc;
+ }
+
+ private FriendsRef readFriends(int maxLevel) throws IOException {
+ FriendsRef friendsRef = new FriendsRef(maxLevel);
+ for (int l = maxLevel; l >= 0; l--) {
+ int friendSize = dataIn.readInt();
+ assert friendSize > 0;
+ int[] frineds = new int[friendSize];
+ int friendId = dataIn.readVInt(); // first friend id
+ frineds[0] = friendId;
+ for (int i = 1; i < friendSize; i++) {
+ int delta = dataIn.readVInt();
+ assert delta > 0;
+ friendId += delta;
+ frineds[i] = friendId;
+ }
+ friendsRef.setFriendsAtLevel(l, frineds);
+ }
+ return friendsRef;
+ }
+
+
+ @Override
+ public long cost() {
+ return maxDoc;
+ }
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
new file mode 100644
index 0000000..6c16fdf
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
@@ -0,0 +1,175 @@
+/*
+ * 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.codecs.lucene90;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.lucene.codecs.CodecUtil;
+import org.apache.lucene.codecs.KnnGraphReader;
+import org.apache.lucene.codecs.KnnGraphWriter;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.IndexFileNames;
+import org.apache.lucene.index.KnnGraphValues;
+import org.apache.lucene.index.SegmentWriteState;
+import org.apache.lucene.index.VectorValues;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.IntsRef;
+
+/**
+ * Writes vector values and knn graphs to index segments.
+ */
+public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
+
+ private final IndexOutput meta, vectorData, graphData;
+
+ private boolean finished;
+
+ public Lucene90KnnGraphWriter(SegmentWriteState state) throws IOException {
+ assert state.fieldInfos.hasVectorValues();
+
+ String metaFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, Lucene90KnnGraphFormat.META_EXTENSION);
+ meta = state.directory.createOutput(metaFileName, state.context);
+
+ String vectorDataFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, Lucene90KnnGraphFormat.VECTOR_DATA_EXTENSION);
+ vectorData = state.directory.createOutput(vectorDataFileName, state.context);
+
+ String graphDataFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, Lucene90KnnGraphFormat.GRAPH_DATA_EXTENSION);
+ graphData = state.directory.createOutput(graphDataFileName, state.context);
+
+ try {
+ CodecUtil.writeIndexHeader(meta,
+ Lucene90KnnGraphFormat.META_CODEC_NAME,
+ Lucene90KnnGraphFormat.VERSION_CURRENT,
+ state.segmentInfo.getId(), state.segmentSuffix);
+ CodecUtil.writeIndexHeader(vectorData,
+ Lucene90KnnGraphFormat.VECTOR_DATA_CODEC_NAME,
+ Lucene90KnnGraphFormat.VERSION_CURRENT,
+ state.segmentInfo.getId(), state.segmentSuffix);
+ CodecUtil.writeIndexHeader(graphData,
+ Lucene90KnnGraphFormat.GRAPH_DATA_CODEC_NAME,
+ Lucene90KnnGraphFormat.VERSION_CURRENT,
+ state.segmentInfo.getId(), state.segmentSuffix);
+ } catch (IOException e) {
+ IOUtils.closeWhileHandlingException(this);
+ }
+ }
+
+ @Override
+ public void writeField(FieldInfo fieldInfo, KnnGraphReader values) throws IOException {
+ Map<Integer, Long> docToOffset = new HashMap<>();
+ List<Integer> enterPoints = new ArrayList<>();
+ long vectorDataOffset = vectorData.getFilePointer();
+ long graphDataOffset = graphData.getFilePointer();
+
+ int numDims = fieldInfo.getVectorNumDimensions();
+
+ VectorValues vectors = values.getVectorValues(fieldInfo.name);
+ KnnGraphValues graph = values.getGraphValues(fieldInfo.name);
+
+ for (int docV = vectors.nextDoc(), docG = graph.nextDoc();
+ docV != DocIdSetIterator.NO_MORE_DOCS && docG != DocIdSetIterator.NO_MORE_DOCS;
+ docV = vectors.nextDoc(), docG = graph.nextDoc()) {
+ assert docV == docG; // must be same
+
+ // write vector value
+ byte[] binaryValue = vectors.binaryValue();
+ VectorValues.verifyNumDimensions(binaryValue, numDims);
+ vectorData.writeBytes(binaryValue, binaryValue.length);
+
+ // write knn graph value
+ docToOffset.put(docG, graphData.getFilePointer() - graphDataOffset);
+ if (graph.isEnterPoint()) {
+ enterPoints.add(docG);
+ }
+ graphData.writeInt(graph.getMaxLevel());
+ for (int l = graph.getMaxLevel(); l >= 0; l--) {
+ IntsRef friends = graph.getFriends(l);
+ assert friends.length > 0;
+ graphData.writeInt(friends.length);
+ int stop = friends.offset + friends.length;
+ // sort friend ids
+ Arrays.sort(friends.ints, friends.offset, stop);
+ // write the smallest friend id
+ graphData.writeVInt(friends.ints[friends.offset]);
+ for (int i = friends.offset + 1; i < stop; i++) {
+ // write delta
+ int delta = friends.ints[i] - friends.ints[i - 1];
+ assert delta > 0;
+ graphData.writeVInt(delta);
+ }
+ }
+ }
+ assert vectors.nextDoc() == DocIdSetIterator.NO_MORE_DOCS; // must be exhausted
+ assert graph.nextDoc() == DocIdSetIterator.NO_MORE_DOCS; // must be exhausted
+
+ long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset;
+ long graphDataLength = graphData.getFilePointer() - graphDataOffset;
+
+ writeMeta(fieldInfo, vectorDataOffset, vectorDataLength, graphDataOffset, graphDataLength, graph.getTopLevel(), enterPoints, docToOffset);
+ }
+
+ private void writeMeta(FieldInfo field, long vectorDataOffset, long vectorDataLength, long graphDataOffset, long graphDataLength,
+ int topLevel, List<Integer> eps, Map<Integer, Long> docToOffset) throws IOException {
+ meta.writeInt(field.number);
+ meta.writeVLong(vectorDataOffset);
+ meta.writeVLong(vectorDataLength);
+ meta.writeVLong(graphDataOffset);
+ meta.writeVLong(graphDataLength);
+ meta.writeInt(topLevel);
+ meta.writeInt(eps.size());
+ for (Integer ep : eps) {
+ meta.writeVInt(ep);
+ }
+ meta.writeInt(docToOffset.size());
+ for (Map.Entry<Integer, Long> entry : docToOffset.entrySet()) {
+ meta.writeVInt(entry.getKey());
+ meta.writeVLong(entry.getValue());
+ }
+ }
+
+ @Override
+ public void finish() throws IOException {
+ if (finished) {
+ throw new IllegalStateException("already finished");
+ }
+ finished = true;
+
+ if (meta != null) {
+ meta.writeInt(-1);
+ CodecUtil.writeFooter(meta);
+ }
+ if (vectorData != null) {
+ CodecUtil.writeFooter(vectorData);
+ }
+ if (graphData != null) {
+ CodecUtil.writeFooter(graphData);
+ }
+ }
+
+ @Override
+ public void close() throws IOException {
+ IOUtils.close(meta, vectorData, graphData);
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/document/FieldType.java b/lucene/core/src/java/org/apache/lucene/document/FieldType.java
index 8dd802c..4355dfb 100644
--- a/lucene/core/src/java/org/apache/lucene/document/FieldType.java
+++ b/lucene/core/src/java/org/apache/lucene/document/FieldType.java
@@ -25,6 +25,7 @@ import org.apache.lucene.index.DocValuesType;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.IndexableFieldType;
import org.apache.lucene.index.PointValues;
+import org.apache.lucene.index.VectorValues;
/**
* Describes the properties of a field.
@@ -44,6 +45,8 @@ public class FieldType implements IndexableFieldType {
private int dataDimensionCount;
private int indexDimensionCount;
private int dimensionNumBytes;
+ private int vectorNumDimensions;
+ private VectorValues.DistanceFunction vectorDistFunc = VectorValues.DistanceFunction.NONE;
private Map<String, String> attributes;
/**
@@ -62,6 +65,8 @@ public class FieldType implements IndexableFieldType {
this.dataDimensionCount = ref.pointDataDimensionCount();
this.indexDimensionCount = ref.pointIndexDimensionCount();
this.dimensionNumBytes = ref.pointNumBytes();
+ this.vectorNumDimensions = ref.vectorNumDimensions();
+ this.vectorDistFunc = ref.vectorDistFunc();
if (ref.getAttributes() != null) {
this.attributes = new HashMap<>(ref.getAttributes());
}
@@ -348,6 +353,31 @@ public class FieldType implements IndexableFieldType {
return dimensionNumBytes;
}
+ public void setVectorDimensionsAndDistanceFunction(int numDimensions, VectorValues.DistanceFunction distFunc) {
+ if (numDimensions < 0) {
+ throw new IllegalArgumentException("vector numDimensions must be >= 0; got " + numDimensions);
+ }
+ if (numDimensions > VectorValues.MAX_DIMENSIONS) {
+ throw new IllegalArgumentException("vector numDimensions must be <= VectorValues.MAX_DIMENSIONS (=" + VectorValues.MAX_DIMENSIONS + "); got " + numDimensions);
+ }
+ if (numDimensions == 0 && distFunc != VectorValues.DistanceFunction.NONE) {
+ throw new IllegalArgumentException("vector distFunc must be NONE when the vector numDimensions = 0; got " + distFunc);
+ }
+
+ this.vectorNumDimensions = numDimensions;
+ this.vectorDistFunc = distFunc;
+ }
+
+ @Override
+ public int vectorNumDimensions() {
+ return vectorNumDimensions;
+ }
+
+ @Override
+ public VectorValues.DistanceFunction vectorDistFunc() {
+ return vectorDistFunc;
+ }
+
/**
* Puts an attribute value.
* <p>
diff --git a/lucene/core/src/java/org/apache/lucene/document/VectorField.java b/lucene/core/src/java/org/apache/lucene/document/VectorField.java
new file mode 100644
index 0000000..664e071
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/document/VectorField.java
@@ -0,0 +1,49 @@
+/*
+ * 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.document;
+
+import org.apache.lucene.index.VectorValues;
+
+/**
+ * Per-document vector value; {@code float} array with indexed knn graph for fast approximate nearest neighbor search.
+ */
+public class VectorField extends Field {
+
+ private static FieldType getType(int dimensions, VectorValues.DistanceFunction distFunc) {
+ if (dimensions == 0) {
+ throw new IllegalArgumentException("VectorField does not support 0 dimensions");
+ }
+ if (dimensions > VectorValues.MAX_DIMENSIONS) {
+ throw new IllegalArgumentException("VectorField does not support greater than " + VectorValues.MAX_DIMENSIONS + " dimensions");
+ }
+ FieldType type = new FieldType();
+ type.setVectorDimensionsAndDistanceFunction(dimensions, distFunc);
+ type.freeze();
+ return type;
+ }
+
+ public VectorField(String name, float[] vector, VectorValues.DistanceFunction distFunc) {
+ super(name, VectorValues.encode(vector), getType(vector.length, distFunc));
+ }
+
+ /** Changes the value of this field */
+ public void setVectorValue(float[] vector) {
+ setBytesValue(VectorValues.encode(vector));
+ }
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/index/CodecReader.java b/lucene/core/src/java/org/apache/lucene/index/CodecReader.java
index af90c7b..74c64e9 100644
--- a/lucene/core/src/java/org/apache/lucene/index/CodecReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/CodecReader.java
@@ -26,6 +26,7 @@ import java.util.Objects;
import org.apache.lucene.codecs.DocValuesProducer;
import org.apache.lucene.codecs.FieldsProducer;
+import org.apache.lucene.codecs.KnnGraphReader;
import org.apache.lucene.codecs.NormsProducer;
import org.apache.lucene.codecs.PointsReader;
import org.apache.lucene.codecs.StoredFieldsReader;
@@ -77,6 +78,12 @@ public abstract class CodecReader extends LeafReader implements Accountable {
* @lucene.internal
*/
public abstract PointsReader getPointsReader();
+
+ /**
+ * Expert: retrieve underlying KnnGraphReader
+ * @lucene.internal
+ */
+ public abstract KnnGraphReader getKnnGraphReader();
@Override
public final void document(int docID, StoredFieldVisitor visitor) throws IOException {
@@ -203,6 +210,30 @@ public abstract class CodecReader extends LeafReader implements Accountable {
}
@Override
+ public final VectorValues getVectorValues(String field) throws IOException {
+ ensureOpen();
+ FieldInfo fi = getFieldInfos().fieldInfo(field);
+ if (fi == null || fi.getVectorNumDimensions() == 0) {
+ // Field does not exist or does not index vectors
+ return null;
+ }
+
+ return getKnnGraphReader().getVectorValues(field);
+ }
+
+ @Override
+ public KnnGraphValues getKnnGraphValues(String field) throws IOException {
+ ensureOpen();
+ FieldInfo fi = getFieldInfos().fieldInfo(field);
+ if (fi == null || fi.getVectorNumDimensions() == 0) {
+ // Field does not exist or does not index vectors
+ return null;
+ }
+
+ return getKnnGraphReader().getGraphValues(field);
+ }
+
+ @Override
protected void doClose() throws IOException {
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java b/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java
index 12bfe07..3ce1246 100644
--- a/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java
+++ b/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java
@@ -31,6 +31,8 @@ import java.util.Set;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.codecs.DocValuesConsumer;
import org.apache.lucene.codecs.DocValuesFormat;
+import org.apache.lucene.codecs.KnnGraphFormat;
+import org.apache.lucene.codecs.KnnGraphWriter;
import org.apache.lucene.codecs.NormsConsumer;
import org.apache.lucene.codecs.NormsFormat;
import org.apache.lucene.codecs.NormsProducer;
@@ -145,6 +147,12 @@ final class DefaultIndexingChain extends DocConsumer {
if (docState.infoStream.isEnabled("IW")) {
docState.infoStream.message("IW", ((System.nanoTime()-t0)/1000000) + " msec to write points");
}
+
+ t0 = System.nanoTime();
+ writeVectors(state, sortMap);
+ if (docState.infoStream.isEnabled("IW")) {
+ docState.infoStream.message("IW", ((System.nanoTime()-t0)/1000000) + " msec to write vectors and knn graph");
+ }
// it's possible all docs hit non-aborting exceptions...
t0 = System.nanoTime();
@@ -294,6 +302,51 @@ final class DefaultIndexingChain extends DocConsumer {
}
}
+ /** Writes all buffered vectors and knn graph. */
+ private void writeVectors(SegmentWriteState state, Sorter.DocMap sortMap) throws IOException {
+ //PointsWriter pointsWriter = null;
+ KnnGraphWriter knnGraphWriter = null;
+ boolean success = false;
+ try {
+ for (int i=0;i<fieldHash.length;i++) {
+ PerField perField = fieldHash[i];
+ while (perField != null) {
+ if (perField.knnGraphValuesWriter != null) {
+ if (perField.fieldInfo.getVectorNumDimensions() == 0) {
+ // BUG
+ throw new AssertionError("segment=" + state.segmentInfo + ": field=\"" + perField.fieldInfo.name + "\" has no vectors but wrote them");
+ }
+ if (knnGraphWriter == null) {
+ // lazy init
+ KnnGraphFormat fmt = state.segmentInfo.getCodec().knnGraphFormat();
+ if (fmt == null) {
+ throw new IllegalStateException("field=\"" + perField.fieldInfo.name + "\" was indexed as vectors but codec does not support vectors");
+ }
+ knnGraphWriter = fmt.fieldsWriter(state);
+ }
+
+ perField.knnGraphValuesWriter.flush(state, sortMap, knnGraphWriter);
+ perField.knnGraphValuesWriter = null;
+ } else if (perField.fieldInfo.getVectorNumDimensions() != 0) {
+ // BUG
+ throw new AssertionError("segment=" + state.segmentInfo + ": field=\"" + perField.fieldInfo.name + "\" has vectors but did not write them");
+ }
+ perField = perField.next;
+ }
+ }
+ if (knnGraphWriter != null) {
+ knnGraphWriter.finish();
+ }
+ success = true;
+ } finally {
+ if (success) {
+ IOUtils.close(knnGraphWriter);
+ } else {
+ IOUtils.closeWhileHandlingException(knnGraphWriter);
+ }
+ }
+ }
+
private void writeNorms(SegmentWriteState state, Sorter.DocMap sortMap) throws IOException {
boolean success = false;
NormsConsumer normsConsumer = null;
@@ -484,6 +537,12 @@ final class DefaultIndexingChain extends DocConsumer {
}
indexPoint(fp, field);
}
+ if (fieldType.vectorNumDimensions() != 0) {
+ if (fp == null) {
+ fp = getOrAddField(fieldName, fieldType, false);
+ }
+ indexVector(fp, field);
+ }
return fieldCount;
}
@@ -633,6 +692,24 @@ final class DefaultIndexingChain extends DocConsumer {
}
}
+ /** Called from processDocument to index one field's doc value */
+ private void indexVector(PerField fp, IndexableField field) throws IOException {
+ int numDimensions = field.fieldType().vectorNumDimensions();
+ VectorValues.DistanceFunction distFunc = field.fieldType().vectorDistFunc();
+
+ // Record dimensions and distance function for this field; this setter will throw IllegalArgExc if
+ // the dimensions or distance function were already set to something different:
+ if (fp.fieldInfo.getVectorNumDimensions() == 0) {
+ fieldInfos.globalFieldNumbers.setVectorDimensionsAndDistanceFunction(fp.fieldInfo.number, fp.fieldInfo.name, numDimensions, distFunc);
+ }
+ fp.fieldInfo.setVectorDimensionsAndDistanceFunction(numDimensions, distFunc);
+
+ if (fp.knnGraphValuesWriter == null) {
+ fp.knnGraphValuesWriter = new KnnGraphValuesWriter(docWriter, fp.fieldInfo);
+ }
+ fp.knnGraphValuesWriter.addValue(docState.docID, field.binaryValue());
+ }
+
/** Returns a previously created {@link PerField}, or null
* if this field name wasn't seen yet. */
private PerField getPerField(String name) {
@@ -719,6 +796,9 @@ final class DefaultIndexingChain extends DocConsumer {
// Non-null if this field ever had points in this segment:
PointValuesWriter pointValuesWriter;
+ // Non-null if this field ever had vector values in this segment:
+ KnnGraphValuesWriter knnGraphValuesWriter;
+
/** We use this to know when a PerField is seen for the
* first time in the current document. */
long fieldGen = -1;
diff --git a/lucene/core/src/java/org/apache/lucene/index/FieldInfo.java b/lucene/core/src/java/org/apache/lucene/index/FieldInfo.java
index 5346523..d71f3b4 100644
--- a/lucene/core/src/java/org/apache/lucene/index/FieldInfo.java
+++ b/lucene/core/src/java/org/apache/lucene/index/FieldInfo.java
@@ -54,6 +54,9 @@ public final class FieldInfo {
private int pointIndexDimensionCount;
private int pointNumBytes;
+ private int vectorNumDimensions; // if it is a positive value, it means this field indexes vectors
+ private VectorValues.DistanceFunction vectorDistFunc = VectorValues.DistanceFunction.NONE;
+
// whether this field is used as the soft-deletes field
private final boolean softDeletesField;
@@ -64,7 +67,8 @@ public final class FieldInfo {
*/
public FieldInfo(String name, int number, boolean storeTermVector, boolean omitNorms, boolean storePayloads,
IndexOptions indexOptions, DocValuesType docValues, long dvGen, Map<String,String> attributes,
- int pointDataDimensionCount, int pointIndexDimensionCount, int pointNumBytes, boolean softDeletesField) {
+ int pointDataDimensionCount, int pointIndexDimensionCount, int pointNumBytes,
+ int vectorNumDimensions, VectorValues.DistanceFunction vectorDistFunc, boolean softDeletesField) {
this.name = Objects.requireNonNull(name);
this.number = number;
this.docValuesType = Objects.requireNonNull(docValues, "DocValuesType must not be null (field: \"" + name + "\")");
@@ -83,6 +87,8 @@ public final class FieldInfo {
this.pointDataDimensionCount = pointDataDimensionCount;
this.pointIndexDimensionCount = pointIndexDimensionCount;
this.pointNumBytes = pointNumBytes;
+ this.vectorNumDimensions = vectorNumDimensions;
+ this.vectorDistFunc = vectorDistFunc;
this.softDeletesField = softDeletesField;
assert checkConsistency();
}
@@ -137,6 +143,14 @@ public final class FieldInfo {
throw new IllegalStateException("field '" + name + "' cannot have a docvalues update generation without having docvalues");
}
+ if (vectorNumDimensions < 0) {
+ throw new IllegalStateException("vectorNumDimensions must be >=0; got " + vectorNumDimensions);
+ }
+
+ if (vectorNumDimensions == 0 && vectorDistFunc != VectorValues.DistanceFunction.NONE) {
+ throw new IllegalStateException("vectorDistFunc must be NONE when vectorNumDimensions = 0; got " + vectorDistFunc);
+ }
+
return true;
}
@@ -232,6 +246,40 @@ public final class FieldInfo {
return pointNumBytes;
}
+ /** Record that this field is indexed with vectors, with the specified num of dimensions and distance function */
+ public void setVectorDimensionsAndDistanceFunction(int numDimensions, VectorValues.DistanceFunction distFunc) {
+ if (numDimensions < 0) {
+ throw new IllegalArgumentException("vector numDimensions must be >= 0; got " + numDimensions);
+ }
+ if (numDimensions > VectorValues.MAX_DIMENSIONS) {
+ throw new IllegalArgumentException("vector numDimensions must be <= VectorValues.MAX_DIMENSIONS (=" + VectorValues.MAX_DIMENSIONS + "); got " + numDimensions);
+ }
+ if (numDimensions == 0 && distFunc != VectorValues.DistanceFunction.NONE) {
+ throw new IllegalArgumentException("vector distFunc must be NONE when the vector numDimensions = 0; got " + distFunc);
+ }
+ if (vectorNumDimensions != 0 && vectorNumDimensions != numDimensions) {
+ throw new IllegalArgumentException("cannot change vector numDimensions from " + vectorNumDimensions + " to " + numDimensions + " for field=\"" + name + "\"");
+ }
+ if (vectorDistFunc != VectorValues.DistanceFunction.NONE && vectorDistFunc != distFunc) {
+ throw new IllegalArgumentException("cannot change vector distFunc from " + vectorDistFunc + " to " + vectorDistFunc + " for field=\"" + name + "\"");
+ }
+
+ this.vectorNumDimensions = numDimensions;
+ this.vectorDistFunc = distFunc;
+
+ assert checkConsistency();
+ }
+
+ /** Returns the number of dimensions of the vector value */
+ public int getVectorNumDimensions() {
+ return vectorNumDimensions;
+ }
+
+ /** Returns {@link org.apache.lucene.index.VectorValues.DistanceFunction} for the field */
+ public VectorValues.DistanceFunction getVectorDistFunc() {
+ return vectorDistFunc;
+ }
+
/** Record that this field is indexed with docvalues, with the specified type */
public void setDocValuesType(DocValuesType type) {
if (type == null) {
diff --git a/lucene/core/src/java/org/apache/lucene/index/FieldInfos.java b/lucene/core/src/java/org/apache/lucene/index/FieldInfos.java
index 979a248..a8c1e4c 100644
--- a/lucene/core/src/java/org/apache/lucene/index/FieldInfos.java
+++ b/lucene/core/src/java/org/apache/lucene/index/FieldInfos.java
@@ -48,6 +48,7 @@ public class FieldInfos implements Iterable<FieldInfo> {
private final boolean hasNorms;
private final boolean hasDocValues;
private final boolean hasPointValues;
+ private final boolean hasVectorValues;
private final String softDeletesField;
// used only by fieldInfo(int)
@@ -68,6 +69,7 @@ public class FieldInfos implements Iterable<FieldInfo> {
boolean hasNorms = false;
boolean hasDocValues = false;
boolean hasPointValues = false;
+ boolean hasVectorValues = false;
String softDeletesField = null;
int size = 0; // number of elements in byNumberTemp, number of used array slots
@@ -99,6 +101,7 @@ public class FieldInfos implements Iterable<FieldInfo> {
hasDocValues |= info.getDocValuesType() != DocValuesType.NONE;
hasPayloads |= info.hasPayloads();
hasPointValues |= (info.getPointDataDimensionCount() != 0);
+ hasVectorValues |= (info.getVectorNumDimensions() != 0);
if (info.isSoftDeletesField()) {
if (softDeletesField != null && softDeletesField.equals(info.name) == false) {
throw new IllegalArgumentException("multiple soft-deletes fields [" + info.name + ", " + softDeletesField + "]");
@@ -115,6 +118,7 @@ public class FieldInfos implements Iterable<FieldInfo> {
this.hasNorms = hasNorms;
this.hasDocValues = hasDocValues;
this.hasPointValues = hasPointValues;
+ this.hasVectorValues = hasVectorValues;
this.softDeletesField = softDeletesField;
List<FieldInfo> valuesTemp = new ArrayList<>();
@@ -204,6 +208,11 @@ public class FieldInfos implements Iterable<FieldInfo> {
return hasPointValues;
}
+ /** Returns true if any fields have VectorValues */
+ public boolean hasVectorValues() {
+ return hasVectorValues;
+ }
+
/** Returns the soft-deletes field name if exists; otherwise returns null */
public String getSoftDeletesField() {
return softDeletesField;
@@ -261,6 +270,16 @@ public class FieldInfos implements Iterable<FieldInfo> {
this.dimensionNumBytes = dimensionNumBytes;
}
}
+
+ static final class FieldVectorProperties {
+ final int numDimensions;
+ final VectorValues.DistanceFunction distFunc;
+
+ FieldVectorProperties(int numDimensions, VectorValues.DistanceFunction distFunc) {
+ this.numDimensions = numDimensions;
+ this.distFunc = distFunc;
+ }
+ }
static final class FieldNumbers {
@@ -274,6 +293,8 @@ public class FieldInfos implements Iterable<FieldInfo> {
private final Map<String,FieldDimensions> dimensions;
+ private final Map<String,FieldVectorProperties> vectorProps;
+
// TODO: we should similarly catch an attempt to turn
// norms back on after they were already ommitted; today
// we silently discard the norm but this is badly trappy
@@ -288,6 +309,7 @@ public class FieldInfos implements Iterable<FieldInfo> {
this.indexOptions = new HashMap<>();
this.docValuesType = new HashMap<>();
this.dimensions = new HashMap<>();
+ this.vectorProps = new HashMap<>();
this.softDeletesFieldName = softDeletesFieldName;
}
@@ -297,7 +319,7 @@ public class FieldInfos implements Iterable<FieldInfo> {
* number assigned if possible otherwise the first unassigned field number
* is used as the field number.
*/
- synchronized int addOrGet(String fieldName, int preferredFieldNumber, IndexOptions indexOptions, DocValuesType dvType, int dataDimensionCount, int indexDimensionCount, int dimensionNumBytes, boolean isSoftDeletesField) {
+ synchronized int addOrGet(String fieldName, int preferredFieldNumber, IndexOptions indexOptions, DocValuesType dvType, int dataDimensionCount, int indexDimensionCount, int dimensionNumBytes, int vectorNumDimensions, VectorValues.DistanceFunction distFunc, boolean isSoftDeletesField) {
if (indexOptions != IndexOptions.NONE) {
IndexOptions currentOpts = this.indexOptions.get(fieldName);
if (currentOpts == null) {
@@ -330,6 +352,19 @@ public class FieldInfos implements Iterable<FieldInfo> {
dimensions.put(fieldName, new FieldDimensions(dataDimensionCount, indexDimensionCount, dimensionNumBytes));
}
}
+ if (vectorNumDimensions != 0) {
+ FieldVectorProperties props = vectorProps.get(fieldName);
+ if (props != null) {
+ if (props.numDimensions != vectorNumDimensions) {
+ throw new IllegalArgumentException("cannot change vector dimensions from " + props.numDimensions + " to " + vectorNumDimensions + " for field=\"" + fieldName + "\"");
+ }
+ if (props.distFunc != distFunc) {
+ throw new IllegalArgumentException("cannot change distance function from " + props.distFunc + " to " + distFunc + " for field=\"" + fieldName + "\"");
+ }
+ } else {
+ vectorProps.put(fieldName, new FieldVectorProperties(vectorNumDimensions, distFunc));
+ }
+ }
Integer fieldNumber = nameToNumber.get(fieldName);
if (fieldNumber == null) {
final Integer preferredBoxed = Integer.valueOf(preferredFieldNumber);
@@ -408,6 +443,24 @@ public class FieldInfos implements Iterable<FieldInfo> {
}
}
+ synchronized void verifyConsistentVectorProperties(Integer number, String name, int numDimensions, VectorValues.DistanceFunction distFunc) {
+ if (name.equals(numberToName.get(number)) == false) {
+ throw new IllegalArgumentException("field number " + number + " is already mapped to field name \"" + numberToName.get(number) + "\", not \"" + name + "\"");
+ }
+ if (number.equals(nameToNumber.get(name)) == false) {
+ throw new IllegalArgumentException("field name \"" + name + "\" is already mapped to field number \"" + nameToNumber.get(name) + "\", not \"" + number + "\"");
+ }
+ FieldVectorProperties props = vectorProps.get(name);
+ if (props != null) {
+ if (props.numDimensions != numDimensions) {
+ throw new IllegalArgumentException("cannot change vector dimensions from " + props.numDimensions + " to " + numDimensions + " for field=\"" + name + "\"");
+ }
+ if (props.distFunc != distFunc) {
+ throw new IllegalArgumentException("cannot change distance function from " + props.distFunc + " to " + distFunc + " for field=\"" + name + "\"");
+ }
+ }
+ }
+
/**
* Returns true if the {@code fieldName} exists in the map and is of the
* same {@code dvType}.
@@ -453,6 +506,20 @@ public class FieldInfos implements Iterable<FieldInfo> {
verifyConsistentDimensions(number, name, dataDimensionCount, indexDimensionCount, dimensionNumBytes);
dimensions.put(name, new FieldDimensions(dataDimensionCount, indexDimensionCount, dimensionNumBytes));
}
+
+ synchronized void setVectorDimensionsAndDistanceFunction(int number, String name, int numDimensions, VectorValues.DistanceFunction distFunc) {
+ if (numDimensions < 0) {
+ throw new IllegalArgumentException("vector numDimensions must be >= 0; got " + numDimensions);
+ }
+ if (numDimensions > VectorValues.MAX_DIMENSIONS) {
+ throw new IllegalArgumentException("vector numDimensions must be <= VectorValues.MAX_DIMENSIONS (=" + VectorValues.MAX_DIMENSIONS + "); got " + numDimensions);
+ }
+ if (numDimensions == 0 && distFunc != VectorValues.DistanceFunction.NONE) {
+ throw new IllegalArgumentException("vector distFunc must be NONE when the vector numDimensions = 0; got " + distFunc);
+ }
+ verifyConsistentVectorProperties(number, name, numDimensions, distFunc);
+ vectorProps.put(name, new FieldVectorProperties(numDimensions, distFunc));
+ }
}
static final class Builder {
@@ -486,8 +553,8 @@ public class FieldInfos implements Iterable<FieldInfo> {
// before then we'll get the same name and number,
// else we'll allocate a new one:
final boolean isSoftDeletesField = name.equals(globalFieldNumbers.softDeletesFieldName);
- final int fieldNumber = globalFieldNumbers.addOrGet(name, -1, IndexOptions.NONE, DocValuesType.NONE, 0, 0, 0, isSoftDeletesField);
- fi = new FieldInfo(name, fieldNumber, false, false, false, IndexOptions.NONE, DocValuesType.NONE, -1, new HashMap<>(), 0, 0, 0, isSoftDeletesField);
+ final int fieldNumber = globalFieldNumbers.addOrGet(name, -1, IndexOptions.NONE, DocValuesType.NONE, 0, 0, 0, 0, VectorValues.DistanceFunction.NONE, isSoftDeletesField);
+ fi = new FieldInfo(name, fieldNumber, false, false, false, IndexOptions.NONE, DocValuesType.NONE, -1, new HashMap<>(), 0, 0, 0, 0, VectorValues.DistanceFunction.NONE, isSoftDeletesField);
assert !byName.containsKey(fi.name);
globalFieldNumbers.verifyConsistent(Integer.valueOf(fi.number), fi.name, DocValuesType.NONE);
byName.put(fi.name, fi);
@@ -502,6 +569,7 @@ public class FieldInfos implements Iterable<FieldInfo> {
DocValuesType docValues, long dvGen,
Map<String, String> attributes,
int dataDimensionCount, int indexDimensionCount, int dimensionNumBytes,
+ int vectorNumDimensions, VectorValues.DistanceFunction vectorDistFunc,
boolean isSoftDeletesField) {
assert assertNotFinished();
if (docValues == null) {
@@ -519,8 +587,8 @@ public class FieldInfos implements Iterable<FieldInfo> {
// number for this field. If the field was seen
// before then we'll get the same name and number,
// else we'll allocate a new one:
- final int fieldNumber = globalFieldNumbers.addOrGet(name, preferredFieldNumber, indexOptions, docValues, dataDimensionCount, indexDimensionCount, dimensionNumBytes, isSoftDeletesField);
- fi = new FieldInfo(name, fieldNumber, storeTermVector, omitNorms, storePayloads, indexOptions, docValues, dvGen, attributes, dataDimensionCount, indexDimensionCount, dimensionNumBytes, isSoftDeletesField);
+ final int fieldNumber = globalFieldNumbers.addOrGet(name, preferredFieldNumber, indexOptions, docValues, dataDimensionCount, indexDimensionCount, dimensionNumBytes, vectorNumDimensions, vectorDistFunc, isSoftDeletesField);
+ fi = new FieldInfo(name, fieldNumber, storeTermVector, omitNorms, storePayloads, indexOptions, docValues, dvGen, attributes, dataDimensionCount, indexDimensionCount, dimensionNumBytes, vectorNumDimensions, vectorDistFunc, isSoftDeletesField);
assert !byName.containsKey(fi.name);
globalFieldNumbers.verifyConsistent(Integer.valueOf(fi.number), fi.name, fi.getDocValuesType());
byName.put(fi.name, fi);
@@ -555,6 +623,7 @@ public class FieldInfos implements Iterable<FieldInfo> {
fi.getIndexOptions(), fi.getDocValuesType(), dvGen,
fi.attributes(),
fi.getPointDataDimensionCount(), fi.getPointIndexDimensionCount(), fi.getPointNumBytes(),
+ fi.getVectorNumDimensions(), fi.getVectorDistFunc(),
fi.isSoftDeletesField());
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/FilterCodecReader.java b/lucene/core/src/java/org/apache/lucene/index/FilterCodecReader.java
index ff3ea18..5265e0f 100644
--- a/lucene/core/src/java/org/apache/lucene/index/FilterCodecReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/FilterCodecReader.java
@@ -23,6 +23,7 @@ import java.util.Objects;
import org.apache.lucene.codecs.DocValuesProducer;
import org.apache.lucene.codecs.FieldsProducer;
+import org.apache.lucene.codecs.KnnGraphReader;
import org.apache.lucene.codecs.NormsProducer;
import org.apache.lucene.codecs.PointsReader;
import org.apache.lucene.codecs.StoredFieldsReader;
@@ -102,6 +103,11 @@ public abstract class FilterCodecReader extends CodecReader {
}
@Override
+ public KnnGraphReader getKnnGraphReader() {
+ return in.getKnnGraphReader();
+ }
+
+ @Override
public int numDocs() {
return in.numDocs();
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/FilterLeafReader.java b/lucene/core/src/java/org/apache/lucene/index/FilterLeafReader.java
index 1d26d17..114d779 100644
--- a/lucene/core/src/java/org/apache/lucene/index/FilterLeafReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/FilterLeafReader.java
@@ -331,6 +331,16 @@ public abstract class FilterLeafReader extends LeafReader {
}
@Override
+ public VectorValues getVectorValues(String field) throws IOException {
+ return in.getVectorValues(field);
+ }
+
+ @Override
+ public KnnGraphValues getKnnGraphValues(String field) throws IOException {
+ return in.getKnnGraphValues(field);
+ }
+
+ @Override
public Fields getTermVectors(int docID)
throws IOException {
ensureOpen();
diff --git a/lucene/core/src/java/org/apache/lucene/index/IndexWriter.java b/lucene/core/src/java/org/apache/lucene/index/IndexWriter.java
index 14f96d4..63d0611 100644
--- a/lucene/core/src/java/org/apache/lucene/index/IndexWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/index/IndexWriter.java
@@ -992,7 +992,7 @@ public class IndexWriter implements Closeable, TwoPhaseCommit, Accountable,
for(SegmentCommitInfo info : segmentInfos) {
FieldInfos fis = readFieldInfos(info);
for(FieldInfo fi : fis) {
- map.addOrGet(fi.name, fi.number, fi.getIndexOptions(), fi.getDocValuesType(), fi.getPointDataDimensionCount(), fi.getPointIndexDimensionCount(), fi.getPointNumBytes(), fi.isSoftDeletesField());
+ map.addOrGet(fi.name, fi.number, fi.getIndexOptions(), fi.getDocValuesType(), fi.getPointDataDimensionCount(), fi.getPointIndexDimensionCount(), fi.getPointNumBytes(), fi.getVectorNumDimensions(), fi.getVectorDistFunc(), fi.isSoftDeletesField());
}
}
@@ -1790,7 +1790,7 @@ public class IndexWriter implements Closeable, TwoPhaseCommit, Accountable,
if (globalFieldNumberMap.contains(f.name(), dvType) == false) {
// if this field doesn't exists we try to add it. if it exists and the DV type doesn't match we
// get a consistent error message as if you try to do that during an indexing operation.
- globalFieldNumberMap.addOrGet(f.name(), -1, IndexOptions.NONE, dvType, 0, 0, 0, f.name().equals(config.softDeletesField));
+ globalFieldNumberMap.addOrGet(f.name(), -1, IndexOptions.NONE, dvType, 0, 0, 0, 0, VectorValues.DistanceFunction.NONE, f.name().equals(config.softDeletesField));
assert globalFieldNumberMap.contains(f.name(), dvType);
}
if (config.getIndexSortFields().contains(f.name())) {
@@ -2814,7 +2814,7 @@ public class IndexWriter implements Closeable, TwoPhaseCommit, Accountable,
FieldInfos fis = readFieldInfos(info);
for(FieldInfo fi : fis) {
// This will throw exceptions if any of the incoming fields have an illegal schema change:
- globalFieldNumberMap.addOrGet(fi.name, fi.number, fi.getIndexOptions(), fi.getDocValuesType(), fi.getPointDataDimensionCount(), fi.getPointIndexDimensionCount(), fi.getPointNumBytes(), fi.isSoftDeletesField());
+ globalFieldNumberMap.addOrGet(fi.name, fi.number, fi.getIndexOptions(), fi.getDocValuesType(), fi.getPointDataDimensionCount(), fi.getPointIndexDimensionCount(), fi.getPointNumBytes(), fi.getVectorNumDimensions(), fi.getVectorDistFunc(), fi.isSoftDeletesField());
}
infos.add(copySegmentAsIs(info, newSegName, context));
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/IndexableFieldType.java b/lucene/core/src/java/org/apache/lucene/index/IndexableFieldType.java
index 59c5ab5..67d9c61 100644
--- a/lucene/core/src/java/org/apache/lucene/index/IndexableFieldType.java
+++ b/lucene/core/src/java/org/apache/lucene/index/IndexableFieldType.java
@@ -115,6 +115,16 @@ public interface IndexableFieldType {
public int pointNumBytes();
/**
+ * The number of dimensions of the field's vector value
+ */
+ public int vectorNumDimensions();
+
+ /**
+ * The {@link org.apache.lucene.index.VectorValues.DistanceFunction} of the field's vector value
+ */
+ public VectorValues.DistanceFunction vectorDistFunc();
+
+ /**
* Attributes for the field type.
*
* Attributes are not thread-safe, user must not add attributes while other threads are indexing documents
diff --git a/lucene/core/src/java/org/apache/lucene/index/KnnGraphValues.java b/lucene/core/src/java/org/apache/lucene/index/KnnGraphValues.java
new file mode 100644
index 0000000..f4d4659
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/index/KnnGraphValues.java
@@ -0,0 +1,116 @@
+/*
+ * 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.index;
+
+import java.io.IOException;
+
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.util.IntsRef;
+
+/**
+ * Access to per-document friends lists in a (hierarchical) knn search graph.
+ */
+public abstract class KnnGraphValues extends DocIdSetIterator {
+
+ /** Sole constructor */
+ protected KnnGraphValues() {}
+
+ /**
+ * Returns the top level of the entire graph.
+ * @return top level
+ */
+ public abstract int getTopLevel();
+
+ /**
+ * Returns true if the current document ID is the enter point of the graph;
+ * this means {@link #getMaxLevel()} is equal to {@link #getTopLevel()}.
+ * It is illegal to call this method after {@link #advanceExact(int)}
+ * returned {@code false}.
+ */
+ public boolean isEnterPoint() {
+ return getMaxLevel() == getTopLevel();
+ }
+
+ /**
+ * Returns the max level (height) for the current document ID.
+ * It is illegal to call this method after {@link #advanceExact(int)}
+ * returned {@code false}.
+ * @return max level
+ */
+ public abstract int getMaxLevel();
+
+ /**
+ * Returns the friend list (doc ID list) at the graph {@code level}.
+ * Each doc ID in the friend list is connected to the current document ID at this {@code level}.
+ * {@code level} must be a valid level, ie. ≥ 0 and ≤ {@link #getMaxLevel()}.
+ * It is illegal to call this method after {@link #advanceExact(int)}
+ * returned {@code false}.
+ * @param level the graph level
+ * @return friend list
+ */
+ public abstract IntsRef getFriends(int level);
+
+ /** Move the pointer to exactly {@code target} and return whether
+ * {@code target} has friends lists.
+ * {@code target} must be a valid doc ID, ie. ≥ 0 and < {@code maxDoc}.
+ * After this method returns, {@link #docID()} retuns {@code target}. */
+ public abstract boolean advanceExact(int target) throws IOException;
+
+ /** Empty graph value */
+ public static KnnGraphValues EMPTY = new KnnGraphValues() {
+ @Override
+ public int getTopLevel() {
+ return -1;
+ }
+
+ @Override
+ public int getMaxLevel() {
+ return -1;
+ }
+
+ @Override
+ public IntsRef getFriends(int level) {
+ return new IntsRef();
+ }
+
+ @Override
+ public boolean advanceExact(int target) throws IOException {
+ return false;
+ }
+
+ @Override
+ public int docID() {
+ return -1;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ return -1;
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ return NO_MORE_DOCS;
+ }
+
+ @Override
+ public long cost() {
+ return 0;
+ }
+ };
+}
diff --git a/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java b/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java
new file mode 100644
index 0000000..0b657dc
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java
@@ -0,0 +1,333 @@
+/*
+ * 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.index;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import org.apache.lucene.codecs.KnnGraphReader;
+import org.apache.lucene.codecs.KnnGraphWriter;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.store.ByteBuffersDataInput;
+import org.apache.lucene.store.ByteBuffersDataOutput;
+import org.apache.lucene.store.DataInput;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.Counter;
+import org.apache.lucene.util.IntsRef;
+import org.apache.lucene.util.IntsRefBuilder;
+import org.apache.lucene.util.RamUsageEstimator;
+import org.apache.lucene.util.hnsw.HNSWGraph;
+import org.apache.lucene.util.hnsw.HNSWGraphBuilder;
+
+/** Buffers up pending vector value(s) and knn graph per doc, then flushes when segment flushes. */
+public class KnnGraphValuesWriter implements Accountable {
+
+ private final FieldInfo fieldInfo;
+ private final Counter iwBytesUsed;
+ private final DocsWithFieldSet docsWithFieldVec;
+ private final DocsWithFieldSet docsWithFieldGrp;
+ private final IntsRefBuilder docsRef;
+ private final List<float[]> buffer;
+ private final ByteBuffersDataOutput bufferOut;
+ private final int numDimensions;
+ private final VectorValues.DistanceFunction distFunc;
+ private final HNSWGraphBuilder hnswBuilder;
+
+ private int addedDocs = 0;
+ private int lastDocID = -1;
+
+ private long bytesUsed = 0L;
+
+ public KnnGraphValuesWriter(DocumentsWriterPerThread docWriter, FieldInfo fieldInfo) {
+ this.fieldInfo = fieldInfo;
+ this.iwBytesUsed = docWriter.bytesUsed;
+ this.docsWithFieldVec = new DocsWithFieldSet();
+ this.docsWithFieldGrp = new DocsWithFieldSet();
+ this.docsRef = new IntsRefBuilder();
+ this.buffer = new ArrayList<>();
+ this.bufferOut = new ByteBuffersDataOutput();
+ this.numDimensions = fieldInfo.getVectorNumDimensions();
+ this.distFunc = fieldInfo.getVectorDistFunc();
+ this.hnswBuilder = new HNSWGraphBuilder(distFunc);
+ this.bytesUsed = docsWithFieldGrp.ramBytesUsed() + docsWithFieldGrp.ramBytesUsed() + RamUsageEstimator.shallowSizeOf(hnswBuilder);
+ iwBytesUsed.addAndGet(bytesUsed);
+ }
+
+ public void addValue(int docID, BytesRef binaryValue) throws IOException {
+ if (docID <= lastDocID) {
+ throw new IllegalArgumentException("VectorValuesField \"" + fieldInfo.name + "\" appears more than once in this document (only one value is allowed per field)");
+ }
+ if (binaryValue == null) {
+ throw new IllegalArgumentException("field=\"" + fieldInfo.name + "\": null value not allowed");
+ }
+
+ // write the vector
+ bufferOut.writeBytes(binaryValue.bytes);
+ docsWithFieldVec.add(docID);
+ docsRef.grow(docID + 1);
+ docsRef.setIntAt(docID, addedDocs++);
+ docsRef.setLength(docID + 1);
+ if (docID > lastDocID + 1) {
+ Arrays.fill(docsRef.ints(), lastDocID + 1, docID, -1);
+ }
+
+ // write the graph
+ buffer.add(VectorValues.decode(binaryValue.bytes, numDimensions));
+ if (distFunc != VectorValues.DistanceFunction.NONE) {
+ float[] query = VectorValues.decode(binaryValue.bytes, numDimensions);
+ hnswBuilder.insert(docID, query, getVectorValues());
+ docsWithFieldGrp.add(docID);
+ }
+
+ updateBytesUsed();
+
+ lastDocID = docID;
+ }
+
+ private void updateBytesUsed() {
+ final long newBytesUsed = docsWithFieldVec.ramBytesUsed() + docsWithFieldGrp.ramBytesUsed() +
+ bufferOut.ramBytesUsed() +
+ RamUsageEstimator.sizeOf(docsRef.ints()) +
+ RamUsageEstimator.shallowSizeOf(buffer) +
+ RamUsageEstimator.shallowSizeOf(hnswBuilder);
+ iwBytesUsed.addAndGet(newBytesUsed - bytesUsed);
+ bytesUsed = newBytesUsed;
+ }
+
+ public void flush(SegmentWriteState state, Sorter.DocMap sortMap, KnnGraphWriter graphWriter) throws IOException {
+ VectorValues vectors = new BufferedVectorValues(numDimensions, bufferOut.toDataInput(), docsWithFieldVec.iterator());
+ KnnGraphValues graph = new BufferedKnnGraphValues(hnswBuilder.build(), docsWithFieldGrp.iterator());
+
+ final VectorValues vectorValues;
+ final KnnGraphValues graphValues;
+ if (sortMap == null) {
+ vectorValues = vectors;
+ graphValues = graph;
+ } else {
+ // TODO
+ vectorValues = vectors;
+ graphValues = graph;
+ }
+
+ graphWriter.writeField(fieldInfo,
+ new KnnGraphReader() {
+ @Override
+ public void checkIntegrity() throws IOException {
+ }
+
+ @Override
+ public VectorValues getVectorValues(String field) throws IOException {
+ return vectorValues;
+ }
+
+ @Override
+ public KnnGraphValues getGraphValues(String field) throws IOException {
+ return graphValues;
+ }
+
+ @Override
+ public void close() throws IOException {
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return 0L;
+ }
+ });
+ }
+
+ private VectorValues getVectorValues() {
+ return new VectorValues() {
+ int docID = -1;
+ float[] value = new float[0];
+
+ @Override
+ public float[] vectorValue() throws IOException {
+ return value;
+ }
+
+ @Override
+ public boolean seek(int target) throws IOException {
+ if (target < 0) {
+ throw new IllegalArgumentException("target must be a positive integer: " + target);
+ }
+ if (target >= KnnGraphValuesWriter.this.docsRef.length()) {
+ docID = NO_MORE_DOCS;
+ return false;
+ }
+ docID = target;
+
+ int position = docsRef.ints()[target];
+ if (position < 0) {
+ throw new IllegalArgumentException("no vector value for doc: " + target);
+ }
+ value = KnnGraphValuesWriter.this.buffer.get(position);
+ return true;
+ }
+
+ @Override
+ public int docID() {
+ return docID;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public long cost() {
+ return KnnGraphValuesWriter.this.docsRef.length();
+ }
+ };
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return bytesUsed;
+ }
+
+ private static class BufferedVectorValues extends VectorValues {
+
+ final DocIdSetIterator docsWithField;
+ final DataInput bytesIterator;
+ final byte[] value;
+ final int numDims;
+
+ BufferedVectorValues(int numDims, DataInput bytesIterator, DocIdSetIterator docsWithField) {
+ this.numDims = numDims;
+ this.value = new byte[Float.BYTES * numDims];
+ this.bytesIterator = bytesIterator;
+ this.docsWithField = docsWithField;
+ }
+
+ @Override
+ public float[] vectorValue() throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public byte[] binaryValue() throws IOException {
+ return value;
+ }
+
+ @Override
+ public boolean seek(int target) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int docID() {
+ return docsWithField.docID();
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ int docID = docsWithField.nextDoc();
+ if (docID != NO_MORE_DOCS) {
+ bytesIterator.readBytes(value, 0, value.length);
+ }
+ return docID;
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public long cost() {
+ return docsWithField.cost();
+ }
+ }
+
+ private static class BufferedKnnGraphValues extends KnnGraphValues {
+
+ final HNSWGraph hnswGraph;
+ final DocIdSetIterator docsWithField;
+
+ private int maxLevel = -1;
+
+ BufferedKnnGraphValues(HNSWGraph hnsw, DocIdSetIterator docsWithField) {
+ this.hnswGraph = hnsw;
+ this.docsWithField = docsWithField;
+ }
+
+ @Override
+ public int getTopLevel() {
+ return hnswGraph.topLevel();
+ }
+
+ @Override
+ public int getMaxLevel() {
+ return maxLevel;
+ }
+
+ @Override
+ public IntsRef getFriends(int level) {
+ return hnswGraph.getFriends(level, docID());
+ }
+
+ @Override
+ public boolean advanceExact(int target) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int docID() {
+ return docsWithField.docID();
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ int docId = docsWithField.nextDoc();
+ if (docId != NO_MORE_DOCS) {
+ boolean levelFound = false;
+ for (int l = hnswGraph.topLevel(); l >= 0; l--) {
+ if (hnswGraph.hasFriends(l, docId)) {
+ maxLevel = l;
+ levelFound = true;
+ break;
+ }
+ }
+ assert levelFound;
+ }
+ return docId;
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public long cost() {
+ return docsWithField.cost();
+ }
+ }
+
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/index/LeafReader.java b/lucene/core/src/java/org/apache/lucene/index/LeafReader.java
index 1d09742..f863a65 100644
--- a/lucene/core/src/java/org/apache/lucene/index/LeafReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/LeafReader.java
@@ -203,6 +203,14 @@ public abstract class LeafReader extends IndexReader {
* used by a single thread. */
public abstract NumericDocValues getNormValues(String field) throws IOException;
+ /** Returns {@link VectorValues} for this field, or null if no {@link VectorValues} were indexed.
+ * The returned instance should only be used by a single thread. */
+ public abstract VectorValues getVectorValues(String field) throws IOException;
+
+ /** Returns {@link KnnGraphValues} for this field, or null if no {@link KnnGraphValues} were indexed.
+ * The returned instance should only be used by a single thread. */
+ public abstract KnnGraphValues getKnnGraphValues(String field) throws IOException;
+
/**
* Get the {@link FieldInfos} describing all fields in
* this reader.
diff --git a/lucene/core/src/java/org/apache/lucene/index/MergeReaderWrapper.java b/lucene/core/src/java/org/apache/lucene/index/MergeReaderWrapper.java
index afa2612..0c80ea1 100644
--- a/lucene/core/src/java/org/apache/lucene/index/MergeReaderWrapper.java
+++ b/lucene/core/src/java/org/apache/lucene/index/MergeReaderWrapper.java
@@ -195,6 +195,16 @@ class MergeReaderWrapper extends LeafReader {
}
@Override
+ public VectorValues getVectorValues(String fieldName) throws IOException {
+ return in.getVectorValues(fieldName);
+ }
+
+ @Override
+ public KnnGraphValues getKnnGraphValues(String fieldName) throws IOException {
+ return in.getKnnGraphValues(fieldName);
+ }
+
+ @Override
public int numDocs() {
return in.numDocs();
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/ParallelLeafReader.java b/lucene/core/src/java/org/apache/lucene/index/ParallelLeafReader.java
index 25f200a..f303a6d 100644
--- a/lucene/core/src/java/org/apache/lucene/index/ParallelLeafReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/ParallelLeafReader.java
@@ -370,6 +370,20 @@ public class ParallelLeafReader extends LeafReader {
}
@Override
+ public VectorValues getVectorValues(String fieldName) throws IOException {
+ ensureOpen();
+ LeafReader reader = fieldToReader.get(fieldName);
+ return reader == null ? null : reader.getVectorValues(fieldName);
+ }
+
+ @Override
+ public KnnGraphValues getKnnGraphValues(String fieldName) throws IOException {
+ ensureOpen();
+ LeafReader reader = fieldToReader.get(fieldName);
+ return reader == null ? null : reader.getKnnGraphValues(fieldName);
+ }
+
+ @Override
public void checkIntegrity() throws IOException {
ensureOpen();
for (LeafReader reader : completeReaderSet) {
diff --git a/lucene/core/src/java/org/apache/lucene/index/SegmentCoreReaders.java b/lucene/core/src/java/org/apache/lucene/index/SegmentCoreReaders.java
index 371abc6..9222839 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SegmentCoreReaders.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SegmentCoreReaders.java
@@ -30,6 +30,7 @@ import java.util.concurrent.atomic.AtomicInteger;
import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.FieldsProducer;
+import org.apache.lucene.codecs.KnnGraphReader;
import org.apache.lucene.codecs.NormsProducer;
import org.apache.lucene.codecs.PointsReader;
import org.apache.lucene.codecs.PostingsFormat;
@@ -61,6 +62,7 @@ final class SegmentCoreReaders {
final StoredFieldsReader fieldsReaderOrig;
final TermVectorsReader termVectorsReaderOrig;
final PointsReader pointsReader;
+ final KnnGraphReader knnGraphReader;
final Directory cfsReader;
final String segment;
/**
@@ -137,6 +139,13 @@ final class SegmentCoreReaders {
} else {
pointsReader = null;
}
+
+ if (coreFieldInfos.hasVectorValues()) {
+ knnGraphReader = codec.knnGraphFormat().fieldsReader(segmentReadState);
+ } else {
+ knnGraphReader = null;
+ }
+
success = true;
} catch (EOFException | FileNotFoundException e) {
throw new CorruptIndexException("Problem reading index from " + dir, dir.toString(), e);
diff --git a/lucene/core/src/java/org/apache/lucene/index/SegmentReader.java b/lucene/core/src/java/org/apache/lucene/index/SegmentReader.java
index d71d384..68f5dee 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SegmentReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SegmentReader.java
@@ -27,6 +27,7 @@ import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.DocValuesProducer;
import org.apache.lucene.codecs.FieldInfosFormat;
import org.apache.lucene.codecs.FieldsProducer;
+import org.apache.lucene.codecs.KnnGraphReader;
import org.apache.lucene.codecs.NormsProducer;
import org.apache.lucene.codecs.PointsReader;
import org.apache.lucene.codecs.StoredFieldsReader;
@@ -261,6 +262,11 @@ public final class SegmentReader extends CodecReader {
}
@Override
+ public KnnGraphReader getKnnGraphReader() {
+ return core.knnGraphReader;
+ }
+
+ @Override
public FieldsProducer getPostingsReader() {
ensureOpen();
return core.fields;
diff --git a/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java b/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java
index 5332557..e6dde79 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java
@@ -24,6 +24,7 @@ import java.util.Iterator;
import org.apache.lucene.codecs.DocValuesProducer;
import org.apache.lucene.codecs.FieldsProducer;
+import org.apache.lucene.codecs.KnnGraphReader;
import org.apache.lucene.codecs.NormsProducer;
import org.apache.lucene.codecs.PointsReader;
import org.apache.lucene.codecs.StoredFieldsReader;
@@ -79,6 +80,12 @@ public final class SlowCodecReaderWrapper {
}
@Override
+ public KnnGraphReader getKnnGraphReader() {
+ reader.ensureOpen();
+ return readerToKnnGraphReader(reader);
+ }
+
+ @Override
public FieldsProducer getPostingsReader() {
reader.ensureOpen();
try {
@@ -160,6 +167,33 @@ public final class SlowCodecReaderWrapper {
};
}
+
+ private static KnnGraphReader readerToKnnGraphReader(LeafReader reader) {
+ return new KnnGraphReader() {
+ @Override
+ public void checkIntegrity() throws IOException {
+ }
+
+ @Override
+ public VectorValues getVectorValues(String field) throws IOException {
+ return reader.getVectorValues(field);
+ }
+
+ @Override
+ public KnnGraphValues getGraphValues(String field) throws IOException {
+ return reader.getKnnGraphValues(field);
+ }
+
+ @Override
+ public void close() throws IOException {
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return 0L;
+ }
+ };
+ }
private static NormsProducer readerToNormsProducer(final LeafReader reader) {
return new NormsProducer() {
diff --git a/lucene/core/src/java/org/apache/lucene/index/VectorValues.java b/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
new file mode 100644
index 0000000..97ddcca
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
@@ -0,0 +1,230 @@
+/*
+ * 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.index;
+
+import java.io.IOException;
+import java.util.Arrays;
+
+import org.apache.lucene.search.DocIdSetIterator;
+
+/**
+ * Access to per-document vector value.
+ */
+public abstract class VectorValues extends DocIdSetIterator {
+
+ public static int MAX_DIMENSIONS = 1024;
+
+ /** Sole constructor */
+ protected VectorValues() {}
+
+ /**
+ * Returns the vector value for the current document ID.
+ * It is illegal to call this method after {@link #seek(int)}
+ * returned {@code false}.
+ * @return vector value
+ */
+ public abstract float[] vectorValue() throws IOException;
+
+ /**
+ * Returns the binary encoded vector value for the current document ID.
+ * It is illegal to call this method after {@link #seek(int)}
+ * returned {@code false}.
+ * @return binary value
+ */
+ public byte[] binaryValue() throws IOException {
+ return encode(vectorValue());
+ }
+
+ /** Move the pointer to exactly {@code target} and return whether
+ * {@code target} has a value.
+ * {@code target} must be a valid doc ID, ie. ≥ 0 and < {@code maxDoc}.
+ * After this method returns, {@link #docID()} retuns {@code target}. */
+ public abstract boolean seek(int target) throws IOException;
+
+ /**
+ * Calculates the distance between the two vectors with specified distance function.
+ */
+ public static float distance(float[] v1, float[] v2, DistanceFunction distFunc) {
+ if (v1.length != v2.length) {
+ throw new IllegalArgumentException("Incompatible number of dimensions: " + v1.length + ", " + v2.length);
+ }
+ return distFunc.distance(v1, v2);
+ }
+
+ /**
+ * Encodes float array to byte array.
+ */
+ public static byte[] encode(float[] value) {
+ byte[] bytes = new byte[Float.BYTES * value.length];
+ for (int i = 0; i < value.length; i++) {
+ int bits = Float.floatToIntBits(value[i]);
+ bytes[i * Float.BYTES] = (byte)(bits >> 24);
+ bytes[i * Float.BYTES + 1] = (byte)(bits >> 16);
+ bytes[i * Float.BYTES + 2] = (byte)(bits >> 8);
+ bytes[i * Float.BYTES + 3] = (byte)(bits);
+ }
+ return bytes;
+ }
+
+ public static boolean verifyNumDimensions(byte[] bytes, int numDims) {
+ if (bytes.length % Float.BYTES != 0) {
+ throw new IllegalArgumentException("Cannot decode bytes array to float array. Reason: invalid length: " + bytes.length);
+ }
+ int dims = bytes.length / Float.BYTES;
+ if (dims != numDims) {
+ throw new IllegalArgumentException("Invalid dimensions: " + dims + " (the number of dimensions should be " + numDims + ")");
+ }
+ return true;
+ }
+
+ /**
+ * Decodes float array from byte array.
+ */
+ public static float[] decode(byte[] bytes, int numDims) {
+ verifyNumDimensions(bytes, numDims);
+ float[] value = new float[numDims];
+ for (int i = 0; i < numDims; i++) {
+ int bits = ((bytes[i * Float.BYTES] & 0xFF) << 24) |
+ ((bytes[i * Float.BYTES + 1] & 0xFF) << 16) |
+ ((bytes[i * Float.BYTES + 2] & 0xFF) << 8) |
+ (bytes[i * Float.BYTES + 3] & 0xFF);
+ value[i] = Float.intBitsToFloat(bits);
+ }
+ return value;
+ }
+
+ /**
+ * Distance function. This is used when both of indexing knn graph and searching.
+ */
+ public enum DistanceFunction {
+ /** No distance function is used.
+ * Note: knn graph is not indexed for the field. */
+ NONE(0),
+
+ /** Manhattan distance */
+ MANHATTAN(1) {
+ @Override
+ float distance(float[] v1, float[] v2) {
+ assert v1.length == v2.length;
+ if (Arrays.equals(v1, v2)) {
+ return 0.0f;
+ }
+ float d = 0.0f;
+ int dim = v1.length;
+ for (int i = 0; i < dim; i++) {
+ d += Math.abs(v1[i] - v2[i]);
+ }
+ return d;
+ }
+ },
+
+ /** Euclidean distance */
+ EUCLIDEAN(2) {
+ @Override
+ float distance(float[] v1, float[] v2) {
+ assert v1.length == v2.length;
+ if (Arrays.equals(v1, v2)) {
+ return 0.0f;
+ }
+ float squareSum = 0.0f;
+ int dim = v1.length;
+ for (int i = 0; i < dim; i++) {
+ float diff = v1[i] - v2[i];
+ squareSum += diff * diff;
+ }
+ return (float)Math.sqrt(squareSum);
+ }
+ },
+
+ /** Cosine distance */
+ COSINE(3) {
+ @Override
+ float distance(float[] v1, float[] v2) {
+ // TODO
+ return 0.0f;
+ }
+ };
+
+ /** ID for each enum value; this is supposed to be persisted to the index and cannot be changed after indexing. */
+ final int id;
+
+ /** Sole constructor */
+ DistanceFunction(int id) {
+ this.id = id;
+ }
+
+ /**
+ * Returns the id of the distance function.
+ * @return id
+ */
+ public int getId() {
+ return id;
+ }
+
+ /**
+ * Calculates the distance between the specified two vectors.
+ */
+ float distance(float[] v1, float[] v2) {
+ throw new UnsupportedOperationException();
+ }
+
+ /**
+ * Returns the distance function that is specified by the id.
+ */
+ public static DistanceFunction fromId(int id) {
+ for (DistanceFunction d : DistanceFunction.values()) {
+ if (d.id == id) {
+ return d;
+ }
+ }
+ throw new IllegalArgumentException("no such distance function with id " + id);
+ }
+ }
+
+ public static VectorValues EMPTY = new VectorValues() {
+ @Override
+ public float[] vectorValue() throws IOException {
+ return new float[0];
+ }
+
+ @Override
+ public boolean seek(int target) throws IOException {
+ return false;
+ }
+
+ @Override
+ public int docID() {
+ return -1;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ return -1;
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ return -1;
+ }
+
+ @Override
+ public long cost() {
+ return 0;
+ }
+ };
+}
diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java b/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
new file mode 100644
index 0000000..b2ca26b
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
@@ -0,0 +1,124 @@
+/*
+ * 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.search;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Objects;
+
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.RamUsageEstimator;
+import org.apache.lucene.util.hnsw.HNSWGraphReader;
+
+/**
+ * Approximate nearest neighbor search query for high dimensional vector values.
+ */
+public class KnnGraphQuery extends Query implements Accountable {
+
+ public static final Integer DEFAULT_EF = 50;
+
+ private final String field;
+ private final float[] queryVector;
+ private final int ef;
+
+ private long bytesUsed = 0L;
+
+ /**
+ * Creates an nearest neighbor search query with default {@code ef} parameter (={@link #DEFAULT_EF}).
+ * @param field field name
+ * @param queryVector query vector. must has same number of dimensions to the indexed vectors.
+ */
+ public KnnGraphQuery(String field, float[] queryVector) {
+ this(field, queryVector, DEFAULT_EF);
+ this.bytesUsed = RamUsageEstimator.shallowSizeOfInstance(getClass());
+ }
+
+ /**
+ * Creates an nearest neighbor search query.
+ * @param field field name
+ * @param queryVector query vector. must has same number of dimensions to the indexed vectors.
+ * @param ef number of per-segment candidates to be scored/collected. the collector does not return results exceeding {@code ef}.
+ * increasing this value leads higher recall at the expense of the search speed.
+ */
+ public KnnGraphQuery(String field, float[] queryVector, int ef) {
+ this.field = field;
+ this.queryVector = queryVector;
+ this.ef = ef;
+ this.bytesUsed = RamUsageEstimator.shallowSizeOfInstance(getClass());
+ }
+
+ /**
+ * Creates an nearest neighbor search query; this also loads per-segment kNN graphs
+ * before executing queries so that the latency of the initial search will be reduced.
+ * @param field field name
+ * @param queryVector query vector. must has same number of dimensions to the indexed vectors.
+ * @param ef number of per-segment candidates to be scored/collected. the collector does not return results exceeding {@code ef}.
+ * increasing this value leads higher recall at the expense of the search speed.
+ * @param reader index reader
+ * @throws IOException
+ */
+ public KnnGraphQuery(String field, float[] queryVector, int ef, IndexReader reader) throws IOException {
+ this.field = field;
+ this.queryVector = queryVector;
+ this.ef = ef;
+ if (reader != null) {
+ this.bytesUsed = HNSWGraphReader.reloadGraph(field, reader);
+ }
+ }
+
+ @Override
+ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
+ return new KnnScoreWeight(this, boost, scoreMode, field, queryVector, ef);
+ }
+
+ @Override
+ public void visit(QueryVisitor visitor) {
+ if (visitor.acceptField(field)) {
+ visitor.visitLeaf(this);
+ }
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ return sameClassAs(other) && equalsTo(getClass().cast(other));
+ }
+
+ private boolean equalsTo(KnnGraphQuery other) {
+ return Objects.equals(field, other.field) && Arrays.equals(queryVector, other.queryVector);
+ }
+
+ @Override
+ public int hashCode() {
+ int hash = classHash();
+ hash = 31 * hash + field.hashCode();
+ hash = 31 * hash + Arrays.hashCode(queryVector);
+ return hash;
+ }
+
+ @Override
+ public String toString(String field) {
+ // TODO: FIXME
+ return null;
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return bytesUsed;
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java b/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java
new file mode 100644
index 0000000..73c8071
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java
@@ -0,0 +1,151 @@
+/*
+ * 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.search;
+
+import java.io.IOException;
+
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.VectorValues;
+import org.apache.lucene.util.hnsw.HNSWGraphReader;
+import org.apache.lucene.util.hnsw.Neighbor;
+import org.apache.lucene.util.hnsw.Neighbors;
+
+class KnnScoreWeight extends ConstantScoreWeight {
+
+ private final String field;
+ private final ScoreMode scoreMode;
+ private final float[] queryVector;
+ private final int ef;
+
+ KnnScoreWeight(Query query, float score, ScoreMode scoreMode, String field, float[] queryVector, int ef) {
+ super(query, score);
+ this.scoreMode = scoreMode;
+ this.field = field;
+ this.queryVector = queryVector;
+ this.ef = ef;
+ }
+
+ @Override
+ public Scorer scorer(LeafReaderContext context) throws IOException {
+ ScorerSupplier supplier = scorerSupplier(context);
+ if (supplier == null) {
+ return null;
+ }
+ return supplier.get(Long.MAX_VALUE);
+ }
+
+ @Override
+ public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
+ FieldInfo fi = context.reader().getFieldInfos().fieldInfo(field);
+ int numDimensions = fi.getVectorNumDimensions();
+ if (numDimensions != queryVector.length) {
+ throw new IllegalArgumentException("field=\"" + field + "\" was indexed with dimensions=" + numDimensions + "; this is incompatible with query dimensions=" + queryVector.length);
+ }
+
+ final HNSWGraphReader hnswReader = new HNSWGraphReader(field, context);
+ final VectorValues vectorValues = context.reader().getVectorValues(field);
+ if (vectorValues == null) {
+ // No docs in this segment/field indexed any vector values
+ return null;
+ }
+
+ final Weight weight = this;
+ return new ScorerSupplier() {
+ @Override
+ public Scorer get(long leadCost) throws IOException {
+ Neighbors neighbors = hnswReader.searchNeighbors(queryVector, ef, vectorValues);
+ return new Scorer(weight) {
+
+ int doc = -1;
+ float score = 0.0f;
+ int size = neighbors.size();
+ int offset = 0;
+
+ @Override
+ public DocIdSetIterator iterator() {
+ return new DocIdSetIterator() {
+ @Override
+ public int docID() {
+ return doc;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ return advance(offset);
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ if (target > size || neighbors.size() == 0) {
+ doc = NO_MORE_DOCS;
+ score = 0.0f;
+ } else {
+ while (offset < target) {
+ neighbors.pop();
+ offset++;
+ }
+ Neighbor next = neighbors.pop();
+ offset++;
+ if (next == null) {
+ doc = NO_MORE_DOCS;
+ score = 0.0f;
+ } else {
+ doc = next.docId();
+ score = 1.0f / (next.distance() / numDimensions + 0.01f);
+ }
+ }
+ return doc;
+ }
+
+ @Override
+ public long cost() {
+ return size;
+ }
+ };
+ }
+
+ @Override
+ public float getMaxScore(int upTo) throws IOException {
+ return Float.POSITIVE_INFINITY;
+ }
+
+ @Override
+ public float score() throws IOException {
+ return score;
+ }
+
+ @Override
+ public int docID() {
+ return doc;
+ }
+ };
+ }
+
+ @Override
+ public long cost() {
+ return ef;
+ }
+ };
+ }
+
+ @Override
+ public boolean isCacheable(LeafReaderContext ctx) {
+ return true;
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/DeferredNeighbor.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/DeferredNeighbor.java
new file mode 100644
index 0000000..17e9295
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/DeferredNeighbor.java
@@ -0,0 +1,75 @@
+/*
+ * 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.util.Arrays;
+
+import org.apache.lucene.index.VectorValues;
+
+/** Calculates the distance as needed basis */
+final class DeferredNeighbor extends Neighbor {
+
+ final int docId;
+ float[] query;
+ VectorValues vectorValues;
+ VectorValues.DistanceFunction distFunc;
+
+ DeferredNeighbor(int docId) {
+ this.docId = docId;
+ }
+
+ @Override
+ public int docId() {
+ return docId;
+ }
+
+ @Override
+ public boolean isDeferred() {
+ return true;
+ }
+
+ @Override
+ public void prepareQuery(float[] query, VectorValues vectorValues, VectorValues.DistanceFunction distFunc) {
+ this.query = query;
+ this.vectorValues = vectorValues;
+ this.distFunc = distFunc;
+ }
+
+ @Override
+ public float distance() {
+ if (query == null || vectorValues == null || distFunc == null) {
+ throw new IllegalStateException("cannot calculate the distance: query, vector values, or distFunc are not set.");
+ }
+ try {
+ if (vectorValues.seek(docId)) {
+ return VectorValues.distance(query, vectorValues.vectorValue(), distFunc);
+ } else {
+ throw new IllegalStateException("docId=" + docId + " has no vector value");
+ }
+ } catch (IOException e) {
+ throw new RuntimeException("cannot calculate the distance between docId=" + docId + " and query=" + Arrays.toString(query), e);
+ }
+ }
+
+ @Override
+ public int compareTo(Neighbor o) {
+ return this.docId() < o.docId() ? -1 : 1;
+ }
+}
+
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/FurthestNeighbors.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/FurthestNeighbors.java
new file mode 100644
index 0000000..d2a47ff
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/FurthestNeighbors.java
@@ -0,0 +1,30 @@
+/*
+ * 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;
+
+final class FurthestNeighbors extends Neighbors {
+
+ FurthestNeighbors(int maxSize, Neighbor seed) {
+ super(maxSize, seed);
+ }
+
+ @Override
+ protected boolean lessThan(Neighbor a, Neighbor b) {
+ return b.distance() < a.distance();
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
new file mode 100644
index 0000000..a62b862
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
@@ -0,0 +1,253 @@
+/*
+ * 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.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import org.apache.lucene.index.VectorValues;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.IntsRef;
+import org.apache.lucene.util.RamUsageEstimator;
+
+/**
+ * Hierarchical NSW graph that provides efficient approximate nearest neighbor search for high dimensional vectors.
+ * This isn't thread-safe.
+ * See <a href="https://arxiv.org/abs/1603.09320">this paper</a> for details.
+ */
+public final class HNSWGraph implements Accountable {
+
+ private final VectorValues.DistanceFunction distFunc;
+ private final Map<Integer, Layer> layers;
+
+ private boolean freeze = false;
+ private int topLevel = 0;
+
+ public HNSWGraph(VectorValues.DistanceFunction distFunc) {
+ this.distFunc = distFunc;
+ this.layers = new HashMap<>();
+ }
+
+ /**
+ * Searches the nearest neighbors for a specified query at a level.
+ * @param query search query vector
+ * @param ep enter point to fhe level
+ * @param ef the number of nodes to be searched
+ * @param level graph level
+ * @param vectorValues vector values
+ * @return nearest neighbors
+ * @throws IOException
+ */
+ public NearestNeighbors searchLayer(float[] query, Neighbor ep, int ef, int level, VectorValues vectorValues) throws IOException {
+ if (level > topLevel) {
+ throw new IllegalArgumentException("layer does not exist for the level: " + level);
+ }
+
+ Layer layer = layers.get(level);
+ if (!layer.getNodes().contains(ep.docId())) {
+ throw new IllegalArgumentException("enter point " + ep.docId() + "does not exist at layer " + level);
+ }
+
+ Set<Integer> visited = new HashSet<>(ep.docId());
+ NearestNeighbors candidates = new NearestNeighbors(layer.size() + 1, ep);
+ FurthestNeighbors results = new FurthestNeighbors(layer.size() + 1, ep);
+
+ if (ep.isDeferred()) {
+ ep.prepareQuery(query, vectorValues, distFunc);
+ }
+ float upperBound = ep.distance();
+
+ while (candidates.size() > 0) {
+ Neighbor c = candidates.pop();
+ Neighbor f = results.top();
+ if (c.isDeferred()) {
+ c.prepareQuery(query, vectorValues, distFunc);
+ }
+ if (f.isDeferred()) {
+ f.prepareQuery(query, vectorValues, distFunc);
+ }
+ if (c.distance() > f.distance()) {
+ break;
+ }
+
+ for (Neighbor e : layer.getFriends(c.docId())) {
+ if (visited.contains(e.docId())) {
+ continue;
+ }
+ visited.add(e.docId());
+
+ f = results.top();
+ if (f.isDeferred()) {
+ f.prepareQuery(query, vectorValues, distFunc);
+ }
+ float dist = distance(query, e.docId(), vectorValues);
+ if ((dist < upperBound && dist < f.distance()) || results.size() < ef) {
+ Neighbor n = new ImmutableNeighbor(e.docId(), dist);
+ candidates.add(n);
+ results.add(n);
+ if (results.size() > ef) {
+ results.pop();
+ }
+ }
+ }
+ }
+
+ return pickNearestNeighbor(results, ef);
+ }
+
+ private float distance(float[] query, int docId, VectorValues vectorValues) throws IOException {
+ if (!vectorValues.seek(docId)) {
+ throw new IllegalStateException("docId=" + docId + " has no vector value");
+ }
+ float[] other = vectorValues.vectorValue();
+ return VectorValues.distance(query, other, distFunc);
+ }
+
+ private NearestNeighbors pickNearestNeighbor(FurthestNeighbors queue, int ef) {
+ NearestNeighbors nearests = new NearestNeighbors(ef);
+ Set<Integer> addedDocs = new HashSet<>();
+ for (int i = 0; i < ef && queue.size() > 0; i++) {
+ Neighbor c = queue.pop();
+ if (!addedDocs.contains(c.docId())) {
+ nearests.add(c);
+ addedDocs.add(c.docId());
+ }
+ }
+ return nearests;
+ }
+
+ public void ensureLevel(int level) {
+ if (freeze) {
+ throw new IllegalStateException("graph is already freezed!");
+ }
+ if (level < 0) {
+ throw new IllegalArgumentException("level must be a positive integer: " + level);
+ }
+ for (int l = level; l >= 0; l--) {
+ layers.putIfAbsent(l, new Layer(l));
+ }
+ if (level > topLevel) {
+ topLevel = level;
+ }
+ }
+
+ public int topLevel() {
+ return topLevel;
+ }
+
+ public boolean isEmpty() {
+ return layers.isEmpty() || layers.get(0).getNodes().isEmpty();
+ }
+
+ public int getEnterPoint() {
+ if (layers.isEmpty()) {
+ throw new IllegalStateException("the graph has no layers!");
+ }
+ List<Integer> nodesAtMaxLevel = layers.get(topLevel).getNodes();
+ if (nodesAtMaxLevel.isEmpty()) {
+ throw new IllegalStateException("the max level of this graph is empty!");
+ }
+ return nodesAtMaxLevel.get(0);
+ }
+
+ public boolean hasNodes(int level) {
+ Layer layer = layers.get(level);
+ if (layer == null) {
+ throw new IllegalArgumentException("layer does not exist for level: " + level);
+ }
+ return layer.getNodes().size() > 0;
+ }
+
+ public IntsRef getFriends(int level, int node) {
+ Layer layer = layers.get(level);
+ if (layer == null) {
+ throw new IllegalArgumentException("layer does not exist for level: " + level);
+ }
+ int[] friends = layer.getFriends(node).stream().mapToInt(Neighbor::docId).sorted().toArray();
+ return new IntsRef(friends, 0, friends.length);
+ }
+
+ public boolean hasFriends(int level, int node) {
+ Layer layer = layers.get(level);
+ if (layer == null) {
+ throw new IllegalArgumentException("layer does not exist for level: " + level);
+ }
+ return layer.getFriends(node).size() > 0;
+ }
+
+ void addNode(int level, int node) {
+ if (freeze) {
+ throw new IllegalStateException("graph is already freezed!");
+ }
+ Layer layer = layers.get(level);
+ if (layer == null) {
+ throw new IllegalArgumentException("layer does not exist for level: " + level);
+ }
+ layer.addNodeIfAbsent(node);
+ }
+
+ /** Connects two nodes; this is supposed to be called when indexing */
+ public void connectNodes(int level, int node1, int node2, float dist, int maxConnections) {
+ if (freeze) {
+ throw new IllegalStateException("graph is already freezed!");
+ }
+ assert level >= 0;
+ assert node1 >= 0 && node2 >= 0;
+ assert node1 != node2;
+
+ Layer layer = layers.get(level);
+ if (layer == null) {
+ throw new IllegalArgumentException("layer does not exist for level: " + level);
+ }
+ layer.connectNodes(node1, node2, dist);
+ // ensure friends size <= maxConnections
+ if (maxConnections > 0) {
+ layer.shrink(node2, maxConnections);
+ }
+ }
+
+ /** Connects two nodes; this is supposed to be called when searching */
+ public void connectNodes(int level, int node1, int node2) {
+ if (freeze) {
+ throw new IllegalStateException("graph is already freezed!");
+ }
+ assert level >= 0;
+ assert node1 >= 0 && node2 >= 0;
+ assert node1 != node2;
+
+ Layer layer = layers.get(level);
+ if (layer == null) {
+ throw new IllegalArgumentException("layer does not exist for level: " + level);
+ }
+ layer.connectNodes(node1, node2);
+ }
+
+ public void finish() {
+ this.freeze = true;
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return RamUsageEstimator.sizeOfMap(layers);
+ }
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphBuilder.java
new file mode 100644
index 0000000..85e17ed
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphBuilder.java
@@ -0,0 +1,111 @@
+/*
+ * 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 org.apache.lucene.index.VectorValues;
+
+/** Builder of {@link HNSWGraph} */
+public final class HNSWGraphBuilder {
+
+ private static final int MAX_LEVEL_LIMIT = 5;
+ // default random seed for level generation
+ private static final long DEFAULT_RAND_SEED = System.currentTimeMillis();
+ // default max connections per node
+ public static final int DEFAULT_MAX_CONNECTIONS = 6;
+ // default max connections per node at layer zero
+ public static final int DEFAULT_MAX_CONNECTIONS_L0 = DEFAULT_MAX_CONNECTIONS * 2;
+ // default candidate list size
+ public static final int DEFAULT_EF_CONST = 50;
+
+ private final int maxConn;
+ private final int maxConn0;
+ private final int efConst;
+ private final VectorValues.DistanceFunction distFunc;
+ private final HNSWGraph hnsw;
+ private final RandomLevelGenerator levelGenerator;
+
+ /** Construct the builder with default configurations */
+ public HNSWGraphBuilder(VectorValues.DistanceFunction distFunc) {
+ this(DEFAULT_MAX_CONNECTIONS, DEFAULT_MAX_CONNECTIONS_L0, DEFAULT_EF_CONST, DEFAULT_RAND_SEED, distFunc);
+ }
+
+ /** Full constructor */
+ public HNSWGraphBuilder(int maxConn, int maxConn0, int efConst, long seed, VectorValues.DistanceFunction distFunc) {
+ this.maxConn = maxConn;
+ this.maxConn0 = maxConn0;
+ this.efConst = efConst;
+ this.distFunc = distFunc;
+ this.hnsw = new HNSWGraph(distFunc);
+ float ml = (float) (1.0 / Math.log(maxConn));
+ this.levelGenerator = new RandomLevelGenerator(seed, ml, MAX_LEVEL_LIMIT);
+ }
+
+ /** Inserts a doc with vector value to the graph */
+ public void insert(int docId, float[] value, VectorValues vectorValues) throws IOException {
+ if (hnsw.isEmpty()) {
+ // initialize the graph
+ hnsw.ensureLevel(0);
+ hnsw.addNode(0, docId);
+ return;
+ }
+
+ int enterPoint = hnsw.getEnterPoint();
+ if (!vectorValues.seek(enterPoint)) {
+ throw new IllegalStateException("enter point (=" + enterPoint + ") has no vector value");
+ }
+ float distFromEp = VectorValues.distance(value, vectorValues.vectorValue(), distFunc);
+ Neighbor ep = new ImmutableNeighbor(enterPoint, distFromEp);
+
+ int level = levelGenerator.randomLevel();
+
+ // down to the level from the hnsw's top level
+ for (int l = hnsw.topLevel(); l > level; l--) {
+ NearestNeighbors neighbors = hnsw.searchLayer(value, ep, 1, l, vectorValues);
+ ep = neighbors.top();
+ }
+
+ // down to level 0 with placing the doc to each layer
+ hnsw.ensureLevel(level);
+ for (int l = level; l >= 0; l--) {
+ if (!hnsw.hasNodes(l)) {
+ hnsw.addNode(l, docId);
+ continue;
+ }
+
+ NearestNeighbors neighbors = hnsw.searchLayer(value, ep, efConst, l, vectorValues);
+ ep = neighbors.top();
+ int maxConnections = l == 0 ? maxConn0 : maxConn;
+ for (int i = 0; i < maxConnections; i++) {
+ if (neighbors.size() == 0) {
+ break;
+ }
+ Neighbor n = neighbors.pop();
+ hnsw.connectNodes(l, docId, n.docId(), n.distance(), maxConnections);
+ }
+ }
+ }
+
+ /** Returns the built HNSW graph*/
+ public HNSWGraph build() {
+ hnsw.finish();
+ return hnsw;
+ }
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
new file mode 100644
index 0000000..4e6385a
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
@@ -0,0 +1,147 @@
+/*
+ * 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.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.KnnGraphValues;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.VectorValues;
+import org.apache.lucene.search.DocIdSetIterator;
+
+/** Executes approximate nearest neighbor search on per-reader {@link HNSWGraph}.
+ * This also caches built {@link HNSWGraph}s for repeated use. */
+public final class HNSWGraphReader {
+
+ private static final Map<GraphKey, HNSWGraph> cache = new ConcurrentHashMap<>();
+
+ private final String field;
+ private final LeafReaderContext context;
+ private final VectorValues.DistanceFunction distFunc;
+
+ public HNSWGraphReader(String field, LeafReaderContext context) {
+ this.field = field;
+ this.context = context;
+ this.distFunc = context.reader().getFieldInfos().fieldInfo(field).getVectorDistFunc();
+ }
+
+ public Neighbors searchNeighbors(float[] query, int ef, VectorValues vectorValues) throws IOException {
+ HNSWGraph hnsw = get(field, context, false);
+ int enterPoint = hnsw.getEnterPoint();
+ if (!vectorValues.seek(enterPoint)) {
+ throw new IllegalStateException("enterPoint=" + enterPoint + " has no vector value");
+ }
+
+ Neighbors neighbors;
+ Neighbor ep = new ImmutableNeighbor(enterPoint, VectorValues.distance(query, vectorValues.vectorValue(), distFunc));
+ for (int l = hnsw.topLevel(); l > 0; l--) {
+ neighbors = hnsw.searchLayer(query, ep, 1, l, vectorValues);
+ ep = neighbors.top();
+ }
+ return hnsw.searchLayer(query, ep, ef, 0, vectorValues);
+ }
+
+ public static long reloadGraph(String field, IndexReader reader) throws IOException {
+ long bytesUsed = 0L;
+ for (LeafReaderContext ctx : reader.leaves()) {
+ HNSWGraph hnsw = get(field, ctx, true);
+ bytesUsed += hnsw.ramBytesUsed();
+ }
+ return bytesUsed;
+ }
+
+ private static HNSWGraph get(String field, LeafReaderContext context, boolean reload) throws IOException {
+ GraphKey key = new GraphKey(field, context.id());
+ IOException[] exc = new IOException[]{null};
+ if (reload) {
+ cache.put(key, load(field, context));
+ } else {
+ cache.computeIfAbsent(key, (k -> {
+ try {
+ return load(k.field, context);
+ } catch (IOException e) {
+ exc[0] = e;
+ return null;
+ }
+ }));
+ if (exc[0] != null) {
+ throw exc[0];
+ }
+ }
+ return cache.get(key);
+ }
+
+ private static HNSWGraph load(String field, LeafReaderContext context) throws IOException {
+ FieldInfo fi = context.reader().getFieldInfos().fieldInfo(field);
+ int numDimensions = fi.getVectorNumDimensions();
+ if (numDimensions == 0) {
+ // the field has no vector values
+ return null;
+ }
+ VectorValues.DistanceFunction distFunc = fi.getVectorDistFunc();
+
+ KnnGraphValues graphValues = context.reader().getKnnGraphValues(field);
+ HNSWGraph hnsw = new HNSWGraph(distFunc);
+ for (int doc = graphValues.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = graphValues.nextDoc()) {
+ int maxLevel = graphValues.getMaxLevel();
+ hnsw.ensureLevel(maxLevel);
+ for (int l = 0; l <= maxLevel; l++) {
+ for (int friend : graphValues.getFriends(l).ints) {
+ hnsw.connectNodes(l, doc, friend);
+ }
+ }
+ }
+ hnsw.finish();
+
+ return hnsw;
+ }
+
+ private static class GraphKey {
+ final String field;
+ final Object readerId;
+
+ GraphKey(String field, Object readerId) {
+ this.field = field;
+ this.readerId = readerId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ GraphKey graphKey = (GraphKey) o;
+ return Objects.equals(field, graphKey.field) &&
+ Objects.equals(readerId, graphKey.readerId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(field, readerId);
+ }
+
+ public String toString() {
+ return "field=" + field + ", readerId=" + readerId;
+ }
+ }
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/ImmutableNeighbor.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/ImmutableNeighbor.java
new file mode 100644
index 0000000..0b46999
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/ImmutableNeighbor.java
@@ -0,0 +1,55 @@
+/*
+ * 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;
+
+final class ImmutableNeighbor extends Neighbor {
+ final int docId;
+ final float distance;
+
+ ImmutableNeighbor(int docId, float distance) {
+ this.docId = docId;
+ this.distance = distance;
+ }
+
+ @Override
+ public int docId() {
+ return docId;
+ }
+
+ @Override
+ public float distance() {
+ return distance;
+ }
+
+ @Override
+ public boolean isDeferred() {
+ return false;
+ }
+
+ @Override
+ public int compareTo(Neighbor o) {
+ if (this.distance() == o.distance() && this.docId() == o.docId()) {
+ return 0;
+ }
+ if (this.distance() == o.distance()) {
+ return this.docId() < o.docId() ? -1 : 1;
+ }
+ return this.distance() < o.distance() ? -1 : 1;
+ }
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java
new file mode 100644
index 0000000..aa4b145
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java
@@ -0,0 +1,101 @@
+/*
+ * 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.util.Collections;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeSet;
+
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.RamUsageEstimator;
+
+/** A knn graph that consists of connected friends (i.e. neighbors) lists. */
+final class Layer implements Accountable {
+
+ private final int level;
+ private final Map<Integer, TreeSet<Neighbor>> friendsMap;
+
+ /**
+ * Constructs a layer with level.
+ * @param level level of this layer
+ */
+ Layer(int level) {
+ this.level = level;
+ this.friendsMap = new HashMap<>();
+ }
+
+ List<Integer> getNodes() {
+ return List.copyOf(friendsMap.keySet());
+ }
+
+ void addNodeIfAbsent(int node) {
+ friendsMap.putIfAbsent(node, new TreeSet<>());
+ }
+
+ void connectNodes(int node1, int node2, float distance) {
+ addNodeIfAbsent(node1);
+ addNodeIfAbsent(node2);
+ Neighbor neighbor1 = new ImmutableNeighbor(node2, distance);
+ friendsMap.get(node1).add(neighbor1);
+ Neighbor neighbor2 = new ImmutableNeighbor(node1, distance);
+ friendsMap.get(node2).add(neighbor2);
+ }
+
+ void connectNodes(int node1, int node2) {
+ addNodeIfAbsent(node1);
+ addNodeIfAbsent(node2);
+ Neighbor neighbor1 = new DeferredNeighbor(node2);
+ friendsMap.get(node1).add(neighbor1);
+ Neighbor neighbor2 = new DeferredNeighbor(node1);
+ friendsMap.get(node2).add(neighbor2);
+ }
+
+ void shrink(int node, int maxFriends) {
+ TreeSet<Neighbor> friends = friendsMap.get(node);
+ if (friends == null) {
+ throw new IllegalArgumentException("node " + node + " does not exist at this layer: " + level);
+ }
+ if (friends.size() > maxFriends) {
+ TreeSet<Neighbor> newFriends = new TreeSet<>();
+ for (Iterator<Neighbor> itr = friends.iterator(); itr.hasNext() && newFriends.size() < maxFriends; ) {
+ newFriends.add(itr.next());
+ }
+ friendsMap.put(node, newFriends);
+ }
+ }
+
+ List<Neighbor> getFriends(int node) {
+ TreeSet<Neighbor> friends = friendsMap.get(node);
+ if (friends == null) {
+ return Collections.emptyList();
+ }
+ return List.copyOf(friends);
+ }
+
+ int size() {
+ return friendsMap.size();
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return RamUsageEstimator.shallowSizeOfInstance(getClass()) + RamUsageEstimator.sizeOfMap(friendsMap);
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/NearestNeighbors.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/NearestNeighbors.java
new file mode 100644
index 0000000..9923075
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/NearestNeighbors.java
@@ -0,0 +1,37 @@
+/*
+ * 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;
+
+/**
+ * Nearest neighbors of a specific query vector.
+ */
+final class NearestNeighbors extends Neighbors {
+
+ NearestNeighbors(int maxSize) {
+ super(maxSize);
+ }
+
+ NearestNeighbors(int maxSize, Neighbor seed) {
+ super(maxSize, seed);
+ }
+
+ @Override
+ protected boolean lessThan(Neighbor a, Neighbor b) {
+ return a.distance() < b.distance();
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/Neighbor.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/Neighbor.java
new file mode 100644
index 0000000..0cd601a
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/Neighbor.java
@@ -0,0 +1,45 @@
+/*
+ * 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 org.apache.lucene.index.VectorValues;
+
+/** An abstract neighbor document of a specific query vector; a query can be the vector value of a document to be indexed
+ * or an arbitrary query. */
+public abstract class Neighbor implements Comparable<Neighbor> {
+
+ /** Returns document id of this neighbor */
+ public abstract int docId();
+
+ /** Returns distance between the query and this neighbor doc; It is illegal if {@link #isDeferred()} = true
+ * and {@link #prepareQuery(float[], VectorValues, VectorValues.DistanceFunction)} is not called before calling
+ * {@link #distance()} */
+ public abstract float distance();
+
+ /** Returns true if distance calculation is deferred */
+ public abstract boolean isDeferred();
+
+ /** Set attributes for distance calculation. This must be called before {@link #distance()} when
+ * {@link #isDeferred()} = true */
+ public void prepareQuery(float[] query, VectorValues vectorValues, VectorValues.DistanceFunction distFunc) {
+ }
+
+ public String toString() {
+ return "(" + docId() + ", " + distance() + ")";
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/Neighbors.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/Neighbors.java
new file mode 100644
index 0000000..5e3a083
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/Neighbors.java
@@ -0,0 +1,45 @@
+/*
+ * 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 org.apache.lucene.util.PriorityQueue;
+
+/** Abstract neighbors queue; subclasses should implement the priority function
+ * {@link PriorityQueue#lessThan(Object, Object)} */
+public abstract class Neighbors extends PriorityQueue<Neighbor> {
+
+ protected Neighbors(int maxSize) {
+ this(maxSize, null);
+ }
+
+ protected Neighbors(int maxSize, Neighbor seed) {
+ super(maxSize);
+ if (seed != null) {
+ add(seed);
+ }
+ }
+
+ public String toString() {
+ StringBuilder sb = new StringBuilder();
+ sb.append("Neighbors=[");
+ this.iterator().forEachRemaining(sb::append);
+ sb.append("]");
+ return sb.toString();
+ }
+
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/RandomLevelGenerator.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/RandomLevelGenerator.java
new file mode 100644
index 0000000..e65f694
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/RandomLevelGenerator.java
@@ -0,0 +1,42 @@
+/*
+ * 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.util.Random;
+
+/** Generates a HNSW's level where a vector (document) is placed to. */
+class RandomLevelGenerator {
+
+ private final Random random;
+ private final float ml;
+ private final int upperLimit;
+
+ /** Creates a level generator */
+ RandomLevelGenerator(long seed, float ml, int upperLimit) {
+ this.random = new Random(seed);
+ this.ml = ml;
+ this.upperLimit = upperLimit;
+ }
+
+ /** Generates a level */
+ int randomLevel() {
+ int level = (int)(-Math.log(random.nextFloat()) * ml);
+ return Math.min(level, upperLimit);
+ }
+
+}
diff --git a/lucene/core/src/resources/META-INF/services/org.apache.lucene.codecs.Codec b/lucene/core/src/resources/META-INF/services/org.apache.lucene.codecs.Codec
index 611e237..d79fd21 100644
--- a/lucene/core/src/resources/META-INF/services/org.apache.lucene.codecs.Codec
+++ b/lucene/core/src/resources/META-INF/services/org.apache.lucene.codecs.Codec
@@ -14,3 +14,4 @@
# limitations under the License.
org.apache.lucene.codecs.lucene84.Lucene84Codec
+org.apache.lucene.codecs.lucene90.Lucene90Codec
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestIndexableField.java b/lucene/core/src/test/org/apache/lucene/index/TestIndexableField.java
index 991807a..d764215 100644
--- a/lucene/core/src/test/org/apache/lucene/index/TestIndexableField.java
+++ b/lucene/core/src/test/org/apache/lucene/index/TestIndexableField.java
@@ -107,6 +107,16 @@ public class TestIndexableField extends LuceneTestCase {
}
@Override
+ public int vectorNumDimensions() {
+ return 0;
+ }
+
+ @Override
+ public VectorValues.DistanceFunction vectorDistFunc() {
+ return VectorValues.DistanceFunction.NONE;
+ }
+
+ @Override
public Map<String, String> getAttributes() {
return null;
}
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestPendingSoftDeletes.java b/lucene/core/src/test/org/apache/lucene/index/TestPendingSoftDeletes.java
index 666b3c4..62ccead 100644
--- a/lucene/core/src/test/org/apache/lucene/index/TestPendingSoftDeletes.java
+++ b/lucene/core/src/test/org/apache/lucene/index/TestPendingSoftDeletes.java
@@ -37,6 +37,8 @@ import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.StringHelper;
import org.apache.lucene.util.Version;
+import static org.apache.lucene.index.VectorValues.DistanceFunction.NONE;
+
public class TestPendingSoftDeletes extends TestPendingDeletes {
@Override
@@ -164,7 +166,7 @@ public class TestPendingSoftDeletes extends TestPendingDeletes {
deletes.onNewReader(segmentReader, commitInfo);
reader.close();
writer.close();
- FieldInfo fieldInfo = new FieldInfo("_soft_deletes", 1, false, false, false, IndexOptions.NONE, DocValuesType.NUMERIC, 0, Collections.emptyMap(), 0, 0, 0, true);
+ FieldInfo fieldInfo = new FieldInfo("_soft_deletes", 1, false, false, false, IndexOptions.NONE, DocValuesType.NUMERIC, 0, Collections.emptyMap(), 0, 0, 0, 0, NONE, true);
List<Integer> docsDeleted = Arrays.asList(1, 3, 7, 8, DocIdSetIterator.NO_MORE_DOCS);
List<DocValuesFieldUpdates> updates = Arrays.asList(singleUpdate(docsDeleted, 10, true));
for (DocValuesFieldUpdates update : updates) {
@@ -185,7 +187,7 @@ public class TestPendingSoftDeletes extends TestPendingDeletes {
docsDeleted = Arrays.asList(1, 2, DocIdSetIterator.NO_MORE_DOCS);
updates = Arrays.asList(singleUpdate(docsDeleted, 10, true));
- fieldInfo = new FieldInfo("_soft_deletes", 1, false, false, false, IndexOptions.NONE, DocValuesType.NUMERIC, 1, Collections.emptyMap(), 0, 0, 0, true);
+ fieldInfo = new FieldInfo("_soft_deletes", 1, false, false, false, IndexOptions.NONE, DocValuesType.NUMERIC, 1, Collections.emptyMap(), 0, 0, 0, 0, NONE, true);
for (DocValuesFieldUpdates update : updates) {
deletes.onDocValuesUpdate(fieldInfo, update.iterator());
}
@@ -228,7 +230,7 @@ public class TestPendingSoftDeletes extends TestPendingDeletes {
SegmentCommitInfo segmentInfo = segmentReader.getSegmentInfo();
PendingDeletes deletes = newPendingDeletes(segmentInfo);
deletes.onNewReader(segmentReader, segmentInfo);
- FieldInfo fieldInfo = new FieldInfo("_soft_deletes", 1, false, false, false, IndexOptions.NONE, DocValuesType.NUMERIC, segmentInfo.getNextDocValuesGen(), Collections.emptyMap(), 0, 0, 0, true);
+ FieldInfo fieldInfo = new FieldInfo("_soft_deletes", 1, false, false, false, IndexOptions.NONE, DocValuesType.NUMERIC, segmentInfo.getNextDocValuesGen(), Collections.emptyMap(), 0, 0, 0, 0, NONE, true);
List<Integer> docsDeleted = Arrays.asList(1, DocIdSetIterator.NO_MORE_DOCS);
List<DocValuesFieldUpdates> updates = Arrays.asList(singleUpdate(docsDeleted, 3, true));
for (DocValuesFieldUpdates update : updates) {
@@ -276,7 +278,7 @@ public class TestPendingSoftDeletes extends TestPendingDeletes {
SegmentCommitInfo segmentInfo = segmentReader.getSegmentInfo();
PendingDeletes deletes = newPendingDeletes(segmentInfo);
deletes.onNewReader(segmentReader, segmentInfo);
- FieldInfo fieldInfo = new FieldInfo("_soft_deletes", 1, false, false, false, IndexOptions.NONE, DocValuesType.NUMERIC, segmentInfo.getNextDocValuesGen(), Collections.emptyMap(), 0, 0, 0, true);
+ FieldInfo fieldInfo = new FieldInfo("_soft_deletes", 1, false, false, false, IndexOptions.NONE, DocValuesType.NUMERIC, segmentInfo.getNextDocValuesGen(), Collections.emptyMap(), 0, 0, 0, 0, NONE, true);
List<DocValuesFieldUpdates> updates = Arrays.asList(singleUpdate(Arrays.asList(0, 1, DocIdSetIterator.NO_MORE_DOCS), 3, false));
for (DocValuesFieldUpdates update : updates) {
deletes.onDocValuesUpdate(fieldInfo, update.iterator());
@@ -295,7 +297,7 @@ public class TestPendingSoftDeletes extends TestPendingDeletes {
assertEquals(0, deletes.numPendingDeletes());
segmentInfo.advanceDocValuesGen();
- fieldInfo = new FieldInfo("_soft_deletes", 1, false, false, false, IndexOptions.NONE, DocValuesType.NUMERIC, segmentInfo.getNextDocValuesGen(), Collections.emptyMap(), 0, 0, 0, true);
+ fieldInfo = new FieldInfo("_soft_deletes", 1, false, false, false, IndexOptions.NONE, DocValuesType.NUMERIC, segmentInfo.getNextDocValuesGen(), Collections.emptyMap(), 0, 0, 0, 0, NONE, true);
updates = Arrays.asList(singleUpdate(Arrays.asList(1, DocIdSetIterator.NO_MORE_DOCS), 3, true));
for (DocValuesFieldUpdates update : updates) {
deletes.onDocValuesUpdate(fieldInfo, update.iterator());
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java b/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java
index 5ea8f83..b88c6be 100644
--- a/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java
+++ b/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java
@@ -107,6 +107,16 @@ public class TestSegmentToThreadMapping extends LuceneTestCase {
}
@Override
+ public VectorValues getVectorValues(String field) {
+ return null;
+ }
+
+ @Override
+ public KnnGraphValues getKnnGraphValues(String field) {
+ return null;
+ }
+
+ @Override
protected void doClose() {
}
diff --git a/lucene/highlighter/src/java/org/apache/lucene/search/highlight/TermVectorLeafReader.java b/lucene/highlighter/src/java/org/apache/lucene/search/highlight/TermVectorLeafReader.java
index 42f7aec..621219d 100644
--- a/lucene/highlighter/src/java/org/apache/lucene/search/highlight/TermVectorLeafReader.java
+++ b/lucene/highlighter/src/java/org/apache/lucene/search/highlight/TermVectorLeafReader.java
@@ -26,6 +26,7 @@ import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.FieldInfos;
import org.apache.lucene.index.Fields;
import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.index.KnnGraphValues;
import org.apache.lucene.index.LeafMetaData;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.NumericDocValues;
@@ -35,6 +36,7 @@ import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.index.SortedSetDocValues;
import org.apache.lucene.index.StoredFieldVisitor;
import org.apache.lucene.index.Terms;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.Version;
@@ -81,7 +83,7 @@ public class TermVectorLeafReader extends LeafReader {
}
FieldInfo fieldInfo = new FieldInfo(field, 0,
true, true, terms.hasPayloads(),
- indexOptions, DocValuesType.NONE, -1, Collections.emptyMap(), 0, 0, 0, false);
+ indexOptions, DocValuesType.NONE, -1, Collections.emptyMap(), 0, 0, 0, 0, VectorValues.DistanceFunction.NONE, false);
fieldInfos = new FieldInfos(new FieldInfo[]{fieldInfo});
}
@@ -140,6 +142,16 @@ public class TermVectorLeafReader extends LeafReader {
}
@Override
+ public VectorValues getVectorValues(String fieldName) {
+ return null;
+ }
+
+ @Override
+ public KnnGraphValues getKnnGraphValues(String fieldName) {
+ return null;
+ }
+
+ @Override
public void checkIntegrity() throws IOException {
}
diff --git a/lucene/memory/src/java/org/apache/lucene/index/memory/MemoryIndex.java b/lucene/memory/src/java/org/apache/lucene/index/memory/MemoryIndex.java
index ca9ccd7..92a3167 100644
--- a/lucene/memory/src/java/org/apache/lucene/index/memory/MemoryIndex.java
+++ b/lucene/memory/src/java/org/apache/lucene/index/memory/MemoryIndex.java
@@ -501,7 +501,8 @@ public class MemoryIndex {
IndexOptions indexOptions = storeOffsets ? IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS : IndexOptions.DOCS_AND_FREQS_AND_POSITIONS;
return new FieldInfo(fieldName, ord, fieldType.storeTermVectors(), fieldType.omitNorms(), storePayloads,
indexOptions, fieldType.docValuesType(), -1, Collections.emptyMap(),
- fieldType.pointDataDimensionCount(), fieldType.pointIndexDimensionCount(), fieldType.pointNumBytes(), false);
+ fieldType.pointDataDimensionCount(), fieldType.pointIndexDimensionCount(), fieldType.pointNumBytes(),
+ fieldType.vectorNumDimensions(), fieldType.vectorDistFunc(), false);
}
private void storePointValues(Info info, BytesRef pointValue) {
@@ -521,6 +522,7 @@ public class MemoryIndex {
info.fieldInfo.name, info.fieldInfo.number, info.fieldInfo.hasVectors(), info.fieldInfo.hasPayloads(),
info.fieldInfo.hasPayloads(), info.fieldInfo.getIndexOptions(), docValuesType, -1, info.fieldInfo.attributes(),
info.fieldInfo.getPointDataDimensionCount(), info.fieldInfo.getPointIndexDimensionCount(), info.fieldInfo.getPointNumBytes(),
+ info.fieldInfo.getVectorNumDimensions(), info.fieldInfo.getVectorDistFunc(),
info.fieldInfo.isSoftDeletesField()
);
} else if (existingDocValuesType != docValuesType) {
@@ -1241,6 +1243,18 @@ public class MemoryIndex {
}
@Override
+ public VectorValues getVectorValues(String fieldName) {
+ // TODO: FIXME
+ return VectorValues.EMPTY;
+ }
+
+ @Override
+ public KnnGraphValues getKnnGraphValues(String fieldName) {
+ // TODO: FIXME
+ return KnnGraphValues.EMPTY;
+ }
+
+ @Override
public void checkIntegrity() throws IOException {
// no-op
}
diff --git a/lucene/test-framework/src/java/org/apache/lucene/index/BaseIndexFileFormatTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/index/BaseIndexFileFormatTestCase.java
index aea5bd5..ea7dd66 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/index/BaseIndexFileFormatTestCase.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/index/BaseIndexFileFormatTestCase.java
@@ -347,7 +347,8 @@ abstract class BaseIndexFileFormatTestCase extends LuceneTestCase {
FieldInfo proto = oneDocReader.getFieldInfos().fieldInfo("field");
FieldInfo field = new FieldInfo(proto.name, proto.number, proto.hasVectors(), proto.omitsNorms(), proto.hasPayloads(),
proto.getIndexOptions(), proto.getDocValuesType(), proto.getDocValuesGen(), new HashMap<>(),
- proto.getPointDataDimensionCount(), proto.getPointIndexDimensionCount(), proto.getPointNumBytes(), proto.isSoftDeletesField());
+ proto.getPointDataDimensionCount(), proto.getPointIndexDimensionCount(), proto.getPointNumBytes(),
+ proto.getVectorNumDimensions(), proto.getVectorDistFunc(), proto.isSoftDeletesField());
FieldInfos fieldInfos = new FieldInfos(new FieldInfo[] { field } );
diff --git a/lucene/test-framework/src/java/org/apache/lucene/index/MismatchedLeafReader.java b/lucene/test-framework/src/java/org/apache/lucene/index/MismatchedLeafReader.java
index e16ce4b..337307b 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/index/MismatchedLeafReader.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/index/MismatchedLeafReader.java
@@ -80,6 +80,8 @@ public class MismatchedLeafReader extends FilterLeafReader {
oldInfo.getPointDataDimensionCount(), // data dimension count
oldInfo.getPointIndexDimensionCount(), // index dimension count
oldInfo.getPointNumBytes(), // dimension numBytes
+ oldInfo.getVectorNumDimensions(), // number of dimensions of the field's vector
+ oldInfo.getVectorDistFunc(), // distance function for calculating similarity of the field's vector
oldInfo.isSoftDeletesField()); // used as soft-deletes field
shuffled.set(i, newInfo);
}
diff --git a/lucene/test-framework/src/java/org/apache/lucene/index/RandomPostingsTester.java b/lucene/test-framework/src/java/org/apache/lucene/index/RandomPostingsTester.java
index 410635b..6f84bb6 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/index/RandomPostingsTester.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/index/RandomPostingsTester.java
@@ -130,7 +130,7 @@ public class RandomPostingsTester {
fieldInfoArray[fieldUpto] = new FieldInfo(field, fieldUpto, false, false, true,
IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS,
DocValuesType.NONE, -1, new HashMap<>(),
- 0, 0, 0, false);
+ 0, 0, 0, 0, VectorValues.DistanceFunction.NONE, false);
fieldUpto++;
SortedMap<BytesRef,SeedAndOrd> postings = new TreeMap<>();
@@ -651,7 +651,7 @@ public class RandomPostingsTester {
DocValuesType.NONE,
-1,
new HashMap<>(),
- 0, 0, 0, false);
+ 0, 0, 0, 0, VectorValues.DistanceFunction.NONE, false);
}
FieldInfos newFieldInfos = new FieldInfos(newFieldInfoArray);
diff --git a/lucene/test-framework/src/java/org/apache/lucene/search/QueryUtils.java b/lucene/test-framework/src/java/org/apache/lucene/search/QueryUtils.java
index cdbac77..90d655e 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/search/QueryUtils.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/search/QueryUtils.java
@@ -25,6 +25,7 @@ import org.apache.lucene.index.BinaryDocValues;
import org.apache.lucene.index.FieldInfos;
import org.apache.lucene.index.Fields;
import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.KnnGraphValues;
import org.apache.lucene.index.LeafMetaData;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
@@ -36,6 +37,7 @@ import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.index.SortedSetDocValues;
import org.apache.lucene.index.StoredFieldVisitor;
import org.apache.lucene.index.Terms;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.LuceneTestCase;
import org.apache.lucene.util.Version;
@@ -208,6 +210,16 @@ public class QueryUtils {
}
@Override
+ public VectorValues getVectorValues(String field) throws IOException {
+ return null;
+ }
+
+ @Override
+ public KnnGraphValues getKnnGraphValues(String field) throws IOException {
+ return null;
+ }
+
+ @Override
public FieldInfos getFieldInfos() {
return FieldInfos.EMPTY;
}
diff --git a/solr/core/src/java/org/apache/solr/index/SlowCompositeReaderWrapper.java b/solr/core/src/java/org/apache/solr/index/SlowCompositeReaderWrapper.java
index 4f1b56c..bd77073 100644
--- a/solr/core/src/java/org/apache/solr/index/SlowCompositeReaderWrapper.java
+++ b/solr/core/src/java/org/apache/solr/index/SlowCompositeReaderWrapper.java
@@ -277,6 +277,16 @@ public final class SlowCompositeReaderWrapper extends LeafReader {
}
@Override
+ public VectorValues getVectorValues(String field) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public KnnGraphValues getKnnGraphValues(String field) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
public FieldInfos getFieldInfos() {
return fieldInfos;
}
diff --git a/solr/core/src/java/org/apache/solr/schema/SchemaField.java b/solr/core/src/java/org/apache/solr/schema/SchemaField.java
index 100a963..594b53e 100644
--- a/solr/core/src/java/org/apache/solr/schema/SchemaField.java
+++ b/solr/core/src/java/org/apache/solr/schema/SchemaField.java
@@ -26,6 +26,7 @@ import org.apache.lucene.index.DocValuesType;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.IndexableField;
import org.apache.lucene.index.IndexableFieldType;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.search.SortField;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.util.SimpleOrderedMap;
@@ -432,6 +433,16 @@ public final class SchemaField extends FieldProperties implements IndexableField
}
@Override
+ public int vectorNumDimensions() {
+ return 0;
+ }
+
+ @Override
+ public VectorValues.DistanceFunction vectorDistFunc() {
+ return VectorValues.DistanceFunction.NONE;
+ }
+
+ @Override
public Map<String, String> getAttributes() {
return null;
}
diff --git a/solr/core/src/java/org/apache/solr/search/CollapsingQParserPlugin.java b/solr/core/src/java/org/apache/solr/search/CollapsingQParserPlugin.java
index b63c119..f991972 100644
--- a/solr/core/src/java/org/apache/solr/search/CollapsingQParserPlugin.java
+++ b/solr/core/src/java/org/apache/solr/search/CollapsingQParserPlugin.java
@@ -439,6 +439,8 @@ public class CollapsingQParserPlugin extends QParserPlugin {
fieldInfo.getPointDataDimensionCount(),
fieldInfo.getPointIndexDimensionCount(),
fieldInfo.getPointNumBytes(),
+ fieldInfo.getVectorNumDimensions(),
+ fieldInfo.getVectorDistFunc(),
fieldInfo.isSoftDeletesField());
newInfos.add(f);
} else {
diff --git a/solr/core/src/java/org/apache/solr/search/Insanity.java b/solr/core/src/java/org/apache/solr/search/Insanity.java
index 7da7a5c..2a36ad1 100644
--- a/solr/core/src/java/org/apache/solr/search/Insanity.java
+++ b/solr/core/src/java/org/apache/solr/search/Insanity.java
@@ -67,7 +67,8 @@ public class Insanity {
if (fi.name.equals(insaneField)) {
filteredInfos.add(new FieldInfo(fi.name, fi.number, fi.hasVectors(), fi.omitsNorms(),
fi.hasPayloads(), fi.getIndexOptions(), DocValuesType.NONE, -1, Collections.emptyMap(),
- fi.getPointDataDimensionCount(), fi.getPointIndexDimensionCount(), fi.getPointNumBytes(), fi.isSoftDeletesField()));
+ fi.getPointDataDimensionCount(), fi.getPointIndexDimensionCount(), fi.getPointNumBytes(),
+ fi.getVectorNumDimensions(), fi.getVectorDistFunc(), fi.isSoftDeletesField()));
} else {
filteredInfos.add(fi);
}
diff --git a/solr/core/src/java/org/apache/solr/uninverting/UninvertingReader.java b/solr/core/src/java/org/apache/solr/uninverting/UninvertingReader.java
index e804635..1bae5d5 100644
--- a/solr/core/src/java/org/apache/solr/uninverting/UninvertingReader.java
+++ b/solr/core/src/java/org/apache/solr/uninverting/UninvertingReader.java
@@ -284,7 +284,8 @@ public class UninvertingReader extends FilterLeafReader {
wrap = true;
newFieldInfos.add(new FieldInfo(fi.name, fi.number, fi.hasVectors(), fi.omitsNorms(),
fi.hasPayloads(), fi.getIndexOptions(), type, fi.getDocValuesGen(), fi.attributes(),
- fi.getPointDataDimensionCount(), fi.getPointIndexDimensionCount(), fi.getPointNumBytes(), fi.isSoftDeletesField()));
+ fi.getPointDataDimensionCount(), fi.getPointIndexDimensionCount(), fi.getPointNumBytes(),
+ fi.getVectorNumDimensions(), fi.getVectorDistFunc(), fi.isSoftDeletesField()));
} else {
newFieldInfos.add(fi);
}
diff --git a/solr/core/src/test/org/apache/solr/search/TestDocSet.java b/solr/core/src/test/org/apache/solr/search/TestDocSet.java
index 20879ea..db70818 100644
--- a/solr/core/src/test/org/apache/solr/search/TestDocSet.java
+++ b/solr/core/src/test/org/apache/solr/search/TestDocSet.java
@@ -26,6 +26,7 @@ import org.apache.lucene.index.FieldInfos;
import org.apache.lucene.index.Fields;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexReaderContext;
+import org.apache.lucene.index.KnnGraphValues;
import org.apache.lucene.index.LeafMetaData;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
@@ -37,6 +38,7 @@ import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.index.SortedSetDocValues;
import org.apache.lucene.index.StoredFieldVisitor;
import org.apache.lucene.index.Terms;
+import org.apache.lucene.index.VectorValues;
import org.apache.lucene.search.DocIdSet;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.BitSetIterator;
@@ -444,6 +446,16 @@ public class TestDocSet extends SolrTestCase {
}
@Override
+ public VectorValues getVectorValues(String field) {
+ return null;
+ }
+
+ @Override
+ public KnnGraphValues getKnnGraphValues(String field) {
+ return null;
+ }
+
+ @Override
protected void doClose() {
}
[lucene-solr] 04/10: Merge pull request #1 from
msokolov/mocobeta/LUCENE-9004-2
Posted by to...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a commit to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git
commit b005dcbe5934a44e0bf92c4df46596f0cb282252
Merge: f22b4b5 444c1d3
Author: Tomoko Uchida <to...@gmail.com>
AuthorDate: Sat Dec 14 00:23:43 2019 +0900
Merge pull request #1 from msokolov/mocobeta/LUCENE-9004-2
Implemented TestKnnGraph unit tests
.../org/apache/lucene/codecs/KnnGraphFormat.java | 2 +-
.../org/apache/lucene/codecs/KnnGraphWriter.java | 9 +-
.../codecs/lucene90/Lucene90FieldInfosFormat.java | 3 +-
.../codecs/lucene90/Lucene90KnnGraphReader.java | 39 +--
.../codecs/lucene90/Lucene90KnnGraphWriter.java | 248 +++++++++++++++--
.../java/org/apache/lucene/index/CodecReader.java | 2 +-
.../apache/lucene/index/DefaultIndexingChain.java | 4 +-
.../java/org/apache/lucene/index/FieldInfo.java | 7 +
.../apache/lucene/index/KnnGraphValuesWriter.java | 17 +-
.../java/org/apache/lucene/index/MergeState.java | 2 +-
.../apache/lucene/index/SegmentCoreReaders.java | 2 +-
.../org/apache/lucene/index/SegmentMerger.java | 147 +++++-----
.../java/org/apache/lucene/index/VectorValues.java | 32 ++-
.../org/apache/lucene/search/KnnGraphQuery.java | 1 -
.../org/apache/lucene/util/hnsw/HNSWGraph.java | 61 +++--
.../apache/lucene/util/hnsw/HNSWGraphWriter.java | 6 +-
.../java/org/apache/lucene/util/hnsw/Layer.java | 18 +-
.../test/org/apache/lucene/index/TestKnnGraph.java | 297 +++++++++++++++++++++
18 files changed, 703 insertions(+), 194 deletions(-)
[lucene-solr] 08/10: Mocobeta/lucene 9004 2 (#3)
Posted by to...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a commit to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git
commit b4a8979e5cb338a17a798525b05eb7ab06cbd613
Author: Michael Sokolov <so...@falutin.net>
AuthorDate: Fri Jan 10 10:52:16 2020 -0500
Mocobeta/lucene 9004 2 (#3)
* Added KnnGraphTester; fixed some small bugs and measured recall
* use FloatBuffer for encode/decode of vector values. Share a single float[] when fetching values. 25% latency reduction
* Don't try to precompute nearest neighbors when generating test data in KnnGraphTester
* added visitedCounter for debugging/optimization
restored shrinking to limit fanout of graph
fixed an upside-down priority queue
* don't terminate exploration of graph until we have a full result queue
* fixing up some small bugs I had introduced in earlier push
---
.../codecs/lucene90/Lucene90KnnGraphReader.java | 27 +++-
.../java/org/apache/lucene/index/VectorValues.java | 12 +-
.../org/apache/lucene/search/KnnGraphQuery.java | 25 +++-
.../org/apache/lucene/search/KnnScoreWeight.java | 12 ++
.../org/apache/lucene/util/hnsw/HNSWGraph.java | 27 ++--
.../apache/lucene/util/hnsw/HNSWGraphReader.java | 17 ++-
.../apache/lucene/util/hnsw/HNSWGraphWriter.java | 12 +-
.../apache/lucene/util/hnsw/KnnGraphTester.java | 165 +++++++++++++++++----
8 files changed, 220 insertions(+), 77 deletions(-)
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
index 1083fd4..c7e615e 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
@@ -18,6 +18,8 @@
package org.apache.lucene.codecs.lucene90;
import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.FloatBuffer;
import java.util.HashMap;
import java.util.Map;
@@ -258,18 +260,29 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
final int maxDoc;
final int numDims;
+ final int byteSize;
final KnnGraphEntry entry;
final IndexInput dataIn;
+ final BytesRef binaryValue;
+ final ByteBuffer byteBuffer;
+ final FloatBuffer floatBuffer;
+ final float[] value;
int doc = -1;
- BytesRef binaryValue;
RandomAccessVectorValuesReader(int maxDoc, int numDims, KnnGraphEntry entry, IndexInput dataIn) {
this.maxDoc = maxDoc;
this.numDims = numDims;
this.entry = entry;
this.dataIn = dataIn;
- this.binaryValue = new BytesRef(new byte[Float.BYTES * numDims]);
+ // TODO: if we had more direct access to file system (a FileChannel) we could use a
+ // MappedByteBuffer here. EG if we know that dataIn is a ByteBufferIndexInput, we could get a
+ // FloatBuffer from that and provide access to it
+ byteSize = Float.BYTES * numDims;
+ byteBuffer = ByteBuffer.allocate(byteSize);
+ floatBuffer = byteBuffer.asFloatBuffer();
+ value = new float[numDims];
+ binaryValue = new BytesRef(byteBuffer.array(), byteBuffer.arrayOffset(), byteSize);
}
@Override
@@ -277,7 +290,10 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
if (doc == NO_MORE_DOCS) {
return null;
}
- return VectorValues.decode(binaryValue, numDims);
+ // return VectorValues.decode(binaryValue, numDims);
+ floatBuffer.position(0);
+ floatBuffer.get(value, 0, numDims);
+ return value;
}
@Override
@@ -296,10 +312,11 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
if (ord == null) {
return false;
}
- int offset = Float.BYTES * numDims * ord;
+ int offset = ord * byteSize;
assert offset >= 0;
dataIn.seek(offset);
- dataIn.readBytes(binaryValue.bytes, binaryValue.offset, binaryValue.length);
+ //dataIn.readBytes(binaryValue.bytes, binaryValue.offset, binaryValue.length);
+ dataIn.readBytes(byteBuffer.array(), byteBuffer.arrayOffset(), byteSize);
return true;
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/VectorValues.java b/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
index e83d174..9e0ce1d 100644
--- a/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
+++ b/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
@@ -73,15 +73,9 @@ public abstract class VectorValues extends DocIdSetIterator {
* Encodes float array to byte array.
*/
public static BytesRef encode(float[] value) {
- byte[] bytes = new byte[Float.BYTES * value.length];
- for (int i = 0; i < value.length; i++) {
- int bits = Float.floatToIntBits(value[i]);
- bytes[i * Float.BYTES] = (byte)(bits >> 24);
- bytes[i * Float.BYTES + 1] = (byte)(bits >> 16);
- bytes[i * Float.BYTES + 2] = (byte)(bits >> 8);
- bytes[i * Float.BYTES + 3] = (byte)(bits);
- }
- return new BytesRef(bytes);
+ ByteBuffer buffer = ByteBuffer.allocate(Float.BYTES * value.length);
+ buffer.asFloatBuffer().put(value);
+ return new BytesRef(buffer.array());
}
public static boolean verifyNumDimensions(int numBytes, int numDims) {
diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java b/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
index b1c14f6..0f30037 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
@@ -21,6 +21,7 @@ import java.io.IOException;
import java.util.Arrays;
import java.util.Objects;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.FieldInfo;
@@ -44,8 +45,8 @@ public class KnnGraphQuery extends Query implements Accountable {
private final String field;
private final float[] queryVector;
private final int ef;
-
- private long bytesUsed = 0L;
+ private final long bytesUsed;
+ private final AtomicLong visitedCounter;
/**
* Creates an nearest neighbor search query with default {@code ef} parameter (={@link #DEFAULT_EF}).
@@ -54,7 +55,6 @@ public class KnnGraphQuery extends Query implements Accountable {
*/
public KnnGraphQuery(String field, float[] queryVector) {
this(field, queryVector, DEFAULT_EF);
- this.bytesUsed = RamUsageEstimator.shallowSizeOfInstance(getClass());
}
/**
@@ -68,7 +68,8 @@ public class KnnGraphQuery extends Query implements Accountable {
this.field = field;
this.queryVector = queryVector;
this.ef = ef;
- this.bytesUsed = RamUsageEstimator.shallowSizeOfInstance(getClass());
+ visitedCounter = new AtomicLong();
+ bytesUsed = RamUsageEstimator.shallowSizeOfInstance(getClass());
}
/**
@@ -85,8 +86,11 @@ public class KnnGraphQuery extends Query implements Accountable {
this.field = field;
this.queryVector = queryVector;
this.ef = ef;
+ visitedCounter = new AtomicLong();
if (reader != null) {
- this.bytesUsed = HNSWGraphReader.loadGraphs(field, reader, forceReload);
+ bytesUsed = HNSWGraphReader.loadGraphs(field, reader, forceReload);
+ } else {
+ bytesUsed = 0;
}
}
@@ -106,7 +110,9 @@ public class KnnGraphQuery extends Query implements Accountable {
@Override
public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
- return new KnnScoreWeight(this, boost, scoreMode, field, queryVector, ef);
+ KnnScoreWeight weight = new KnnScoreWeight(this, boost, scoreMode, field, queryVector, ef);
+ weight.setVisitedCounter(visitedCounter);
+ return weight;
}
@Override
@@ -143,4 +149,11 @@ public class KnnGraphQuery extends Query implements Accountable {
public long ramBytesUsed() {
return bytesUsed;
}
+
+ /**
+ * @return this total the number of documents visited by this query
+ */
+ public long getVisitedCount() {
+ return visitedCounter.get();
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java b/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java
index 02b6a5c..a3847ad 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnScoreWeight.java
@@ -18,6 +18,7 @@
package org.apache.lucene.search;
import java.io.IOException;
+import java.util.concurrent.atomic.AtomicLong;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.LeafReaderContext;
@@ -32,6 +33,7 @@ class KnnScoreWeight extends ConstantScoreWeight {
private final ScoreMode scoreMode;
private final float[] queryVector;
private final int ef;
+ private AtomicLong visitedCounter;
KnnScoreWeight(Query query, float score, ScoreMode scoreMode, String field, float[] queryVector, int ef) {
super(query, score);
@@ -70,6 +72,7 @@ class KnnScoreWeight extends ConstantScoreWeight {
@Override
public Scorer get(long leadCost) throws IOException {
Neighbors neighbors = hnswReader.searchNeighbors(queryVector, ef, vectorValues);
+ visitedCounter.addAndGet(hnswReader.getVisitedCount());
return new Scorer(weight) {
int doc = -1;
@@ -110,6 +113,7 @@ class KnnScoreWeight extends ConstantScoreWeight {
switch (fi.getVectorDistFunc()) {
case MANHATTAN:
case EUCLIDEAN:
+ // is it necessary to normalize these scores?
score = 1.0f / (next.distance() / numDimensions + 0.01f);
break;
case COSINE:
@@ -158,4 +162,12 @@ class KnnScoreWeight extends ConstantScoreWeight {
public boolean isCacheable(LeafReaderContext ctx) {
return true;
}
+
+ public void setVisitedCounter(AtomicLong counter) {
+ visitedCounter = counter;
+ }
+
+ public long getVisitedCount() {
+ return visitedCounter.get();
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
index c4d8332..37c7338 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
@@ -55,29 +55,27 @@ public final class HNSWGraph implements Accountable {
* @param ef the number of nodes to be searched
* @param level graph level
* @param vectorValues vector values
- * @return nearest neighbors
+ * @return number of candidates visited
*/
- void searchLayer(float[] query, FurthestNeighbors results, int ef, int level, VectorValues vectorValues) throws IOException {
+ int searchLayer(float[] query, FurthestNeighbors results, int ef, int level, VectorValues vectorValues) throws IOException {
if (level >= layers.size()) {
throw new IllegalArgumentException("layer does not exist for the level: " + level);
}
Layer layer = layers.get(level);
TreeSet<Neighbor> candidates = new TreeSet<>();
+ // set of docids that have been visited by search on this layer, used to avoid backtracking
+ Set<Integer> visited = new HashSet<>();
for (Neighbor n : results) {
- // TODO: candidates should get neighbors of neighbors
candidates.add(n);
+ visited.add(n.docId());
}
- // set of docids that have been visited by search on this layer, used to avoid backtracking
- Set<Integer> visited = new HashSet<>();
- // We want to efficiently pop the best (nearest, least distance) candidate, so use NearestNeighbors,
- // but we don't want to overflow the heap and lose the best candidate!
+ Neighbor f = results.top();
while (candidates.size() > 0) {
Neighbor c = candidates.pollFirst();
- Neighbor f = results.top();
assert c.isDeferred() == false;
assert f.isDeferred() == false;
- if (c.distance() > f.distance()) {
+ if (c.distance() > f.distance() && results.size() >= ef) {
break;
}
for (Neighbor e : layer.getFriends(c.docId())) {
@@ -85,21 +83,19 @@ public final class HNSWGraph implements Accountable {
continue;
}
visited.add(e.docId());
- assert f.isDeferred() == false;
float dist = distance(query, e.docId(), vectorValues);
if (dist < f.distance() || results.size() < ef) {
Neighbor n = new ImmutableNeighbor(e.docId(), dist);
candidates.add(n);
- Neighbor popped = results.insertWithOverflow(n);
- if (popped != null && popped != n) {
- f = results.top();
- }
+ results.insertWithOverflow(n);
+ f = results.top();
}
}
}
//System.out.println("level=" + level + ", visited nodes=" + visited.size());
//return pickNearestNeighbor(results);
+ return visited.size();
}
private float distance(float[] query, int docId, VectorValues vectorValues) throws IOException {
@@ -216,13 +212,10 @@ public final class HNSWGraph implements Accountable {
throw new IllegalArgumentException("layer does not exist for level: " + level);
}
layer.connectNodes(node1, node2, dist);
- /*
// ensure friends size <= maxConnections
- //
if (maxConnections > 0) {
layer.shrink(node2, maxConnections);
}
- */
}
/** Connects two nodes; this is supposed to be called when searching */
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
index ef11386..28abd8e 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
@@ -35,12 +35,16 @@ import org.apache.lucene.search.DocIdSetIterator;
* This also caches built {@link HNSWGraph}s for repeated use. */
public final class HNSWGraphReader {
+ // TODO: evict from this cache when a reader's segment is merged. Probably should hold the cache
+ // somewhere that has a similar lifecycle, like Lucene90KnnGraphReader
private static final Map<GraphKey, HNSWGraph> cache = new ConcurrentHashMap<>();
private final String field;
private final LeafReaderContext context;
private final VectorValues.DistanceFunction distFunc;
+ private long visitedCount;
+
public HNSWGraphReader(String field, LeafReaderContext context) {
this.field = field;
this.context = context;
@@ -58,9 +62,9 @@ public final class HNSWGraphReader {
Neighbor ep = new ImmutableNeighbor(enterPoint, VectorValues.distance(query, vectorValues.vectorValue(), distFunc));
FurthestNeighbors neighbors = new FurthestNeighbors(ef, ep);
for (int l = hnsw.topLevel(); l > 0; l--) {
- hnsw.searchLayer(query, neighbors, 1, l, vectorValues);
+ visitedCount += hnsw.searchLayer(query, neighbors, 1, l, vectorValues);
}
- hnsw.searchLayer(query, neighbors, ef, 0, vectorValues);
+ visitedCount += hnsw.searchLayer(query, neighbors, ef, 0, vectorValues);
return neighbors;
}
@@ -122,6 +126,15 @@ public final class HNSWGraphReader {
return hnsw;
}
+ /**
+ * @return the number of documents visited by this reader. For each visited document, the reader
+ * computed the distance to a target vector. This count accumulates over the lifetime of the
+ * reader.
+ */
+ public long getVisitedCount() {
+ return visitedCount;
+ }
+
private static class GraphKey {
final String field;
final Object readerId;
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
index 2760481..a00cfd3 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
@@ -64,7 +64,7 @@ public final class HNSWGraphWriter implements Accountable {
}
/** Full constructor */
- public HNSWGraphWriter(int maxConn, int maxConn0, int efConst, long seed, int numDimensions, VectorValues.DistanceFunction distFunc) {
+ HNSWGraphWriter(int maxConn, int maxConn0, int efConst, long seed, int numDimensions, VectorValues.DistanceFunction distFunc) {
this.maxConn = maxConn;
this.maxConn0 = maxConn0;
this.efConst = efConst;
@@ -115,7 +115,7 @@ public final class HNSWGraphWriter implements Accountable {
FurthestNeighbors results = new FurthestNeighbors(efConst, ep);
// down to the level from the hnsw's top level
for (int l = hnsw.topLevel(); l > level; l--) {
- hnsw.searchLayer(value, results, efConst, l, vectorValues);
+ hnsw.searchLayer(value, results, 1, l, vectorValues);
}
// down to level 0 with placing the doc to each layer
@@ -128,12 +128,10 @@ public final class HNSWGraphWriter implements Accountable {
hnsw.searchLayer(value, results, efConst, l, vectorValues);
int maxConnections = l == 0 ? maxConn0 : maxConn;
- NearestNeighbors neighbors = new NearestNeighbors(maxConnections, results.top());
- for (Neighbor n : results) {
- neighbors.insertWithOverflow(n);
+ while (results.size() > maxConnections) {
+ results.pop();
}
- for (Neighbor n : neighbors) {
- // TODO: limit *total* num connections by pruning (shrinking)
+ for (Neighbor n : results) {
hnsw.connectNodes(l, docId, n.docId(), n.distance(), maxConnections);
}
}
diff --git a/lucene/core/src/test/org/apache/lucene/util/hnsw/KnnGraphTester.java b/lucene/core/src/test/org/apache/lucene/util/hnsw/KnnGraphTester.java
index ad21eca..f731d15 100644
--- a/lucene/core/src/test/org/apache/lucene/util/hnsw/KnnGraphTester.java
+++ b/lucene/core/src/test/org/apache/lucene/util/hnsw/KnnGraphTester.java
@@ -46,6 +46,8 @@ import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.KnnGraphValues;
+import org.apache.lucene.index.NoMergePolicy;
+import org.apache.lucene.index.ReaderUtil;
import org.apache.lucene.index.VectorValues;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.KnnGraphQuery;
@@ -53,6 +55,7 @@ import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
+import org.apache.lucene.util.IntsRef;
import org.apache.lucene.util.PriorityQueue;
/** For testing indexing and search performance of a knn-graph using random vectors */
@@ -65,51 +68,93 @@ public class KnnGraphTester {
private int numDocs;
private int dim;
private int topK;
- private int numProbe;
+ private int fanout;
private float[] vectors;
private int[] nabors;
+ private Path indexPath;
+
public KnnGraphTester() {
// set defaults
numDocs = 10_000;
dim = 256;
topK = 10;
- numProbe = 500;
+ fanout = 50;
random = new Random();
+ indexPath = Paths.get("knn_test_index");
}
public static void main(String... args) throws Exception {
- if (args.length != 2) {
+ new KnnGraphTester().run(args);
+ }
+
+ private void run(String... args) throws Exception {
+ String operation = null, path = null;
+ for (int iarg = 0; iarg < args.length; iarg++) {
+ String arg = args[iarg];
+ switch(arg) {
+ case "-generate":
+ case "-search":
+ case "-check":
+ case "-stats":
+ if (operation != null) {
+ throw new IllegalArgumentException("Specify only one operation, not both " + arg + " and " + operation);
+ }
+ if (iarg == args.length - 1) {
+ throw new IllegalArgumentException("Operation " + arg + " requires a following pathname");
+ }
+ operation = arg;
+ path = args[++iarg];
+ break;
+ case "-fanout":
+ if (iarg == args.length - 1) {
+ throw new IllegalArgumentException("-fanout requires a following number");
+ }
+ fanout = Integer.parseInt(args[++iarg]);
+ break;
+ case "-ndoc":
+ if (iarg == args.length - 1) {
+ throw new IllegalArgumentException("-ndoc requires a following number");
+ }
+ numDocs = Integer.parseInt(args[++iarg]);
+ break;
+ default:
+ usage();
+ }
+ }
+ if (operation == null || path == null) {
usage();
}
- switch (args[0]) {
+ switch (operation) {
case "-generate":
- new KnnGraphTester().create(args[1]);
+ create(path);
break;
case "-search":
- new KnnGraphTester().search(args[1]);
+ search(path);
+ break;
+ case "-check":
+ check(path);
break;
case "-stats":
- new KnnGraphTester().stats(args[1]);
+ stats(path);
break;
- default:
- usage();
}
}
private void search(String dataFile) throws IOException {
readDataFile(dataFile);
- Path indexPath = Paths.get("knn_test_index");
createIndex(indexPath);
- // topK = 25;
testSearch(indexPath, 1000);
- //GraphSearch.VERBOSE = true;
- //testSearch(indexPath, 1);
+ }
+
+ private void check(String dataFile) throws IOException {
+ readDataFile(dataFile);
+ createIndex(indexPath);
+ checkIndex(indexPath, 1000);
}
private void stats(String dataFile) throws IOException {
readDataFile(dataFile);
- Path indexPath = Paths.get("knn_test_index");
createIndex(indexPath);
printFanoutHist(indexPath);
}
@@ -176,6 +221,45 @@ public class KnnGraphTester {
System.out.println("");
}
+ private void checkIndex(Path indexPath, int numIters) throws IOException {
+ float[][] targets = new float[numIters][];
+ TopDocs[] results = new TopDocs[numIters];
+ System.out.println("checking " + numIters + " documents");
+ long start = System.nanoTime();
+ long minVisited = Long.MAX_VALUE, maxVisited = 0, totalVisited = 0;
+ try (Directory dir = FSDirectory.open(indexPath);
+ DirectoryReader reader = DirectoryReader.open(dir)) {
+ int leafCount = reader.leaves().size();
+ for (int i = 0; i < numIters; i++) {
+ targets[i] = new float[dim];
+ int docid = random.nextInt(reader.maxDoc());
+ int readerOrd = ReaderUtil.subIndex(docid, reader.leaves());
+ LeafReaderContext leafReaderContext = reader.leaves().get(readerOrd);
+ int leafDocid = docid - leafReaderContext.docBase;
+ KnnGraphValues graphValues = leafReaderContext.reader().getKnnGraphValues(KNN_FIELD);
+ VectorValues vectorValues = leafReaderContext.reader().getVectorValues(KNN_FIELD);
+ boolean ok = graphValues.advanceExact(leafDocid);
+ if (ok == false) {
+ throw new IllegalStateException("KnnGraphValues failed to advance to docid=" + docid);
+ }
+ IntsRef friends = graphValues.getFriends(0);
+ // check all the supposed friends for overlap with the true list of closest documents
+ int advanced = vectorValues.advance(leafDocid);
+ if (advanced != leafDocid) {
+ throw new IllegalStateException("VectorValues failed to advance to docid=" + docid
+ + "; in leaf=" + readerOrd + " having maxDoc=" + leafReaderContext.reader().maxDoc()
+ + ", advance(" + leafDocid + ") returned " + advanced);
+ }
+ System.arraycopy(vectorValues.vectorValue(), 0, targets[i], 0, dim);
+ results[i] = new TopDocs(null, new ScoreDoc[friends.length]);
+ for (int j = 0; j < friends.length; j++) {
+ results[i].scoreDocs[j] = new ScoreDoc(leafReaderContext.docBase + friends.ints[friends.offset + j], 0);
+ }
+ }
+ }
+ checkResults(targets, results);
+ }
+
private void testSearch(Path indexPath, int numIters) throws IOException {
float[][] targets = new float[numIters][];
TopDocs[] results = new TopDocs[numIters];
@@ -183,15 +267,19 @@ public class KnnGraphTester {
targets[i] = new float[dim];
randomVector(targets[i]);
}
- System.out.println("running " + numIters + " targets; topK=" + topK + ", numProbe=" + numProbe);
+ System.out.println("running " + numIters + " targets; topK=" + topK + ", fanout=" + fanout);
long start = System.nanoTime();
+ long minVisited = Long.MAX_VALUE, maxVisited = 0, totalVisited = 0;
try (Directory dir = FSDirectory.open(indexPath);
DirectoryReader reader = DirectoryReader.open(dir)) {
IndexSearcher searcher = new IndexSearcher(reader);
- int result = 0;
for (int i = 0; i < numIters; i++) {
- KnnGraphQuery query = new KnnGraphQuery(KNN_FIELD, targets[i], numProbe);
+ KnnGraphQuery query = new KnnGraphQuery(KNN_FIELD, targets[i], fanout);
results[i] = searcher.search(query, topK);
+ long numVisited = query.getVisitedCount();
+ minVisited = Math.min(minVisited, numVisited);
+ maxVisited = Math.max(maxVisited, numVisited);
+ totalVisited += numVisited;
for (ScoreDoc scoreDoc : results[i].scoreDocs) {
int id = searcher.doc(scoreDoc.doc).getFields().get(0).numericValue().intValue();
scoreDoc.doc = id;
@@ -200,6 +288,7 @@ public class KnnGraphTester {
}
long elapsed = (System.nanoTime() - start) / 1_000_000; // ns -> ms
System.out.println("completed " + numIters + " searches in " + elapsed + " ms: " + (1000 * numIters / elapsed) + " QPS");
+ System.out.println(" visited documents avg=" + (totalVisited / 1000) + " min=" + minVisited + " max=" + maxVisited);
System.out.println("checking results");
checkResults(targets, results);
}
@@ -207,21 +296,28 @@ public class KnnGraphTester {
private void checkResults(float[][] targets, TopDocs[] results) {
int[] expected = new int[topK];
int totalMatches = 0;
+ int totalResults = 0;
for (int i = 0; i < results.length; i++) {
+ /*
if (results[i].scoreDocs.length != topK) {
System.err.println("search " + i + " got " + results[i].scoreDocs.length + " results, expecting " + topK);
}
- getActualNN(targets[i], 0, expected, 0);
+ */
+ int n = results[i].scoreDocs.length;
+ totalResults += n;
+ if (expected.length < n) {
+ expected = new int[n];
+ }
+ getActualNN(targets[i], 0, expected, 0, n);
int matched = compareNN(expected, results[i]);
totalMatches += matched;
}
- System.out.println("total matches = " + totalMatches);
- System.out.println("Average overlap = " + (100.0 * totalMatches / (results.length * topK)) + "%");
+ System.out.println("total matches = " + totalMatches + " out of " + totalResults);
+ System.out.printf("Average overlap = %.2f%%\n", ((100.0 * totalMatches) / totalResults));
}
int compareNN(int[] expected, TopDocs results) {
int matched = 0;
- int i = 0;
/*
System.out.print("expected=");
for (int j = 0; j < expected.length; j++) {
@@ -236,8 +332,8 @@ public class KnnGraphTester {
System.out.print('\n');
*/
Set<Integer> expectedSet = new HashSet<>();
- for (int doc : expected) {
- expectedSet.add(doc);
+ for (int i = 0; i < results.scoreDocs.length; i++) {
+ expectedSet.add(expected[i]);
}
for (ScoreDoc scoreDoc : results.scoreDocs) {
if (expectedSet.contains(scoreDoc.doc)) {
@@ -253,10 +349,11 @@ public class KnnGraphTester {
* @param targetOffset offset of the target vector in the target parameter
* @param nn an array in which to write the resulting nearest neighbor vector indexes
* @param nnOfset the offset in the nearest neighbor array at which to start writing
+ * @param n how many results to find
*/
- void getActualNN(float[] targetArray, int targetOffset, int[] nn, int nnOffset) {
- final ScoreDocQueue queue = new ScoreDocQueue(topK);
- assert queue.size() == topK : " queue.size()=" + queue.size();
+ void getActualNN(float[] targetArray, int targetOffset, int[] nn, int nnOffset, int n) {
+ final ScoreDocQueue queue = new ScoreDocQueue(n);
+ assert queue.size() == n : " queue.size()=" + queue.size();
int vectorOffset = 0;
ScoreDoc bottom = queue.top();
for (int j = 0; j < numDocs; j++) {
@@ -273,8 +370,9 @@ public class KnnGraphTester {
vectorOffset += dim;
}
assert queue.size() == topK;
- nnOffset += topK;
- for (int k = 1; k <= topK; k++) {
+ nnOffset += n;
+ for (int k = 1; k <= n; k++) {
+ // we don't really need these in score order, but we might?
ScoreDoc scoreDoc = queue.pop();
//System.out.println("" + scoreDoc.doc + ":" + scoreDoc.score);
nn[nnOffset - k] = scoreDoc.doc;
@@ -290,6 +388,7 @@ public class KnnGraphTester {
private void createIndex(Path indexPath) throws IOException {
IndexWriterConfig iwc = new IndexWriterConfig()
.setOpenMode(IndexWriterConfig.OpenMode.CREATE);
+ // iwc.setMergePolicy(NoMergePolicy.INSTANCE);
System.out.println("creating index in " + indexPath);
long start = System.nanoTime();
try (FSDirectory dir = FSDirectory.open(indexPath);
@@ -310,7 +409,7 @@ public class KnnGraphTester {
private void create(String dataFile) throws IOException {
generateRandomVectors(dim * numDocs);
System.out.println("Generated " + numDocs + " random vectors");
- computeNearest();
+ //computeNearest();
writeDataFile(dataFile);
}
@@ -326,9 +425,11 @@ public class KnnGraphTester {
vectors[i] = din.readFloat();
}
nabors = new int[numDocs * topK];
+ /*
for (int i = 0; i < nabors.length; i++) {
nabors[i] = din.readInt();
}
+ */
}
}
@@ -342,9 +443,11 @@ public class KnnGraphTester {
for (int i = 0; i < vectors.length; i++) {
dout.writeFloat(vectors[i]);
}
+ /*
for (int i = 0; i < nabors.length; i++) {
dout.writeInt(nabors[i]);
}
+ */
}
}
@@ -361,7 +464,7 @@ public class KnnGraphTester {
if (i % 1000 == 1) {
System.out.println(" " + (i - 1));
}
- getActualNN(vectors, i * dim, nabors, i * topK);
+ getActualNN(vectors, i * dim, nabors, i * topK, topK);
}
}
@@ -379,7 +482,7 @@ public class KnnGraphTester {
}
private static void usage() {
- String error = "Usage: TestKnnGraph -generate|-search {datafile}";
+ String error = "Usage: TestKnnGraph -generate|-search|-stats|-check {datafile} [-fanout N]";
System.err.println(error);
System.exit(1);
}
[lucene-solr] 09/10: fix single node scenario (#4)
Posted by to...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a commit to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git
commit e4d348b8a5ce4e2e257fa2da43db35fef56ced2e
Author: Xin-Chun Zhang <ir...@qq.com>
AuthorDate: Sun Jan 12 13:49:21 2020 +0800
fix single node scenario (#4)
---
.../apache/lucene/util/hnsw/HNSWGraphReader.java | 11 +++++--
.../test/org/apache/lucene/index/TestKnnGraph.java | 35 ++++++++++++++++++++++
2 files changed, 44 insertions(+), 2 deletions(-)
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
index 28abd8e..0a5c4b2 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
@@ -30,6 +30,7 @@ import org.apache.lucene.index.KnnGraphValues;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.VectorValues;
import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.util.IntsRef;
/** Executes approximate nearest neighbor search on per-reader {@link HNSWGraph}.
* This also caches built {@link HNSWGraph}s for repeated use. */
@@ -117,8 +118,14 @@ public final class HNSWGraphReader {
int maxLevel = graphValues.getMaxLevel();
hnsw.ensureLevel(maxLevel);
for (int l = 0; l <= maxLevel; l++) {
- for (int friend : graphValues.getFriends(l).ints) {
- hnsw.connectNodes(l, doc, friend);
+ final IntsRef friends = graphValues.getFriends(l);
+ if (friends.length > 0) {
+ for (int friend : friends.ints) {
+ hnsw.connectNodes(l, doc, friend);
+ }
+ } else {
+ /// if there's single one node (thus no friends), the node should be added to all layers
+ hnsw.addNode(l, doc);
}
}
}
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java b/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java
index 21de79a..fe84c22 100644
--- a/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java
+++ b/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java
@@ -24,6 +24,7 @@ import java.util.LinkedList;
import java.util.List;
import java.util.Set;
+import org.apache.lucene.codecs.Codec;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.VectorField;
@@ -81,6 +82,21 @@ public class TestKnnGraph extends LuceneTestCase {
}
}
+ public void testSingleDocRecall() throws Exception {
+ try (Directory dir = newDirectory();
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null)
+ .setCodec(Codec.forName("Lucene90")))) {
+ float[][] values = new float[][]{new float[]{0, 1, 2}};
+ add(iw, 0, values[0]);
+ assertConsistentGraph(iw, values);
+
+ iw.commit();
+ assertConsistentGraph(iw, values);
+
+ assertRecall(dir, 0, values[0]);
+ }
+ }
+
/**
* Verify that the graph properties are preserved when merging
*/
@@ -294,4 +310,23 @@ public class TestKnnGraph extends LuceneTestCase {
iw.addDocument(doc);
}
+ private void assertRecall(Directory dir, int expectDocId, float[] value) throws IOException {
+ try (IndexReader reader = DirectoryReader.open(dir)) {
+ IndexSearcher searcher = new IndexSearcher(reader);
+ KnnGraphQuery query = new KnnGraphQuery(KNN_GRAPH_FIELD, value);
+ TopDocs result = searcher.search(query, 1);
+
+ int recallCnt = 0;
+ for (LeafReaderContext ctx : reader.leaves()) {
+ VectorValues vector = ctx.reader().getVectorValues(KNN_GRAPH_FIELD);
+ if (vector.seek(result.scoreDocs[0].doc)) {
+ ++recallCnt;
+ assertEquals(expectDocId, vector.docID());
+ assertEquals(0, Arrays.compare(value, vector.vectorValue()));
+ }
+ }
+ assertEquals(1, recallCnt);
+ }
+ }
+
}
[lucene-solr] 03/10: Implemented TestKnnGraph unit tests
Posted by to...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a commit to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git
commit 444c1d3b8db14592ca31253e1272dba57178266a
Author: Michael Sokolov <so...@amazon.com>
AuthorDate: Fri Dec 6 10:55:24 2019 -0500
Implemented TestKnnGraph unit tests
---
.../org/apache/lucene/codecs/KnnGraphFormat.java | 2 +-
.../org/apache/lucene/codecs/KnnGraphWriter.java | 9 +-
.../codecs/lucene90/Lucene90FieldInfosFormat.java | 3 +-
.../codecs/lucene90/Lucene90KnnGraphReader.java | 39 +--
.../codecs/lucene90/Lucene90KnnGraphWriter.java | 248 +++++++++++++++--
.../java/org/apache/lucene/index/CodecReader.java | 2 +-
.../apache/lucene/index/DefaultIndexingChain.java | 4 +-
.../java/org/apache/lucene/index/FieldInfo.java | 7 +
.../apache/lucene/index/KnnGraphValuesWriter.java | 17 +-
.../java/org/apache/lucene/index/MergeState.java | 2 +-
.../apache/lucene/index/SegmentCoreReaders.java | 2 +-
.../org/apache/lucene/index/SegmentMerger.java | 147 +++++-----
.../java/org/apache/lucene/index/VectorValues.java | 32 ++-
.../org/apache/lucene/search/KnnGraphQuery.java | 1 -
.../org/apache/lucene/util/hnsw/HNSWGraph.java | 61 +++--
.../apache/lucene/util/hnsw/HNSWGraphWriter.java | 6 +-
.../java/org/apache/lucene/util/hnsw/Layer.java | 18 +-
.../test/org/apache/lucene/index/TestKnnGraph.java | 297 +++++++++++++++++++++
18 files changed, 703 insertions(+), 194 deletions(-)
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphFormat.java
index d053301..9893337 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphFormat.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphFormat.java
@@ -45,7 +45,7 @@ public abstract class KnnGraphFormat {
public static KnnGraphFormat EMPTY = new KnnGraphFormat() {
@Override
public KnnGraphWriter fieldsWriter(SegmentWriteState state) throws IOException {
- throw new UnsupportedOperationException();
+ throw new UnsupportedOperationException("Attempt to write EMPTY KnnGraphValues: maybe you forgot to use codec=Lucene90");
}
@Override
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java
index a8d9834..3732003 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java
@@ -54,7 +54,9 @@ public abstract class KnnGraphWriter implements Closeable {
}
for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) {
- mergeOneField(fieldInfo, mergeState);
+ if (fieldInfo.hasVectorValues()) {
+ mergeOneField(fieldInfo, mergeState);
+ }
}
finish();
}
@@ -86,6 +88,8 @@ public abstract class KnnGraphWriter implements Closeable {
}
}
}
+ assert vectorSubs.isEmpty() == false;
+ assert graphSubs.isEmpty() == false;
final DocIDMerger<VectorValuesSub> vectorDocIdMerger = DocIDMerger.of(vectorSubs, mergeState.needsIndexSort);
final DocIDMerger<KnnGraphValuesSub> graphDocIdMerger = DocIDMerger.of(graphSubs, mergeState.needsIndexSort);
@@ -195,6 +199,9 @@ public abstract class KnnGraphWriter implements Closeable {
}
// make links to other graphs
if (enterPoints.stream().anyMatch(ep -> ep.docId == docID)) {
+ // this docid is a mapped enter point (mapped from the first enter points of each segment's graph).
+ // add all the other mapped enter points as its friends. The result is that segments' graphs are linked all-to-all
+ // by N*(N+1)/2 links
enterPoints.stream().filter(ep -> ep.docId != docID && ep.topLevel >= level).forEach(ep -> {
mappedFriends.append(ep.docId);
});
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java
index b846bc0..e91ab88 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90FieldInfosFormat.java
@@ -240,8 +240,9 @@ public final class Lucene90FieldInfosFormat extends FieldInfosFormat {
switch (distFunc) {
case NONE:
case MANHATTAN:
- case COSINE:
+ case EUCLIDEAN:
return (byte)distFunc.getId();
+ case COSINE:
default:
// BUG
throw new AssertionError("unhandled DistanceFunction: " + distFunc);
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
index 34a6f3d..1083fd4 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
@@ -33,6 +33,7 @@ import org.apache.lucene.index.VectorValues;
import org.apache.lucene.store.ChecksumIndexInput;
import org.apache.lucene.store.IndexInput;
import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.IntsRef;
import org.apache.lucene.util.RamUsageEstimator;
@@ -261,14 +262,14 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
final IndexInput dataIn;
int doc = -1;
- byte[] binaryValue;
+ BytesRef binaryValue;
RandomAccessVectorValuesReader(int maxDoc, int numDims, KnnGraphEntry entry, IndexInput dataIn) {
this.maxDoc = maxDoc;
this.numDims = numDims;
this.entry = entry;
this.dataIn = dataIn;
- this.binaryValue = new byte[Float.BYTES * numDims];
+ this.binaryValue = new BytesRef(new byte[Float.BYTES * numDims]);
}
@Override
@@ -280,6 +281,11 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
}
@Override
+ public BytesRef binaryValue() {
+ return binaryValue;
+ }
+
+ @Override
public boolean seek(int target) throws IOException {
if (target > maxDoc) {
doc = NO_MORE_DOCS;
@@ -293,7 +299,7 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
int offset = Float.BYTES * numDims * ord;
assert offset >= 0;
dataIn.seek(offset);
- dataIn.readBytes(binaryValue, 0, binaryValue.length);
+ dataIn.readBytes(binaryValue.bytes, binaryValue.offset, binaryValue.length);
return true;
}
@@ -382,7 +388,8 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
@Override
public int advance(int target) throws IOException {
- if (doc > target || target > maxDoc) {
+ // enabled random access. Also I think maxDoc really means maxDoc + 1
+ if (target >= maxDoc) {
return doc = NO_MORE_DOCS;
}
doc = target - 1;
@@ -397,7 +404,7 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
}
dataIn.seek(offset);
int maxLevel = dataIn.readInt();
- assert maxLevel >= 0;
+ assert maxLevel >= 0 || entry.enterPoints.length == 0;
this.maxLevel = maxLevel;
this.friendsRef = readFriends(maxLevel);
return doc;
@@ -407,17 +414,19 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
FriendsRef friendsRef = new FriendsRef(maxLevel);
for (int l = maxLevel; l >= 0; l--) {
int friendSize = dataIn.readInt();
- assert friendSize > 0;
- int[] frineds = new int[friendSize];
- int friendId = dataIn.readVInt(); // first friend id
- frineds[0] = friendId;
- for (int i = 1; i < friendSize; i++) {
- int delta = dataIn.readVInt();
- assert delta > 0;
- friendId += delta;
- frineds[i] = friendId;
+ // assert friendSize > 0;
+ int[] friends = new int[friendSize];
+ if (friendSize > 0) {
+ int friendId = dataIn.readVInt(); // first friend id
+ friends[0] = friendId;
+ for (int i = 1; i < friendSize; i++) {
+ int delta = dataIn.readVInt();
+ assert delta > 0;
+ friendId += delta;
+ friends[i] = friendId;
+ }
}
- friendsRef.setFriendsAtLevel(l, frineds);
+ friendsRef.setFriendsAtLevel(l, friends);
}
return friendsRef;
}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
index 4c7f4fd..8644c50 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
@@ -18,23 +18,31 @@
package org.apache.lucene.codecs.lucene90;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.Arrays;
-import java.util.HashMap;
+import java.util.List;
import java.util.Map;
+import java.util.TreeMap;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnGraphReader;
import org.apache.lucene.codecs.KnnGraphWriter;
+import org.apache.lucene.index.DocIDMerger;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.IndexFileNames;
import org.apache.lucene.index.KnnGraphValues;
+import org.apache.lucene.index.KnnGraphValuesWriter;
+import org.apache.lucene.index.MergeState;
import org.apache.lucene.index.SegmentWriteState;
import org.apache.lucene.index.VectorValues;
-import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.Counter;
import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.IntsRef;
+import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
+
/**
* Writes vector values and knn graphs to index segments.
*/
@@ -44,7 +52,7 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
private boolean finished;
- public Lucene90KnnGraphWriter(SegmentWriteState state) throws IOException {
+ Lucene90KnnGraphWriter(SegmentWriteState state) throws IOException {
assert state.fieldInfos.hasVectorValues();
String metaFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, Lucene90KnnGraphFormat.META_EXTENSION);
@@ -76,7 +84,10 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
@Override
public void writeField(FieldInfo fieldInfo, KnnGraphReader values) throws IOException {
- Map<Integer, Long> docToOffset = new HashMap<>();
+ // we require a TreeMap so that we can iterate over docid keys in order. Instead we should move
+ // to the IndexedDISI encoding used by BinaryDocValues
+ Map<Integer, Long> docToOffset = new TreeMap<>();
+ List<Integer> enterPoints = new ArrayList<>();
long vectorDataOffset = vectorData.getFilePointer();
long graphDataOffset = graphData.getFilePointer();
@@ -85,15 +96,13 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
VectorValues vectors = values.getVectorValues(fieldInfo.name);
KnnGraphValues graph = values.getGraphValues(fieldInfo.name);
- for (int docV = vectors.nextDoc(), docG = graph.nextDoc();
- docV != DocIdSetIterator.NO_MORE_DOCS && docG != DocIdSetIterator.NO_MORE_DOCS;
+ int docV, docG;
+ for (docV = vectors.nextDoc(), docG = graph.nextDoc();
+ docV != NO_MORE_DOCS && docG != NO_MORE_DOCS;
docV = vectors.nextDoc(), docG = graph.nextDoc()) {
assert docV == docG; // must be same
- // write vector value
- byte[] binaryValue = vectors.binaryValue();
- VectorValues.verifyNumDimensions(binaryValue, numDims);
- vectorData.writeBytes(binaryValue, binaryValue.length);
+ writeVectorValue(numDims, vectors);
// write knn graph value
docToOffset.put(docG, graphData.getFilePointer() - graphDataOffset);
@@ -103,21 +112,25 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
graphData.writeInt(maxLevel);
for (int l = graph.getMaxLevel(); l >= 0; l--) {
IntsRef friends = graph.getFriends(l);
- assert friends.length > 0 : "doc " + docG + " has empty friends list at level=" + l;
+ // assert friends.length > 0 : "doc " + docG + " has empty friends list at level=" + l;
graphData.writeInt(friends.length);
- int stop = friends.offset + friends.length;
- // sort friend ids
- Arrays.sort(friends.ints, friends.offset, stop);
- // write the smallest friend id
- graphData.writeVInt(friends.ints[friends.offset]);
- for (int i = friends.offset + 1; i < stop; i++) {
- // write delta
- int delta = friends.ints[i] - friends.ints[i - 1];
- assert delta > 0;
- graphData.writeVInt(delta);
+ if (friends.length > 0) {
+ int stop = friends.offset + friends.length;
+ // sort friend ids
+ Arrays.sort(friends.ints, friends.offset, stop);
+ // write the smallest friend id
+ graphData.writeVInt(friends.ints[friends.offset]);
+ for (int i = friends.offset + 1; i < stop; i++) {
+ // write delta
+ int delta = friends.ints[i] - friends.ints[i - 1];
+ assert delta > 0;
+ graphData.writeVInt(delta);
+ }
}
}
}
+ assert docV == NO_MORE_DOCS; // must be exhausted
+ assert docG == NO_MORE_DOCS; // must be exhausted
long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset;
long graphDataLength = graphData.getFilePointer() - graphDataOffset;
@@ -125,6 +138,13 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
writeMeta(fieldInfo, vectorDataOffset, vectorDataLength, graphDataOffset, graphDataLength, graph.getTopLevel(), graph.getEnterPoints(), docToOffset);
}
+ private void writeVectorValue(int numDims, VectorValues vectors) throws IOException {
+ // write vector value
+ BytesRef binaryValue = vectors.binaryValue();
+ VectorValues.verifyNumDimensions(binaryValue.length, numDims);
+ vectorData.writeBytes(binaryValue.bytes, binaryValue.offset, binaryValue.length);
+ }
+
private void writeMeta(FieldInfo field, long vectorDataOffset, long vectorDataLength, long graphDataOffset, long graphDataLength,
int topLevel, int[] eps, Map<Integer, Long> docToOffset) throws IOException {
meta.writeInt(field.number);
@@ -139,11 +159,197 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
}
meta.writeInt(docToOffset.size());
for (Map.Entry<Integer, Long> entry : docToOffset.entrySet()) {
+ // these are not in sorted order, yet we write the vectors in order by docid
meta.writeVInt(entry.getKey());
meta.writeVLong(entry.getValue());
}
}
+ /*
+ @Override
+ public void merge(MergeState mergeState) throws IOException {
+ for (KnnGraphReader reader : mergeState.knnGraphReaders) {
+ if (reader != null) {
+ reader.checkIntegrity();
+ }
+ }
+ for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) {
+ if (fieldInfo.hasVectorValues()) {
+ mergeKnnGraph(fieldInfo, mergeState);
+ }
+ }
+ finish();
+ }
+ */
+
+ /**
+ * Merges the segment HNSW graphs by constructing a new merged graph using HNSWGraph
+ */
+ private void mergeKnnGraph(FieldInfo mergeFieldInfo, final MergeState mergeState) throws IOException {
+ // We must compute the entire merged field in memory since each document's values depend on its neighbors
+ //mergeState.infoStream.message("ReferenceDocValues", "merging " + mergeState.segmentInfo);
+ List<VectorValuesSub> subs = new ArrayList<>();
+ int dimension = -1;
+ for (int i = 0 ; i < mergeState.knnGraphReaders.length; i++) {
+ KnnGraphReader graphReader = mergeState.knnGraphReaders[i];
+ if (graphReader != null) {
+ if (mergeFieldInfo != null && mergeFieldInfo.hasVectorValues()) {
+ int segmentDimension = mergeFieldInfo.getVectorNumDimensions();
+ if (dimension == -1) {
+ dimension = segmentDimension;
+ } else if (dimension != segmentDimension) {
+ throw new IllegalStateException("Varying dimensions for vector-valued field " + mergeFieldInfo.name
+ + ": " + dimension + "!=" + segmentDimension);
+ }
+ VectorValues values = graphReader.getVectorValues(mergeFieldInfo.name);
+ subs.add(new VectorValuesSub(i, mergeState.docMaps[i], values));
+ }
+ }
+ }
+ // Create a new KnnGraphValues by iterating over the vectors, searching for
+ // its nearest neighbor vectors in the newly merged segments' vectors, mapping the resulting
+ // docids using docMaps in the mergeState.
+ MultiVectorValues multiVectors = new MultiVectorValues(subs, mergeState.maxDocs);
+ KnnGraphValuesWriter valuesWriter = new KnnGraphValuesWriter(mergeFieldInfo, Counter.newCounter());
+ for (int i = 0; i < subs.size(); i++) {
+ VectorValuesSub sub = subs.get(i);
+ MergeState.DocMap docMap = mergeState.docMaps[sub.segmentIndex];
+ // nocommit: test sorted index and test index with deletions
+ int docid;
+ while ((docid = sub.nextDoc()) != NO_MORE_DOCS) {
+ int mappedDocId = docMap.get(docid);
+ assert sub.values.docID() == docid;
+ assert docid == multiVectors.unmap(mappedDocId) : "unmap mismatch " + docid + " != " + multiVectors.unmap(mappedDocId);
+ valuesWriter.addValue(mappedDocId, sub.values.binaryValue());
+ }
+ }
+ valuesWriter.flush(null, this);
+ //writeField(mergeFieldInfo, valuesWriter.getGraphValues(), valuesWriter.getVectorValues());
+ //mergeState.infoStream.message("ReferenceDocValues", " mergeReferenceField done: " + mergeState.segmentInfo);
+ }
+
+ /** Tracks state of one binary sub-reader that we are merging */
+ private static class VectorValuesSub extends DocIDMerger.Sub {
+
+ final VectorValues values;
+ final int segmentIndex;
+
+ VectorValuesSub(int segmentIndex, MergeState.DocMap docMap, VectorValues values) {
+ super(docMap);
+ this.values = values;
+ this.segmentIndex = segmentIndex;
+ assert values.docID() == -1;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ return values.nextDoc();
+ }
+ }
+
+ // provides a view over multiple VectorValues by concatenating their docid spaces
+ private static class MultiVectorValues extends VectorValues {
+ private final VectorValuesSub[] subValues;
+ private final int[] docBase;
+ private final int[] segmentMaxDocs;
+ private final int cost;
+
+ private int whichSub;
+
+ MultiVectorValues(List<VectorValuesSub> subs, int[] maxDocs) throws IOException {
+ this.subValues = new VectorValuesSub[subs.size()];
+ // TODO: this complicated logic needs its own test
+ // maxDocs actually says *how many* docs there are, not what the number of the max doc is
+ int maxDoc = -1;
+ int lastMaxDoc = -1;
+ segmentMaxDocs = new int[subs.size() - 1];
+ docBase = new int[subs.size()];
+ for (int i = 0, j = 0; j < subs.size(); i++) {
+ lastMaxDoc = maxDoc;
+ maxDoc += maxDocs[i];
+ if (i == subs.get(j).segmentIndex) {
+ // we may skip some segments if they have no docs with values for this field
+ if (j > 0) {
+ segmentMaxDocs[j - 1] = lastMaxDoc;
+ }
+ docBase[j] = lastMaxDoc + 1;
+ ++j;
+ }
+ }
+ int i = 0;
+ int totalCost = 0;
+ for (VectorValuesSub sub : subs) {
+ totalCost += sub.values.cost();
+ this.subValues[i++] = sub;
+ }
+ cost = totalCost;
+ whichSub = 0;
+ }
+
+ @Override
+ public long cost() {
+ return cost;
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ int rebased = unmapSettingWhich(target);
+ if (rebased < 0) {
+ rebased = 0;
+ }
+ int segmentDocId = subValues[whichSub].values.advance(rebased);
+ if (segmentDocId == NO_MORE_DOCS) {
+ if (++whichSub < subValues.length) {
+ // Get the first document in the next segment; Note that all segments have values.
+ segmentDocId = subValues[whichSub].values.advance(0);
+ } else {
+ return NO_MORE_DOCS;
+ }
+ }
+ return docBase[whichSub] + segmentDocId;
+ }
+
+ @Override
+ public boolean seek(int target) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public float[] vectorValue() throws IOException {
+ return subValues[whichSub].values.vectorValue();
+ }
+
+ int unmap(int docid) {
+ // map from global (merged) to segment-local (unmerged)
+ // like mapDocid but no side effect - used for assertion
+ return docid - docBase[findSegment(docid)];
+ }
+
+ private int unmapSettingWhich(int target) {
+ whichSub = findSegment(target);
+ return target - docBase[whichSub];
+ }
+
+ private int findSegment(int docid) {
+ int segment = Arrays.binarySearch(segmentMaxDocs, docid);
+ if (segment < 0) {
+ return -1 - segment;
+ } else {
+ return segment;
+ }
+ }
+
+ @Override
+ public int docID() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ throw new UnsupportedOperationException();
+ }
+ }
+
@Override
public void finish() throws IOException {
if (finished) {
diff --git a/lucene/core/src/java/org/apache/lucene/index/CodecReader.java b/lucene/core/src/java/org/apache/lucene/index/CodecReader.java
index 74c64e9..0f838ba 100644
--- a/lucene/core/src/java/org/apache/lucene/index/CodecReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/CodecReader.java
@@ -222,7 +222,7 @@ public abstract class CodecReader extends LeafReader implements Accountable {
}
@Override
- public KnnGraphValues getKnnGraphValues(String field) throws IOException {
+ public final KnnGraphValues getKnnGraphValues(String field) throws IOException {
ensureOpen();
FieldInfo fi = getFieldInfos().fieldInfo(field);
if (fi == null || fi.getVectorNumDimensions() == 0) {
diff --git a/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java b/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java
index 4d1de5d..eed8339 100644
--- a/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java
+++ b/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java
@@ -308,7 +308,7 @@ final class DefaultIndexingChain extends DocConsumer {
KnnGraphWriter knnGraphWriter = null;
boolean success = false;
try {
- for (int i=0;i<fieldHash.length;i++) {
+ for (int i = 0; i<fieldHash.length; i++) {
PerField perField = fieldHash[i];
while (perField != null) {
if (perField.knnGraphValuesWriter != null) {
@@ -325,7 +325,7 @@ final class DefaultIndexingChain extends DocConsumer {
knnGraphWriter = fmt.fieldsWriter(state);
}
- perField.knnGraphValuesWriter.flush(state, sortMap, knnGraphWriter);
+ perField.knnGraphValuesWriter.flush(sortMap, knnGraphWriter);
perField.knnGraphValuesWriter = null;
} else if (perField.fieldInfo.getVectorNumDimensions() != 0) {
// BUG
diff --git a/lucene/core/src/java/org/apache/lucene/index/FieldInfo.java b/lucene/core/src/java/org/apache/lucene/index/FieldInfo.java
index d71f3b4..79d2c19 100644
--- a/lucene/core/src/java/org/apache/lucene/index/FieldInfo.java
+++ b/lucene/core/src/java/org/apache/lucene/index/FieldInfo.java
@@ -382,6 +382,13 @@ public final class FieldInfo {
public boolean hasVectors() {
return storeTermVector;
}
+
+ /**
+ * @return true if any (numeric) vector values exist for this field
+ */
+ public boolean hasVectorValues() {
+ return vectorNumDimensions > 0;
+ }
/**
* Get a codec attribute value, or null if it does not exist
diff --git a/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java b/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java
index 0a6f491..10c834c 100644
--- a/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java
@@ -43,10 +43,6 @@ public class KnnGraphValuesWriter implements Accountable {
private long bytesUsed = 0L;
- public KnnGraphValuesWriter(FieldInfo fieldInfo) {
- this(fieldInfo, null);
- }
-
public KnnGraphValuesWriter(FieldInfo fieldInfo, Counter iwBytesUsed) {
this.fieldInfo = fieldInfo;
this.iwBytesUsed = iwBytesUsed;
@@ -84,7 +80,7 @@ public class KnnGraphValuesWriter implements Accountable {
bytesUsed = newBytesUsed;
}
- public void flush(SegmentWriteState state, Sorter.DocMap sortMap, KnnGraphWriter graphWriter) throws IOException {
+ public void flush(Sorter.DocMap sortMap, KnnGraphWriter graphWriter) throws IOException {
hnswGraphWriter.finish();
VectorValues vectors = new BufferedVectorValues(docsWithFieldVec.iterator(), hnswGraphWriter.rawVectorsArray());
KnnGraphValues graph = new BufferedKnnGraphValues(docsWithFieldGrp.iterator(), hnswGraphWriter.hnswGraph());
@@ -137,8 +133,8 @@ public class KnnGraphValuesWriter implements Accountable {
final DocIdSetIterator docsWithField;
final float[][] vectorsArray;
- float[] value = new float[0];
int bufferPos = 0;
+ float[] value;
BufferedVectorValues(DocIdSetIterator docsWithField, float[][] vectorsArray) {
this.docsWithField = docsWithField;
@@ -226,15 +222,14 @@ public class KnnGraphValuesWriter implements Accountable {
public int nextDoc() throws IOException {
int docId = docsWithField.nextDoc();
if (docId != NO_MORE_DOCS) {
- boolean levelFound = false;
- for (int l = hnswGraph.topLevel(); l >= 0; l--) {
+ for (int l = hnswGraph.topLevel(); l > 0; l--) {
if (hnswGraph.hasFriends(l, docId)) {
maxLevel = l;
- levelFound = true;
- break;
+ return docId;
}
}
- assert levelFound;
+ // A node in a singleton graph has no friends
+ maxLevel = 0;
}
return docId;
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/MergeState.java b/lucene/core/src/java/org/apache/lucene/index/MergeState.java
index 30763fc..8fb23ca 100644
--- a/lucene/core/src/java/org/apache/lucene/index/MergeState.java
+++ b/lucene/core/src/java/org/apache/lucene/index/MergeState.java
@@ -78,7 +78,7 @@ public class MergeState {
/** Point readers to merge */
public final PointsReader[] pointsReaders;
- /** Knn graph readers to merge */
+ /** KNN Graph readers to merge */
public final KnnGraphReader[] knnGraphReaders;
/** Max docs per reader */
diff --git a/lucene/core/src/java/org/apache/lucene/index/SegmentCoreReaders.java b/lucene/core/src/java/org/apache/lucene/index/SegmentCoreReaders.java
index 9222839..1526cbb 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SegmentCoreReaders.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SegmentCoreReaders.java
@@ -178,7 +178,7 @@ final class SegmentCoreReaders {
Throwable th = null;
try (Closeable finalizer = this::notifyCoreClosedListeners){
IOUtils.close(termVectorsLocal, fieldsReaderLocal, fields, termVectorsReaderOrig, fieldsReaderOrig,
- cfsReader, normsProducer, pointsReader);
+ cfsReader, normsProducer, pointsReader, knnGraphReader);
}
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/SegmentMerger.java b/lucene/core/src/java/org/apache/lucene/index/SegmentMerger.java
index 974212a..7d94e0a 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SegmentMerger.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SegmentMerger.java
@@ -20,6 +20,7 @@ package org.apache.lucene.index;
import java.io.IOException;
import java.util.Collections;
import java.util.List;
+import java.util.function.Supplier;
import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.DocValuesConsumer;
@@ -100,15 +101,8 @@ final class SegmentMerger {
throw new IllegalStateException("Merge would result in 0 document segment");
}
mergeFieldInfos();
- long t0 = 0;
- if (mergeState.infoStream.isEnabled("SM")) {
- t0 = System.nanoTime();
- }
- int numMerged = mergeFields();
- if (mergeState.infoStream.isEnabled("SM")) {
- long t1 = System.nanoTime();
- mergeState.infoStream.message("SM", ((t1-t0)/1000000) + " msec to merge stored fields [" + numMerged + " docs]");
- }
+
+ int numMerged = mergeWithLogging(this::mergeFields, "stored fields");
assert numMerged == mergeState.segmentInfo.maxDoc(): "numMerged=" + numMerged + " vs mergeState.segmentInfo.maxDoc()=" + mergeState.segmentInfo.maxDoc();
final SegmentWriteState segmentWriteState = new SegmentWriteState(mergeState.infoStream, directory, mergeState.segmentInfo,
@@ -117,88 +111,29 @@ final class SegmentMerger {
true, IOContext.READ, segmentWriteState.segmentSuffix, Collections.emptyMap());
if (mergeState.mergeFieldInfos.hasNorms()) {
- if (mergeState.infoStream.isEnabled("SM")) {
- t0 = System.nanoTime();
- }
- mergeNorms(segmentWriteState);
- if (mergeState.infoStream.isEnabled("SM")) {
- long t1 = System.nanoTime();
- mergeState.infoStream.message("SM", ((t1-t0)/1000000) + " msec to merge norms [" + numMerged + " docs]");
- }
+ mergeWithLogging(() -> mergeNorms(segmentWriteState), "norms", numMerged);
}
- if (mergeState.infoStream.isEnabled("SM")) {
- t0 = System.nanoTime();
- }
- try (NormsProducer norms = mergeState.mergeFieldInfos.hasNorms()
- ? codec.normsFormat().normsProducer(segmentReadState)
- : null) {
- NormsProducer normsMergeInstance = null;
- if (norms != null) {
- // Use the merge instance in order to reuse the same IndexInput for all terms
- normsMergeInstance = norms.getMergeInstance();
- }
- mergeTerms(segmentWriteState, normsMergeInstance);
- }
- if (mergeState.infoStream.isEnabled("SM")) {
- long t1 = System.nanoTime();
- mergeState.infoStream.message("SM", ((t1-t0)/1000000) + " msec to merge postings [" + numMerged + " docs]");
- }
+ mergeWithLogging(() -> mergeTerms(segmentWriteState, segmentReadState), "postings", numMerged);
- if (mergeState.infoStream.isEnabled("SM")) {
- t0 = System.nanoTime();
- }
if (mergeState.mergeFieldInfos.hasDocValues()) {
- mergeDocValues(segmentWriteState);
- }
- if (mergeState.infoStream.isEnabled("SM")) {
- long t1 = System.nanoTime();
- mergeState.infoStream.message("SM", ((t1-t0)/1000000) + " msec to merge doc values [" + numMerged + " docs]");
+ mergeWithLogging(() -> mergeDocValues(segmentWriteState), "doc values", numMerged);
}
- if (mergeState.infoStream.isEnabled("SM")) {
- t0 = System.nanoTime();
- }
if (mergeState.mergeFieldInfos.hasPointValues()) {
- mergePoints(segmentWriteState);
- }
- if (mergeState.infoStream.isEnabled("SM")) {
- long t1 = System.nanoTime();
- mergeState.infoStream.message("SM", ((t1-t0)/1000000) + " msec to merge points [" + numMerged + " docs]");
+ mergeWithLogging(() -> mergePoints(segmentWriteState), "points", numMerged);
}
- if (mergeState.infoStream.isEnabled("SM")) {
- t0 = System.nanoTime();
- }
if (mergeState.mergeFieldInfos.hasVectorValues()) {
- mergeKnnGraphValues(segmentWriteState);
- }
- if (mergeState.infoStream.isEnabled("SM")) {
- long t1 = System.nanoTime();
- mergeState.infoStream.message("SM", ((t1-t0)/1000000) + " msec to merge knn graph [" + numMerged + " docs]");
+ mergeWithLogging(() -> mergeKnnGraphValues(segmentWriteState), "knn graph", numMerged);
}
if (mergeState.mergeFieldInfos.hasVectors()) {
- if (mergeState.infoStream.isEnabled("SM")) {
- t0 = System.nanoTime();
- }
- numMerged = mergeVectors();
- if (mergeState.infoStream.isEnabled("SM")) {
- long t1 = System.nanoTime();
- mergeState.infoStream.message("SM", ((t1-t0)/1000000) + " msec to merge vectors [" + numMerged + " docs]");
- }
- assert numMerged == mergeState.segmentInfo.maxDoc();
+ mergeWithLogging(this::mergeVectors, "term vectors");
}
-
+
// write the merged infos
- if (mergeState.infoStream.isEnabled("SM")) {
- t0 = System.nanoTime();
- }
- codec.fieldInfosFormat().write(directory, mergeState.segmentInfo, "", mergeState.mergeFieldInfos, context);
- if (mergeState.infoStream.isEnabled("SM")) {
- long t1 = System.nanoTime();
- mergeState.infoStream.message("SM", ((t1-t0)/1000000) + " msec to write field infos [" + numMerged + " docs]");
- }
+ mergeWithLogging(() -> codec.fieldInfosFormat().write(directory, mergeState.segmentInfo, "", mergeState.mergeFieldInfos, context), "field infos", numMerged);
return mergeState;
}
@@ -220,7 +155,22 @@ final class SegmentMerger {
consumer.merge(mergeState);
}
}
-
+
+ private void mergeTerms(SegmentWriteState segmentWriteState, SegmentReadState segmentReadState) throws IOException {
+ try (NormsProducer norms = mergeState.mergeFieldInfos.hasNorms()
+ ? codec.normsFormat().normsProducer(segmentReadState)
+ : null) {
+ NormsProducer normsMergeInstance = null;
+ if (norms != null) {
+ // Use the merge instance in order to reuse the same IndexInput for all terms
+ normsMergeInstance = norms.getMergeInstance();
+ }
+ try (FieldsConsumer consumer = codec.postingsFormat().fieldsConsumer(segmentWriteState)) {
+ consumer.merge(mergeState, normsMergeInstance);
+ }
+ }
+ }
+
public void mergeFieldInfos() throws IOException {
for (FieldInfos readerFieldInfos : mergeState.fieldInfos) {
for (FieldInfo fi : readerFieldInfos) {
@@ -248,13 +198,9 @@ final class SegmentMerger {
*/
private int mergeVectors() throws IOException {
try (TermVectorsWriter termVectorsWriter = codec.termVectorsFormat().vectorsWriter(directory, mergeState.segmentInfo, context)) {
- return termVectorsWriter.merge(mergeState);
- }
- }
-
- private void mergeTerms(SegmentWriteState segmentWriteState, NormsProducer norms) throws IOException {
- try (FieldsConsumer consumer = codec.postingsFormat().fieldsConsumer(segmentWriteState)) {
- consumer.merge(mergeState, norms);
+ int numMerged = termVectorsWriter.merge(mergeState);
+ assert numMerged == mergeState.segmentInfo.maxDoc();
+ return numMerged;
}
}
@@ -264,4 +210,37 @@ final class SegmentMerger {
}
}
+ private interface Merger {
+ int merge() throws IOException;
+ }
+
+ private interface VoidMerger {
+ void merge() throws IOException;
+ }
+
+ private int mergeWithLogging(Merger merger, String formatName) throws IOException {
+ long t0 = 0;
+ if (mergeState.infoStream.isEnabled("SM")) {
+ t0 = System.nanoTime();
+ }
+ int numMerged = merger.merge();
+ if (mergeState.infoStream.isEnabled("SM")) {
+ long t1 = System.nanoTime();
+ mergeState.infoStream.message("SM", ((t1 - t0) / 1000000) + " msec to merge vector values [" + numMerged + " docs]");
+ }
+ return numMerged;
+ }
+
+ private void mergeWithLogging(VoidMerger merger, String formatName, int numMerged) throws IOException {
+ long t0 = 0;
+ if (mergeState.infoStream.isEnabled("SM")) {
+ t0 = System.nanoTime();
+ }
+ merger.merge();
+ if (mergeState.infoStream.isEnabled("SM")) {
+ long t1 = System.nanoTime();
+ mergeState.infoStream.message("SM", ((t1 - t0) / 1000000) + " msec to merge vector values [" + numMerged + " docs]");
+ }
+ }
+
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/VectorValues.java b/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
index 97ddcca..b0aed22 100644
--- a/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
+++ b/lucene/core/src/java/org/apache/lucene/index/VectorValues.java
@@ -18,9 +18,12 @@
package org.apache.lucene.index;
import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.FloatBuffer;
import java.util.Arrays;
import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.util.BytesRef;
/**
* Access to per-document vector value.
@@ -46,7 +49,7 @@ public abstract class VectorValues extends DocIdSetIterator {
* returned {@code false}.
* @return binary value
*/
- public byte[] binaryValue() throws IOException {
+ public BytesRef binaryValue() throws IOException {
return encode(vectorValue());
}
@@ -69,7 +72,7 @@ public abstract class VectorValues extends DocIdSetIterator {
/**
* Encodes float array to byte array.
*/
- public static byte[] encode(float[] value) {
+ public static BytesRef encode(float[] value) {
byte[] bytes = new byte[Float.BYTES * value.length];
for (int i = 0; i < value.length; i++) {
int bits = Float.floatToIntBits(value[i]);
@@ -78,14 +81,14 @@ public abstract class VectorValues extends DocIdSetIterator {
bytes[i * Float.BYTES + 2] = (byte)(bits >> 8);
bytes[i * Float.BYTES + 3] = (byte)(bits);
}
- return bytes;
+ return new BytesRef(bytes);
}
- public static boolean verifyNumDimensions(byte[] bytes, int numDims) {
- if (bytes.length % Float.BYTES != 0) {
- throw new IllegalArgumentException("Cannot decode bytes array to float array. Reason: invalid length: " + bytes.length);
+ public static boolean verifyNumDimensions(int numBytes, int numDims) {
+ if (numBytes % Float.BYTES != 0) {
+ throw new IllegalArgumentException("Cannot decode bytes array to float array. Reason: invalid length: " + numBytes);
}
- int dims = bytes.length / Float.BYTES;
+ int dims = numBytes / Float.BYTES;
if (dims != numDims) {
throw new IllegalArgumentException("Invalid dimensions: " + dims + " (the number of dimensions should be " + numDims + ")");
}
@@ -93,18 +96,13 @@ public abstract class VectorValues extends DocIdSetIterator {
}
/**
- * Decodes float array from byte array.
+ * Decodes float array from byte array. TODO: allow caller to supply the array so they control allocation.
*/
- public static float[] decode(byte[] bytes, int numDims) {
- verifyNumDimensions(bytes, numDims);
+ public static float[] decode(BytesRef bytes, int numDims) {
+ verifyNumDimensions(bytes.length, numDims);
float[] value = new float[numDims];
- for (int i = 0; i < numDims; i++) {
- int bits = ((bytes[i * Float.BYTES] & 0xFF) << 24) |
- ((bytes[i * Float.BYTES + 1] & 0xFF) << 16) |
- ((bytes[i * Float.BYTES + 2] & 0xFF) << 8) |
- (bytes[i * Float.BYTES + 3] & 0xFF);
- value[i] = Float.intBitsToFloat(bits);
- }
+ ByteBuffer buffer = ByteBuffer.wrap(bytes.bytes, bytes.offset, bytes.length);
+ buffer.asFloatBuffer().get(value);
return value;
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java b/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
index b2ca26b..0ee5ab2 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnGraphQuery.java
@@ -71,7 +71,6 @@ public class KnnGraphQuery extends Query implements Accountable {
* @param ef number of per-segment candidates to be scored/collected. the collector does not return results exceeding {@code ef}.
* increasing this value leads higher recall at the expense of the search speed.
* @param reader index reader
- * @throws IOException
*/
public KnnGraphQuery(String field, float[] queryVector, int ef, IndexReader reader) throws IOException {
this.field = field;
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
index d2acbef..6722e76 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
@@ -18,10 +18,10 @@
package org.apache.lucene.util.hnsw;
import java.io.IOException;
-import java.util.HashMap;
+import java.util.ArrayList;
+import java.util.Collections;
import java.util.HashSet;
import java.util.List;
-import java.util.Map;
import java.util.Set;
import org.apache.lucene.index.VectorValues;
@@ -37,14 +37,13 @@ import org.apache.lucene.util.RamUsageEstimator;
public final class HNSWGraph implements Accountable {
private final VectorValues.DistanceFunction distFunc;
- private final Map<Integer, Layer> layers;
+ private final List<Layer> layers;
- private boolean freeze = false;
- private int topLevel = 0;
+ private boolean frozen = false;
public HNSWGraph(VectorValues.DistanceFunction distFunc) {
this.distFunc = distFunc;
- this.layers = new HashMap<>();
+ this.layers = new ArrayList<>();
}
/**
@@ -55,10 +54,9 @@ public final class HNSWGraph implements Accountable {
* @param level graph level
* @param vectorValues vector values
* @return nearest neighbors
- * @throws IOException
*/
public NearestNeighbors searchLayer(float[] query, Neighbor ep, int ef, int level, VectorValues vectorValues) throws IOException {
- if (level > topLevel) {
+ if (level >= layers.size()) {
throw new IllegalArgumentException("layer does not exist for the level: " + level);
}
@@ -87,7 +85,6 @@ public final class HNSWGraph implements Accountable {
if (c.distance() > f.distance()) {
break;
}
-
for (Neighbor e : layer.getFriends(c.docId())) {
if (visited.contains(e.docId())) {
continue;
@@ -100,12 +97,12 @@ public final class HNSWGraph implements Accountable {
}
float dist = distance(query, e.docId(), vectorValues);
if (dist < f.distance() || results.size() < ef) {
- if (results.size() == ef) {
- results.pop();
- }
Neighbor n = new ImmutableNeighbor(e.docId(), dist);
candidates.insertWithOverflow(n);
- results.add(n);
+ Neighbor popped = results.insertWithOverflow(n);
+ if (popped != null && popped != n) {
+ f = results.top();
+ }
}
}
}
@@ -125,7 +122,8 @@ public final class HNSWGraph implements Accountable {
private NearestNeighbors pickNearestNeighbor(FurthestNeighbors queue) {
NearestNeighbors nearests = new NearestNeighbors(queue.size());
Set<Integer> addedDocs = new HashSet<>();
- while (queue.size() > 0) {
+ int ef = queue.size();
+ while (addedDocs.size() < ef && queue.size() > 0) {
Neighbor c = queue.pop();
if (!addedDocs.contains(c.docId())) {
nearests.add(c);
@@ -136,22 +134,19 @@ public final class HNSWGraph implements Accountable {
}
public void ensureLevel(int level) {
- if (freeze) {
+ if (frozen) {
throw new IllegalStateException("graph is already freezed!");
}
if (level < 0) {
throw new IllegalArgumentException("level must be a positive integer: " + level);
}
- for (int l = level; l >= 0; l--) {
- layers.putIfAbsent(l, new Layer(l));
- }
- if (level > topLevel) {
- topLevel = level;
+ for (int l = layers.size(); l <= level; l++) {
+ layers.add(new Layer(l));
}
}
public int topLevel() {
- return topLevel;
+ return layers.size() - 1;
}
public boolean isEmpty() {
@@ -162,7 +157,7 @@ public final class HNSWGraph implements Accountable {
if (layers.isEmpty()) {
throw new IllegalStateException("the graph has no layers!");
}
- List<Integer> nodesAtMaxLevel = layers.get(topLevel).getNodes();
+ List<Integer> nodesAtMaxLevel = layers.get(layers.size() - 1).getNodes();
if (nodesAtMaxLevel.isEmpty()) {
throw new IllegalStateException("the max level of this graph is empty!");
}
@@ -173,7 +168,7 @@ public final class HNSWGraph implements Accountable {
if (layers.isEmpty()) {
throw new IllegalStateException("the graph has no layers!");
}
- List<Integer> nodesAtMaxLevel = layers.get(topLevel).getNodes();
+ List<Integer> nodesAtMaxLevel = layers.get(layers.size() - 1).getNodes();
if (nodesAtMaxLevel.isEmpty()) {
throw new IllegalStateException("the max level of this graph is empty!");
}
@@ -202,11 +197,11 @@ public final class HNSWGraph implements Accountable {
if (layer == null) {
throw new IllegalArgumentException("layer does not exist for level: " + level);
}
- return layer.getFriends(node).size() > 0;
+ return layer.getFriends(node) != Layer.NO_FRIENDS;
}
void addNode(int level, int node) {
- if (freeze) {
+ if (frozen) {
throw new IllegalStateException("graph is already freezed!");
}
Layer layer = layers.get(level);
@@ -218,7 +213,7 @@ public final class HNSWGraph implements Accountable {
/** Connects two nodes; this is supposed to be called when indexing */
public void connectNodes(int level, int node1, int node2, float dist, int maxConnections) {
- if (freeze) {
+ if (frozen) {
throw new IllegalStateException("graph is already freezed!");
}
assert level >= 0;
@@ -230,15 +225,19 @@ public final class HNSWGraph implements Accountable {
throw new IllegalArgumentException("layer does not exist for level: " + level);
}
layer.connectNodes(node1, node2, dist);
+ layer.connectNodes(node2, node1, dist);
+ /*
// ensure friends size <= maxConnections
+ //
if (maxConnections > 0) {
layer.shrink(node2, maxConnections);
}
+ */
}
/** Connects two nodes; this is supposed to be called when searching */
public void connectNodes(int level, int node1, int node2) {
- if (freeze) {
+ if (frozen) {
throw new IllegalStateException("graph is already freezed!");
}
assert level >= 0;
@@ -253,12 +252,16 @@ public final class HNSWGraph implements Accountable {
}
public void finish() {
- this.freeze = true;
+ while (layers.isEmpty() == false && layers.get(layers.size() - 1).size() == 0) {
+ // remove empty top layers
+ layers.remove(layers.size() - 1);
+ }
+ this.frozen = true;
}
@Override
public long ramBytesUsed() {
- return RamUsageEstimator.sizeOfMap(layers);
+ return RamUsageEstimator.sizeOfCollection(layers);
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
index 70f7b4e..dafbb05 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
@@ -34,6 +34,8 @@ public final class HNSWGraphWriter implements Accountable {
private static final int MAX_LEVEL_LIMIT = 5;
// default random seed for level generation
private static final long DEFAULT_RAND_SEED = System.currentTimeMillis();
+ // expose for testing. TODO: make a better way to initialize this
+ public static long RAND_SEED = DEFAULT_RAND_SEED;
// default max connections per node
public static final int DEFAULT_MAX_CONNECTIONS = 6;
// default max connections per node at layer zero
@@ -56,7 +58,7 @@ public final class HNSWGraphWriter implements Accountable {
/** Construct the builder with default configurations */
public HNSWGraphWriter(int numDimensions, VectorValues.DistanceFunction distFunc) {
- this(DEFAULT_MAX_CONNECTIONS, DEFAULT_MAX_CONNECTIONS_L0, DEFAULT_EF_CONST, DEFAULT_RAND_SEED, numDimensions, distFunc);
+ this(DEFAULT_MAX_CONNECTIONS, DEFAULT_MAX_CONNECTIONS_L0, DEFAULT_EF_CONST, RAND_SEED, numDimensions, distFunc);
}
/** Full constructor */
@@ -76,7 +78,7 @@ public final class HNSWGraphWriter implements Accountable {
/** Inserts a doc with vector value to the graph */
public void insert(int docId, BytesRef binaryValue) throws IOException {
// add the vector value
- float[] value = VectorValues.decode(binaryValue.bytes, numDimensions);
+ float[] value = VectorValues.decode(binaryValue, numDimensions);
rawVectors.add(value);
docsRef.grow(docId + 1);
docsRef.setIntAt(docId, addedDocs++);
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java
index 18befd3..018dd47 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java
@@ -17,9 +17,10 @@
package org.apache.lucene.util.hnsw;
-import java.util.ArrayList;
+import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
@@ -31,6 +32,10 @@ import org.apache.lucene.util.RamUsageEstimator;
/** A knn graph that consists of connected friends (i.e. neighbors) lists. */
final class Layer implements Accountable {
+ // Sentinel that indicates a document has no friends *ie does not participate in the graph*
+ // This is distinct from a document in a single-node graph, which has no friends *yet* (sad).
+ public static final Collection<Neighbor> NO_FRIENDS = Collections.unmodifiableSet(new HashSet<>());
+
private final int level;
private final Map<Integer, TreeSet<Neighbor>> friendsMap;
@@ -89,12 +94,13 @@ final class Layer implements Accountable {
}
}
- List<Neighbor> getFriends(int node) {
- TreeSet<Neighbor> friends = friendsMap.get(node);
- if (friends == null) {
- return Collections.emptyList();
+ Collection<Neighbor> getFriends(int node) {
+ Collection<Neighbor> friends = friendsMap.get(node);
+ if (friends != null) {
+ return friends;
+ } else {
+ return NO_FRIENDS;
}
- return List.copyOf(friends);
}
int size() {
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java b/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java
new file mode 100644
index 0000000..21de79a
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java
@@ -0,0 +1,297 @@
+/*
+ * 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.index;
+
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Set;
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.VectorField;
+import org.apache.lucene.document.StringField;
+import org.apache.lucene.search.KnnGraphQuery;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.util.IntsRef;
+import org.apache.lucene.util.LuceneTestCase;
+import org.junit.Before;
+
+import static org.apache.lucene.util.hnsw.HNSWGraphWriter.RAND_SEED;
+
+/** Tests indexing of a knn-graph */
+public class TestKnnGraph extends LuceneTestCase {
+
+ private static final String KNN_GRAPH_FIELD = "vector";
+
+ @Before
+ public void setup() {
+ RAND_SEED = random().nextLong();
+ }
+
+ /**
+ * Basic test of creating documents in a graph
+ */
+ public void testBasic() throws Exception {
+ try (Directory dir = newDirectory();
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
+ int numDoc = atLeast(10);
+ int dimension = atLeast(3);
+ float[][] values = new float[numDoc][];
+ for (int i = 0; i < numDoc; i++) {
+ if (random().nextBoolean()) {
+ values[i] = new float[dimension];
+ for (int j = 0; j < dimension; j++) {
+ values[i][j] = random().nextFloat();
+ }
+ }
+ add(iw, i, values[i]);
+ }
+ assertConsistentGraph(iw, values);
+ }
+ }
+
+ public void testSingleDocument() throws Exception {
+ try (Directory dir = newDirectory();
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
+ float[][] values = new float[][]{new float[]{0, 1, 2}};
+ add(iw, 0, values[0]);
+ assertConsistentGraph(iw, values);
+ iw.commit();
+ assertConsistentGraph(iw, values);
+ }
+ }
+
+ /**
+ * Verify that the graph properties are preserved when merging
+ */
+ public void testMerge() throws Exception {
+ try (Directory dir = newDirectory();
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
+ int numDoc = atLeast(100);
+ int dimension = atLeast(10);
+ float[][] values = new float[numDoc][];
+ for (int i = 0; i < numDoc; i++) {
+ if (random().nextBoolean()) {
+ values[i] = new float[dimension];
+ for (int j = 0; j < dimension; j++) {
+ values[i][j] = random().nextFloat();
+ }
+ }
+ add(iw, i, values[i]);
+ if (random().nextInt(10) == 3) {
+ //System.out.println("commit @" + i);
+ iw.commit();
+ }
+ }
+ if (random().nextBoolean()) {
+ iw.forceMerge(1);
+ }
+ assertConsistentGraph(iw, values);
+ }
+ }
+
+ // TODO: testSorted
+ // TODO: testDeletions
+
+ /**
+ * Verify that searching does something reasonable
+ */
+ public void testSearch() throws Exception {
+ try (Directory dir = newDirectory();
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
+ // Add a document for every cartesian point in an NxN square so we can
+ // easily know which are the nearest neighbors to every point. Insert by iterating
+ // using a prime number that is not a divisor of N*N so that we will hit each point once,
+ // and chosen so that points will be inserted in a deterministic
+ // but somewhat distributed pattern
+ int n = 5, stepSize = 17;
+ float[][] values = new float[n * n][];
+ int index = 0;
+ for (int i = 0; i < values.length; i++) {
+ // System.out.printf("%d: (%d, %d)\n", i, index % n, index / n);
+ values[i] = new float[]{index % n, index / n};
+ index = (index + stepSize) % (n * n);
+ add(iw, i, values[i]);
+ if (i == 13) {
+ // create 2 segments
+ iw.commit();
+ }
+ }
+ //System.out.println("");
+ if (random().nextBoolean()) {
+ iw.forceMerge(1);
+ }
+ assertConsistentGraph(iw, values);
+ try (DirectoryReader dr = DirectoryReader.open(iw)) {
+ IndexSearcher searcher = new IndexSearcher(dr);
+ // results are ordered by distance (descending) and docid (ascending);
+ // This is the docid ordering:
+ // column major, origin at upper left
+ // 0 15 5 20 10
+ // 3 18 8 23 13
+ // 6 21 11 1 16
+ // 9 24 14 4 19
+ // 12 2 17 7 22
+
+ // For this small graph it seems we can always get exact results with 2 probes
+ assertGraphSearch(new int[]{0, 15, 3, 18, 5}, new float[]{0f, 0.1f}, searcher);
+ assertGraphSearch(new int[]{11, 1, 8, 14, 21}, new float[]{2, 2}, searcher);
+ assertGraphSearch(new int[]{15, 18, 0, 3, 5},new float[]{0.3f, 0.8f}, searcher);
+ }
+ }
+ }
+
+ private void assertGraphSearch(int[] expected, float[] vector, IndexSearcher searcher) throws IOException {
+ TopDocs results = searcher.search(new KnnGraphQuery(KNN_GRAPH_FIELD, vector), 5);
+ assertResults(expected, results);
+ }
+
+ private void assertResults(int[] expected, TopDocs topDocs) {
+ String resultString = Arrays.asList(topDocs.scoreDocs).toString();
+ assertEquals(resultString, expected.length, topDocs.scoreDocs.length);
+ for (int i = 0; i < expected.length; i++) {
+ assertEquals(resultString, expected[i], topDocs.scoreDocs[i].doc);
+ }
+ }
+
+ private void assertConsistentGraph(IndexWriter iw, float[][] values) throws IOException {
+ int totalGraphDocs = 0;
+ try (DirectoryReader dr = DirectoryReader.open(iw)) {
+ for (LeafReaderContext ctx: dr.leaves()) {
+ LeafReader reader = ctx.reader();
+ VectorValues vectorValues = reader.getVectorValues(KNN_GRAPH_FIELD);
+ KnnGraphValues neighbors = reader.getKnnGraphValues(KNN_GRAPH_FIELD);
+ assertTrue((vectorValues == null) == (neighbors == null));
+ if (vectorValues == null) {
+ continue;
+ }
+ int[][] graph = new int[reader.maxDoc()][];
+ boolean singleNodeGraph = false;
+ int graphSize = 0;
+ for (int i = 0; i < reader.maxDoc(); i++) {
+ int id = Integer.parseInt(reader.document(i).get("id"));
+ if (values[id] == null) {
+ // documents without KnnGraphValues have no vectors or neighbors
+ assertNotEquals("document " + id + " was not expected to have values", i, vectorValues.advance(i));
+ assertNotEquals(i, neighbors.advance(i));
+ } else {
+ ++graphSize;
+ // documents with KnnGraphValues have the expected vectors
+ int doc = vectorValues.advance(i);
+ assertEquals("doc " + i + " with id=" + id + " has no vector value", i, doc);
+ float[] scratch = vectorValues.vectorValue();
+ assertArrayEquals("vector did not match for doc " + i + ", id=" + id + ": " + Arrays.toString(scratch),
+ values[id], scratch, 0f);
+ // We collect neighbors for analysis below
+ int nextWithNeighbors = neighbors.advance(i);
+ if (i == nextWithNeighbors) {
+ // FIXME: check every level, not only level 0
+ IntsRef friends = neighbors.getFriends(0);
+ if (friends.length == 0) {
+ //System.out.printf("knngraph @%d is singleton (advance returns %d)\n", i, nextWithNeighbors);
+ singleNodeGraph = true;
+ } else {
+ graph[i] = new int[friends.length];
+ System.arraycopy(friends.ints, friends.offset, graph[i], 0, friends.length);
+ //System.out.printf("knngraph @%d => %s\n", i, Arrays.toString(graph[i]));
+ }
+ } else {
+ // graph must have a single node
+ //System.out.printf("knngraph @%d is singleton (advance returns %d)\n", i, nextWithNeighbors);
+ singleNodeGraph = true;
+ }
+ }
+ }
+ if (singleNodeGraph) {
+ assertEquals("graph is not fully connected", 1, graphSize);
+ } else {
+ assertTrue("Graph has " + graphSize + " nodes, but one of them has no neighbors", graphSize > 1);
+ }
+ // assert that the graph in each leaf is connected and undirected (ie links are reciprocated)
+ assertReciprocal(graph);
+ assertConnected(graph);
+ totalGraphDocs += graphSize;
+ }
+ }
+ int expectedCount = 0;
+ for (float[] friends : values) {
+ if (friends != null) {
+ ++expectedCount;
+ }
+ }
+ assertEquals(expectedCount, totalGraphDocs);
+ }
+
+ private void assertReciprocal(int[][] graph) {
+ // The graph is undirected: if a -> b then b -> a.
+ for (int i = 0; i < graph.length; i++) {
+ if (graph[i] != null) {
+ for (int j = 0; j < graph[i].length; j++) {
+ int k = graph[i][j];
+ assertTrue("" + i + "->" + k + " is not reciprocated", Arrays.binarySearch(graph[k], i) >= 0);
+ }
+ }
+ }
+ }
+
+ private void assertConnected(int[][] graph) {
+ // every node in the graph is reachable from every other node
+ Set<Integer> visited = new HashSet<>();
+ List<Integer> queue = new LinkedList<>();
+ int count = 0;
+ for (int[] entry : graph) {
+ if (entry != null) {
+ if (queue.isEmpty()) {
+ queue.add(entry[0]); // start from any node
+ //System.out.println("start at " + entry[0]);
+ }
+ ++count;
+ }
+ }
+ while(queue.isEmpty() == false) {
+ int i = queue.remove(0);
+ assertNotNull("expected neighbors of " + i, graph[i]);
+ visited.add(i);
+ for (int j : graph[i]) {
+ if (visited.contains(j) == false) {
+ //System.out.println(" ... " + j);
+ queue.add(j);
+ }
+ }
+ }
+ // we visited each node exactly once
+ assertEquals("Attempted to walk entire graph but only visited " + visited.size(), count, visited.size());
+ }
+
+
+ private void add(IndexWriter iw, int id, float[] vector) throws IOException {
+ Document doc = new Document();
+ if (vector != null) {
+ doc.add(new VectorField(KNN_GRAPH_FIELD, vector, VectorValues.DistanceFunction.EUCLIDEAN));
+ }
+ doc.add(new StringField("id", Integer.toString(id), Field.Store.YES));
+ //System.out.println("add " + id + " " + Arrays.toString(vector));
+ iw.addDocument(doc);
+ }
+
+}
[lucene-solr] 02/10: Support segment merging
Posted by to...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tomoko pushed a commit to branch jira/lucene-9004-aknn-2
in repository https://gitbox.apache.org/repos/asf/lucene-solr.git
commit f22b4b5e2201f32c7fe69dc108499da227411855
Author: Tomoko Uchida <to...@apache.org>
AuthorDate: Sun Dec 8 16:39:31 2019 +0900
Support segment merging
---
.../org/apache/lucene/codecs/KnnGraphReader.java | 8 +
.../org/apache/lucene/codecs/KnnGraphWriter.java | 242 +++++++++++++++++++++
.../codecs/lucene90/Lucene90KnnGraphReader.java | 7 +-
.../codecs/lucene90/Lucene90KnnGraphWriter.java | 21 +-
.../apache/lucene/index/DefaultIndexingChain.java | 4 +-
.../org/apache/lucene/index/KnnGraphValues.java | 15 +-
.../apache/lucene/index/KnnGraphValuesWriter.java | 167 ++++----------
.../java/org/apache/lucene/index/MergeState.java | 11 +
.../org/apache/lucene/index/SegmentMerger.java | 19 ++
.../org/apache/lucene/util/hnsw/HNSWGraph.java | 37 ++--
.../apache/lucene/util/hnsw/HNSWGraphBuilder.java | 111 ----------
.../apache/lucene/util/hnsw/HNSWGraphReader.java | 7 +-
.../apache/lucene/util/hnsw/HNSWGraphWriter.java | 229 +++++++++++++++++++
.../java/org/apache/lucene/util/hnsw/Layer.java | 11 +-
14 files changed, 615 insertions(+), 274 deletions(-)
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphReader.java b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphReader.java
index fbf5bfd..58b0904 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphReader.java
@@ -47,4 +47,12 @@ public abstract class KnnGraphReader implements Closeable, Accountable {
/** Returns the {@link KnnGraphValues} for the given {@code field} */
public abstract KnnGraphValues getGraphValues(String field) throws IOException;
+ /**
+ * Returns an instance optimized for merging. This instance may only be
+ * consumed in the thread that called {@link #getMergeInstance()}.
+ * <p>
+ * The default implementation returns {@code this} */
+ public KnnGraphReader getMergeInstance() {
+ return this;
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java
index c0b5638..a8d9834 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/KnnGraphWriter.java
@@ -19,8 +19,18 @@ package org.apache.lucene.codecs;
import java.io.Closeable;
import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import org.apache.lucene.index.DocIDMerger;
import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.KnnGraphValues;
+import org.apache.lucene.index.MergeState;
+import org.apache.lucene.index.VectorValues;
+import org.apache.lucene.util.IntsRef;
+import org.apache.lucene.util.IntsRefBuilder;
/**
* Writes vectors and knn-graph to an index.
@@ -36,4 +46,236 @@ public abstract class KnnGraphWriter implements Closeable {
/** Called once at the end before close */
public abstract void finish() throws IOException;
+ public void merge(MergeState mergeState) throws IOException {
+ for (KnnGraphReader reader : mergeState.knnGraphReaders) {
+ if (reader != null) {
+ reader.checkIntegrity();
+ }
+ }
+
+ for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) {
+ mergeOneField(fieldInfo, mergeState);
+ }
+ finish();
+ }
+
+ private void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
+ List<VectorValuesSub> vectorSubs = new ArrayList<>();
+ List<KnnGraphValuesSub> graphSubs = new ArrayList<>();
+ int topLevel = 0;
+ final Set<EnterPoint> mappedEnterPoints = new HashSet<>();
+
+ long cost = 0;
+ for (int i = 0; i < mergeState.knnGraphReaders.length; i++) {
+ KnnGraphReader knnGraphReader = mergeState.knnGraphReaders[i];
+ if (knnGraphReader != null) {
+ FieldInfo readerFieldInfo = mergeState.fieldInfos[i].fieldInfo(fieldInfo.name);
+ if (readerFieldInfo != null && readerFieldInfo.getVectorNumDimensions() != 0) {
+ VectorValues vectorValues = knnGraphReader.getVectorValues(readerFieldInfo.name);
+ KnnGraphValues graphValues = knnGraphReader.getGraphValues(readerFieldInfo.name);
+ if (vectorValues != null && graphValues != null && graphValues.getTopLevel() >= 0) {
+ cost += vectorValues.cost();
+ vectorSubs.add(new VectorValuesSub(mergeState.docMaps[i], vectorValues));
+ graphSubs.add(new KnnGraphValuesSub(mergeState.docMaps[i], graphValues));
+ if (graphValues.getTopLevel() > topLevel) {
+ topLevel = graphValues.getTopLevel();
+ }
+ int enterPoint = graphValues.getEnterPoints()[0];
+ mappedEnterPoints.add(new EnterPoint(mergeState.docMaps[i].get(enterPoint), graphValues.getTopLevel()));
+ }
+ }
+ }
+ }
+
+ final DocIDMerger<VectorValuesSub> vectorDocIdMerger = DocIDMerger.of(vectorSubs, mergeState.needsIndexSort);
+ final DocIDMerger<KnnGraphValuesSub> graphDocIdMerger = DocIDMerger.of(graphSubs, mergeState.needsIndexSort);
+ final long finalCost = cost;
+ final int finalTopLevel = topLevel;
+
+ final VectorValues meregdVectorValues = mergedVectorValues(vectorDocIdMerger, finalCost);
+ final KnnGraphValues mergedGraphValues = mergedGraphValues(graphDocIdMerger, finalCost, finalTopLevel, mappedEnterPoints);
+
+ writeField(fieldInfo, new KnnGraphReader() {
+ @Override
+ public void checkIntegrity() throws IOException {
+ }
+
+ @Override
+ public VectorValues getVectorValues(String field) throws IOException {
+ return meregdVectorValues;
+ }
+
+ @Override
+ public KnnGraphValues getGraphValues(String field) throws IOException {
+ return mergedGraphValues;
+ }
+
+ @Override
+ public void close() throws IOException {
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return 0L;
+ }
+ });
+ }
+
+ private VectorValues mergedVectorValues(DocIDMerger<VectorValuesSub> docIDMerger, long finalCost) {
+ return new VectorValues() {
+ private VectorValuesSub current;
+ private int docID = -1;
+
+ @Override
+ public float[] vectorValue() throws IOException {
+ return current.values.vectorValue();
+ }
+
+ @Override
+ public boolean seek(int target) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int docID() {
+ return docID;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ current = docIDMerger.next();
+ if (current == null) {
+ docID = NO_MORE_DOCS;
+ } else {
+ docID = current.mappedDocID;
+ }
+ return docID;
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ return 0;
+ }
+
+ @Override
+ public long cost() {
+ return finalCost;
+ }
+ };
+ }
+
+ private KnnGraphValues mergedGraphValues(DocIDMerger<KnnGraphValuesSub> docIDMerger, long finalCost, int finalTopLevel, Set<EnterPoint> enterPoints) throws IOException{
+ return new KnnGraphValues() {
+ private KnnGraphValuesSub current;
+ private int docID = -1;
+
+ @Override
+ public int getTopLevel() {
+ return finalTopLevel;
+ }
+
+ @Override
+ public int[] getEnterPoints() {
+ return enterPoints.stream().mapToInt(ep -> ep.docId).toArray();
+ }
+
+ @Override
+ public int getMaxLevel() {
+ return current.values.getMaxLevel();
+ }
+
+ @Override
+ public IntsRef getFriends(int level) {
+ // re-map friend ids
+ IntsRef friends = current.values.getFriends(level);
+ IntsRefBuilder mappedFriends = new IntsRefBuilder();
+ mappedFriends.copyInts(friends);
+ for (int i = 0; i < friends.length; i++) {
+ mappedFriends.setIntAt(i, current.docMap.get(friends.ints[i]));
+ }
+ // make links to other graphs
+ if (enterPoints.stream().anyMatch(ep -> ep.docId == docID)) {
+ enterPoints.stream().filter(ep -> ep.docId != docID && ep.topLevel >= level).forEach(ep -> {
+ mappedFriends.append(ep.docId);
+ });
+ }
+ return mappedFriends.toIntsRef();
+ }
+
+ @Override
+ public boolean advanceExact(int target) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int docID() {
+ return docID;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ current = docIDMerger.next();
+ if (current == null) {
+ docID = NO_MORE_DOCS;
+ } else {
+ docID = current.mappedDocID;
+ }
+ return docID;
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public long cost() {
+ return finalCost;
+ }
+ };
+ }
+
+ private static class VectorValuesSub extends DocIDMerger.Sub {
+
+ private final VectorValues values;
+
+ VectorValuesSub(MergeState.DocMap docMap, VectorValues values) {
+ super(docMap);
+ this.values = values;
+ assert values.docID() == -1;
+ }
+
+
+ @Override
+ public int nextDoc() throws IOException {
+ return values.nextDoc();
+ }
+ }
+
+ private class KnnGraphValuesSub extends DocIDMerger.Sub {
+
+ private final KnnGraphValues values;
+ private final MergeState.DocMap docMap;
+
+ KnnGraphValuesSub(MergeState.DocMap docMap, KnnGraphValues values) {
+ super(docMap);
+ this.docMap = docMap;
+ this.values = values;
+ assert values.docID() == -1;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ return values.nextDoc();
+ }
+ }
+
+ private class EnterPoint {
+ final int docId;
+ final int topLevel;
+ EnterPoint(int docId, int topLevel) {
+ this.docId = docId;
+ this.topLevel = topLevel;
+ }
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
index e53e260..34a6f3d 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphReader.java
@@ -346,6 +346,11 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
}
@Override
+ public int[] getEnterPoints() {
+ return entry.enterPoints;
+ }
+
+ @Override
public int getMaxLevel() {
return maxLevel;
}
@@ -392,7 +397,7 @@ public final class Lucene90KnnGraphReader extends KnnGraphReader {
}
dataIn.seek(offset);
int maxLevel = dataIn.readInt();
- assert maxLevel > 0;
+ assert maxLevel >= 0;
this.maxLevel = maxLevel;
this.friendsRef = readFriends(maxLevel);
return doc;
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
index 6c16fdf..4c7f4fd 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90KnnGraphWriter.java
@@ -18,10 +18,8 @@
package org.apache.lucene.codecs.lucene90;
import java.io.IOException;
-import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
-import java.util.List;
import java.util.Map;
import org.apache.lucene.codecs.CodecUtil;
@@ -79,7 +77,6 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
@Override
public void writeField(FieldInfo fieldInfo, KnnGraphReader values) throws IOException {
Map<Integer, Long> docToOffset = new HashMap<>();
- List<Integer> enterPoints = new ArrayList<>();
long vectorDataOffset = vectorData.getFilePointer();
long graphDataOffset = graphData.getFilePointer();
@@ -100,13 +97,13 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
// write knn graph value
docToOffset.put(docG, graphData.getFilePointer() - graphDataOffset);
- if (graph.isEnterPoint()) {
- enterPoints.add(docG);
- }
- graphData.writeInt(graph.getMaxLevel());
+
+ int maxLevel = graph.getMaxLevel();
+ assert maxLevel >= 0;
+ graphData.writeInt(maxLevel);
for (int l = graph.getMaxLevel(); l >= 0; l--) {
IntsRef friends = graph.getFriends(l);
- assert friends.length > 0;
+ assert friends.length > 0 : "doc " + docG + " has empty friends list at level=" + l;
graphData.writeInt(friends.length);
int stop = friends.offset + friends.length;
// sort friend ids
@@ -121,24 +118,22 @@ public final class Lucene90KnnGraphWriter extends KnnGraphWriter {
}
}
}
- assert vectors.nextDoc() == DocIdSetIterator.NO_MORE_DOCS; // must be exhausted
- assert graph.nextDoc() == DocIdSetIterator.NO_MORE_DOCS; // must be exhausted
long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset;
long graphDataLength = graphData.getFilePointer() - graphDataOffset;
- writeMeta(fieldInfo, vectorDataOffset, vectorDataLength, graphDataOffset, graphDataLength, graph.getTopLevel(), enterPoints, docToOffset);
+ writeMeta(fieldInfo, vectorDataOffset, vectorDataLength, graphDataOffset, graphDataLength, graph.getTopLevel(), graph.getEnterPoints(), docToOffset);
}
private void writeMeta(FieldInfo field, long vectorDataOffset, long vectorDataLength, long graphDataOffset, long graphDataLength,
- int topLevel, List<Integer> eps, Map<Integer, Long> docToOffset) throws IOException {
+ int topLevel, int[] eps, Map<Integer, Long> docToOffset) throws IOException {
meta.writeInt(field.number);
meta.writeVLong(vectorDataOffset);
meta.writeVLong(vectorDataLength);
meta.writeVLong(graphDataOffset);
meta.writeVLong(graphDataLength);
meta.writeInt(topLevel);
- meta.writeInt(eps.size());
+ meta.writeInt(eps.length);
for (Integer ep : eps) {
meta.writeVInt(ep);
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java b/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java
index 3ce1246..4d1de5d 100644
--- a/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java
+++ b/lucene/core/src/java/org/apache/lucene/index/DefaultIndexingChain.java
@@ -692,7 +692,7 @@ final class DefaultIndexingChain extends DocConsumer {
}
}
- /** Called from processDocument to index one field's doc value */
+ /** Called from processDocument to index one field's vector value */
private void indexVector(PerField fp, IndexableField field) throws IOException {
int numDimensions = field.fieldType().vectorNumDimensions();
VectorValues.DistanceFunction distFunc = field.fieldType().vectorDistFunc();
@@ -705,7 +705,7 @@ final class DefaultIndexingChain extends DocConsumer {
fp.fieldInfo.setVectorDimensionsAndDistanceFunction(numDimensions, distFunc);
if (fp.knnGraphValuesWriter == null) {
- fp.knnGraphValuesWriter = new KnnGraphValuesWriter(docWriter, fp.fieldInfo);
+ fp.knnGraphValuesWriter = new KnnGraphValuesWriter(fp.fieldInfo, docWriter.bytesUsed);
}
fp.knnGraphValuesWriter.addValue(docState.docID, field.binaryValue());
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/KnnGraphValues.java b/lucene/core/src/java/org/apache/lucene/index/KnnGraphValues.java
index f4d4659..55c41ce 100644
--- a/lucene/core/src/java/org/apache/lucene/index/KnnGraphValues.java
+++ b/lucene/core/src/java/org/apache/lucene/index/KnnGraphValues.java
@@ -37,14 +37,10 @@ public abstract class KnnGraphValues extends DocIdSetIterator {
public abstract int getTopLevel();
/**
- * Returns true if the current document ID is the enter point of the graph;
- * this means {@link #getMaxLevel()} is equal to {@link #getTopLevel()}.
- * It is illegal to call this method after {@link #advanceExact(int)}
- * returned {@code false}.
+ * Returns the enter points of the graph.
+ * @return enter points
*/
- public boolean isEnterPoint() {
- return getMaxLevel() == getTopLevel();
- }
+ public abstract int[] getEnterPoints();
/**
* Returns the max level (height) for the current document ID.
@@ -79,6 +75,11 @@ public abstract class KnnGraphValues extends DocIdSetIterator {
}
@Override
+ public int[] getEnterPoints() {
+ return new int[0];
+ }
+
+ @Override
public int getMaxLevel() {
return -1;
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java b/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java
index 0b657dc..0a6f491 100644
--- a/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/index/KnnGraphValuesWriter.java
@@ -18,24 +18,17 @@
package org.apache.lucene.index;
import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
import org.apache.lucene.codecs.KnnGraphReader;
import org.apache.lucene.codecs.KnnGraphWriter;
import org.apache.lucene.search.DocIdSetIterator;
-import org.apache.lucene.store.ByteBuffersDataInput;
-import org.apache.lucene.store.ByteBuffersDataOutput;
-import org.apache.lucene.store.DataInput;
import org.apache.lucene.util.Accountable;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.Counter;
import org.apache.lucene.util.IntsRef;
-import org.apache.lucene.util.IntsRefBuilder;
import org.apache.lucene.util.RamUsageEstimator;
import org.apache.lucene.util.hnsw.HNSWGraph;
-import org.apache.lucene.util.hnsw.HNSWGraphBuilder;
+import org.apache.lucene.util.hnsw.HNSWGraphWriter;
/** Buffers up pending vector value(s) and knn graph per doc, then flushes when segment flushes. */
public class KnnGraphValuesWriter implements Accountable {
@@ -44,31 +37,26 @@ public class KnnGraphValuesWriter implements Accountable {
private final Counter iwBytesUsed;
private final DocsWithFieldSet docsWithFieldVec;
private final DocsWithFieldSet docsWithFieldGrp;
- private final IntsRefBuilder docsRef;
- private final List<float[]> buffer;
- private final ByteBuffersDataOutput bufferOut;
- private final int numDimensions;
- private final VectorValues.DistanceFunction distFunc;
- private final HNSWGraphBuilder hnswBuilder;
-
- private int addedDocs = 0;
+ private final HNSWGraphWriter hnswGraphWriter;
+
private int lastDocID = -1;
private long bytesUsed = 0L;
- public KnnGraphValuesWriter(DocumentsWriterPerThread docWriter, FieldInfo fieldInfo) {
+ public KnnGraphValuesWriter(FieldInfo fieldInfo) {
+ this(fieldInfo, null);
+ }
+
+ public KnnGraphValuesWriter(FieldInfo fieldInfo, Counter iwBytesUsed) {
this.fieldInfo = fieldInfo;
- this.iwBytesUsed = docWriter.bytesUsed;
+ this.iwBytesUsed = iwBytesUsed;
this.docsWithFieldVec = new DocsWithFieldSet();
this.docsWithFieldGrp = new DocsWithFieldSet();
- this.docsRef = new IntsRefBuilder();
- this.buffer = new ArrayList<>();
- this.bufferOut = new ByteBuffersDataOutput();
- this.numDimensions = fieldInfo.getVectorNumDimensions();
- this.distFunc = fieldInfo.getVectorDistFunc();
- this.hnswBuilder = new HNSWGraphBuilder(distFunc);
- this.bytesUsed = docsWithFieldGrp.ramBytesUsed() + docsWithFieldGrp.ramBytesUsed() + RamUsageEstimator.shallowSizeOf(hnswBuilder);
- iwBytesUsed.addAndGet(bytesUsed);
+ this.hnswGraphWriter = new HNSWGraphWriter(fieldInfo.getVectorNumDimensions(), fieldInfo.getVectorDistFunc());
+ this.bytesUsed = docsWithFieldGrp.ramBytesUsed() + docsWithFieldGrp.ramBytesUsed() + RamUsageEstimator.shallowSizeOf(hnswGraphWriter);
+ if (iwBytesUsed != null) {
+ iwBytesUsed.addAndGet(bytesUsed);
+ }
}
public void addValue(int docID, BytesRef binaryValue) throws IOException {
@@ -79,23 +67,9 @@ public class KnnGraphValuesWriter implements Accountable {
throw new IllegalArgumentException("field=\"" + fieldInfo.name + "\": null value not allowed");
}
- // write the vector
- bufferOut.writeBytes(binaryValue.bytes);
+ hnswGraphWriter.insert(docID, binaryValue);
docsWithFieldVec.add(docID);
- docsRef.grow(docID + 1);
- docsRef.setIntAt(docID, addedDocs++);
- docsRef.setLength(docID + 1);
- if (docID > lastDocID + 1) {
- Arrays.fill(docsRef.ints(), lastDocID + 1, docID, -1);
- }
-
- // write the graph
- buffer.add(VectorValues.decode(binaryValue.bytes, numDimensions));
- if (distFunc != VectorValues.DistanceFunction.NONE) {
- float[] query = VectorValues.decode(binaryValue.bytes, numDimensions);
- hnswBuilder.insert(docID, query, getVectorValues());
- docsWithFieldGrp.add(docID);
- }
+ docsWithFieldGrp.add(docID);
updateBytesUsed();
@@ -103,18 +77,17 @@ public class KnnGraphValuesWriter implements Accountable {
}
private void updateBytesUsed() {
- final long newBytesUsed = docsWithFieldVec.ramBytesUsed() + docsWithFieldGrp.ramBytesUsed() +
- bufferOut.ramBytesUsed() +
- RamUsageEstimator.sizeOf(docsRef.ints()) +
- RamUsageEstimator.shallowSizeOf(buffer) +
- RamUsageEstimator.shallowSizeOf(hnswBuilder);
- iwBytesUsed.addAndGet(newBytesUsed - bytesUsed);
+ final long newBytesUsed = docsWithFieldVec.ramBytesUsed() + docsWithFieldGrp.ramBytesUsed() + hnswGraphWriter.ramBytesUsed();
+ if (iwBytesUsed != null) {
+ iwBytesUsed.addAndGet(newBytesUsed - bytesUsed);
+ }
bytesUsed = newBytesUsed;
}
public void flush(SegmentWriteState state, Sorter.DocMap sortMap, KnnGraphWriter graphWriter) throws IOException {
- VectorValues vectors = new BufferedVectorValues(numDimensions, bufferOut.toDataInput(), docsWithFieldVec.iterator());
- KnnGraphValues graph = new BufferedKnnGraphValues(hnswBuilder.build(), docsWithFieldGrp.iterator());
+ hnswGraphWriter.finish();
+ VectorValues vectors = new BufferedVectorValues(docsWithFieldVec.iterator(), hnswGraphWriter.rawVectorsArray());
+ KnnGraphValues graph = new BufferedKnnGraphValues(docsWithFieldGrp.iterator(), hnswGraphWriter.hnswGraph());
final VectorValues vectorValues;
final KnnGraphValues graphValues;
@@ -154,57 +127,6 @@ public class KnnGraphValuesWriter implements Accountable {
});
}
- private VectorValues getVectorValues() {
- return new VectorValues() {
- int docID = -1;
- float[] value = new float[0];
-
- @Override
- public float[] vectorValue() throws IOException {
- return value;
- }
-
- @Override
- public boolean seek(int target) throws IOException {
- if (target < 0) {
- throw new IllegalArgumentException("target must be a positive integer: " + target);
- }
- if (target >= KnnGraphValuesWriter.this.docsRef.length()) {
- docID = NO_MORE_DOCS;
- return false;
- }
- docID = target;
-
- int position = docsRef.ints()[target];
- if (position < 0) {
- throw new IllegalArgumentException("no vector value for doc: " + target);
- }
- value = KnnGraphValuesWriter.this.buffer.get(position);
- return true;
- }
-
- @Override
- public int docID() {
- return docID;
- }
-
- @Override
- public int nextDoc() throws IOException {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public int advance(int target) throws IOException {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public long cost() {
- return KnnGraphValuesWriter.this.docsRef.length();
- }
- };
- }
-
@Override
public long ramBytesUsed() {
return bytesUsed;
@@ -213,33 +135,22 @@ public class KnnGraphValuesWriter implements Accountable {
private static class BufferedVectorValues extends VectorValues {
final DocIdSetIterator docsWithField;
- final DataInput bytesIterator;
- final byte[] value;
- final int numDims;
-
- BufferedVectorValues(int numDims, DataInput bytesIterator, DocIdSetIterator docsWithField) {
- this.numDims = numDims;
- this.value = new byte[Float.BYTES * numDims];
- this.bytesIterator = bytesIterator;
+ final float[][] vectorsArray;
+
+ float[] value = new float[0];
+ int bufferPos = 0;
+
+ BufferedVectorValues(DocIdSetIterator docsWithField, float[][] vectorsArray) {
this.docsWithField = docsWithField;
+ this.vectorsArray = vectorsArray;
}
@Override
public float[] vectorValue() throws IOException {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public byte[] binaryValue() throws IOException {
return value;
}
@Override
- public boolean seek(int target) throws IOException {
- throw new UnsupportedOperationException();
- }
-
- @Override
public int docID() {
return docsWithField.docID();
}
@@ -248,7 +159,7 @@ public class KnnGraphValuesWriter implements Accountable {
public int nextDoc() throws IOException {
int docID = docsWithField.nextDoc();
if (docID != NO_MORE_DOCS) {
- bytesIterator.readBytes(value, 0, value.length);
+ value = vectorsArray[bufferPos++];
}
return docID;
}
@@ -259,6 +170,11 @@ public class KnnGraphValuesWriter implements Accountable {
}
@Override
+ public boolean seek(int target) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
public long cost() {
return docsWithField.cost();
}
@@ -266,14 +182,14 @@ public class KnnGraphValuesWriter implements Accountable {
private static class BufferedKnnGraphValues extends KnnGraphValues {
- final HNSWGraph hnswGraph;
final DocIdSetIterator docsWithField;
+ final HNSWGraph hnswGraph;
private int maxLevel = -1;
- BufferedKnnGraphValues(HNSWGraph hnsw, DocIdSetIterator docsWithField) {
- this.hnswGraph = hnsw;
+ BufferedKnnGraphValues(DocIdSetIterator docsWithField, HNSWGraph hnsw) {
this.docsWithField = docsWithField;
+ this.hnswGraph = hnsw;
}
@Override
@@ -282,6 +198,11 @@ public class KnnGraphValuesWriter implements Accountable {
}
@Override
+ public int[] getEnterPoints() {
+ return hnswGraph.getEnterPoints().stream().mapToInt(Integer::intValue).toArray();
+ }
+
+ @Override
public int getMaxLevel() {
return maxLevel;
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/MergeState.java b/lucene/core/src/java/org/apache/lucene/index/MergeState.java
index 0e47a5a..30763fc 100644
--- a/lucene/core/src/java/org/apache/lucene/index/MergeState.java
+++ b/lucene/core/src/java/org/apache/lucene/index/MergeState.java
@@ -24,6 +24,7 @@ import java.util.Locale;
import org.apache.lucene.codecs.DocValuesProducer;
import org.apache.lucene.codecs.FieldsProducer;
+import org.apache.lucene.codecs.KnnGraphReader;
import org.apache.lucene.codecs.NormsProducer;
import org.apache.lucene.codecs.PointsReader;
import org.apache.lucene.codecs.StoredFieldsReader;
@@ -77,6 +78,9 @@ public class MergeState {
/** Point readers to merge */
public final PointsReader[] pointsReaders;
+ /** Knn graph readers to merge */
+ public final KnnGraphReader[] knnGraphReaders;
+
/** Max docs per reader */
public final int[] maxDocs;
@@ -103,6 +107,7 @@ public class MergeState {
termVectorsReaders = new TermVectorsReader[numReaders];
docValuesProducers = new DocValuesProducer[numReaders];
pointsReaders = new PointsReader[numReaders];
+ knnGraphReaders = new KnnGraphReader[numReaders];
fieldInfos = new FieldInfos[numReaders];
liveDocs = new Bits[numReaders];
@@ -139,6 +144,12 @@ public class MergeState {
if (pointsReaders[i] != null) {
pointsReaders[i] = pointsReaders[i].getMergeInstance();
}
+
+ knnGraphReaders[i] = reader.getKnnGraphReader();
+ if (knnGraphReaders[i] != null) {
+ knnGraphReaders[i] = knnGraphReaders[i].getMergeInstance();
+ }
+
numDocs += reader.numDocs();
}
diff --git a/lucene/core/src/java/org/apache/lucene/index/SegmentMerger.java b/lucene/core/src/java/org/apache/lucene/index/SegmentMerger.java
index c2677c9..974212a 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SegmentMerger.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SegmentMerger.java
@@ -24,6 +24,7 @@ import java.util.List;
import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.DocValuesConsumer;
import org.apache.lucene.codecs.FieldsConsumer;
+import org.apache.lucene.codecs.KnnGraphWriter;
import org.apache.lucene.codecs.NormsConsumer;
import org.apache.lucene.codecs.NormsProducer;
import org.apache.lucene.codecs.PointsWriter;
@@ -166,6 +167,17 @@ final class SegmentMerger {
mergeState.infoStream.message("SM", ((t1-t0)/1000000) + " msec to merge points [" + numMerged + " docs]");
}
+ if (mergeState.infoStream.isEnabled("SM")) {
+ t0 = System.nanoTime();
+ }
+ if (mergeState.mergeFieldInfos.hasVectorValues()) {
+ mergeKnnGraphValues(segmentWriteState);
+ }
+ if (mergeState.infoStream.isEnabled("SM")) {
+ long t1 = System.nanoTime();
+ mergeState.infoStream.message("SM", ((t1-t0)/1000000) + " msec to merge knn graph [" + numMerged + " docs]");
+ }
+
if (mergeState.mergeFieldInfos.hasVectors()) {
if (mergeState.infoStream.isEnabled("SM")) {
t0 = System.nanoTime();
@@ -245,4 +257,11 @@ final class SegmentMerger {
consumer.merge(mergeState, norms);
}
}
+
+ private void mergeKnnGraphValues(SegmentWriteState segmentWriteState) throws IOException {
+ try (KnnGraphWriter writer = codec.knnGraphFormat().fieldsWriter(segmentWriteState)) {
+ writer.merge(mergeState);
+ }
+ }
+
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
index a62b862..d2acbef 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraph.java
@@ -68,13 +68,12 @@ public final class HNSWGraph implements Accountable {
}
Set<Integer> visited = new HashSet<>(ep.docId());
- NearestNeighbors candidates = new NearestNeighbors(layer.size() + 1, ep);
- FurthestNeighbors results = new FurthestNeighbors(layer.size() + 1, ep);
+ NearestNeighbors candidates = new NearestNeighbors(ef, ep);
+ FurthestNeighbors results = new FurthestNeighbors(ef, ep);
if (ep.isDeferred()) {
ep.prepareQuery(query, vectorValues, distFunc);
}
- float upperBound = ep.distance();
while (candidates.size() > 0) {
Neighbor c = candidates.pop();
@@ -100,18 +99,19 @@ public final class HNSWGraph implements Accountable {
f.prepareQuery(query, vectorValues, distFunc);
}
float dist = distance(query, e.docId(), vectorValues);
- if ((dist < upperBound && dist < f.distance()) || results.size() < ef) {
- Neighbor n = new ImmutableNeighbor(e.docId(), dist);
- candidates.add(n);
- results.add(n);
- if (results.size() > ef) {
+ if (dist < f.distance() || results.size() < ef) {
+ if (results.size() == ef) {
results.pop();
}
+ Neighbor n = new ImmutableNeighbor(e.docId(), dist);
+ candidates.insertWithOverflow(n);
+ results.add(n);
}
}
}
- return pickNearestNeighbor(results, ef);
+ //System.out.println("level=" + level + ", visited nodes=" + visited.size());
+ return pickNearestNeighbor(results);
}
private float distance(float[] query, int docId, VectorValues vectorValues) throws IOException {
@@ -122,10 +122,10 @@ public final class HNSWGraph implements Accountable {
return VectorValues.distance(query, other, distFunc);
}
- private NearestNeighbors pickNearestNeighbor(FurthestNeighbors queue, int ef) {
- NearestNeighbors nearests = new NearestNeighbors(ef);
+ private NearestNeighbors pickNearestNeighbor(FurthestNeighbors queue) {
+ NearestNeighbors nearests = new NearestNeighbors(queue.size());
Set<Integer> addedDocs = new HashSet<>();
- for (int i = 0; i < ef && queue.size() > 0; i++) {
+ while (queue.size() > 0) {
Neighbor c = queue.pop();
if (!addedDocs.contains(c.docId())) {
nearests.add(c);
@@ -158,7 +158,7 @@ public final class HNSWGraph implements Accountable {
return layers.isEmpty() || layers.get(0).getNodes().isEmpty();
}
- public int getEnterPoint() {
+ public int getFirstEnterPoint() {
if (layers.isEmpty()) {
throw new IllegalStateException("the graph has no layers!");
}
@@ -169,6 +169,17 @@ public final class HNSWGraph implements Accountable {
return nodesAtMaxLevel.get(0);
}
+ public List<Integer> getEnterPoints() {
+ if (layers.isEmpty()) {
+ throw new IllegalStateException("the graph has no layers!");
+ }
+ List<Integer> nodesAtMaxLevel = layers.get(topLevel).getNodes();
+ if (nodesAtMaxLevel.isEmpty()) {
+ throw new IllegalStateException("the max level of this graph is empty!");
+ }
+ return List.copyOf(nodesAtMaxLevel);
+ }
+
public boolean hasNodes(int level) {
Layer layer = layers.get(level);
if (layer == null) {
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphBuilder.java
deleted file mode 100644
index 85e17ed..0000000
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphBuilder.java
+++ /dev/null
@@ -1,111 +0,0 @@
-/*
- * 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 org.apache.lucene.index.VectorValues;
-
-/** Builder of {@link HNSWGraph} */
-public final class HNSWGraphBuilder {
-
- private static final int MAX_LEVEL_LIMIT = 5;
- // default random seed for level generation
- private static final long DEFAULT_RAND_SEED = System.currentTimeMillis();
- // default max connections per node
- public static final int DEFAULT_MAX_CONNECTIONS = 6;
- // default max connections per node at layer zero
- public static final int DEFAULT_MAX_CONNECTIONS_L0 = DEFAULT_MAX_CONNECTIONS * 2;
- // default candidate list size
- public static final int DEFAULT_EF_CONST = 50;
-
- private final int maxConn;
- private final int maxConn0;
- private final int efConst;
- private final VectorValues.DistanceFunction distFunc;
- private final HNSWGraph hnsw;
- private final RandomLevelGenerator levelGenerator;
-
- /** Construct the builder with default configurations */
- public HNSWGraphBuilder(VectorValues.DistanceFunction distFunc) {
- this(DEFAULT_MAX_CONNECTIONS, DEFAULT_MAX_CONNECTIONS_L0, DEFAULT_EF_CONST, DEFAULT_RAND_SEED, distFunc);
- }
-
- /** Full constructor */
- public HNSWGraphBuilder(int maxConn, int maxConn0, int efConst, long seed, VectorValues.DistanceFunction distFunc) {
- this.maxConn = maxConn;
- this.maxConn0 = maxConn0;
- this.efConst = efConst;
- this.distFunc = distFunc;
- this.hnsw = new HNSWGraph(distFunc);
- float ml = (float) (1.0 / Math.log(maxConn));
- this.levelGenerator = new RandomLevelGenerator(seed, ml, MAX_LEVEL_LIMIT);
- }
-
- /** Inserts a doc with vector value to the graph */
- public void insert(int docId, float[] value, VectorValues vectorValues) throws IOException {
- if (hnsw.isEmpty()) {
- // initialize the graph
- hnsw.ensureLevel(0);
- hnsw.addNode(0, docId);
- return;
- }
-
- int enterPoint = hnsw.getEnterPoint();
- if (!vectorValues.seek(enterPoint)) {
- throw new IllegalStateException("enter point (=" + enterPoint + ") has no vector value");
- }
- float distFromEp = VectorValues.distance(value, vectorValues.vectorValue(), distFunc);
- Neighbor ep = new ImmutableNeighbor(enterPoint, distFromEp);
-
- int level = levelGenerator.randomLevel();
-
- // down to the level from the hnsw's top level
- for (int l = hnsw.topLevel(); l > level; l--) {
- NearestNeighbors neighbors = hnsw.searchLayer(value, ep, 1, l, vectorValues);
- ep = neighbors.top();
- }
-
- // down to level 0 with placing the doc to each layer
- hnsw.ensureLevel(level);
- for (int l = level; l >= 0; l--) {
- if (!hnsw.hasNodes(l)) {
- hnsw.addNode(l, docId);
- continue;
- }
-
- NearestNeighbors neighbors = hnsw.searchLayer(value, ep, efConst, l, vectorValues);
- ep = neighbors.top();
- int maxConnections = l == 0 ? maxConn0 : maxConn;
- for (int i = 0; i < maxConnections; i++) {
- if (neighbors.size() == 0) {
- break;
- }
- Neighbor n = neighbors.pop();
- hnsw.connectNodes(l, docId, n.docId(), n.distance(), maxConnections);
- }
- }
- }
-
- /** Returns the built HNSW graph*/
- public HNSWGraph build() {
- hnsw.finish();
- return hnsw;
- }
-
-}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
index 4e6385a..8a9d0dd 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
@@ -47,7 +47,7 @@ public final class HNSWGraphReader {
public Neighbors searchNeighbors(float[] query, int ef, VectorValues vectorValues) throws IOException {
HNSWGraph hnsw = get(field, context, false);
- int enterPoint = hnsw.getEnterPoint();
+ int enterPoint = hnsw.getFirstEnterPoint();
if (!vectorValues.seek(enterPoint)) {
throw new IllegalStateException("enterPoint=" + enterPoint + " has no vector value");
}
@@ -101,6 +101,10 @@ public final class HNSWGraphReader {
VectorValues.DistanceFunction distFunc = fi.getVectorDistFunc();
KnnGraphValues graphValues = context.reader().getKnnGraphValues(field);
+ return load(field, distFunc, graphValues);
+ }
+
+ public static HNSWGraph load(String field, VectorValues.DistanceFunction distFunc, KnnGraphValues graphValues) throws IOException {
HNSWGraph hnsw = new HNSWGraph(distFunc);
for (int doc = graphValues.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = graphValues.nextDoc()) {
int maxLevel = graphValues.getMaxLevel();
@@ -112,7 +116,6 @@ public final class HNSWGraphReader {
}
}
hnsw.finish();
-
return hnsw;
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
new file mode 100644
index 0000000..70f7b4e
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphWriter.java
@@ -0,0 +1,229 @@
+/*
+ * 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.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import org.apache.lucene.index.VectorValues;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.IntsRef;
+import org.apache.lucene.util.IntsRefBuilder;
+import org.apache.lucene.util.RamUsageEstimator;
+
+public final class HNSWGraphWriter implements Accountable {
+
+ private static final int MAX_LEVEL_LIMIT = 5;
+ // default random seed for level generation
+ private static final long DEFAULT_RAND_SEED = System.currentTimeMillis();
+ // default max connections per node
+ public static final int DEFAULT_MAX_CONNECTIONS = 6;
+ // default max connections per node at layer zero
+ public static final int DEFAULT_MAX_CONNECTIONS_L0 = DEFAULT_MAX_CONNECTIONS * 2;
+ // default candidate list size
+ public static final int DEFAULT_EF_CONST = 50;
+
+ private final int maxConn;
+ private final int maxConn0;
+ private final int efConst;
+ private final IntsRefBuilder docsRef;
+ private final List<float[]> rawVectors;
+ private final int numDimensions;
+ private final VectorValues.DistanceFunction distFunc;
+ private final HNSWGraph hnsw;
+ private final RandomLevelGenerator levelGenerator;
+
+ private int addedDocs = 0;
+ private int lastDocID = -1;
+
+ /** Construct the builder with default configurations */
+ public HNSWGraphWriter(int numDimensions, VectorValues.DistanceFunction distFunc) {
+ this(DEFAULT_MAX_CONNECTIONS, DEFAULT_MAX_CONNECTIONS_L0, DEFAULT_EF_CONST, DEFAULT_RAND_SEED, numDimensions, distFunc);
+ }
+
+ /** Full constructor */
+ public HNSWGraphWriter(int maxConn, int maxConn0, int efConst, long seed, int numDimensions, VectorValues.DistanceFunction distFunc) {
+ this.maxConn = maxConn;
+ this.maxConn0 = maxConn0;
+ this.efConst = efConst;
+ this.numDimensions = numDimensions;
+ this.distFunc = distFunc;
+ this.docsRef = new IntsRefBuilder();
+ this.rawVectors = new ArrayList<>();
+ this.hnsw = new HNSWGraph(distFunc);
+ float ml = (float) (1.0 / Math.log(maxConn));
+ this.levelGenerator = new RandomLevelGenerator(seed, ml, MAX_LEVEL_LIMIT);
+ }
+
+ /** Inserts a doc with vector value to the graph */
+ public void insert(int docId, BytesRef binaryValue) throws IOException {
+ // add the vector value
+ float[] value = VectorValues.decode(binaryValue.bytes, numDimensions);
+ rawVectors.add(value);
+ docsRef.grow(docId + 1);
+ docsRef.setIntAt(docId, addedDocs++);
+ docsRef.setLength(docId + 1);
+ if (docId > lastDocID + 1) {
+ Arrays.fill(docsRef.ints(), lastDocID + 1, docId, -1);
+ }
+ lastDocID = docId;
+
+ if (distFunc == VectorValues.DistanceFunction.NONE) {
+ return;
+ }
+
+ // add the graph node
+ if (hnsw.isEmpty()) {
+ // initialize the graph
+ hnsw.ensureLevel(0);
+ hnsw.addNode(0, docId);
+ return;
+ }
+
+ VectorValues vectorValues = getVectorValues();
+
+ int enterPoint = hnsw.getFirstEnterPoint();
+ if (!vectorValues.seek(enterPoint)) {
+ throw new IllegalStateException("enter point (=" + enterPoint + ") has no vector value");
+ }
+ float distFromEp = VectorValues.distance(value, vectorValues.vectorValue(), distFunc);
+ Neighbor ep = new ImmutableNeighbor(enterPoint, distFromEp);
+
+ int level = levelGenerator.randomLevel();
+
+ // down to the level from the hnsw's top level
+ for (int l = hnsw.topLevel(); l > level; l--) {
+ NearestNeighbors neighbors = hnsw.searchLayer(value, ep, 1, l, vectorValues);
+ ep = neighbors.top();
+ }
+
+ // down to level 0 with placing the doc to each layer
+ hnsw.ensureLevel(level);
+ for (int l = level; l >= 0; l--) {
+ if (!hnsw.hasNodes(l)) {
+ hnsw.addNode(l, docId);
+ continue;
+ }
+
+ NearestNeighbors neighbors = hnsw.searchLayer(value, ep, efConst, l, vectorValues);
+ ep = neighbors.top();
+ int maxConnections = l == 0 ? maxConn0 : maxConn;
+ for (int i = 0; i < maxConnections; i++) {
+ if (neighbors.size() == 0) {
+ break;
+ }
+ Neighbor n = neighbors.pop();
+ hnsw.connectNodes(l, docId, n.docId(), n.distance(), maxConnections);
+ }
+ }
+ }
+
+ public void finish() {
+ hnsw.finish();
+ }
+
+ public float[][] rawVectorsArray() {
+ return rawVectors.toArray(new float[0][]);
+ }
+
+ /** Returns the built HNSW graph*/
+ public HNSWGraph hnswGraph() {
+ return hnsw;
+ }
+
+ private VectorValues getVectorValues() {
+ return new RandomAccessVectorValues(docsRef.toIntsRef(), rawVectors);
+ }
+
+ private static class RandomAccessVectorValues extends VectorValues {
+
+ final IntsRef docsRef;
+ final List<float[]> buffer;
+
+ int docID = -1;
+ float[] value = new float[0];
+
+ RandomAccessVectorValues(IntsRef docsRef, List<float[]> buffer) {
+ this.docsRef = docsRef;
+ this.buffer = buffer;
+ }
+
+ @Override
+ public float[] vectorValue() throws IOException {
+ return value;
+ }
+
+ @Override
+ public boolean seek(int target) throws IOException {
+ if (target < 0) {
+ throw new IllegalArgumentException("target must be a positive integer: " + target);
+ }
+ if (target >= docsRef.length) {
+ docID = NO_MORE_DOCS;
+ return false;
+ }
+ docID = target;
+
+ int position = docsRef.ints[target];
+ if (position < 0) {
+ throw new IllegalArgumentException("no vector value for doc: " + target);
+ }
+ value = buffer.get(position);
+ return true;
+ }
+
+ @Override
+ public int docID() {
+ return docID;
+ }
+
+ @Override
+ public int nextDoc() throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public long cost() {
+ return docsRef.length;
+ }
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ // calculating the exact ram usage is time consuming so we make rough estimation here
+ long bytesUsed = RamUsageEstimator.sizeOf(docsRef.ints()) +
+ Float.BYTES * numDimensions * rawVectors.size() +
+ RamUsageEstimator.shallowSizeOfInstance(HNSWGraph.class) +
+ RamUsageEstimator.shallowSizeOfInstance(Layer.class) +
+ Integer.BYTES * 2 * DEFAULT_MAX_CONNECTIONS_L0 * addedDocs;
+ if (hnsw.topLevel() > 0) {
+ for (int l = hnsw.topLevel(); l > 0; l--) {
+ bytesUsed += Integer.BYTES * 2 * DEFAULT_MAX_CONNECTIONS * addedDocs;
+ }
+ }
+ return bytesUsed;
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java
index aa4b145..18befd3 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/Layer.java
@@ -17,6 +17,7 @@
package org.apache.lucene.util.hnsw;
+import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
@@ -75,8 +76,14 @@ final class Layer implements Accountable {
}
if (friends.size() > maxFriends) {
TreeSet<Neighbor> newFriends = new TreeSet<>();
- for (Iterator<Neighbor> itr = friends.iterator(); itr.hasNext() && newFriends.size() < maxFriends; ) {
- newFriends.add(itr.next());
+ Iterator<Neighbor> itr = friends.iterator();
+ while(itr.hasNext()) {
+ Neighbor n = itr.next();
+ if (newFriends.size() < maxFriends || friendsMap.get(n.docId()).size() == 1) {
+ newFriends.add(n);
+ } else {
+ friendsMap.get(n.docId()).removeIf(n2 -> n2.docId() == node);
+ }
}
friendsMap.put(node, newFriends);
}