Added new exception type to signify expected warmup behavior (#3070)

Signed-off-by: Andrew Klepchick <aklepchi@amazon.com>
(cherry picked from commit 29420a7)

Author: MrFlap

CHANGELOG.md

@@ -7,3 +7,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 3.2](https://github.com/opensearch-project/k-NN/compare/main...HEAD)
 * Memory optimized search warmup ([#2954](https://github.com/opensearch-project/k-NN/pull/2954))
 * Changed warmup seek to use long instead of int to avoid overflow [#3067](https://github.com/opensearch-project/k-NN/pull/3067)
+* 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

@@ -28,7 +28,6 @@
 import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
 import org.opensearch.knn.index.memory.NativeMemoryEntryContext;
 import org.opensearch.knn.index.memory.NativeMemoryLoadStrategy;
-import org.opensearch.knn.index.query.SegmentLevelQuantizationInfo;
 import org.opensearch.knn.index.warmup.MemoryOptimizedSearchWarmup;
 
 import java.io.IOException;
@@ -90,7 +89,6 @@ public void warmup() throws IOException {
 
         final MapperService mapperService = indexShard.mapperService();
         final String indexName = indexShard.shardId().getIndexName();
-        final Directory directory = indexShard.store().directory();
 
         try (Engine.Searcher searcher = indexShard.acquireSearcher("knn-warmup-mem")) {
             final Directory directory = indexShard.store().directory();

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

@@ -36,6 +36,7 @@
 import org.opensearch.knn.index.quantizationservice.QuantizationService;
 import org.opensearch.knn.memoryoptsearch.VectorSearcher;
 import org.opensearch.knn.memoryoptsearch.VectorSearcherFactory;
+import org.opensearch.knn.memoryoptsearch.faiss.FaissMemoryOptimizedSearcher;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationState;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCacheManager;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateReadConfig;
@@ -245,6 +246,10 @@ private boolean trySearchWithMemoryOptimizedSearch(
         // Try with memory optimized searcher
         final VectorSearcher memoryOptimizedSearcher = loadMemoryOptimizedSearcherIfRequired(field);
 
+        if (target == null) {
+            throw new FaissMemoryOptimizedSearcher.WarmupInitializationException("Null vector supplied for warmup");
+        }
+
         if (memoryOptimizedSearcher != null) {
             if (isFloatVector) {
                 memoryOptimizedSearcher.search((float[]) target, knnCollector, acceptDocs);

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

@@ -19,12 +19,12 @@
 import org.opensearch.common.lucene.Lucene;
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.mapper.MapperService;
-import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.knn.index.codec.util.KNNCodecUtil;
 import org.opensearch.knn.index.engine.KNNEngine;
 import org.opensearch.knn.index.engine.MemoryOptimizedSearchSupportSpec;
 import org.opensearch.knn.index.mapper.KNNVectorFieldMapper;
 import org.opensearch.knn.index.mapper.KNNVectorFieldType;
+import org.opensearch.knn.memoryoptsearch.faiss.FaissMemoryOptimizedSearcher;
 
 import java.io.IOException;
 import java.nio.file.Path;
@@ -33,7 +33,6 @@
 import java.util.List;
 
 import static org.opensearch.knn.common.FieldInfoExtractor.extractKNNEngine;
-import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD;
 
 @Log4j2
 public class MemoryOptimizedSearchWarmup {
@@ -96,18 +95,10 @@ private boolean warmUpField(FieldInfo field, SegmentReader segmentReader, Direct
         }
 
         try {
-            final String dataTypeStr = field.getAttribute(VECTOR_DATA_TYPE_FIELD);
-            if (dataTypeStr == null) {
-                return true;
-            }
-            final VectorDataType vectorDataType = VectorDataType.get(dataTypeStr);
-            // We expect a NPE here:
-            if (vectorDataType == VectorDataType.FLOAT) {
-                segmentReader.getVectorReader().search(field.getName(), (float[]) null, null, null);
-            } else {
-                segmentReader.getVectorReader().search(field.getName(), (byte[]) null, null, null);
-            }
-        } catch (NullPointerException ignored) {}
+            segmentReader.getVectorReader().search(field.getName(), (float[]) null, null, null); // codecov[ignore]
+        } catch (FaissMemoryOptimizedSearcher.WarmupInitializationException ignored) {
+            // Expected during warmup initialization
+        }
 
         return true;
     }

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

@@ -28,6 +28,18 @@
  * This searcher directly reads FAISS index file via the provided {@link IndexInput} then perform vector search on it.
  */
 public class FaissMemoryOptimizedSearcher implements VectorSearcher {
+
+    /**
+     * Exception thrown during warmup initialization when the search cannot proceed.
+     * This encapsulates expected exceptions like NullPointerException (from null target vectors)
+     * and UnsupportedOperationException (from vector encoding mismatches with quantized indices).
+     */
+    public static class WarmupInitializationException extends RuntimeException {
+        public WarmupInitializationException(String message) {
+            super(message);
+        }
+    }
+
     private final IndexInput indexInput;
     private final FaissIndex faissIndex;
     private final FlatVectorsScorer flatVectorsScorer;