@@ -59,85 +59,27 @@ public class NativeEngines990KnnVectorsReader extends KnnVectorsReader {
5959
6060 private final FlatVectorsReader flatVectorsReader ;
6161 private Map <String , String > quantizationStateCacheKeyPerField ;
62- private SegmentReadState segmentReadState ;
62+ private final SegmentReadState segmentReadState ;
6363 private final List <String > cacheKeys ;
6464 private volatile Map <String , VectorSearcher > vectorSearchers ;
6565
6666 public NativeEngines990KnnVectorsReader (final SegmentReadState state , final FlatVectorsReader flatVectorsReader ) {
67+ this (state , flatVectorsReader , false );
68+ }
69+
70+ public NativeEngines990KnnVectorsReader (
71+ final SegmentReadState state ,
72+ final FlatVectorsReader flatVectorsReader ,
73+ final boolean memoryOptimizedSearchEnabled
74+ ) {
6775 this .flatVectorsReader = flatVectorsReader ;
6876 this .segmentReadState = state ;
6977 this .cacheKeys = getVectorCacheKeysFromSegmentReaderState (state );
70- this .vectorSearchers = null ;
7178
7279 loadCacheKeyMap ();
7380
74- // Memory optimized searcher will be ONLY used for searching, we don't need this during merging.
75- // TODO(KDY) : Enable this based on setting flag value. e.g. index.knn.memory_optimized_search = True
76- if (state .context .context () != IOContext .Context .MERGE ) {
77- loadMemoryOptimizedSearcher ();
78- }
79- }
80-
81- private IOSupplier <VectorSearcher > getIndexFileNameIfMemoryOptimizedSearchSupported (final FieldInfo fieldInfo ) {
82- // Skip non-knn fields.
83- final Map <String , String > attributes = fieldInfo .attributes ();
84- if (attributes == null || attributes .containsKey (KNN_FIELD ) == false ) {
85- return null ;
86- }
87-
88- // Get engine
89- final String engineName = attributes .getOrDefault (KNN_ENGINE , KNNEngine .DEFAULT .getName ());
90- final KNNEngine knnEngine = KNNEngine .getEngine (engineName );
91-
92- // Get memory optimized searcher from engine
93- final VectorSearcherFactory searcherFactory = knnEngine .getVectorSearcherFactory ();
94- if (searcherFactory == null ) {
95- // It's not supported
96- return null ;
97- }
98-
99- // Start creating searcher
100- final String fileName = KNNCodecUtil .getNativeEngineFileFromFieldInfo (fieldInfo , segmentReadState .segmentInfo );
101- if (fileName != null ) {
102- return () -> searcherFactory .createVectorSearcher (segmentReadState .directory , fileName );
103- }
104-
105- // Not supported
106- return null ;
107- }
108-
109- private synchronized void loadMemoryOptimizedSearcher () {
110- if (vectorSearchers != null ) {
111- return ;
112- }
113-
114- final Map <String , VectorSearcher > vectorSearcherPerField = new HashMap <>(
115- RESERVE_TWICE_SPACE * segmentReadState .fieldInfos .size (),
116- SUFFICIENT_LOAD_FACTOR
117- );
118-
119- try {
120- for (FieldInfo fieldInfo : segmentReadState .fieldInfos ) {
121- final IOSupplier <VectorSearcher > searcherSupplier = getIndexFileNameIfMemoryOptimizedSearchSupported (fieldInfo );
122- if (searcherSupplier != null ) {
123- final VectorSearcher searcher = searcherSupplier .get ();
124- if (searcher != null ) {
125- // It's supported. There can be a case where a certain index type underlying is not yet supported while KNNEngine
126- // itself supports memory optimized searching.
127- vectorSearcherPerField .put (fieldInfo .getName (), searcher );
128- }
129- }
130- }
131-
132- vectorSearchers = vectorSearcherPerField ;
133- } catch (Exception e ) {
134- // Close opened searchers first, then suppress
135- try {
136- IOUtils .closeWhileHandlingException (vectorSearcherPerField .values ());
137- } catch (Exception closeException ) {
138- log .error (closeException .getMessage (), closeException );
139- }
140- throw new RuntimeException (e );
81+ if (memoryOptimizedSearchEnabled && state .context .context () != IOContext .Context .MERGE ) {
82+ loadMemoryOptimizedSearcherIfRequired ();
14183 }
14284 }
14385
@@ -308,12 +250,7 @@ private boolean trySearchWithMemoryOptimizedSearch(
308250 Bits acceptDocs ,
309251 boolean isFloatVector
310252 ) throws IOException {
311- if (vectorSearchers == null ) {
312- // Null `vectorSearchers` indicates that this reader was instantiated during merge.
313- // In this case, we load the searcher on demand for searching.
314- // This will not likely happen, by the time searching, this old segment will be merged away anyway.
315- loadMemoryOptimizedSearcher ();
316- }
253+ loadMemoryOptimizedSearcherIfRequired ();
317254
318255 // Try with memory optimized searcher
319256 final VectorSearcher memoryOptimizedSearcher = vectorSearchers .get (field );
@@ -350,4 +287,76 @@ private static List<String> getVectorCacheKeysFromSegmentReaderState(SegmentRead
350287
351288 return cacheKeys ;
352289 }
290+
291+ private void loadMemoryOptimizedSearcherIfRequired () {
292+ if (vectorSearchers != null ) {
293+ return ;
294+ }
295+
296+ synchronized (this ) {
297+ if (vectorSearchers != null ) {
298+ return ;
299+ }
300+
301+ // We need sufficient memory space for this table as it will be queried for every single search.
302+ // Hence, having larger space to approximate a perfect hash here.
303+ final Map <String , VectorSearcher > vectorSearcherPerField = new HashMap <>(
304+ RESERVE_TWICE_SPACE * segmentReadState .fieldInfos .size (),
305+ SUFFICIENT_LOAD_FACTOR
306+ );
307+
308+ try {
309+ for (FieldInfo fieldInfo : segmentReadState .fieldInfos ) {
310+ final IOSupplier <VectorSearcher > searcherSupplier = getIndexFileNameIfMemoryOptimizedSearchSupported (fieldInfo );
311+ if (searcherSupplier != null ) {
312+ final VectorSearcher searcher = searcherSupplier .get ();
313+ if (searcher != null ) {
314+ // It's supported. There can be a case where a certain index type underlying is not yet supported while
315+ // KNNEngine
316+ // itself supports memory optimized searching.
317+ vectorSearcherPerField .put (fieldInfo .getName (), searcher );
318+ }
319+ }
320+ }
321+
322+ vectorSearchers = vectorSearcherPerField ;
323+ } catch (Exception e ) {
324+ // Close opened searchers first, then suppress
325+ try {
326+ IOUtils .closeWhileHandlingException (vectorSearcherPerField .values ());
327+ } catch (Exception closeException ) {
328+ log .error (closeException .getMessage (), closeException );
329+ }
330+ throw new RuntimeException (e );
331+ }
332+ }
333+ }
334+
335+ private IOSupplier <VectorSearcher > getIndexFileNameIfMemoryOptimizedSearchSupported (final FieldInfo fieldInfo ) {
336+ // Skip non-knn fields.
337+ final Map <String , String > attributes = fieldInfo .attributes ();
338+ if (attributes == null || attributes .containsKey (KNN_FIELD ) == false ) {
339+ return null ;
340+ }
341+
342+ // Get engine
343+ final String engineName = attributes .getOrDefault (KNN_ENGINE , KNNEngine .DEFAULT .getName ());
344+ final KNNEngine knnEngine = KNNEngine .getEngine (engineName );
345+
346+ // Get memory optimized searcher from engine
347+ final VectorSearcherFactory searcherFactory = knnEngine .getVectorSearcherFactory ();
348+ if (searcherFactory == null ) {
349+ // It's not supported
350+ return null ;
351+ }
352+
353+ // Start creating searcher
354+ final String fileName = KNNCodecUtil .getNativeEngineFileFromFieldInfo (fieldInfo , segmentReadState .segmentInfo );
355+ if (fileName != null ) {
356+ return () -> searcherFactory .createVectorSearcher (segmentReadState .directory , fileName );
357+ }
358+
359+ // Not supported
360+ return null ;
361+ }
353362}
0 commit comments