kNN vector rescoring for quantized vectors #116663
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
|
There was a problem hiding this comment. lets add an There was a problem hiding this comment. Sounds good - b0c2221 |
||
| 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) | ||
carlosdelest marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| && vectorSimilarityFunction == that.vectorSimilarityFunction; | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| return "VectorSimilarityFloatValueSource(" + field + ", " + Arrays.toString(target) + ", " + vectorSimilarityFunction + ")"; | ||
carlosdelest marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| } | ||
|
|
||
| @Override | ||
| public boolean isCacheable(LeafReaderContext ctx) { | ||
| return false; | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We don't need this.
bytevalues cannot be quantized.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🤦 Thanks for the catch - it's just
floats that get quantized. Removing it in 2379596