Skip to content

Commit 17e6c0e

Browse files
john-wagsterjohn-wagster
authored
[Preconditioning] Fixes for Tests That Reuse Writers (elastic#141372)
* fixed test failures related to precond usage in tests
1 parent 85e501a commit 17e6c0e

File tree

6 files changed

+64
-51
lines changed

6 files changed

+64
-51
lines changed

‎muted-tests.yml‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -327,12 +327,6 @@ tests:
327327
- class: org.elasticsearch.xpack.logsdb.StandardToLogsDbIndexModeRollingUpgradeIT
328328
method: testIndexing
329329
issue: https://github.com/elastic/elasticsearch/issues/141235
330-
- class: org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormatTests
331-
method: testSparseVectors
332-
issue: https://github.com/elastic/elasticsearch/issues/141270
333-
- class: org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQBFloat16VectorsFormatTests
334-
method: testSparseVectors
335-
issue: https://github.com/elastic/elasticsearch/issues/141271
336330
- class: org.elasticsearch.snapshots.SnapshotStatusApisIT
337331
method: testFailedSnapshot
338332
issue: https://github.com/elastic/elasticsearch/issues/141279

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

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import java.nio.ByteOrder;
3838
import java.util.Arrays;
3939
import java.util.List;
40+
import java.util.function.Consumer;
4041

4142
import static org.elasticsearch.index.codec.vectors.cluster.HierarchicalKMeans.NO_SOAR_ASSIGNMENT;
4243
import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.BULK_SIZE;
@@ -383,28 +384,31 @@ public CentroidSupplier createCentroidSupplier(
383384
}
384385

385386
@Override
386-
protected void inheritPreconditioner(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
387+
protected Preconditioner inheritPreconditioner(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
387388
// no-op
389+
return null;
388390
}
389391

390392
@Override
391-
protected void createPreconditioner(int dimension) {
393+
protected Preconditioner createPreconditioner(int dimension) {
392394
// no-op
395+
return null;
393396
}
394397

395398
@Override
396-
protected FloatVectorValues preconditionVectors(FloatVectorValues vectors) {
399+
protected FloatVectorValues preconditionVectors(Preconditioner Preconditioner, FloatVectorValues vectors) {
397400
// no-op
398401
return vectors;
399402
}
400403

401404
@Override
402-
protected void preconditionVectors(List<float[]> vectors) {
405+
protected Consumer<List<float[]>> preconditionVectors(Preconditioner preconditioner) {
403406
// no-op
407+
return (vectors) -> {};
404408
}
405409

406410
@Override
407-
protected void writePreconditioner(IndexOutput out) throws IOException {
411+
protected void writePreconditioner(Preconditioner Preconditioner, IndexOutput out) throws IOException {
408412
// no-op
409413
}
410414

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

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -191,23 +191,23 @@ public abstract CentroidSupplier createCentroidSupplier(
191191
float[] globalCentroid
192192
) throws IOException;
193193

194-
protected abstract void inheritPreconditioner(FieldInfo fieldInfo, MergeState mergeState) throws IOException;
194+
protected abstract Preconditioner inheritPreconditioner(FieldInfo fieldInfo, MergeState mergeState) throws IOException;
195195

196-
protected abstract void createPreconditioner(int dimension);
196+
protected abstract Preconditioner createPreconditioner(int dimension);
197197

198-
protected abstract void writePreconditioner(IndexOutput out) throws IOException;
198+
protected abstract void writePreconditioner(Preconditioner precondtioner, IndexOutput out) throws IOException;
199199

200-
protected abstract FloatVectorValues preconditionVectors(FloatVectorValues vectors);
200+
protected abstract FloatVectorValues preconditionVectors(Preconditioner precondtioner, FloatVectorValues vectors);
201201

202-
protected abstract void preconditionVectors(List<float[]> vectors);
202+
protected abstract Consumer<List<float[]>> preconditionVectors(Preconditioner preconditioner);
203203

204204
@Override
205205
public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
206206
rawVectorDelegate.flush(maxDoc, sortMap);
207207
for (FieldWriter fieldWriter : fieldWriters) {
208208
// build preconditioner if necessary, only need one given that this writer is tied to a format that has a fixed dim & block dim
209209
// write preconditioner subsequently in the centroids file
210-
createPreconditioner(fieldWriter.fieldInfo().getVectorDimension());
210+
Preconditioner preconditioner = createPreconditioner(fieldWriter.fieldInfo().getVectorDimension());
211211
if (fieldWriter.delegate == null) {
212212
// field is not float, we just write meta information
213213
writeMeta(fieldWriter.fieldInfo, 0, 0, 0, 0, 0, null, 0, 0);
@@ -218,7 +218,7 @@ public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
218218
fieldWriter.fieldInfo,
219219
fieldWriter.delegate,
220220
maxDoc,
221-
this::preconditionVectors
221+
preconditionVectors(preconditioner)
222222
);
223223

224224
// build centroids
@@ -253,7 +253,7 @@ public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
253253
);
254254
final long centroidLength = ivfCentroids.getFilePointer() - centroidOffset;
255255
long preconditionerOffset = ivfCentroids.getFilePointer();
256-
writePreconditioner(ivfCentroids);
256+
writePreconditioner(preconditioner, ivfCentroids);
257257
long preconditionerLength = ivfCentroids.getFilePointer() - preconditionerOffset;
258258
// write meta file
259259
writeMeta(
@@ -349,15 +349,16 @@ private void mergeOneFieldIVF(FieldInfo fieldInfo, MergeState mergeState) throws
349349
String docsFileName = null;
350350
// build a float vector values with random access. In order to do that we dump the vectors to
351351
// a temporary file and if the segment is not dense, the docs to another file/
352+
Preconditioner preconditioner;
352353
try (
353354
IndexOutput vectorsOut = mergeState.segmentInfo.dir.createTempOutput(mergeState.segmentInfo.name, "ivfvec_", IOContext.DEFAULT)
354355
) {
355356
tempRawVectorsFileName = vectorsOut.getName();
356357
FloatVectorValues mergedFloatVectorValues = MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState);
357358

358359
// TODO: we only want to write this once but we'll wind up doing it for every field with the same dim and blockdim
359-
inheritPreconditioner(fieldInfo, mergeState);
360-
mergedFloatVectorValues = preconditionVectors(mergedFloatVectorValues);
360+
preconditioner = inheritPreconditioner(fieldInfo, mergeState);
361+
mergedFloatVectorValues = preconditionVectors(preconditioner, mergedFloatVectorValues);
361362

362363
// if the segment is dense, we don't need to do anything with docIds.
363364
boolean dense = mergedFloatVectorValues.size() == mergeState.segmentInfo.maxDoc();
@@ -480,7 +481,7 @@ private void mergeOneFieldIVF(FieldInfo fieldInfo, MergeState mergeState) throws
480481
);
481482
centroidLength = ivfCentroids.getFilePointer() - centroidOffset;
482483
long preconditionerOffset = ivfCentroids.getFilePointer();
483-
writePreconditioner(ivfCentroids);
484+
writePreconditioner(preconditioner, ivfCentroids);
484485
long preconditionerLength = ivfCentroids.getFilePointer() - preconditionerOffset;
485486
// write meta
486487
writeMeta(

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ private static void randomFill(Random random, float[][] m) {
103103
}
104104

105105
private static float[][][] generateRandomOrthogonalMatrix(int dim, int blockDim, Random random) {
106-
blockDim = Math.min(dim, blockDim);
106+
assert blockDim <= dim;
107107
int nBlocks = dim / blockDim;
108108
int rem = dim % blockDim;
109109

@@ -127,8 +127,8 @@ private static float[][][] generateRandomOrthogonalMatrix(int dim, int blockDim,
127127
}
128128

129129
private static void matrixVectorMultiply(float[][] m, float[] x, float[] out) {
130-
assert m.length == x.length;
131130
assert m.length == out.length;
131+
assert m.length > 0 && m[0].length == x.length;
132132
int dim = out.length;
133133
// TODO: write Panama version of this to do all multiplications in one pass
134134
for (int i = 0; i < dim; i++) {
@@ -190,6 +190,10 @@ public void write(IndexOutput out) throws IOException {
190190
}
191191
}
192192

193+
// TODO: cache these preconditioners based on vectorDimension and blockDimension
194+
// need something thread safe and a way to clear the cache when done indexing (after flush or merge ... but that defeats the point)
195+
// maybe not possible or we limit it to a fixed number of cached preconditioners
196+
// maybe use setExpireAfterAccess in CacheBuilder; to be fair this code is not a hot path though
193197
public static Preconditioner createPreconditioner(int vectorDimension, int blockDimension) {
194198
if (blockDimension <= 0) {
195199
throw new IllegalArgumentException("block dimension must be positive but was [" + blockDimension + "]");
@@ -198,6 +202,7 @@ public static Preconditioner createPreconditioner(int vectorDimension, int block
198202
throw new IllegalArgumentException("vector dimension must be positive but was [" + vectorDimension + "]");
199203
}
200204
Random random = new Random(42L);
205+
blockDimension = Math.min(vectorDimension, blockDimension);
201206
float[][][] blocks = Preconditioner.generateRandomOrthogonalMatrix(vectorDimension, blockDimension, random);
202207
int[] dimBlocks = new int[blocks.length];
203208
for (int i = 0; i < blocks.length; i++) {
@@ -232,5 +237,4 @@ public static Preconditioner read(IndexInput input) throws IOException {
232237

233238
return new Preconditioner(blockDim, permutationMatrix, blocks);
234239
}
235-
236240
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,10 @@ protected FieldEntry doReadField(
213213
@Override
214214
public Preconditioner getPreconditioner(FieldInfo fieldInfo) throws IOException {
215215
final FieldEntry fieldEntry = fields.get(fieldInfo.number);
216+
// only seems possible in tests
217+
if (fieldEntry == null) {
218+
return null;
219+
}
216220
long preconditionerOffset = ((NextFieldEntry) fieldEntry).preconditionerOffset();
217221
long preconditionerLength = ((NextFieldEntry) fieldEntry).preconditionerLength();
218222
if (preconditionerLength > 0) {

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

Lines changed: 32 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
import java.nio.ByteOrder;
5050
import java.util.Arrays;
5151
import java.util.List;
52+
import java.util.function.Consumer;
5253
import java.util.function.IntUnaryOperator;
5354

5455
import static org.elasticsearch.index.codec.vectors.cluster.HierarchicalKMeans.NO_SOAR_ASSIGNMENT;
@@ -69,7 +70,6 @@ public class ESNextDiskBBQVectorsWriter extends IVFVectorsWriter {
6970
private final int numMergeWorkers;
7071
private final int blockDimension;
7172
private final boolean doPrecondition;
72-
private Preconditioner preconditioner;
7373

7474
public ESNextDiskBBQVectorsWriter(
7575
SegmentWriteState state,
@@ -95,52 +95,58 @@ public ESNextDiskBBQVectorsWriter(
9595
}
9696

9797
@Override
98-
protected void inheritPreconditioner(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
99-
if (doPrecondition && this.preconditioner == null) {
98+
protected Preconditioner inheritPreconditioner(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
99+
if (doPrecondition) {
100100
for (KnnVectorsReader reader : mergeState.knnVectorsReaders) {
101101
if (reader instanceof VectorPreconditioner) {
102-
this.preconditioner = ((VectorPreconditioner) reader).getPreconditioner(fieldInfo);
103-
break;
102+
Preconditioner preconditioner = ((VectorPreconditioner) reader).getPreconditioner(fieldInfo);
103+
if (preconditioner != null) {
104+
return preconditioner;
105+
}
104106
}
105107
}
106-
if (this.preconditioner == null) {
107-
createPreconditioner(fieldInfo.getVectorDimension());
108-
}
108+
// else
109+
return createPreconditioner(fieldInfo.getVectorDimension());
109110
}
111+
return null;
110112
}
111113

112114
@Override
113-
protected void createPreconditioner(int dimension) {
114-
if (doPrecondition && this.preconditioner == null) {
115-
this.preconditioner = Preconditioner.createPreconditioner(dimension, blockDimension);
115+
protected Preconditioner createPreconditioner(int dimension) {
116+
if (doPrecondition) {
117+
return Preconditioner.createPreconditioner(dimension, blockDimension);
118+
} else {
119+
return null;
116120
}
117121
}
118122

119123
@Override
120-
protected void writePreconditioner(IndexOutput out) throws IOException {
124+
protected void writePreconditioner(Preconditioner preconditioner, IndexOutput out) throws IOException {
121125
if (preconditioner != null) {
122126
preconditioner.write(out);
123127
}
124128
}
125129

126130
@Override
127-
protected void preconditionVectors(List<float[]> vectors) {
128-
if (doPrecondition == false || vectors.isEmpty()) {
129-
return;
130-
}
131-
if (preconditioner == null) {
132-
throw new IllegalStateException("preconditioner was not created but should be first");
133-
}
134-
float[] out = new float[vectors.getFirst().length];
135-
for (int i = 0; i < vectors.size(); i++) {
136-
float[] vector = vectors.get(i);
137-
preconditioner.applyTransform(vector, out);
138-
System.arraycopy(out, 0, vector, 0, vector.length);
139-
}
131+
protected Consumer<List<float[]>> preconditionVectors(Preconditioner preconditioner) {
132+
return (vectors) -> {
133+
if (doPrecondition == false || vectors.isEmpty()) {
134+
return;
135+
}
136+
if (preconditioner == null) {
137+
throw new IllegalStateException("preconditioner was not created but should be first");
138+
}
139+
float[] out = new float[vectors.getFirst().length];
140+
for (int i = 0; i < vectors.size(); i++) {
141+
float[] vector = vectors.get(i);
142+
preconditioner.applyTransform(vector, out);
143+
System.arraycopy(out, 0, vector, 0, vector.length);
144+
}
145+
};
140146
}
141147

142148
@Override
143-
protected FloatVectorValues preconditionVectors(FloatVectorValues vectors) {
149+
protected FloatVectorValues preconditionVectors(Preconditioner preconditioner, FloatVectorValues vectors) {
144150
if (doPrecondition == false) {
145151
return vectors;
146152
}

0 commit comments

Comments
 (0)