Skip to content

Commit 1e09bf0

Browse files
0ctopus13prime0ctopus13prime
committed
Enable optimistic search for LuceneOnFaiss.
Signed-off-by: Dooyong Kim <[email protected]>
1 parent 1523ccc commit 1e09bf0

20 files changed

+1616
-57
lines changed

‎build.gradle‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -495,6 +495,7 @@ def commonIntegTest(RestIntegTestTask task, project, integTestDependOnJniLib, op
495495
task.systemProperty 'cluster.debug', isDebuggingCluster
496496
// Set number of nodes system property to be used in tests
497497
task.systemProperty 'cluster.number_of_nodes', "${_numNodes}"
498+
498499
// There seems to be an issue when running multi node run or integ tasks with unicast_hosts
499500
// not being written, the waitForAllConditions ensures it's written
500501
task.getClusters().forEach { cluster ->
@@ -554,6 +555,7 @@ def commonIntegTestClusters(OpenSearchCluster cluster, _numNodes){
554555
debugPort += 1
555556
}
556557
}
558+
557559
cluster.systemProperty("java.library.path", "$rootDir/jni/build/release")
558560
final testSnapshotFolder = file("${buildDir}/testSnapshotFolder")
559561
testSnapshotFolder.mkdirs()
@@ -563,6 +565,8 @@ def commonIntegTestClusters(OpenSearchCluster cluster, _numNodes){
563565

564566
testClusters.integTest {
565567
commonIntegTestClusters(it, _numNodes)
568+
// Forcing optimistic search for testing
569+
systemProperty 'mem_opt_srch.force_reenter', 'true'
566570
}
567571

568572
testClusters.integTestRemoteIndexBuild {
@@ -575,6 +579,8 @@ testClusters.integTestRemoteIndexBuild {
575579
keystore 's3.client.default.access_key', "${System.getProperty("access_key")}"
576580
keystore 's3.client.default.secret_key', "${System.getProperty("secret_key")}"
577581
keystore 's3.client.default.session_token', "${System.getProperty("session_token")}"
582+
// Forcing optimistic search for testing
583+
systemProperty 'mem_opt_srch.force_reenter', 'true'
578584
}
579585

580586
task integTestRemote(type: RestIntegTestTask) {

‎src/main/java/org/opensearch/knn/index/query/KNNQueryFactory.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,8 @@ public static Query create(CreateQueryRequest createQueryRequest) {
117117
.build();
118118
}
119119

120-
if (createQueryRequest.getRescoreContext().isPresent()
120+
if (memoryOptimizedSearchEnabled
121+
|| createQueryRequest.getRescoreContext().isPresent()
121122
|| (ENGINES_SUPPORTING_NESTED_FIELDS.contains(createQueryRequest.getKnnEngine()) && expandNested)) {
122123
return new NativeEngineKnnVectorQuery(knnQuery, QueryUtils.getInstance(), expandNested);
123124
}

‎src/main/java/org/opensearch/knn/index/query/KNNWeight.java‎

Lines changed: 58 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -302,8 +302,9 @@ public PerLeafResult searchLeaf(LeafReaderContext context, int k) throws IOExcep
302302
final BitSet filterBitSet = getFilteredDocsBitSet(context);
303303
stopStopWatchAndLog(stopWatch, "FilterBitSet creation", segmentName);
304304

305-
final int maxDoc = context.reader().maxDoc();
306-
int filterCardinality = filterBitSet.cardinality();
305+
// Save its cardinality, as the cardinality calculation is expensive.
306+
final int filterCardinality = filterBitSet.cardinality();
307+
307308
// We don't need to go to JNI layer if no documents are found which satisfy the filters
308309
// We should give this condition a deeper look that where it should be placed. For now I feel this is a good
309310
// place,
@@ -320,18 +321,17 @@ public PerLeafResult searchLeaf(LeafReaderContext context, int k) throws IOExcep
320321
* This improves the recall.
321322
*/
322323
if (isFilteredExactSearchPreferred(filterCardinality)) {
323-
TopDocs result = doExactSearch(context, new BitSetIterator(filterBitSet, filterCardinality), filterCardinality, k);
324-
return new PerLeafResult(filterWeight == null ? null : filterBitSet, result);
324+
final TopDocs result = doExactSearch(context, new BitSetIterator(filterBitSet, filterCardinality), filterCardinality, k);
325+
return new PerLeafResult(
326+
filterWeight == null ? null : filterBitSet,
327+
filterCardinality,
328+
result,
329+
PerLeafResult.SearchMode.EXACT_SEARCH
330+
);
325331
}
326332

327-
/*
328-
* If filters match all docs in this segment, then null should be passed as filterBitSet
329-
* so that it will not do a bitset look up in bottom search layer.
330-
*/
331-
final BitSet annFilter = (filterWeight != null && filterCardinality == maxDoc) ? null : filterBitSet;
332-
333333
StopWatch annStopWatch = startStopWatch();
334-
final TopDocs topDocs = approximateSearch(context, annFilter, filterCardinality, k);
334+
final TopDocs topDocs = approximateSearch(context, filterBitSet, filterCardinality, k);
335335
stopStopWatchAndLog(annStopWatch, "ANN search", segmentName);
336336
if (knnQuery.isExplain()) {
337337
knnExplanation.addLeafResult(context.id(), topDocs.scoreDocs.length);
@@ -341,10 +341,21 @@ public PerLeafResult searchLeaf(LeafReaderContext context, int k) throws IOExcep
341341
// results less than K, though we have more than k filtered docs
342342
if (isExactSearchRequire(context, filterCardinality, topDocs.scoreDocs.length)) {
343343
final BitSetIterator docs = filterWeight != null ? new BitSetIterator(filterBitSet, filterCardinality) : null;
344-
TopDocs result = doExactSearch(context, docs, filterCardinality, k);
345-
return new PerLeafResult(filterWeight == null ? null : filterBitSet, result);
344+
final TopDocs result = doExactSearch(context, docs, filterCardinality, k);
345+
return new PerLeafResult(
346+
filterWeight == null ? null : filterBitSet,
347+
filterCardinality,
348+
result,
349+
PerLeafResult.SearchMode.EXACT_SEARCH
350+
);
346351
}
347-
return new PerLeafResult(filterWeight == null ? null : filterBitSet, topDocs);
352+
353+
return new PerLeafResult(
354+
filterWeight == null ? null : filterBitSet,
355+
filterCardinality,
356+
topDocs,
357+
PerLeafResult.SearchMode.APPROXIMATE_SEARCH
358+
);
348359
}
349360

350361
private void stopStopWatchAndLog(@Nullable final StopWatch stopWatch, final String prefixMessage, String segmentName) {
@@ -413,9 +424,33 @@ private TopDocs doExactSearch(
413424
return exactSearch(context, exactSearcherContextBuilder.build());
414425
}
415426

416-
protected TopDocs approximateSearch(final LeafReaderContext context, final BitSet filterIdsBitSet, final int cardinality, final int k)
417-
throws IOException {
427+
/**
428+
* Performs an approximate nearest neighbor (ANN) search on the provided index segment.
429+
* <p>
430+
* This method prepares all necessary query metadata before triggering the actual ANN search.
431+
* It extracts the {@code model_id} from field-level attributes if required, retrieves any
432+
* quantization or auxiliary metadata associated with the vector field, and applies quantization
433+
* to the query vector when applicable. After these preprocessing steps, it invokes
434+
* {@code doANNSearch(LeafReaderContext, BitSet, int, int)} to execute the approximate search
435+
* and obtain the top results.
436+
*
437+
* @param context the {@link LeafReaderContext} representing the current index segment
438+
* @param filterIdsBitSet an optional {@link BitSet} indicating document IDs to include in the search;
439+
* may be {@code null} if no filtering is required
440+
* @param filterCardinality the number of documents included in {@code filterIdsBitSet};
441+
* used to optimize search filtering
442+
* @param k the number of nearest neighbors to retrieve
443+
* @return a {@link TopDocs} object containing the top {@code k} approximate search results
444+
* @throws IOException if an error occurs while reading index data or accessing vector fields
445+
*/
446+
public TopDocs approximateSearch(
447+
final LeafReaderContext context,
448+
final BitSet filterIdsBitSet,
449+
final int filterCardinality,
450+
final int k
451+
) throws IOException {
418452
final SegmentReader reader = Lucene.segmentReader(context.reader());
453+
419454
FieldInfo fieldInfo = FieldInfoExtractor.getFieldInfo(reader, knnQuery.getField());
420455

421456
if (fieldInfo == null) {
@@ -465,6 +500,11 @@ protected TopDocs approximateSearch(final LeafReaderContext context, final BitSe
465500
// TODO: Change type of vector once more quantization methods are supported
466501
byte[] quantizedVector = maybeQuantizeVector(segmentLevelQuantizationInfo);
467502
float[] transformedVector = maybeTransformVector(segmentLevelQuantizationInfo, spaceType);
503+
/*
504+
* If filters match all docs in this segment, then null should be passed as filterBitSet
505+
* so that it will not do a bitset look up in bottom search layer.
506+
*/
507+
final BitSet annFilter = filterCardinality == context.reader().maxDoc() ? null : filterIdsBitSet;
468508

469509
KNNCounter.GRAPH_QUERY_REQUESTS.increment();
470510
final TopDocs results = doANNSearch(
@@ -477,8 +517,8 @@ protected TopDocs approximateSearch(final LeafReaderContext context, final BitSe
477517
quantizedVector,
478518
transformedVector,
479519
modelId,
480-
filterIdsBitSet,
481-
cardinality,
520+
annFilter,
521+
filterCardinality,
482522
k
483523
);
484524

‎src/main/java/org/opensearch/knn/index/query/PerLeafResult.java‎

Lines changed: 173 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,21 +7,189 @@
77

88
import lombok.Getter;
99
import lombok.Setter;
10+
import org.apache.lucene.index.LeafReaderContext;
1011
import org.apache.lucene.search.TopDocs;
1112
import org.apache.lucene.search.TopDocsCollector;
12-
import org.apache.lucene.util.Bits;
13+
import org.apache.lucene.util.BitSet;
1314
import org.opensearch.common.Nullable;
1415

16+
/**
17+
* Represents the per-segment (leaf-level) result of a vector search operation.
18+
* <p>
19+
* This class encapsulates the intermediate search state and results produced from
20+
* a single {@link LeafReaderContext} during approximate or exact vector search.
21+
* It stores the active filter bitset (if any), its cardinality, the top document
22+
* results for that segment, and the search mode used.
23+
* <p>
24+
* Instances of this class are typically aggregated at a higher level to produce
25+
* the global {@code TopDocs} result set.
26+
*/
1527
@Getter
1628
public class PerLeafResult {
17-
public static final PerLeafResult EMPTY_RESULT = new PerLeafResult(new Bits.MatchNoBits(0), TopDocsCollector.EMPTY_TOPDOCS);
29+
/**
30+
* An immutable, empty {@link BitSet} implementation used to represent
31+
* the absence of filter bits without incurring null checks or allocations.
32+
*/
33+
public static final BitSet MATCH_ALL_BIT_SET = new BitSet() {
34+
@Override
35+
public void set(int i) {
36+
throw new UnsupportedOperationException();
37+
}
38+
39+
@Override
40+
public boolean getAndSet(int i) {
41+
throw new UnsupportedOperationException();
42+
}
43+
44+
@Override
45+
public void clear(int i) {
46+
throw new UnsupportedOperationException();
47+
}
48+
49+
@Override
50+
public void clear(int startIndex, int endIndex) {
51+
throw new UnsupportedOperationException();
52+
}
53+
54+
@Override
55+
public int cardinality() {
56+
throw new UnsupportedOperationException();
57+
}
58+
59+
@Override
60+
public int approximateCardinality() {
61+
throw new UnsupportedOperationException();
62+
}
63+
64+
@Override
65+
public int prevSetBit(int index) {
66+
throw new UnsupportedOperationException();
67+
}
68+
69+
@Override
70+
public int nextSetBit(int start, int end) {
71+
throw new UnsupportedOperationException();
72+
}
73+
74+
@Override
75+
public long ramBytesUsed() {
76+
throw new UnsupportedOperationException();
77+
}
78+
79+
@Override
80+
public boolean get(int i) {
81+
return true;
82+
}
83+
84+
@Override
85+
public int length() {
86+
throw new UnsupportedOperationException();
87+
}
88+
};
89+
90+
/**
91+
* An immutable, empty {@link BitSet} implementation used to represent
92+
* the absence of filter bits without incurring null checks or allocations.
93+
*/
94+
private static final BitSet MATCH_NO_BIT_SET = new BitSet() {
95+
@Override
96+
public void set(int i) {
97+
throw new UnsupportedOperationException();
98+
}
99+
100+
@Override
101+
public boolean getAndSet(int i) {
102+
throw new UnsupportedOperationException();
103+
}
104+
105+
@Override
106+
public void clear(int i) {
107+
throw new UnsupportedOperationException();
108+
}
109+
110+
@Override
111+
public void clear(int startIndex, int endIndex) {
112+
throw new UnsupportedOperationException();
113+
}
114+
115+
@Override
116+
public int cardinality() {
117+
throw new UnsupportedOperationException();
118+
}
119+
120+
@Override
121+
public int approximateCardinality() {
122+
throw new UnsupportedOperationException();
123+
}
124+
125+
@Override
126+
public int prevSetBit(int index) {
127+
throw new UnsupportedOperationException();
128+
}
129+
130+
@Override
131+
public int nextSetBit(int start, int end) {
132+
throw new UnsupportedOperationException();
133+
}
134+
135+
@Override
136+
public long ramBytesUsed() {
137+
throw new UnsupportedOperationException();
138+
}
139+
140+
@Override
141+
public boolean get(int i) {
142+
return false;
143+
}
144+
145+
@Override
146+
public int length() {
147+
throw new UnsupportedOperationException();
148+
}
149+
};
150+
151+
// A statically defined empty {@code PerLeafResult} used as a lightweight placeholder when a segment produces no hits.
152+
public static final PerLeafResult EMPTY_RESULT = new PerLeafResult(
153+
MATCH_NO_BIT_SET,
154+
0,
155+
TopDocsCollector.EMPTY_TOPDOCS,
156+
SearchMode.EXACT_SEARCH
157+
);
158+
159+
/**
160+
* Indicates the search mode applied within a segment. Either exact or approximate nearest neighbor (ANN) search.
161+
*/
162+
public enum SearchMode {
163+
EXACT_SEARCH,
164+
APPROXIMATE_SEARCH,
165+
}
166+
167+
// Active filter bitset limiting document candidates in this leaf (may be empty).
18168
@Nullable
19-
private final Bits filterBits;
169+
private final BitSet filterBits;
170+
171+
// Cardinality of {@link #filterBits}, used for filtering optimizations.
172+
private final int filterBitsCardinality;
173+
174+
// Top document results for this leaf segment.
20175
@Setter
21176
private TopDocs result;
22177

23-
public PerLeafResult(final Bits filterBits, final TopDocs result) {
24-
this.filterBits = filterBits == null ? new Bits.MatchAllBits(0) : filterBits;
178+
// Indicates whether this result was produced via exact or approximate search.
179+
private final SearchMode searchMode;
180+
181+
/**
182+
* Constructs a new {@code PerLeafResult}.
183+
*
184+
* @param filterBits the document filter bitset for this leaf, or {@code null} if none
185+
* @param filterBitsCardinality the number of bits set in {@code filterBits}
186+
* @param result the top document results for this leaf
187+
* @param searchMode the search mode (exact or approximate) used
188+
*/
189+
public PerLeafResult(final BitSet filterBits, final int filterBitsCardinality, final TopDocs result, final SearchMode searchMode) {
190+
this.filterBits = filterBits == null ? MATCH_ALL_BIT_SET : filterBits;
191+
this.filterBitsCardinality = filterBitsCardinality;
25192
this.result = result;
193+
this.searchMode = searchMode;
26194
}
27195
}

0 commit comments

Comments
 (0)