Skip to content

Commit 2639d28

Browse files
weizijunweizijun
committed
add max_search_ef check
1 parent 1f30eaa commit 2639d28

File tree

4 files changed

+143
-32
lines changed

4 files changed

+143
-32
lines changed

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

Lines changed: 64 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@
9494
import static org.elasticsearch.common.Strings.format;
9595
import static org.elasticsearch.common.xcontent.XContentParserUtils.ensureExpectedToken;
9696
import static org.elasticsearch.index.IndexVersions.DEFAULT_DENSE_VECTOR_TO_INT8_HNSW;
97+
import static org.elasticsearch.search.vectors.KnnSearchBuilder.NUM_CANDS_FIELD;
9798

9899
/**
99100
* A {@link FieldMapper} for indexing a dense vector of floats.
@@ -102,6 +103,7 @@ public class DenseVectorFieldMapper extends FieldMapper {
102103
public static final String COSINE_MAGNITUDE_FIELD_SUFFIX = "._magnitude";
103104
private static final float EPS = 1e-3f;
104105
public static final int BBQ_MIN_DIMS = 64;
106+
public static final int NUM_CANDS_LIMIT = 10_000;
105107

106108
public static boolean isNotUnitVector(float magnitude) {
107109
return Math.abs(magnitude - 1.0f) > EPS;
@@ -120,7 +122,6 @@ public static boolean isNotUnitVector(float magnitude) {
120122
public static final short MIN_DIMS_FOR_DYNAMIC_FLOAT_MAPPING = 128; // minimum number of dims for floats to be dynamically mapped to
121123
// vector
122124
public static final int MAGNITUDE_BYTES = 4;
123-
public static final int OVERSAMPLE_LIMIT = 10_000; // Max oversample allowed
124125

125126
private static DenseVectorFieldMapper toType(FieldMapper in) {
126127
return (DenseVectorFieldMapper) in;
@@ -210,6 +211,7 @@ public Builder(String name, IndexVersion indexVersionCreated) {
210211
? new Int8HnswIndexOptions(
211212
Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN,
212213
Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH,
214+
NUM_CANDS_LIMIT,
213215
null
214216
)
215217
: null,
@@ -1236,6 +1238,14 @@ public void validateDimension(int dim) {
12361238
throw new IllegalArgumentException(type.name + " only supports even dimensions; provided=" + dim);
12371239
}
12381240

1241+
public void validateNumCandidates(int numCands) {
1242+
1243+
}
1244+
1245+
public int maxSearchEf() {
1246+
return Integer.MAX_VALUE;
1247+
}
1248+
12391249
@Override
12401250
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
12411251
builder.startObject();
@@ -1272,25 +1282,32 @@ public final int hashCode() {
12721282
abstract static class AbstractHnswIndexOptions extends IndexOptions {
12731283
protected final int m;
12741284
protected final int efConstruction;
1285+
protected final int maxSearchEf;
12751286

1276-
AbstractHnswIndexOptions(VectorIndexType type, int m, int efConstruction) {
1287+
AbstractHnswIndexOptions(VectorIndexType type, int m, int efConstruction, int maxSearchEf) {
12771288
super(type);
12781289
this.m = m;
12791290
this.efConstruction = efConstruction;
1291+
this.maxSearchEf = maxSearchEf;
12801292
}
12811293

1282-
public int getM() {
1283-
return m;
1294+
@Override
1295+
public void validateNumCandidates(int numCands) {
1296+
if (numCands > maxSearchEf) {
1297+
throw new IllegalArgumentException("[" + NUM_CANDS_FIELD.getPreferredName() + "] cannot exceed [" + maxSearchEf + "]");
1298+
}
12841299
}
12851300

1286-
public int getEfConstruction() {
1287-
return efConstruction;
1301+
@Override
1302+
public int maxSearchEf() {
1303+
return maxSearchEf;
12881304
}
12891305

12901306
@Override
12911307
public XContentBuilder innerXContent(XContentBuilder builder, Params params) throws IOException {
12921308
builder.field("m", m);
12931309
builder.field("ef_construction", efConstruction);
1310+
builder.field("max_search_ef", maxSearchEf);
12941311
innerHnswXContent(builder, params);
12951312
return builder;
12961313
}
@@ -1301,18 +1318,18 @@ public XContentBuilder innerXContent(XContentBuilder builder, Params params) thr
13011318
public boolean doEquals(IndexOptions o) {
13021319
if (this == o) return true;
13031320
if (o == null || getClass() != o.getClass()) return false;
1304-
HnswIndexOptions that = (HnswIndexOptions) o;
1305-
return m == that.m && efConstruction == that.efConstruction;
1321+
AbstractHnswIndexOptions that = (AbstractHnswIndexOptions) o;
1322+
return m == that.m && efConstruction == that.efConstruction && maxSearchEf == that.maxSearchEf;
13061323
}
13071324

13081325
@Override
13091326
public int doHashCode() {
1310-
return Objects.hash(m, efConstruction);
1327+
return Objects.hash(m, efConstruction, maxSearchEf);
13111328
}
13121329

13131330
@Override
13141331
public String toString() {
1315-
return "{type=" + type + ", m=" + m + ", ef_construction=" + efConstruction + "}";
1332+
return "{type=" + type + ", m=" + m + ", ef_construction=" + efConstruction + ", max_search_ef=" + "}";
13161333
}
13171334
}
13181335

@@ -1322,16 +1339,22 @@ public enum VectorIndexType {
13221339
public IndexOptions parseIndexOptions(String fieldName, Map<String, ?> indexOptionsMap) {
13231340
Object mNode = indexOptionsMap.remove("m");
13241341
Object efConstructionNode = indexOptionsMap.remove("ef_construction");
1342+
Object maxSearchEfNode = indexOptionsMap.remove("max_search_ef");
13251343
if (mNode == null) {
13261344
mNode = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
13271345
}
13281346
if (efConstructionNode == null) {
13291347
efConstructionNode = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
13301348
}
1349+
if (maxSearchEfNode == null) {
1350+
maxSearchEfNode = NUM_CANDS_LIMIT;
1351+
}
13311352
int m = XContentMapValues.nodeIntegerValue(mNode);
13321353
int efConstruction = XContentMapValues.nodeIntegerValue(efConstructionNode);
1354+
int maxSearchEf = XContentMapValues.nodeIntegerValue(maxSearchEfNode);
1355+
13331356
MappingParser.checkNoRemainingFields(fieldName, indexOptionsMap);
1334-
return new HnswIndexOptions(m, efConstruction);
1357+
return new HnswIndexOptions(m, efConstruction, maxSearchEf);
13351358
}
13361359

13371360
@Override
@@ -1349,21 +1372,26 @@ public boolean supportsDimension(int dims) {
13491372
public IndexOptions parseIndexOptions(String fieldName, Map<String, ?> indexOptionsMap) {
13501373
Object mNode = indexOptionsMap.remove("m");
13511374
Object efConstructionNode = indexOptionsMap.remove("ef_construction");
1375+
Object maxSearchEfNode = indexOptionsMap.remove("max_search_ef");
13521376
Object confidenceIntervalNode = indexOptionsMap.remove("confidence_interval");
13531377
if (mNode == null) {
13541378
mNode = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
13551379
}
13561380
if (efConstructionNode == null) {
13571381
efConstructionNode = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
13581382
}
1383+
if (maxSearchEfNode == null) {
1384+
maxSearchEfNode = NUM_CANDS_LIMIT;
1385+
}
13591386
int m = XContentMapValues.nodeIntegerValue(mNode);
13601387
int efConstruction = XContentMapValues.nodeIntegerValue(efConstructionNode);
1388+
int maxSearchEf = XContentMapValues.nodeIntegerValue(maxSearchEfNode);
13611389
Float confidenceInterval = null;
13621390
if (confidenceIntervalNode != null) {
13631391
confidenceInterval = (float) XContentMapValues.nodeDoubleValue(confidenceIntervalNode);
13641392
}
13651393
MappingParser.checkNoRemainingFields(fieldName, indexOptionsMap);
1366-
return new Int8HnswIndexOptions(m, efConstruction, confidenceInterval);
1394+
return new Int8HnswIndexOptions(m, efConstruction, maxSearchEf, confidenceInterval);
13671395
}
13681396

13691397
@Override
@@ -1380,21 +1408,26 @@ public boolean supportsDimension(int dims) {
13801408
public IndexOptions parseIndexOptions(String fieldName, Map<String, ?> indexOptionsMap) {
13811409
Object mNode = indexOptionsMap.remove("m");
13821410
Object efConstructionNode = indexOptionsMap.remove("ef_construction");
1411+
Object maxSearchEfNode = indexOptionsMap.remove("max_search_ef");
13831412
Object confidenceIntervalNode = indexOptionsMap.remove("confidence_interval");
13841413
if (mNode == null) {
13851414
mNode = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
13861415
}
13871416
if (efConstructionNode == null) {
13881417
efConstructionNode = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
13891418
}
1419+
if (maxSearchEfNode == null) {
1420+
maxSearchEfNode = NUM_CANDS_LIMIT;
1421+
}
13901422
int m = XContentMapValues.nodeIntegerValue(mNode);
13911423
int efConstruction = XContentMapValues.nodeIntegerValue(efConstructionNode);
1424+
int maxSearchEf = XContentMapValues.nodeIntegerValue(maxSearchEfNode);
13921425
Float confidenceInterval = null;
13931426
if (confidenceIntervalNode != null) {
13941427
confidenceInterval = (float) XContentMapValues.nodeDoubleValue(confidenceIntervalNode);
13951428
}
13961429
MappingParser.checkNoRemainingFields(fieldName, indexOptionsMap);
1397-
return new Int4HnswIndexOptions(m, efConstruction, confidenceInterval);
1430+
return new Int4HnswIndexOptions(m, efConstruction, maxSearchEf, confidenceInterval);
13981431
}
13991432

14001433
@Override
@@ -1473,16 +1506,21 @@ public boolean supportsDimension(int dims) {
14731506
public IndexOptions parseIndexOptions(String fieldName, Map<String, ?> indexOptionsMap) {
14741507
Object mNode = indexOptionsMap.remove("m");
14751508
Object efConstructionNode = indexOptionsMap.remove("ef_construction");
1509+
Object maxSearchEfNode = indexOptionsMap.remove("max_search_ef");
14761510
if (mNode == null) {
14771511
mNode = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
14781512
}
14791513
if (efConstructionNode == null) {
14801514
efConstructionNode = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
14811515
}
1516+
if (maxSearchEfNode == null) {
1517+
maxSearchEfNode = NUM_CANDS_LIMIT;
1518+
}
14821519
int m = XContentMapValues.nodeIntegerValue(mNode);
14831520
int efConstruction = XContentMapValues.nodeIntegerValue(efConstructionNode);
1521+
int maxSearchEf = XContentMapValues.nodeIntegerValue(maxSearchEfNode);
14841522
MappingParser.checkNoRemainingFields(fieldName, indexOptionsMap);
1485-
return new BBQHnswIndexOptions(m, efConstruction);
1523+
return new BBQHnswIndexOptions(m, efConstruction, maxSearchEf);
14861524
}
14871525

14881526
@Override
@@ -1622,8 +1660,8 @@ public int doHashCode() {
16221660
static class Int4HnswIndexOptions extends AbstractHnswIndexOptions {
16231661
private final float confidenceInterval;
16241662

1625-
Int4HnswIndexOptions(int m, int efConstruction, Float confidenceInterval) {
1626-
super(VectorIndexType.INT4_HNSW, m, efConstruction);
1663+
Int4HnswIndexOptions(int m, int efConstruction, int maxSearchEf, Float confidenceInterval) {
1664+
super(VectorIndexType.INT4_HNSW, m, efConstruction, maxSearchEf);
16271665
// The default confidence interval for int4 is dynamic quantiles, this provides the best relevancy and is
16281666
// effectively required for int4 to behave well across a wide range of data.
16291667
this.confidenceInterval = confidenceInterval == null ? 0f : confidenceInterval;
@@ -1731,8 +1769,8 @@ boolean updatableTo(IndexOptions update) {
17311769
static class Int8HnswIndexOptions extends AbstractHnswIndexOptions {
17321770
private final Float confidenceInterval;
17331771

1734-
Int8HnswIndexOptions(int m, int efConstruction, Float confidenceInterval) {
1735-
super(VectorIndexType.INT8_HNSW, m, efConstruction);
1772+
Int8HnswIndexOptions(int m, int efConstruction, int maxSearchEf, Float confidenceInterval) {
1773+
super(VectorIndexType.INT8_HNSW, m, efConstruction, maxSearchEf);
17361774
this.confidenceInterval = confidenceInterval;
17371775
}
17381776

@@ -1793,8 +1831,8 @@ boolean updatableTo(IndexOptions update) {
17931831

17941832
static class HnswIndexOptions extends AbstractHnswIndexOptions {
17951833

1796-
HnswIndexOptions(int m, int efConstruction) {
1797-
super(VectorIndexType.HNSW, m, efConstruction);
1834+
HnswIndexOptions(int m, int efConstruction, int maxSearchEf) {
1835+
super(VectorIndexType.HNSW, m, efConstruction, maxSearchEf);
17981836
}
17991837

18001838
@Override
@@ -1826,8 +1864,8 @@ public XContentBuilder innerHnswXContent(XContentBuilder builder, Params params)
18261864

18271865
static class BBQHnswIndexOptions extends AbstractHnswIndexOptions {
18281866

1829-
BBQHnswIndexOptions(int m, int efConstruction) {
1830-
super(VectorIndexType.BBQ_HNSW, m, efConstruction);
1867+
BBQHnswIndexOptions(int m, int efConstruction, int maxSearchEf) {
1868+
super(VectorIndexType.BBQ_HNSW, m, efConstruction, maxSearchEf);
18311869
}
18321870

18331871
@Override
@@ -2036,6 +2074,9 @@ public Query createKnnQuery(
20362074
"to perform knn search on field [" + name() + "], its mapping must have [index] set to [true]"
20372075
);
20382076
}
2077+
2078+
indexOptions.validateNumCandidates(numCands);
2079+
20392080
return switch (getElementType()) {
20402081
case BYTE -> createKnnByteQuery(queryVector.asByteVector(), k, numCands, filter, similarityThreshold, parentFilter);
20412082
case FLOAT -> createKnnFloatQuery(
@@ -2137,7 +2178,7 @@ && isNotUnitVector(squaredMagnitude)) {
21372178
boolean rescore = needsRescore(oversample);
21382179
if (rescore) {
21392180
// Will get k * oversample for rescoring, and get the top k
2140-
adjustedK = Math.min((int) Math.ceil(k * oversample), OVERSAMPLE_LIMIT);
2181+
adjustedK = Math.min((int) Math.ceil(k * oversample), indexOptions.maxSearchEf());
21412182
numCands = Math.max(adjustedK, numCands);
21422183
}
21432184
Query knnQuery = parentFilter != null

‎server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -937,6 +937,7 @@ public void testMergeDims() throws IOException {
937937
.field("type", "int8_hnsw")
938938
.field("m", 16)
939939
.field("ef_construction", 100)
940+
.field("max_search_ef", DenseVectorFieldMapper.NUM_CANDS_LIMIT)
940941
.endObject();
941942
b.endObject();
942943
});

0 commit comments

Comments
 (0)