| /* |
| * Licensed to the Apache Software Foundation (ASF) under one or more |
| * contributor license agreements. See the NOTICE file distributed with |
| * this work for additional information regarding copyright ownership. |
| * The ASF licenses this file to You under the Apache License, Version 2.0 |
| * (the "License"); you may not use this file except in compliance with |
| * the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| package org.apache.lucene.util.hnsw; |
| |
| import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; |
| |
| import java.io.Closeable; |
| import java.io.IOException; |
| import java.io.OutputStream; |
| import java.lang.management.ManagementFactory; |
| import java.lang.management.ThreadMXBean; |
| import java.nio.ByteBuffer; |
| import java.nio.ByteOrder; |
| import java.nio.FloatBuffer; |
| import java.nio.IntBuffer; |
| import java.nio.channels.FileChannel; |
| import java.nio.file.Files; |
| import java.nio.file.Path; |
| import java.nio.file.Paths; |
| import java.nio.file.attribute.FileTime; |
| import java.util.Arrays; |
| import java.util.HashSet; |
| import java.util.Locale; |
| import java.util.Objects; |
| import java.util.Set; |
| import org.apache.lucene.codecs.KnnVectorsFormat; |
| import org.apache.lucene.codecs.KnnVectorsReader; |
| import org.apache.lucene.codecs.lucene94.Lucene94Codec; |
| import org.apache.lucene.codecs.lucene94.Lucene94HnswVectorsFormat; |
| import org.apache.lucene.codecs.lucene94.Lucene94HnswVectorsReader; |
| import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; |
| import org.apache.lucene.document.Document; |
| import org.apache.lucene.document.FieldType; |
| import org.apache.lucene.document.KnnVectorField; |
| import org.apache.lucene.document.StoredField; |
| import org.apache.lucene.index.CodecReader; |
| import org.apache.lucene.index.DirectoryReader; |
| 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.RandomAccessVectorValues; |
| import org.apache.lucene.index.RandomAccessVectorValuesProducer; |
| import org.apache.lucene.index.VectorEncoding; |
| import org.apache.lucene.index.VectorSimilarityFunction; |
| import org.apache.lucene.search.ConstantScoreScorer; |
| import org.apache.lucene.search.ConstantScoreWeight; |
| import org.apache.lucene.search.IndexSearcher; |
| import org.apache.lucene.search.KnnVectorQuery; |
| import org.apache.lucene.search.Query; |
| import org.apache.lucene.search.QueryVisitor; |
| import org.apache.lucene.search.ScoreDoc; |
| import org.apache.lucene.search.ScoreMode; |
| import org.apache.lucene.search.Scorer; |
| import org.apache.lucene.search.TopDocs; |
| import org.apache.lucene.search.Weight; |
| import org.apache.lucene.store.Directory; |
| import org.apache.lucene.store.FSDirectory; |
| import org.apache.lucene.util.BitSetIterator; |
| import org.apache.lucene.util.BytesRef; |
| import org.apache.lucene.util.FixedBitSet; |
| import org.apache.lucene.util.IntroSorter; |
| import org.apache.lucene.util.PrintStreamInfoStream; |
| import org.apache.lucene.util.SuppressForbidden; |
| |
| /** |
| * For testing indexing and search performance of a knn-graph |
| * |
| * <p>java -cp .../lib/*.jar org.apache.lucene.util.hnsw.KnnGraphTester -ndoc 1000000 -search |
| * .../vectors.bin |
| */ |
| public class KnnGraphTester { |
| |
| private static final String KNN_FIELD = "knn"; |
| private static final String ID_FIELD = "id"; |
| |
| private int numDocs; |
| private int dim; |
| private int topK; |
| private int numIters; |
| private int fanout; |
| private Path indexPath; |
| private boolean quiet; |
| private boolean reindex; |
| private boolean forceMerge; |
| private int reindexTimeMsec; |
| private int beamWidth; |
| private int maxConn; |
| private VectorSimilarityFunction similarityFunction; |
| private VectorEncoding vectorEncoding; |
| private FixedBitSet matchDocs; |
| private float selectivity; |
| private boolean prefilter; |
| |
| private KnnGraphTester() { |
| // set defaults |
| numDocs = 1000; |
| numIters = 1000; |
| dim = 256; |
| topK = 100; |
| fanout = topK; |
| similarityFunction = VectorSimilarityFunction.DOT_PRODUCT; |
| vectorEncoding = VectorEncoding.FLOAT32; |
| selectivity = 1f; |
| prefilter = false; |
| } |
| |
| public static void main(String... args) throws Exception { |
| new KnnGraphTester().run(args); |
| } |
| |
| private void run(String... args) throws Exception { |
| String operation = null; |
| Path docVectorsPath = null, queryPath = null, outputPath = null; |
| for (int iarg = 0; iarg < args.length; iarg++) { |
| String arg = args[iarg]; |
| switch (arg) { |
| case "-search": |
| case "-check": |
| case "-stats": |
| case "-dump": |
| if (operation != null) { |
| throw new IllegalArgumentException( |
| "Specify only one operation, not both " + arg + " and " + operation); |
| } |
| operation = arg; |
| if (operation.equals("-search")) { |
| if (iarg == args.length - 1) { |
| throw new IllegalArgumentException( |
| "Operation " + arg + " requires a following pathname"); |
| } |
| queryPath = Paths.get(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 "-beamWidthIndex": |
| if (iarg == args.length - 1) { |
| throw new IllegalArgumentException("-beamWidthIndex requires a following number"); |
| } |
| beamWidth = Integer.parseInt(args[++iarg]); |
| break; |
| case "-maxConn": |
| if (iarg == args.length - 1) { |
| throw new IllegalArgumentException("-maxConn requires a following number"); |
| } |
| maxConn = Integer.parseInt(args[++iarg]); |
| break; |
| case "-dim": |
| if (iarg == args.length - 1) { |
| throw new IllegalArgumentException("-dim requires a following number"); |
| } |
| dim = 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; |
| case "-niter": |
| if (iarg == args.length - 1) { |
| throw new IllegalArgumentException("-niter requires a following number"); |
| } |
| numIters = Integer.parseInt(args[++iarg]); |
| break; |
| case "-reindex": |
| reindex = true; |
| break; |
| case "-topK": |
| if (iarg == args.length - 1) { |
| throw new IllegalArgumentException("-topK requires a following number"); |
| } |
| topK = Integer.parseInt(args[++iarg]); |
| break; |
| case "-out": |
| outputPath = Paths.get(args[++iarg]); |
| break; |
| case "-docs": |
| docVectorsPath = Paths.get(args[++iarg]); |
| break; |
| case "-encoding": |
| String encoding = args[++iarg]; |
| switch (encoding) { |
| case "byte": |
| vectorEncoding = VectorEncoding.BYTE; |
| break; |
| case "float32": |
| vectorEncoding = VectorEncoding.FLOAT32; |
| break; |
| default: |
| throw new IllegalArgumentException("-encoding can be 'byte' or 'float32' only"); |
| } |
| break; |
| case "-metric": |
| String metric = args[++iarg]; |
| switch (metric) { |
| case "euclidean": |
| similarityFunction = VectorSimilarityFunction.EUCLIDEAN; |
| break; |
| case "angular": |
| similarityFunction = VectorSimilarityFunction.DOT_PRODUCT; |
| break; |
| default: |
| throw new IllegalArgumentException("-metric can be 'angular' or 'euclidean' only"); |
| } |
| break; |
| case "-forceMerge": |
| forceMerge = true; |
| break; |
| case "-prefilter": |
| prefilter = true; |
| break; |
| case "-filterSelectivity": |
| if (iarg == args.length - 1) { |
| throw new IllegalArgumentException("-filterSelectivity requires a following float"); |
| } |
| selectivity = Float.parseFloat(args[++iarg]); |
| if (selectivity <= 0 || selectivity >= 1) { |
| throw new IllegalArgumentException("-filterSelectivity must be between 0 and 1"); |
| } |
| break; |
| case "-quiet": |
| quiet = true; |
| break; |
| default: |
| throw new IllegalArgumentException("unknown argument " + arg); |
| // usage(); |
| } |
| } |
| if (operation == null && reindex == false) { |
| usage(); |
| } |
| if (prefilter && selectivity == 1f) { |
| throw new IllegalArgumentException("-prefilter requires filterSelectivity between 0 and 1"); |
| } |
| indexPath = Paths.get(formatIndexPath(docVectorsPath)); |
| if (reindex) { |
| if (docVectorsPath == null) { |
| throw new IllegalArgumentException("-docs argument is required when indexing"); |
| } |
| reindexTimeMsec = createIndex(docVectorsPath, indexPath); |
| if (forceMerge) { |
| forceMerge(); |
| } |
| } |
| if (operation != null) { |
| switch (operation) { |
| case "-search": |
| if (docVectorsPath == null) { |
| throw new IllegalArgumentException("missing -docs arg"); |
| } |
| if (selectivity < 1) { |
| matchDocs = generateRandomBitSet(numDocs, selectivity); |
| } |
| if (outputPath != null) { |
| testSearch(indexPath, queryPath, outputPath, null); |
| } else { |
| testSearch(indexPath, queryPath, null, getNN(docVectorsPath, queryPath)); |
| } |
| break; |
| case "-dump": |
| dumpGraph(docVectorsPath); |
| break; |
| case "-stats": |
| printFanoutHist(indexPath); |
| break; |
| } |
| } |
| } |
| |
| private String formatIndexPath(Path docsPath) { |
| return docsPath.getFileName() + "-" + maxConn + "-" + beamWidth + ".index"; |
| } |
| |
| @SuppressForbidden(reason = "Prints stuff") |
| private void printFanoutHist(Path indexPath) throws IOException { |
| try (Directory dir = FSDirectory.open(indexPath); |
| DirectoryReader reader = DirectoryReader.open(dir)) { |
| for (LeafReaderContext context : reader.leaves()) { |
| LeafReader leafReader = context.reader(); |
| KnnVectorsReader vectorsReader = |
| ((PerFieldKnnVectorsFormat.FieldsReader) ((CodecReader) leafReader).getVectorReader()) |
| .getFieldReader(KNN_FIELD); |
| HnswGraph knnValues = ((Lucene94HnswVectorsReader) vectorsReader).getGraph(KNN_FIELD); |
| System.out.printf("Leaf %d has %d documents\n", context.ord, leafReader.maxDoc()); |
| printGraphFanout(knnValues, leafReader.maxDoc()); |
| } |
| } |
| } |
| |
| @SuppressWarnings("unchecked") |
| private void dumpGraph(Path docsPath) throws IOException { |
| try (BinaryFileVectors vectors = new BinaryFileVectors(docsPath)) { |
| RandomAccessVectorValues values = vectors.randomAccess(); |
| HnswGraphBuilder<float[]> builder = |
| (HnswGraphBuilder<float[]>) |
| HnswGraphBuilder.create( |
| vectors, vectorEncoding, similarityFunction, maxConn, beamWidth, 0); |
| // start at node 1 |
| for (int i = 1; i < numDocs; i++) { |
| builder.addGraphNode(i, values); |
| System.out.println("\nITERATION " + i); |
| dumpGraph(builder.hnsw); |
| } |
| } |
| } |
| |
| private void dumpGraph(OnHeapHnswGraph hnsw) { |
| for (int i = 0; i < hnsw.size(); i++) { |
| NeighborArray neighbors = hnsw.getNeighbors(0, i); |
| System.out.printf(Locale.ROOT, "%5d", i); |
| NeighborArray sorted = new NeighborArray(neighbors.size(), true); |
| for (int j = 0; j < neighbors.size(); j++) { |
| int node = neighbors.node[j]; |
| float score = neighbors.score[j]; |
| sorted.add(node, score); |
| } |
| new NeighborArraySorter(sorted).sort(0, sorted.size()); |
| for (int j = 0; j < sorted.size(); j++) { |
| System.out.printf(Locale.ROOT, " [%d, %.4f]", sorted.node[j], sorted.score[j]); |
| } |
| System.out.println(); |
| } |
| } |
| |
| @SuppressForbidden(reason = "Prints stuff") |
| private void forceMerge() throws IOException { |
| IndexWriterConfig iwc = new IndexWriterConfig().setOpenMode(IndexWriterConfig.OpenMode.APPEND); |
| iwc.setInfoStream(new PrintStreamInfoStream(System.out)); |
| System.out.println("Force merge index in " + indexPath); |
| try (IndexWriter iw = new IndexWriter(FSDirectory.open(indexPath), iwc)) { |
| iw.forceMerge(1); |
| } |
| } |
| |
| @SuppressForbidden(reason = "Prints stuff") |
| private void printGraphFanout(HnswGraph knnValues, int numDocs) throws IOException { |
| int min = Integer.MAX_VALUE, max = 0, total = 0; |
| int count = 0; |
| int[] leafHist = new int[numDocs]; |
| for (int node = 0; node < numDocs; node++) { |
| knnValues.seek(0, node); |
| int n = 0; |
| while (knnValues.nextNeighbor() != NO_MORE_DOCS) { |
| ++n; |
| } |
| ++leafHist[n]; |
| max = Math.max(max, n); |
| min = Math.min(min, n); |
| if (n > 0) { |
| ++count; |
| total += n; |
| } |
| } |
| System.out.printf( |
| "Graph size=%d, Fanout min=%d, mean=%.2f, max=%d\n", |
| count, min, total / (float) count, max); |
| printHist(leafHist, max, count, 10); |
| } |
| |
| @SuppressForbidden(reason = "Prints stuff") |
| 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]; |
| while (total >= count * ibucket / nbuckets) { |
| System.out.printf("%4d", i); |
| ++ibucket; |
| } |
| } |
| System.out.println(); |
| } |
| |
| @SuppressForbidden(reason = "Prints stuff") |
| private void testSearch(Path indexPath, Path queryPath, Path outputPath, int[][] nn) |
| throws IOException { |
| TopDocs[] results = new TopDocs[numIters]; |
| long elapsed, totalCpuTime, totalVisited = 0; |
| try (FileChannel input = FileChannel.open(queryPath)) { |
| VectorReader targetReader = VectorReader.create(input, dim, vectorEncoding, numIters); |
| if (quiet == false) { |
| System.out.println("running " + numIters + " targets; topK=" + topK + ", fanout=" + fanout); |
| } |
| long start; |
| ThreadMXBean bean = ManagementFactory.getThreadMXBean(); |
| long cpuTimeStartNs; |
| try (Directory dir = FSDirectory.open(indexPath); |
| DirectoryReader reader = DirectoryReader.open(dir)) { |
| IndexSearcher searcher = new IndexSearcher(reader); |
| numDocs = reader.maxDoc(); |
| Query bitSetQuery = prefilter ? new BitSetQuery(matchDocs) : null; |
| for (int i = 0; i < numIters; i++) { |
| // warm up |
| float[] target = targetReader.next(); |
| if (prefilter) { |
| doKnnVectorQuery(searcher, KNN_FIELD, target, topK, fanout, bitSetQuery); |
| } else { |
| doKnnVectorQuery(searcher, KNN_FIELD, target, (int) (topK / selectivity), fanout, null); |
| } |
| } |
| targetReader.reset(); |
| start = System.nanoTime(); |
| cpuTimeStartNs = bean.getCurrentThreadCpuTime(); |
| for (int i = 0; i < numIters; i++) { |
| float[] target = targetReader.next(); |
| if (prefilter) { |
| results[i] = doKnnVectorQuery(searcher, KNN_FIELD, target, topK, fanout, bitSetQuery); |
| } else { |
| results[i] = |
| doKnnVectorQuery( |
| searcher, KNN_FIELD, target, (int) (topK / selectivity), fanout, null); |
| |
| if (matchDocs != null) { |
| results[i].scoreDocs = |
| Arrays.stream(results[i].scoreDocs) |
| .filter(scoreDoc -> matchDocs.get(scoreDoc.doc)) |
| .toArray(ScoreDoc[]::new); |
| } |
| } |
| } |
| totalCpuTime = (bean.getCurrentThreadCpuTime() - cpuTimeStartNs) / 1_000_000; |
| elapsed = (System.nanoTime() - start) / 1_000_000; // ns -> ms |
| for (int i = 0; i < numIters; i++) { |
| totalVisited += results[i].totalHits.value; |
| for (ScoreDoc doc : results[i].scoreDocs) { |
| if (doc.doc != NO_MORE_DOCS) { |
| // there is a bug somewhere that can result in doc=NO_MORE_DOCS! I think it happens |
| // in some degenerate case (like input query has NaN in it?) that causes no results to |
| // be returned from HNSW search? |
| doc.doc = Integer.parseInt(reader.document(doc.doc).get("id")); |
| } else { |
| System.out.println("NO_MORE_DOCS!"); |
| } |
| } |
| } |
| } |
| if (quiet == false) { |
| System.out.println( |
| "completed " |
| + numIters |
| + " searches in " |
| + elapsed |
| + " ms: " |
| + ((1000 * numIters) / elapsed) |
| + " QPS " |
| + "CPU time=" |
| + totalCpuTime |
| + "ms"); |
| } |
| } |
| if (outputPath != null) { |
| ByteBuffer buf = ByteBuffer.allocate(4); |
| IntBuffer ibuf = buf.order(ByteOrder.LITTLE_ENDIAN).asIntBuffer(); |
| try (OutputStream out = Files.newOutputStream(outputPath)) { |
| for (int i = 0; i < numIters; i++) { |
| for (ScoreDoc doc : results[i].scoreDocs) { |
| ibuf.position(0); |
| ibuf.put(doc.doc); |
| out.write(buf.array()); |
| } |
| } |
| } |
| } else { |
| if (quiet == false) { |
| System.out.println("checking results"); |
| } |
| float recall = checkResults(results, nn); |
| totalVisited /= numIters; |
| System.out.printf( |
| Locale.ROOT, |
| "%5.3f\t%5.2f\t%d\t%d\t%d\t%d\t%d\t%d\t%.2f\t%s\n", |
| recall, |
| totalCpuTime / (float) numIters, |
| numDocs, |
| fanout, |
| maxConn, |
| beamWidth, |
| totalVisited, |
| reindexTimeMsec, |
| selectivity, |
| prefilter ? "pre-filter" : "post-filter"); |
| } |
| } |
| |
| private abstract static class VectorReader { |
| final float[] target; |
| final ByteBuffer bytes; |
| |
| static VectorReader create(FileChannel input, int dim, VectorEncoding vectorEncoding, int n) |
| throws IOException { |
| int bufferSize = n * dim * vectorEncoding.byteSize; |
| return switch (vectorEncoding) { |
| case BYTE -> new VectorReaderByte(input, dim, bufferSize); |
| case FLOAT32 -> new VectorReaderFloat32(input, dim, bufferSize); |
| }; |
| } |
| |
| VectorReader(FileChannel input, int dim, int bufferSize) throws IOException { |
| bytes = |
| input.map(FileChannel.MapMode.READ_ONLY, 0, bufferSize).order(ByteOrder.LITTLE_ENDIAN); |
| target = new float[dim]; |
| } |
| |
| void reset() { |
| bytes.position(0); |
| } |
| |
| abstract float[] next(); |
| } |
| |
| private static class VectorReaderFloat32 extends VectorReader { |
| private final FloatBuffer floats; |
| |
| VectorReaderFloat32(FileChannel input, int dim, int bufferSize) throws IOException { |
| super(input, dim, bufferSize); |
| floats = bytes.asFloatBuffer(); |
| } |
| |
| @Override |
| void reset() { |
| super.reset(); |
| floats.position(0); |
| } |
| |
| @Override |
| float[] next() { |
| floats.get(target); |
| return target; |
| } |
| } |
| |
| private static class VectorReaderByte extends VectorReader { |
| private byte[] scratch; |
| private BytesRef bytesRef; |
| |
| VectorReaderByte(FileChannel input, int dim, int bufferSize) throws IOException { |
| super(input, dim, bufferSize); |
| scratch = new byte[dim]; |
| bytesRef = new BytesRef(scratch); |
| } |
| |
| @Override |
| float[] next() { |
| bytes.get(scratch); |
| for (int i = 0; i < scratch.length; i++) { |
| target[i] = scratch[i]; |
| } |
| return target; |
| } |
| |
| BytesRef nextBytes() { |
| bytes.get(scratch); |
| return bytesRef; |
| } |
| } |
| |
| private static TopDocs doKnnVectorQuery( |
| IndexSearcher searcher, String field, float[] vector, int k, int fanout, Query filter) |
| throws IOException { |
| return searcher.search(new KnnVectorQuery(field, vector, k + fanout, filter), k); |
| } |
| |
| private float checkResults(TopDocs[] results, int[][] nn) { |
| int totalMatches = 0; |
| int totalResults = results.length * topK; |
| for (int i = 0; i < results.length; i++) { |
| // System.out.println(Arrays.toString(nn[i])); |
| // System.out.println(Arrays.toString(results[i].scoreDocs)); |
| totalMatches += compareNN(nn[i], results[i]); |
| } |
| return totalMatches / (float) totalResults; |
| } |
| |
| private int compareNN(int[] expected, TopDocs results) { |
| int matched = 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 i = 0; i < topK; i++) { |
| expectedSet.add(expected[i]); |
| } |
| for (ScoreDoc scoreDoc : results.scoreDocs) { |
| if (expectedSet.contains(scoreDoc.doc)) { |
| ++matched; |
| } |
| } |
| return matched; |
| } |
| |
| private int[][] getNN(Path docPath, Path queryPath) throws IOException { |
| // look in working directory for cached nn file |
| String hash = Integer.toString(Objects.hash(docPath, queryPath, numDocs, numIters, topK), 36); |
| String nnFileName = "nn-" + hash + ".bin"; |
| Path nnPath = Paths.get(nnFileName); |
| if (Files.exists(nnPath) && isNewer(nnPath, docPath, queryPath) && selectivity == 1f) { |
| return readNN(nnPath); |
| } else { |
| // TODO: enable computing NN from high precision vectors when |
| // checking low-precision recall |
| int[][] nn = computeNN(docPath, queryPath, vectorEncoding); |
| if (selectivity == 1f) { |
| writeNN(nn, nnPath); |
| } |
| return nn; |
| } |
| } |
| |
| private boolean isNewer(Path path, Path... others) throws IOException { |
| FileTime modified = Files.getLastModifiedTime(path); |
| for (Path other : others) { |
| if (Files.getLastModifiedTime(other).compareTo(modified) >= 0) { |
| return false; |
| } |
| } |
| return true; |
| } |
| |
| private int[][] readNN(Path nnPath) throws IOException { |
| int[][] result = new int[numIters][]; |
| try (FileChannel in = FileChannel.open(nnPath)) { |
| IntBuffer intBuffer = |
| in.map(FileChannel.MapMode.READ_ONLY, 0, numIters * topK * Integer.BYTES) |
| .order(ByteOrder.LITTLE_ENDIAN) |
| .asIntBuffer(); |
| for (int i = 0; i < numIters; i++) { |
| result[i] = new int[topK]; |
| intBuffer.get(result[i]); |
| } |
| } |
| return result; |
| } |
| |
| private void writeNN(int[][] nn, Path nnPath) throws IOException { |
| if (quiet == false) { |
| System.out.println("writing true nearest neighbors to " + nnPath); |
| } |
| ByteBuffer tmp = |
| ByteBuffer.allocate(nn[0].length * Integer.BYTES).order(ByteOrder.LITTLE_ENDIAN); |
| try (OutputStream out = Files.newOutputStream(nnPath)) { |
| for (int i = 0; i < numIters; i++) { |
| tmp.asIntBuffer().put(nn[i]); |
| out.write(tmp.array()); |
| } |
| } |
| } |
| |
| @SuppressForbidden(reason = "Uses random()") |
| private static FixedBitSet generateRandomBitSet(int size, float selectivity) { |
| FixedBitSet bitSet = new FixedBitSet(size); |
| for (int i = 0; i < size; i++) { |
| if (Math.random() < selectivity) { |
| bitSet.set(i); |
| } else { |
| bitSet.clear(i); |
| } |
| } |
| return bitSet; |
| } |
| |
| private int[][] computeNN(Path docPath, Path queryPath, VectorEncoding encoding) |
| throws IOException { |
| int[][] result = new int[numIters][]; |
| if (quiet == false) { |
| System.out.println("computing true nearest neighbors of " + numIters + " target vectors"); |
| } |
| try (FileChannel in = FileChannel.open(docPath); |
| FileChannel qIn = FileChannel.open(queryPath)) { |
| VectorReader docReader = VectorReader.create(in, dim, encoding, numDocs); |
| VectorReader queryReader = VectorReader.create(qIn, dim, encoding, numIters); |
| for (int i = 0; i < numIters; i++) { |
| float[] query = queryReader.next(); |
| NeighborQueue queue = new NeighborQueue(topK, false); |
| for (int j = 0; j < numDocs; j++) { |
| float[] doc = docReader.next(); |
| float d = similarityFunction.compare(query, doc); |
| if (matchDocs == null || matchDocs.get(j)) { |
| queue.insertWithOverflow(j, d); |
| } |
| } |
| docReader.reset(); |
| result[i] = new int[topK]; |
| for (int k = topK - 1; k >= 0; k--) { |
| result[i][k] = queue.topNode(); |
| queue.pop(); |
| // System.out.print(" " + n); |
| } |
| if (quiet == false && (i + 1) % 10 == 0) { |
| System.out.print(" " + (i + 1)); |
| System.out.flush(); |
| } |
| } |
| } |
| return result; |
| } |
| |
| private int createIndex(Path docsPath, Path indexPath) throws IOException { |
| IndexWriterConfig iwc = new IndexWriterConfig().setOpenMode(IndexWriterConfig.OpenMode.CREATE); |
| iwc.setCodec( |
| new Lucene94Codec() { |
| @Override |
| public KnnVectorsFormat getKnnVectorsFormatForField(String field) { |
| return new Lucene94HnswVectorsFormat(maxConn, beamWidth); |
| } |
| }); |
| // iwc.setMergePolicy(NoMergePolicy.INSTANCE); |
| iwc.setRAMBufferSizeMB(1994d); |
| iwc.setUseCompoundFile(false); |
| // iwc.setMaxBufferedDocs(10000); |
| |
| FieldType fieldType = KnnVectorField.createFieldType(dim, vectorEncoding, similarityFunction); |
| if (quiet == false) { |
| iwc.setInfoStream(new PrintStreamInfoStream(System.out)); |
| System.out.println("creating index in " + indexPath); |
| } |
| long start = System.nanoTime(); |
| try (FSDirectory dir = FSDirectory.open(indexPath); |
| IndexWriter iw = new IndexWriter(dir, iwc)) { |
| try (FileChannel in = FileChannel.open(docsPath)) { |
| VectorReader vectorReader = VectorReader.create(in, dim, vectorEncoding, numDocs); |
| for (int i = 0; i < numDocs; i++) { |
| Document doc = new Document(); |
| switch (vectorEncoding) { |
| case BYTE -> doc.add( |
| new KnnVectorField( |
| KNN_FIELD, ((VectorReaderByte) vectorReader).nextBytes(), fieldType)); |
| case FLOAT32 -> doc.add(new KnnVectorField(KNN_FIELD, vectorReader.next(), fieldType)); |
| } |
| doc.add(new StoredField(ID_FIELD, i)); |
| iw.addDocument(doc); |
| } |
| if (quiet == false) { |
| System.out.println("Done indexing " + numDocs + " documents; now flush"); |
| } |
| } |
| } |
| long elapsed = System.nanoTime() - start; |
| if (quiet == false) { |
| System.out.println("Indexed " + numDocs + " documents in " + elapsed / 1_000_000_000 + "s"); |
| } |
| return (int) (elapsed / 1_000_000); |
| } |
| |
| private static void usage() { |
| String error = |
| "Usage: TestKnnGraph [-reindex] [-search {queryfile}|-stats|-check] [-docs {datafile}] [-niter N] [-fanout N] [-maxConn N] [-beamWidth N] [-filterSelectivity N] [-prefilter]"; |
| System.err.println(error); |
| System.exit(1); |
| } |
| |
| class BinaryFileVectors implements RandomAccessVectorValuesProducer, Closeable { |
| |
| private final int size; |
| private final FileChannel in; |
| private final FloatBuffer mmap; |
| |
| BinaryFileVectors(Path filePath) throws IOException { |
| in = FileChannel.open(filePath); |
| long totalBytes = (long) numDocs * dim * Float.BYTES; |
| if (totalBytes > Integer.MAX_VALUE) { |
| throw new IllegalArgumentException("input over 2GB not supported"); |
| } |
| int vectorByteSize = dim * Float.BYTES; |
| size = (int) (totalBytes / vectorByteSize); |
| mmap = |
| in.map(FileChannel.MapMode.READ_ONLY, 0, totalBytes) |
| .order(ByteOrder.LITTLE_ENDIAN) |
| .asFloatBuffer(); |
| } |
| |
| @Override |
| public void close() throws IOException { |
| in.close(); |
| } |
| |
| @Override |
| public RandomAccessVectorValues randomAccess() { |
| return new Values(); |
| } |
| |
| class Values implements RandomAccessVectorValues { |
| |
| float[] vector = new float[dim]; |
| FloatBuffer source = mmap.slice(); |
| |
| @Override |
| public int size() { |
| return size; |
| } |
| |
| @Override |
| public int dimension() { |
| return dim; |
| } |
| |
| @Override |
| public float[] vectorValue(int targetOrd) { |
| int pos = targetOrd * dim; |
| source.position(pos); |
| source.get(vector); |
| return vector; |
| } |
| |
| @Override |
| public BytesRef binaryValue(int targetOrd) { |
| throw new UnsupportedOperationException(); |
| } |
| } |
| } |
| |
| static class NeighborArraySorter extends IntroSorter { |
| private final int[] node; |
| private final float[] score; |
| |
| NeighborArraySorter(NeighborArray neighbors) { |
| node = neighbors.node; |
| score = neighbors.score; |
| } |
| |
| int pivot; |
| |
| @Override |
| protected void swap(int i, int j) { |
| int tmpNode = node[i]; |
| float tmpScore = score[i]; |
| node[i] = node[j]; |
| score[i] = score[j]; |
| node[j] = tmpNode; |
| score[j] = tmpScore; |
| } |
| |
| @Override |
| protected void setPivot(int i) { |
| pivot = i; |
| } |
| |
| @Override |
| protected int comparePivot(int j) { |
| return Float.compare(score[pivot], score[j]); |
| } |
| } |
| |
| private static class BitSetQuery extends Query { |
| |
| private final FixedBitSet docs; |
| private final int cardinality; |
| |
| BitSetQuery(FixedBitSet docs) { |
| this.docs = docs; |
| this.cardinality = docs.cardinality(); |
| } |
| |
| @Override |
| public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) |
| throws IOException { |
| return new ConstantScoreWeight(this, boost) { |
| @Override |
| public Scorer scorer(LeafReaderContext context) throws IOException { |
| return new ConstantScoreScorer( |
| this, score(), scoreMode, new BitSetIterator(docs, cardinality)); |
| } |
| |
| @Override |
| public boolean isCacheable(LeafReaderContext ctx) { |
| return false; |
| } |
| }; |
| } |
| |
| @Override |
| public void visit(QueryVisitor visitor) {} |
| |
| @Override |
| public String toString(String field) { |
| return "BitSetQuery"; |
| } |
| |
| @Override |
| public boolean equals(Object other) { |
| return sameClassAs(other) && docs.equals(((BitSetQuery) other).docs); |
| } |
| |
| @Override |
| public int hashCode() { |
| return 31 * classHash() + docs.hashCode(); |
| } |
| } |
| } |
