Preconditioning

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

@@ -59,7 +59,9 @@ record CmdLineArgs(
     double writerBufferSizeInMb,
     int writerMaxBufferedDocs,
     int forceMergeMaxNumSegments,
-    boolean onDiskRescore
+    boolean onDiskRescore,
+    boolean doPrecondition,
+    int preconditioningDims
 ) implements ToXContentObject {
 
     static final ParseField DOC_VECTORS_FIELD = new ParseField("doc_vectors");
@@ -92,6 +94,8 @@ record CmdLineArgs(
     static final ParseField WRITER_BUFFER_MB_FIELD = new ParseField("writer_buffer_mb");
     static final ParseField WRITER_BUFFER_DOCS_FIELD = new ParseField("writer_buffer_docs");
     static final ParseField ON_DISK_RESCORE_FIELD = new ParseField("on_disk_rescore");
+    static final ParseField DO_PRECONDITION = new ParseField("do_precondition");
+    static final ParseField PRECONDITIONING_DIMS = new ParseField("preconditioning_dims");
 
     /** By default, in ES the default writer buffer size is 10% of the heap space
      * (see {@code IndexingMemoryController.INDEX_BUFFER_SIZE_SETTING}).
@@ -138,6 +142,8 @@ static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         PARSER.declareInt(Builder::setWriterMaxBufferedDocs, WRITER_BUFFER_DOCS_FIELD);
         PARSER.declareInt(Builder::setForceMergeMaxNumSegments, FORCE_MERGE_MAX_NUM_SEGMENTS_FIELD);
         PARSER.declareBoolean(Builder::setOnDiskRescore, ON_DISK_RESCORE_FIELD);
+        PARSER.declareBoolean(Builder::setDoPrecondition, DO_PRECONDITION);
+        PARSER.declareInt(Builder::setPreconditioningDims, PRECONDITIONING_DIMS);
     }
 
     @Override
@@ -213,6 +219,8 @@ static class Builder {
         private KnnIndexTester.MergePolicyType mergePolicy = null;
         private double writerBufferSizeInMb = DEFAULT_WRITER_BUFFER_MB;
         private boolean onDiskRescore = false;
+        private boolean doPrecondition = false;
+        private int preconditioningDims = 64;
 
         /**
          * Elasticsearch does not set this explicitly, and in Lucene this setting is
@@ -369,6 +377,16 @@ public Builder setOnDiskRescore(boolean onDiskRescore) {
             return this;
         }
 
+        public Builder setDoPrecondition(boolean doPrecondition) {
+            this.doPrecondition = doPrecondition;
+            return this;
+        }
+
+        public Builder setPreconditioningDims(int preconditioningDims) {
+            this.preconditioningDims = preconditioningDims;
+            return this;
+        }
+
         public CmdLineArgs build() {
             if (docVectors == null) {
                 throw new IllegalArgumentException("Document vectors path must be provided");
@@ -407,7 +425,9 @@ public CmdLineArgs build() {
                 writerBufferSizeInMb,
                 writerMaxBufferedDocs,
                 forceMergeMaxNumSegments,
-                onDiskRescore
+                onDiskRescore,
+                doPrecondition,
+                preconditioningDims
             );
         }
     }

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

@@ -133,7 +133,9 @@ static Codec createCodec(CmdLineArgs args) {
                 args.ivfClusterSize(),
                 ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER,
                 DenseVectorFieldMapper.ElementType.FLOAT,
-                args.onDiskRescore()
+                args.onDiskRescore(),
+                args.doPrecondition(),
+                args.preconditioningDims()
             );
         } else if (args.indexType() == IndexType.GPU_HNSW) {
             if (quantizeBits == 32) {

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

@@ -174,6 +174,12 @@ protected FieldEntry doReadField(
         );
     }
 
+    @Override
+    protected float[] preconditionVector(FieldInfo fieldInfo, float[] vector) {
+        // no-op
+        return vector;
+    }
+
     private static CentroidIterator getCentroidIteratorNoParent(
         FieldInfo fieldInfo,
         IndexInput centroids,

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

@@ -378,6 +378,12 @@ public CentroidSupplier createCentroidSupplier(
         return new OffHeapCentroidSupplier(centroidsInput, numCentroids, fieldInfo);
     }
 
+    @Override
+    public FloatVectorValues preconditionVectors(FloatVectorValues floatVectorValues) throws IOException {
+        // no-op
+        return floatVectorValues;
+    }
+
     @Override
     public void writeCentroids(
         FieldInfo fieldInfo,

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

@@ -275,6 +275,8 @@ public final ByteVectorValues getByteVectorValues(String field) throws IOExcepti
         return getReaderForField(field).getByteVectorValues(field);
     }
 
+    protected abstract float[] preconditionVector(FieldInfo fieldInfo, float[] vector);
+
     @Override
     public final void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException {
         final FieldInfo fieldInfo = state.fieldInfos.fieldInfo(field);
@@ -317,6 +319,9 @@ public final void search(String field, float[] target, KnnCollector knnCollector
             // clip so we visit at least one vector
             visitRatio = estimated / numVectors;
         }
+        // precondition the query vector if necessary
+        target = preconditionVector(fieldInfo, target);
+
         // we account for soar vectors here. We can potentially visit a vector twice so we multiply by 2 here.
         long maxVectorVisited = (long) (2.0 * visitRatio * numVectors);
         IndexInput postListSlice = entry.postingListSlice(ivfClusters);
@@ -332,6 +337,7 @@ public final void search(String field, float[] target, KnnCollector knnCollector
             visitRatio
         );
         Bits acceptDocsBits = acceptDocs.bits();
+
         PostingVisitor scorer = getPostingVisitor(fieldInfo, postListSlice, target, acceptDocsBits);
         long expectedDocs = 0;
         long actualDocs = 0;

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

@@ -185,6 +185,8 @@ public abstract CentroidSupplier createCentroidSupplier(
         float[] globalCentroid
     ) throws IOException;
 
+    public abstract FloatVectorValues preconditionVectors(FloatVectorValues floatVectorValues) throws IOException;
+
     @Override
     public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
         rawVectorDelegate.flush(maxDoc, sortMap);
@@ -195,7 +197,9 @@ public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
                 continue;
             }
             // build a float vector values with random access
-            final FloatVectorValues floatVectorValues = getFloatVectorValues(fieldWriter.fieldInfo, fieldWriter.delegate, maxDoc);
+            FloatVectorValues floatVectorValues = getFloatVectorValues(fieldWriter.fieldInfo, fieldWriter.delegate, maxDoc);
+            // precondition the vectors if necessary
+            floatVectorValues = preconditionVectors(floatVectorValues);
             // build centroids
             final CentroidAssignments centroidAssignments = calculateCentroids(fieldWriter.fieldInfo, floatVectorValues);
             // wrap centroids with a supplier
@@ -378,7 +382,9 @@ private void mergeOneFieldIVF(FieldInfo fieldInfo, MergeState mergeState) throws
                 ? null
                 : mergeState.segmentInfo.dir.openInput(docsFileName, IOContext.DEFAULT.withHints(DataAccessHint.SEQUENTIAL))
         ) {
-            final FloatVectorValues floatVectorValues = getFloatVectorValues(fieldInfo, docs, vectors, numVectors);
+            FloatVectorValues floatVectorValues = getFloatVectorValues(fieldInfo, docs, vectors, numVectors);
+            // precondition vectors if necessary
+            floatVectorValues = preconditionVectors(floatVectorValues);
 
             final long centroidOffset;
             final long centroidLength;
@@ -394,11 +400,7 @@ private void mergeOneFieldIVF(FieldInfo fieldInfo, MergeState mergeState) throws
             try {
                 centroidTemp = mergeState.segmentInfo.dir.createTempOutput(mergeState.segmentInfo.name, "civf_", IOContext.DEFAULT);
                 centroidTempName = centroidTemp.getName();
-                CentroidAssignments centroidAssignments = calculateCentroids(
-                    fieldInfo,
-                    getFloatVectorValues(fieldInfo, docs, vectors, numVectors),
-                    mergeState
-                );
+                CentroidAssignments centroidAssignments = calculateCentroids(fieldInfo, floatVectorValues, mergeState);
                 // write the centroids to a temporary file so we are not holding them on heap
                 final ByteBuffer buffer = ByteBuffer.allocate(fieldInfo.getVectorDimension() * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN);
                 for (float[] centroid : centroidAssignments.centroids()) {