Skip to content

Commit a47daee

Browse files
benwtrentbenwtrent
authored
Upgrade DiskBBQ algorithm (#143760)
This updates our diskbbq algorithm and format. - It now provides 3x or more better search performance on very restrictive filters (prefilters on centroids) - Provides a way to condition non iid vectors (expert API for now) - Gives more bit options (1, 2, 4, and 7 bits!) - More native code improvements for overall performance
1 parent eee273c commit a47daee

21 files changed

+3287
-137
lines changed

‎docs/changelog/143760.yaml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
area: Vector Search
2+
highlight:
3+
body: |-
4+
This updates our diskbbq algorithm and format.
5+
6+
- It now provides 3x or more better search performance on very restrictive filters (prefilters on centroids)
7+
- Provides a way to condition non iid vectors (expert API for now)
8+
- Gives more bit options (1, 2, 4, and 7 bits!)
9+
- More native code improvements for overall performance
10+
notable: true
11+
title: Upgrade DiskBBQ algorithm
12+
issues: []
13+
pr: 143760
14+
summary: Upgrade DiskBBQ algorithm
15+
type: enhancement

‎server/src/main/java/module-info.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -472,6 +472,7 @@
472472
org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat,
473473
org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat,
474474
org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat,
475+
org.elasticsearch.index.codec.vectors.diskbbq.es94.ES940DiskBBQVectorsFormat,
475476
org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat,
476477
org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat,
477478
org.elasticsearch.index.codec.vectors.es93.ES93HnswVectorsFormat,
@@ -513,6 +514,7 @@
513514
exports org.elasticsearch.inference.telemetry;
514515
exports org.elasticsearch.index.codec.vectors.diskbbq to org.elasticsearch.test.knn, org.elasticsearch.xpack.diskbbq;
515516
exports org.elasticsearch.index.codec.vectors.diskbbq.next to org.elasticsearch.test.knn, org.elasticsearch.xpack.diskbbq;
517+
exports org.elasticsearch.index.codec.vectors.diskbbq.es94 to org.elasticsearch.test.knn, org.elasticsearch.xpack.diskbbq;
516518
exports org.elasticsearch.index.codec.vectors.cluster to org.elasticsearch.test.knn;
517519
exports org.elasticsearch.index.codec.vectors.es93 to org.elasticsearch.test.knn;
518520
exports org.elasticsearch.index.codec.vectors.es94 to org.elasticsearch.test.knn;

‎server/src/main/java/org/elasticsearch/index/IndexVersions.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,7 @@ private static Version parseUnchecked(String version) {
226226
public static final IndexVersion STORE_IGNORED_MALFORMED_IN_BINARY_DOC_VALUES = def(9_073_0_00, Version.LUCENE_10_3_2);
227227
public static final IndexVersion DISABLE_SEQUENCE_NUMBERS = def(9_074_0_00, Version.LUCENE_10_3_2);
228228
public static final IndexVersion UPGRADE_TO_LUCENE_10_4_0 = def(9_075_00_0, Version.LUCENE_10_4_0);
229+
public static final IndexVersion UPGRADE_DISKBBQ_ES940 = def(9_076_00_0, Version.LUCENE_10_4_0);
229230

230231
/*
231232
* STOP! READ THIS FIRST! No, really,

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

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
5757
// offsets contained in cen_ivf, [vector ordinals, actually just docIds](long varint), quantized
5858
// vectors (OSQ bit)
5959
public static final String CLUSTER_EXTENSION = "clivf";
60-
static final String IVF_META_EXTENSION = "mivf";
60+
public static final String IVF_META_EXTENSION = "mivf";
6161

6262
public static final int VERSION_START = 0;
6363
public static final int VERSION_DIRECT_IO = 1;
@@ -193,10 +193,12 @@ public ES920DiskBBQVectorsFormat() {
193193

194194
@Override
195195
public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
196+
final Boolean writeDirectIOReads = Boolean.valueOf(useDirectIO);
197+
validateDirectIOMetadata(writeDirectIOReads, VERSION_CURRENT);
196198
return new ES920DiskBBQVectorsWriter(
197199
state,
198200
rawVectorFormat.getName(),
199-
useDirectIO,
201+
writeDirectIOReads,
200202
rawVectorFormat.fieldsWriter(state),
201203
vectorPerCluster,
202204
centroidsPerParentCluster,
@@ -207,7 +209,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
207209
}
208210

209211
// for testing
210-
KnnVectorsWriter version0FieldsWriter(SegmentWriteState state) throws IOException {
212+
protected KnnVectorsWriter version0FieldsWriter(SegmentWriteState state) throws IOException {
213+
validateDirectIOMetadata(null, VERSION_START);
211214
return new ES920DiskBBQVectorsWriter(
212215
state,
213216
rawVectorFormat.getName(),
@@ -222,6 +225,16 @@ KnnVectorsWriter version0FieldsWriter(SegmentWriteState state) throws IOExceptio
222225
);
223226
}
224227

228+
private static void validateDirectIOMetadata(Boolean useDirectIOReads, int writeVersion) {
229+
final boolean shouldWriteDirectIOReads = writeVersion >= VERSION_DIRECT_IO;
230+
if (shouldWriteDirectIOReads && useDirectIOReads == null) {
231+
throw new IllegalArgumentException("useDirectIOReads must be provided when writing direct-IO metadata");
232+
}
233+
if (shouldWriteDirectIOReads == false && useDirectIOReads != null) {
234+
throw new IllegalArgumentException("useDirectIOReads must be null when direct-IO metadata is not written");
235+
}
236+
}
237+
225238
@Override
226239
public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
227240
return new ES920DiskBBQVectorsReader(state, (f, dio) -> {

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,18 @@ public class ES920DiskBBQVectorsReader extends IVFVectorsReader {
4444
private static final byte QUERY_BITS = 4;
4545

4646
ES920DiskBBQVectorsReader(SegmentReadState state, GenericFlatVectorReaders.LoadFlatVectorsReader getFormatReader) throws IOException {
47-
super(state, getFormatReader);
47+
super(
48+
state,
49+
getFormatReader,
50+
ES920DiskBBQVectorsFormat.NAME,
51+
ES920DiskBBQVectorsFormat.CENTROID_EXTENSION,
52+
ES920DiskBBQVectorsFormat.CLUSTER_EXTENSION,
53+
ES920DiskBBQVectorsFormat.IVF_META_EXTENSION,
54+
ES920DiskBBQVectorsFormat.VERSION_START,
55+
ES920DiskBBQVectorsFormat.VERSION_CURRENT,
56+
ES920DiskBBQVectorsFormat.VERSION_DIRECT_IO,
57+
ES920DiskBBQVectorsFormat.DYNAMIC_VISIT_RATIO
58+
);
4859
}
4960

5061
public CentroidIterator getPostingListPrefetchIterator(CentroidIterator centroidIterator, IndexInput postingListSlice)

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,19 @@ public ES920DiskBBQVectorsWriter(
9191
int numMergeWorkers,
9292
int flatVectorThreshold
9393
) throws IOException {
94-
super(state, rawVectorFormatName, useDirectIOReads, rawVectorDelegate, writeVersion, flatVectorThreshold);
94+
super(
95+
state,
96+
rawVectorFormatName,
97+
useDirectIOReads,
98+
rawVectorDelegate,
99+
writeVersion,
100+
ES920DiskBBQVectorsFormat.NAME,
101+
ES920DiskBBQVectorsFormat.IVF_META_EXTENSION,
102+
ES920DiskBBQVectorsFormat.CENTROID_EXTENSION,
103+
ES920DiskBBQVectorsFormat.CLUSTER_EXTENSION,
104+
writeVersion >= ES920DiskBBQVectorsFormat.VERSION_DIRECT_IO,
105+
flatVectorThreshold
106+
);
95107
this.vectorPerCluster = vectorPerCluster;
96108
this.centroidsPerParentCluster = centroidsPerParentCluster;
97109
this.mergeExec = mergeExec;

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

Lines changed: 35 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -42,11 +42,6 @@
4242
import java.util.Map;
4343

4444
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS;
45-
import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.CENTROID_EXTENSION;
46-
import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.CLUSTER_EXTENSION;
47-
import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.DYNAMIC_VISIT_RATIO;
48-
import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.IVF_META_EXTENSION;
49-
import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.VERSION_DIRECT_IO;
5045

5146
/**
5247
* Reader for IVF vectors. This reader is used to read the IVF vectors from the index.
@@ -58,24 +53,43 @@ public abstract class IVFVectorsReader extends KnnVectorsReader {
5853
private final FieldInfos fieldInfos;
5954
protected final IntObjectHashMap<FieldEntry> fields;
6055
private final GenericFlatVectorReaders genericReaders;
56+
private final String centroidExtension;
57+
private final String clusterExtension;
58+
private final int versionDirectIo;
59+
private final float dynamicVisitRatio;
6160

6261
@SuppressWarnings("this-escape")
63-
protected IVFVectorsReader(SegmentReadState state, GenericFlatVectorReaders.LoadFlatVectorsReader loadReader) throws IOException {
62+
protected IVFVectorsReader(
63+
SegmentReadState state,
64+
GenericFlatVectorReaders.LoadFlatVectorsReader loadReader,
65+
String codecName,
66+
String centroidExtension,
67+
String clusterExtension,
68+
String metaExtension,
69+
int versionStart,
70+
int versionCurrent,
71+
int versionDirectIo,
72+
float dynamicVisitRatio
73+
) throws IOException {
6474
this.state = state;
6575
this.fieldInfos = state.fieldInfos;
6676
this.fields = new IntObjectHashMap<>();
6777
this.genericReaders = new GenericFlatVectorReaders();
68-
String meta = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, IVF_META_EXTENSION);
78+
this.centroidExtension = centroidExtension;
79+
this.clusterExtension = clusterExtension;
80+
this.versionDirectIo = versionDirectIo;
81+
this.dynamicVisitRatio = dynamicVisitRatio;
82+
String meta = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, metaExtension);
6983

7084
int versionMeta = -1;
7185
try (ChecksumIndexInput ivfMeta = state.directory.openChecksumInput(meta)) {
7286
Throwable priorE = null;
7387
try {
7488
versionMeta = CodecUtil.checkIndexHeader(
7589
ivfMeta,
76-
ES920DiskBBQVectorsFormat.NAME,
77-
ES920DiskBBQVectorsFormat.VERSION_START,
78-
ES920DiskBBQVectorsFormat.VERSION_CURRENT,
90+
codecName,
91+
versionStart,
92+
versionCurrent,
7993
state.segmentInfo.getId(),
8094
state.segmentSuffix
8195
);
@@ -85,8 +99,8 @@ protected IVFVectorsReader(SegmentReadState state, GenericFlatVectorReaders.Load
8599
} finally {
86100
CodecUtil.checkFooter(ivfMeta, priorE);
87101
}
88-
ivfCentroids = openDataInput(state, versionMeta, CENTROID_EXTENSION, ES920DiskBBQVectorsFormat.NAME, state.context);
89-
ivfClusters = openDataInput(state, versionMeta, CLUSTER_EXTENSION, ES920DiskBBQVectorsFormat.NAME, state.context);
102+
ivfCentroids = openDataInput(state, versionMeta, centroidExtension, codecName, versionStart, versionCurrent, state.context);
103+
ivfClusters = openDataInput(state, versionMeta, clusterExtension, codecName, versionStart, versionCurrent, state.context);
90104
} catch (Throwable t) {
91105
IOUtils.closeWhileHandlingException(this);
92106
throw t;
@@ -110,6 +124,8 @@ protected static IndexInput openDataInput(
110124
int versionMeta,
111125
String fileExtension,
112126
String codecName,
127+
int versionStart,
128+
int versionCurrent,
113129
IOContext context
114130
) throws IOException {
115131
final String fileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, fileExtension);
@@ -118,8 +134,8 @@ protected static IndexInput openDataInput(
118134
final int versionVectorData = CodecUtil.checkIndexHeader(
119135
in,
120136
codecName,
121-
ES920DiskBBQVectorsFormat.VERSION_START,
122-
ES920DiskBBQVectorsFormat.VERSION_CURRENT,
137+
versionStart,
138+
versionCurrent,
123139
state.segmentInfo.getId(),
124140
state.segmentSuffix
125141
);
@@ -158,7 +174,7 @@ private void readFields(
158174

159175
private FieldEntry readField(IndexInput input, FieldInfo info, int versionMeta) throws IOException {
160176
final String rawVectorFormat = input.readString();
161-
final boolean useDirectIOReads = versionMeta >= VERSION_DIRECT_IO && input.readByte() == 1;
177+
final boolean useDirectIOReads = versionMeta >= versionDirectIo && input.readByte() == 1;
162178
final VectorEncoding vectorEncoding = readVectorEncoding(input);
163179
final VectorSimilarityFunction similarityFunction = readSimilarityFunction(input);
164180
if (similarityFunction != info.getVectorSimilarityFunction()) {
@@ -283,14 +299,14 @@ public final void search(String field, float[] target, KnnCollector knnCollector
283299
: esAcceptDocs instanceof ESAcceptDocs.ESAcceptDocsAll ? numVectors
284300
: esAcceptDocs.approximateCost());
285301
float percentFiltered = Math.max(0f, Math.min(1f, approximateCost / numVectors));
286-
float visitRatio = DYNAMIC_VISIT_RATIO;
302+
float visitRatio = dynamicVisitRatio;
287303
// Search strategy may be null if this is being called from checkIndex (e.g. from a test)
288304
if (knnCollector.getSearchStrategy() instanceof IVFKnnSearchStrategy ivfSearchStrategy) {
289305
visitRatio = ivfSearchStrategy.getVisitRatio();
290306
}
291307

292308
FieldEntry entry = fields.get(fieldInfo.number);
293-
if (visitRatio == DYNAMIC_VISIT_RATIO) {
309+
if (visitRatio == dynamicVisitRatio) {
294310
// empirically based, and a good dynamic to get decent recall while scaling a la "efSearch"
295311
// scaling by the number of vectors vs. the nearest neighbors requested
296312
// not perfect, but a comparative heuristic.
@@ -369,7 +385,7 @@ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
369385
return raw;
370386
}
371387

372-
var centroidsClusters = Map.of(CENTROID_EXTENSION, fe.centroidLength, CLUSTER_EXTENSION, fe.postingListLength);
388+
var centroidsClusters = Map.of(centroidExtension, fe.centroidLength, clusterExtension, fe.postingListLength);
373389
return KnnVectorsReader.mergeOffHeapByteSizeMaps(raw, centroidsClusters);
374390
}
375391

@@ -394,7 +410,7 @@ protected static class FieldEntry implements GenericFlatVectorReaders.Field {
394410
protected final float globalCentroidDp;
395411
protected final int bulkSize;
396412

397-
protected FieldEntry(
413+
public FieldEntry(
398414
String rawVectorFormatName,
399415
boolean useDirectIOReads,
400416
VectorSimilarityFunction similarityFunction,

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

Lines changed: 15 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,10 @@ public abstract class IVFVectorsWriter extends KnnVectorsWriter {
5353
private final IndexOutput ivfCentroids, ivfClusters;
5454
private final IndexOutput ivfMeta;
5555
private final String rawVectorFormatName;
56-
private final int writeVersion;
5756
private final Boolean useDirectIOReads;
5857
private final FlatVectorsWriter rawVectorDelegate;
5958
private final int flatVectorThreshold;
59+
private final boolean shouldWriteDirectIoReads;
6060

6161
@SuppressWarnings("this-escape")
6262
protected IVFVectorsWriter(
@@ -65,58 +65,28 @@ protected IVFVectorsWriter(
6565
Boolean useDirectIOReads,
6666
FlatVectorsWriter rawVectorDelegate,
6767
int writeVersion,
68+
String codecName,
69+
String metaExtension,
70+
String centroidExtension,
71+
String clusterExtension,
72+
boolean shouldWriteDirectIoReads,
6873
int flatVectorThreshold
6974
) throws IOException {
70-
// if version >= VERSION_DIRECT_IO, useDirectIOReads should have a value
71-
if ((writeVersion >= ES920DiskBBQVectorsFormat.VERSION_DIRECT_IO) == (useDirectIOReads == null)) throw new IllegalArgumentException(
72-
"Write version " + writeVersion + " does not match direct IO value " + useDirectIOReads
73-
);
74-
7575
this.rawVectorFormatName = rawVectorFormatName;
76-
this.writeVersion = writeVersion;
7776
this.useDirectIOReads = useDirectIOReads;
7877
this.rawVectorDelegate = rawVectorDelegate;
7978
this.flatVectorThreshold = flatVectorThreshold;
80-
final String metaFileName = IndexFileNames.segmentFileName(
81-
state.segmentInfo.name,
82-
state.segmentSuffix,
83-
ES920DiskBBQVectorsFormat.IVF_META_EXTENSION
84-
);
85-
final String ivfCentroidsFileName = IndexFileNames.segmentFileName(
86-
state.segmentInfo.name,
87-
state.segmentSuffix,
88-
ES920DiskBBQVectorsFormat.CENTROID_EXTENSION
89-
);
90-
final String ivfClustersFileName = IndexFileNames.segmentFileName(
91-
state.segmentInfo.name,
92-
state.segmentSuffix,
93-
ES920DiskBBQVectorsFormat.CLUSTER_EXTENSION
94-
);
79+
this.shouldWriteDirectIoReads = shouldWriteDirectIoReads;
80+
final String metaFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, metaExtension);
81+
final String ivfCentroidsFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, centroidExtension);
82+
final String ivfClustersFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, clusterExtension);
9583
try {
9684
ivfMeta = state.directory.createOutput(metaFileName, state.context);
97-
CodecUtil.writeIndexHeader(
98-
ivfMeta,
99-
ES920DiskBBQVectorsFormat.NAME,
100-
writeVersion,
101-
state.segmentInfo.getId(),
102-
state.segmentSuffix
103-
);
85+
CodecUtil.writeIndexHeader(ivfMeta, codecName, writeVersion, state.segmentInfo.getId(), state.segmentSuffix);
10486
ivfCentroids = state.directory.createOutput(ivfCentroidsFileName, state.context);
105-
CodecUtil.writeIndexHeader(
106-
ivfCentroids,
107-
ES920DiskBBQVectorsFormat.NAME,
108-
writeVersion,
109-
state.segmentInfo.getId(),
110-
state.segmentSuffix
111-
);
87+
CodecUtil.writeIndexHeader(ivfCentroids, codecName, writeVersion, state.segmentInfo.getId(), state.segmentSuffix);
11288
ivfClusters = state.directory.createOutput(ivfClustersFileName, state.context);
113-
CodecUtil.writeIndexHeader(
114-
ivfClusters,
115-
ES920DiskBBQVectorsFormat.NAME,
116-
writeVersion,
117-
state.segmentInfo.getId(),
118-
state.segmentSuffix
119-
);
89+
CodecUtil.writeIndexHeader(ivfClusters, codecName, writeVersion, state.segmentInfo.getId(), state.segmentSuffix);
12090
} catch (Throwable t) {
12191
IOUtils.closeWhileHandlingException(this);
12292
throw t;
@@ -360,7 +330,8 @@ private void writeMeta(
360330
) throws IOException {
361331
ivfMeta.writeInt(field.number);
362332
ivfMeta.writeString(rawVectorFormatName);
363-
if (writeVersion >= ES920DiskBBQVectorsFormat.VERSION_DIRECT_IO) {
333+
if (shouldWriteDirectIoReads) {
334+
assert useDirectIOReads != null : "shouldWriteDirectIoReads is true but useDirectIOReads is null";
364335
ivfMeta.writeByte(useDirectIOReads ? (byte) 1 : 0);
365336
}
366337
ivfMeta.writeInt(field.getVectorEncoding().ordinal());

0 commit comments

Comments
 (0)