Skip to content

Commit 881e79c

Browse files
committed
Added new exception type to signify expected warmup behavior (#3070)
Signed-off-by: Andrew Klepchick <aklepchi@amazon.com> (cherry picked from commit 29420a7)
1 parent 19f6d9a commit 881e79c

File tree

7 files changed

+703
-319
lines changed

7 files changed

+703
-319
lines changed

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,3 +7,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
77
## [Unreleased 3.2](https://github.com/opensearch-project/k-NN/compare/main...HEAD)
88
* Memory optimized search warmup ([#2954](https://github.com/opensearch-project/k-NN/pull/2954))
99
* Changed warmup seek to use long instead of int to avoid overflow [#3067](https://github.com/opensearch-project/k-NN/pull/3067)
10+
* Fix memory-optimized-search reentrant search bug in byte index. [#3071](https://github.com/opensearch-project/k-NN/pull/3071)

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
2929
import org.opensearch.knn.index.memory.NativeMemoryEntryContext;
3030
import org.opensearch.knn.index.memory.NativeMemoryLoadStrategy;
31-
import org.opensearch.knn.index.query.SegmentLevelQuantizationInfo;
3231
import org.opensearch.knn.index.warmup.MemoryOptimizedSearchWarmup;
3332

3433
import java.io.IOException;
@@ -90,7 +89,6 @@ public void warmup() throws IOException {
9089

9190
final MapperService mapperService = indexShard.mapperService();
9291
final String indexName = indexShard.shardId().getIndexName();
93-
final Directory directory = indexShard.store().directory();
9492

9593
try (Engine.Searcher searcher = indexShard.acquireSearcher("knn-warmup-mem")) {
9694
final Directory directory = indexShard.store().directory();

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsReader.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.opensearch.knn.index.quantizationservice.QuantizationService;
3737
import org.opensearch.knn.memoryoptsearch.VectorSearcher;
3838
import org.opensearch.knn.memoryoptsearch.VectorSearcherFactory;
39+
import org.opensearch.knn.memoryoptsearch.faiss.FaissMemoryOptimizedSearcher;
3940
import org.opensearch.knn.quantization.models.quantizationState.QuantizationState;
4041
import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCacheManager;
4142
import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateReadConfig;
@@ -245,6 +246,10 @@ private boolean trySearchWithMemoryOptimizedSearch(
245246
// Try with memory optimized searcher
246247
final VectorSearcher memoryOptimizedSearcher = loadMemoryOptimizedSearcherIfRequired(field);
247248

249+
if (target == null) {
250+
throw new FaissMemoryOptimizedSearcher.WarmupInitializationException("Null vector supplied for warmup");
251+
}
252+
248253
if (memoryOptimizedSearcher != null) {
249254
if (isFloatVector) {
250255
memoryOptimizedSearcher.search((float[]) target, knnCollector, acceptDocs);

src/main/java/org/opensearch/knn/index/warmup/MemoryOptimizedSearchWarmup.java

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,12 @@
1919
import org.opensearch.common.lucene.Lucene;
2020
import org.opensearch.index.mapper.MappedFieldType;
2121
import org.opensearch.index.mapper.MapperService;
22-
import org.opensearch.knn.index.VectorDataType;
2322
import org.opensearch.knn.index.codec.util.KNNCodecUtil;
2423
import org.opensearch.knn.index.engine.KNNEngine;
2524
import org.opensearch.knn.index.engine.MemoryOptimizedSearchSupportSpec;
2625
import org.opensearch.knn.index.mapper.KNNVectorFieldMapper;
2726
import org.opensearch.knn.index.mapper.KNNVectorFieldType;
27+
import org.opensearch.knn.memoryoptsearch.faiss.FaissMemoryOptimizedSearcher;
2828

2929
import java.io.IOException;
3030
import java.nio.file.Path;
@@ -33,7 +33,6 @@
3333
import java.util.List;
3434

3535
import static org.opensearch.knn.common.FieldInfoExtractor.extractKNNEngine;
36-
import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD;
3736

3837
@Log4j2
3938
public class MemoryOptimizedSearchWarmup {
@@ -96,18 +95,10 @@ private boolean warmUpField(FieldInfo field, SegmentReader segmentReader, Direct
9695
}
9796

9897
try {
99-
final String dataTypeStr = field.getAttribute(VECTOR_DATA_TYPE_FIELD);
100-
if (dataTypeStr == null) {
101-
return true;
102-
}
103-
final VectorDataType vectorDataType = VectorDataType.get(dataTypeStr);
104-
// We expect a NPE here:
105-
if (vectorDataType == VectorDataType.FLOAT) {
106-
segmentReader.getVectorReader().search(field.getName(), (float[]) null, null, null);
107-
} else {
108-
segmentReader.getVectorReader().search(field.getName(), (byte[]) null, null, null);
109-
}
110-
} catch (NullPointerException ignored) {}
98+
segmentReader.getVectorReader().search(field.getName(), (float[]) null, null, null); // codecov[ignore]
99+
} catch (FaissMemoryOptimizedSearcher.WarmupInitializationException ignored) {
100+
// Expected during warmup initialization
101+
}
111102

112103
return true;
113104
}

src/main/java/org/opensearch/knn/memoryoptsearch/faiss/FaissMemoryOptimizedSearcher.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,18 @@
2828
* This searcher directly reads FAISS index file via the provided {@link IndexInput} then perform vector search on it.
2929
*/
3030
public class FaissMemoryOptimizedSearcher implements VectorSearcher {
31+
32+
/**
33+
* Exception thrown during warmup initialization when the search cannot proceed.
34+
* This encapsulates expected exceptions like NullPointerException (from null target vectors)
35+
* and UnsupportedOperationException (from vector encoding mismatches with quantized indices).
36+
*/
37+
public static class WarmupInitializationException extends RuntimeException {
38+
public WarmupInitializationException(String message) {
39+
super(message);
40+
}
41+
}
42+
3143
private final IndexInput indexInput;
3244
private final FaissIndex faissIndex;
3345
private final FlatVectorsScorer flatVectorsScorer;

0 commit comments

Comments
 (0)