|
19 | 19 | import java.io.IOException; |
20 | 20 | import java.util.Arrays; |
21 | 21 | import java.util.Objects; |
| 22 | +import org.apache.lucene.document.KnnFloatVectorField; |
22 | 23 | import org.apache.lucene.index.FieldInfo; |
23 | 24 | import org.apache.lucene.index.FloatVectorValues; |
24 | 25 | import org.apache.lucene.index.LeafReaderContext; |
25 | 26 | import org.apache.lucene.search.knn.KnnCollectorManager; |
26 | 27 | import org.apache.lucene.util.ArrayUtil; |
27 | 28 | import org.apache.lucene.util.Bits; |
28 | 29 |
|
| 30 | +/** A subclass of KnnFloatVectorQuery which does oversampling and full-precision reranking. */ |
29 | 31 | public class TwoPhaseKnnVectorQuery extends KnnFloatVectorQuery { |
30 | 32 |
|
31 | 33 | private final int originalK; |
32 | 34 | private final double oversample; |
33 | 35 |
|
| 36 | + /** |
| 37 | + * Find the <code>k</code> nearest documents to the target vector according to the vectors in the |
| 38 | + * given field. <code>target</code> vector. It also over-samples by oversample parameter and does |
| 39 | + * full precision reranking if oversample > 0 |
| 40 | + * |
| 41 | + * @param field a field that has been indexed as a {@link KnnFloatVectorField}. |
| 42 | + * @param target the target of the search |
| 43 | + * @param k the number of documents to find |
| 44 | + * @param oversample the oversampling factor, a value of 0 means no oversampling |
| 45 | + * @param filter a filter applied before the vector search |
| 46 | + * @throws IllegalArgumentException if <code>k</code> is less than 1 |
| 47 | + */ |
34 | 48 | public TwoPhaseKnnVectorQuery( |
35 | 49 | String field, float[] target, int k, double oversample, Query filter) { |
36 | 50 | super(field, target, k + (int) Math.round(k * oversample), filter); |
|
0 commit comments