Use a FunctionScoreQuery to replace scores using a VectorSimilarity based DoubleValueSource

Author: carlosdelest

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -190,6 +190,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion LOGSDB_TELEMETRY_STATS = def(8_785_00_0);
     public static final TransportVersion KQL_QUERY_ADDED = def(8_786_00_0);
     public static final TransportVersion ROLE_MONITOR_STATS = def(8_787_00_0);
+    public static final TransportVersion KNN_QUERY_RESCORE_OVERSAMPLE = def(8_788_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java

@@ -30,6 +30,7 @@
 import org.apache.lucene.index.SegmentWriteState;
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.queries.function.FunctionScoreQuery;
 import org.apache.lucene.search.FieldExistsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.join.BitSetProducer;
@@ -121,6 +122,8 @@ public static boolean isNotUnitVector(float magnitude) {
     public static short MAX_DIMS_COUNT = 4096; // maximum allowed number of dimensions
     public static int MAX_DIMS_COUNT_BIT = 4096 * Byte.SIZE; // maximum allowed number of dimensions
 
+    public static final int OVERSAMPLE_LIMIT = 10_000; // Max oversample allowed for k and num_candidates
+
     public static short MIN_DIMS_FOR_DYNAMIC_FLOAT_MAPPING = 128; // minimum number of dims for floats to be dynamically mapped to vector
     public static final int MAGNITUDE_BYTES = 4;
 
@@ -2000,6 +2003,7 @@ public Query createKnnQuery(
             VectorData queryVector,
             Integer k,
             int numCands,
+            Float rescoreOversample,
             Query filter,
             Float similarityThreshold,
             BitSetProducer parentFilter
@@ -2010,21 +2014,50 @@ public Query createKnnQuery(
                 );
             }
             return switch (getElementType()) {
-                case BYTE -> createKnnByteQuery(queryVector.asByteVector(), k, numCands, filter, similarityThreshold, parentFilter);
-                case FLOAT -> createKnnFloatQuery(queryVector.asFloatVector(), k, numCands, filter, similarityThreshold, parentFilter);
-                case BIT -> createKnnBitQuery(queryVector.asByteVector(), k, numCands, filter, similarityThreshold, parentFilter);
+                case BYTE -> createKnnByteQuery(
+                    queryVector.asByteVector(),
+                    k,
+                    numCands,
+                    filter,
+                    rescoreOversample,
+                    similarityThreshold,
+                    parentFilter
+                );
+                case FLOAT -> createKnnFloatQuery(
+                    queryVector.asFloatVector(),
+                    k,
+                    numCands,
+                    rescoreOversample,
+                    filter,
+                    similarityThreshold,
+                    parentFilter
+                );
+                case BIT -> createKnnBitQuery(
+                    queryVector.asByteVector(),
+                    k,
+                    numCands,
+                    rescoreOversample,
+                    filter,
+                    similarityThreshold,
+                    parentFilter
+                );
             };
         }
 
         private Query createKnnBitQuery(
             byte[] queryVector,
             Integer k,
             int numCands,
+            Float rescoreOversample,
             Query filter,
             Float similarityThreshold,
             BitSetProducer parentFilter
         ) {
             elementType.checkDimensions(dims, queryVector.length);
+            if (similarity == VectorSimilarity.DOT_PRODUCT || similarity == VectorSimilarity.COSINE) {
+                float squaredMagnitude = VectorUtil.dotProduct(queryVector, queryVector);
+                elementType.checkVectorMagnitude(similarity, ElementType.errorByteElementsAppender(queryVector), squaredMagnitude);
+            }
             Query knnQuery = parentFilter != null
                 ? new ESDiversifyingChildrenByteKnnVectorQuery(name(), queryVector, filter, k, numCands, parentFilter)
                 : new ESKnnByteVectorQuery(name(), queryVector, k, numCands, filter);
@@ -2035,6 +2068,17 @@ private Query createKnnBitQuery(
                     similarity.score(similarityThreshold, elementType, dims)
                 );
             }
+            if (rescoreOversample != null) {
+                knnQuery = new FunctionScoreQuery(
+                    knnQuery,
+                    new VectorSimilarityByteValueSource(
+                        name(),
+                        queryVector,
+                        similarity.vectorSimilarityFunction(indexVersionCreated, ElementType.BYTE)
+                    )
+                );
+
+            }
             return knnQuery;
         }
 
@@ -2043,6 +2087,7 @@ private Query createKnnByteQuery(
             Integer k,
             int numCands,
             Query filter,
+            Float rescoreOversample,
             Float similarityThreshold,
             BitSetProducer parentFilter
         ) {
@@ -2052,23 +2097,38 @@ private Query createKnnByteQuery(
                 float squaredMagnitude = VectorUtil.dotProduct(queryVector, queryVector);
                 elementType.checkVectorMagnitude(similarity, ElementType.errorByteElementsAppender(queryVector), squaredMagnitude);
             }
+            int adjustedK = rescoreOversample == null ? k : Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample));
+            int adjustedNumCands = Math.max(adjustedK, numCands);
+
             Query knnQuery = parentFilter != null
-                ? new ESDiversifyingChildrenByteKnnVectorQuery(name(), queryVector, filter, k, numCands, parentFilter)
-                : new ESKnnByteVectorQuery(name(), queryVector, k, numCands, filter);
+                ? new ESDiversifyingChildrenByteKnnVectorQuery(name(), queryVector, filter, adjustedK, adjustedNumCands, parentFilter)
+                : new ESKnnByteVectorQuery(name(), queryVector, adjustedK, adjustedNumCands, filter);
             if (similarityThreshold != null) {
                 knnQuery = new VectorSimilarityQuery(
                     knnQuery,
                     similarityThreshold,
                     similarity.score(similarityThreshold, elementType, dims)
                 );
             }
+            if (rescoreOversample != null) {
+                knnQuery = new FunctionScoreQuery(
+                    knnQuery,
+                    new VectorSimilarityByteValueSource(
+                        name(),
+                        queryVector,
+                        similarity.vectorSimilarityFunction(indexVersionCreated, ElementType.BYTE)
+                    )
+                );
+
+            }
             return knnQuery;
         }
 
         private Query createKnnFloatQuery(
             float[] queryVector,
             Integer k,
             int numCands,
+            Float rescoreOversample,
             Query filter,
             Float similarityThreshold,
             BitSetProducer parentFilter
@@ -2088,16 +2148,30 @@ && isNotUnitVector(squaredMagnitude)) {
                     }
                 }
             }
+
+            int adjustedK = rescoreOversample == null ? k : Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample));
+            int adjustedNumCands = Math.max(adjustedK, numCands);
             Query knnQuery = parentFilter != null
-                ? new ESDiversifyingChildrenFloatKnnVectorQuery(name(), queryVector, filter, k, numCands, parentFilter)
-                : new ESKnnFloatVectorQuery(name(), queryVector, k, numCands, filter);
+                ? new ESDiversifyingChildrenFloatKnnVectorQuery(name(), queryVector, filter, adjustedK, adjustedNumCands, parentFilter)
+                : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, adjustedNumCands, filter);
             if (similarityThreshold != null) {
                 knnQuery = new VectorSimilarityQuery(
                     knnQuery,
                     similarityThreshold,
                     similarity.score(similarityThreshold, elementType, dims)
                 );
             }
+            if (rescoreOversample != null) {
+                knnQuery = new FunctionScoreQuery(
+                    knnQuery,
+                    new VectorSimilarityFloatValueSource(
+                        name(),
+                        queryVector,
+                        similarity.vectorSimilarityFunction(indexVersionCreated, ElementType.FLOAT)
+                    )
+                );
+
+            }
             return knnQuery;
         }
 

server/src/main/java/org/elasticsearch/index/mapper/vectors/VectorSimilarityByteValueSource.java

@@ -0,0 +1,95 @@
+/*
+ * 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.index.mapper.vectors;
+
+import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.KnnVectorValues;
+import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.DoubleValues;
+import org.apache.lucene.search.DoubleValuesSource;
+import org.apache.lucene.search.IndexSearcher;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Objects;
+
+public class VectorSimilarityByteValueSource extends DoubleValuesSource {
+
+    private final String field;
+    private final byte[] target;
+    private final VectorSimilarityFunction vectorSimilarityFunction;
+
+    public VectorSimilarityByteValueSource(String field, byte[] target, VectorSimilarityFunction vectorSimilarityFunction) {
+        this.field = field;
+        this.target = target;
+        this.vectorSimilarityFunction = vectorSimilarityFunction;
+    }
+
+    @Override
+    public DoubleValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
+        final LeafReader reader = ctx.reader();
+
+        ByteVectorValues vectorValues = reader.getByteVectorValues(field);
+        KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator();
+
+        return new DoubleValues() {
+            private int docId = -1;
+
+            @Override
+            public double doubleValue() throws IOException {
+                return vectorSimilarityFunction.compare(target, vectorValues.vectorValue(docId));
+            }
+
+            @Override
+            public boolean advanceExact(int doc) throws IOException {
+                docId = doc;
+                return iterator.advance(docId) != DocIdSetIterator.NO_MORE_DOCS;
+            }
+        };
+    }
+
+    @Override
+    public boolean needsScores() {
+        return false;
+    }
+
+    @Override
+    public DoubleValuesSource rewrite(IndexSearcher reader) throws IOException {
+        return this;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(field, Arrays.hashCode(target), vectorSimilarityFunction);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        VectorSimilarityByteValueSource that = (VectorSimilarityByteValueSource) o;
+        return Objects.equals(field, that.field)
+            && Objects.deepEquals(target, that.target)
+            && vectorSimilarityFunction == that.vectorSimilarityFunction;
+    }
+
+    @Override
+    public String toString() {
+        return "VectorSimilarityByteValueSource(" + field + ", " + Arrays.toString(target) + ", " + vectorSimilarityFunction + ")";
+    }
+
+    @Override
+    public boolean isCacheable(LeafReaderContext ctx) {
+        return false;
+    }
+}

server/src/main/java/org/elasticsearch/index/mapper/vectors/VectorSimilarityFloatValueSource.java

@@ -0,0 +1,95 @@
+/*
+ * 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.index.mapper.vectors;
+
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.KnnVectorValues;
+import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.DoubleValues;
+import org.apache.lucene.search.DoubleValuesSource;
+import org.apache.lucene.search.IndexSearcher;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Objects;
+
+public class VectorSimilarityFloatValueSource extends DoubleValuesSource {
+
+    private final String field;
+    private final float[] target;
+    private final VectorSimilarityFunction vectorSimilarityFunction;
+
+    public VectorSimilarityFloatValueSource(String field, float[] target, VectorSimilarityFunction vectorSimilarityFunction) {
+        this.field = field;
+        this.target = target;
+        this.vectorSimilarityFunction = vectorSimilarityFunction;
+    }
+
+    @Override
+    public DoubleValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
+        final LeafReader reader = ctx.reader();
+
+        FloatVectorValues vectorValues = reader.getFloatVectorValues(field);
+        KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator();
+
+        return new DoubleValues() {
+            private int docId = -1;
+
+            @Override
+            public double doubleValue() throws IOException {
+                return vectorSimilarityFunction.compare(target, vectorValues.vectorValue(docId));
+            }
+
+            @Override
+            public boolean advanceExact(int doc) throws IOException {
+                docId = doc;
+                return iterator.advance(docId) != DocIdSetIterator.NO_MORE_DOCS;
+            }
+        };
+    }
+
+    @Override
+    public boolean needsScores() {
+        return false;
+    }
+
+    @Override
+    public DoubleValuesSource rewrite(IndexSearcher reader) throws IOException {
+        return this;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(field, Arrays.hashCode(target), vectorSimilarityFunction);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        VectorSimilarityFloatValueSource that = (VectorSimilarityFloatValueSource) o;
+        return Objects.equals(field, that.field)
+            && Objects.deepEquals(target, that.target)
+            && vectorSimilarityFunction == that.vectorSimilarityFunction;
+    }
+
+    @Override
+    public String toString() {
+        return "VectorSimilarityFloatValueSource(" + field + ", " + Arrays.toString(target) + ", " + vectorSimilarityFunction + ")";
+    }
+
+    @Override
+    public boolean isCacheable(LeafReaderContext ctx) {
+        return false;
+    }
+}