added preconditioning to help edge cases where quantization causes vectors to become indistinguishable

Author: john-wagster

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnFileUtils.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.test.knn;
+
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FSDirectory;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.MMapDirectory;
+import org.apache.lucene.store.NativeFSLockFactory;
+import org.apache.lucene.store.ReadAdvice;
+import org.elasticsearch.index.StandardIOBehaviorHint;
+import org.elasticsearch.index.store.FsDirectoryFactory;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.Optional;
+import java.util.function.BiFunction;
+
+public class KnnFileUtils {
+
+    static Directory getDirectory(Path indexPath) throws IOException {
+        Directory dir = FSDirectory.open(indexPath);
+        if (dir instanceof MMapDirectory mmapDir) {
+            mmapDir.setReadAdvice(getReadAdviceFunc()); // enable madvise
+            return new FsDirectoryFactory.HybridDirectory(NativeFSLockFactory.INSTANCE, mmapDir, 64);
+        }
+        return dir;
+    }
+
+    private static BiFunction<String, IOContext, Optional<ReadAdvice>> getReadAdviceFunc() {
+        return (name, context) -> {
+            if (context.hints().contains(StandardIOBehaviorHint.INSTANCE) || name.endsWith(".cfs")) {
+                return Optional.of(ReadAdvice.NORMAL);
+            }
+            return MMapDirectory.ADVISE_BY_CONTEXT.apply(name, context);
+        };
+    }
+}

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java

@@ -29,6 +29,7 @@
 import org.apache.lucene.store.FSDirectory;
 import org.elasticsearch.cli.ProcessInfo;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.io.Channels;
 import org.elasticsearch.common.logging.LogConfigurator;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.PathUtils;
@@ -52,6 +53,9 @@
 import java.io.InputStream;
 import java.io.UncheckedIOException;
 import java.lang.management.ThreadInfo;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.channels.FileChannel;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
@@ -116,7 +120,7 @@ private static String formatIndexPath(CmdLineArgs args) {
         return INDEX_DIR + "/" + args.docVectors().get(0).getFileName() + "-" + String.join("-", suffix) + ".index";
     }
 
-    static Codec createCodec(CmdLineArgs args) {
+    static Codec createCodec(CmdLineArgs args, List<Path> docsPaths) throws IOException {
         final KnnVectorsFormat format;
         int quantizeBits = args.quantizeBits();
         if (args.indexType() == IndexType.IVF) {
@@ -128,8 +132,29 @@ static Codec createCodec(CmdLineArgs args) {
                     "IVF index type only supports 1, 2 or 4 bits quantization, but got: " + quantizeBits
                 );
             };
+
+            // get the actual dims
+            int dims = args.dimensions();
+            if (dims == -1) {
+                Path docsPath = docsPaths.getFirst();
+                try (FileChannel in = FileChannel.open(docsPath)) {
+                    ByteBuffer preamble = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN);
+                    int bytesRead = Channels.readFromFileChannel(in, 0, preamble);
+                    if (bytesRead < 4) {
+                        long docsPathSizeInBytes = in.size();
+                        throw new IllegalArgumentException(
+                            "docsPath \"" + docsPath + "\" does not contain a valid dims?  size=" + docsPathSizeInBytes
+                        );
+                    }
+                    dims = preamble.getInt(0);
+                    if (dims <= 0) {
+                        throw new IllegalArgumentException("docsPath \"" + docsPath + "\" has invalid dimension: " + dims);
+                    }
+                }
+            }
+
             format = new ESNextDiskBBQVectorsFormat(
-                args.dimensions(),
+                dims,
                 encoding,
                 args.ivfClusterSize(),
                 ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER,
@@ -265,7 +290,7 @@ public static void main(String[] args) throws Exception {
             }
             logger.info("Running with Java: " + Runtime.version());
             logger.info("Running KNN index tester with arguments: " + cmdLineArgs);
-            Codec codec = createCodec(cmdLineArgs);
+            Codec codec = createCodec(cmdLineArgs, cmdLineArgs.docVectors());
             Path indexPath = PathUtils.get(formatIndexPath(cmdLineArgs));
             MergePolicy mergePolicy = getMergePolicy(cmdLineArgs);
             if (cmdLineArgs.reindex() || cmdLineArgs.forceMerge()) {

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexer.java

@@ -34,15 +34,8 @@
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.store.Directory;
-import org.apache.lucene.store.FSDirectory;
-import org.apache.lucene.store.IOContext;
-import org.apache.lucene.store.MMapDirectory;
-import org.apache.lucene.store.NativeFSLockFactory;
-import org.apache.lucene.store.ReadAdvice;
 import org.apache.lucene.util.PrintStreamInfoStream;
 import org.elasticsearch.common.io.Channels;
-import org.elasticsearch.index.StandardIOBehaviorHint;
-import org.elasticsearch.index.store.FsDirectoryFactory;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -54,15 +47,14 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.BiFunction;
 
+import static org.elasticsearch.test.knn.KnnFileUtils.getDirectory;
 import static org.elasticsearch.test.knn.KnnIndexTester.logger;
 
 class KnnIndexer {
@@ -243,24 +235,6 @@ public boolean isEnabled(String component) {
         results.forceMergeTimeMS = TimeUnit.NANOSECONDS.toMillis(elapsedNSec);
     }
 
-    static Directory getDirectory(Path indexPath) throws IOException {
-        Directory dir = FSDirectory.open(indexPath);
-        if (dir instanceof MMapDirectory mmapDir) {
-            mmapDir.setReadAdvice(getReadAdviceFunc()); // enable madvise
-            return new FsDirectoryFactory.HybridDirectory(NativeFSLockFactory.INSTANCE, mmapDir, 64);
-        }
-        return dir;
-    }
-
-    private static BiFunction<String, IOContext, Optional<ReadAdvice>> getReadAdviceFunc() {
-        return (name, context) -> {
-            if (context.hints().contains(StandardIOBehaviorHint.INSTANCE) || name.endsWith(".cfs")) {
-                return Optional.of(ReadAdvice.NORMAL);
-            }
-            return MMapDirectory.ADVISE_BY_CONTEXT.apply(name, context);
-        };
-    }
-
     static class IndexerThread extends Thread {
         private final IndexWriter iw;
         private final AtomicInteger numDocsIndexed;

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnSearcher.java

@@ -90,6 +90,7 @@
 import java.util.function.IntConsumer;
 
 import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
+import static org.elasticsearch.test.knn.KnnFileUtils.getDirectory;
 import static org.elasticsearch.test.knn.KnnIndexTester.logger;
 import static org.elasticsearch.test.knn.KnnIndexer.ID_FIELD;
 import static org.elasticsearch.test.knn.KnnIndexer.VECTOR_FIELD;
@@ -180,7 +181,7 @@ void runSearch(KnnIndexTester.Results finalResults, boolean earlyTermination) th
             );
             KnnIndexer.VectorReader targetReader = KnnIndexer.VectorReader.create(input, dim, vectorEncoding, offsetByteSize);
             long startNS;
-            try (Directory dir = KnnIndexer.getDirectory(indexPath)) {
+            try (Directory dir = getDirectory(indexPath)) {
                 try (DirectoryReader reader = DirectoryReader.open(dir)) {
                     IndexSearcher searcher = searchThreads > 1 ? new IndexSearcher(reader, executorService) : new IndexSearcher(reader);
                     byte[] targetBytes = new byte[dim];

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsReader.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.index.codec.vectors.diskbbq;
 
 import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorEncoding;
@@ -27,6 +28,7 @@
 import org.elasticsearch.simdvec.ESVectorUtil;
 
 import java.io.IOException;
+import java.util.List;
 
 import static org.apache.lucene.codecs.lucene102.Lucene102BinaryQuantizedVectorsFormat.QUERY_BITS;
 import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
@@ -46,7 +48,7 @@ public class ES920DiskBBQVectorsReader extends IVFVectorsReader {
     }
 
     @Override
-    public void doInitExtraFiles(SegmentReadState state, int version) throws IOException {
+    public void doInitExtraFiles(SegmentReadState state, int version, FieldInfos fieldInfo) throws IOException {
         // no extra files to init
     }
 
@@ -179,11 +181,22 @@ protected FieldEntry doReadField(
     }
 
     @Override
-    protected float[] preconditionVector(float[] vector) {
+    protected void doAdditionalIntegrityChecks() throws IOException {
+        // no-op
+    }
+
+    @Override
+    protected float[] preconditionVector(FieldInfo fieldInfo, float[] vector) {
         // no-op
         return vector;
     }
 
+    @Override
+    protected List<IndexInput> getAdditionalCloseables() {
+        // no-op
+        return List.of();
+    }
+
     private static CentroidIterator getCentroidIteratorNoParent(
         FieldInfo fieldInfo,
         IndexInput centroids,

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsWriter.java

@@ -9,7 +9,6 @@
 
 package org.elasticsearch.index.codec.vectors.diskbbq;
 
-import org.apache.lucene.codecs.CodecUtil;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FloatVectorValues;
@@ -22,7 +21,6 @@
 import org.apache.lucene.util.hnsw.IntToIntFunction;
 import org.apache.lucene.util.packed.PackedInts;
 import org.apache.lucene.util.packed.PackedLongValues;
-import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.index.codec.vectors.BQVectorUtils;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
@@ -39,6 +37,7 @@
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.util.Arrays;
+import java.util.List;
 
 import static org.elasticsearch.index.codec.vectors.cluster.HierarchicalKMeans.NO_SOAR_ASSIGNMENT;
 
@@ -379,12 +378,24 @@ public CentroidSupplier createCentroidSupplier(
         return new OffHeapCentroidSupplier(centroidsInput, numCentroids, fieldInfo);
     }
 
+    @Override
+    public boolean createPreconditioner() throws IOException {
+        // no-op
+        return false;
+    }
+
     @Override
     public FloatVectorValues preconditionVectors(FloatVectorValues floatVectorValues) throws IOException {
         // no-op
         return floatVectorValues;
     }
 
+    @Override
+    public List<float[]> preconditionVectors(List<float[]> vectors) {
+        // no-op
+        return vectors;
+    }
+
     @Override
     public void writeCentroids(
         FieldInfo fieldInfo,