-
Notifications
You must be signed in to change notification settings - Fork 127
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Adds in lazy execution for Lucene kNN queries (#2305)
Signed-off-by: Kunal Kotwani <[email protected]> (cherry picked from commit 2b9a741)
- Loading branch information
1 parent
8a5bef2
commit ce6f32e
Showing
6 changed files
with
191 additions
and
18 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
69 changes: 69 additions & 0 deletions
69
src/main/java/org/opensearch/knn/index/query/lucene/LuceneEngineKnnVectorQuery.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,69 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.index.query.lucene; | ||
|
|
||
| import lombok.AllArgsConstructor; | ||
| import lombok.Getter; | ||
| import lombok.extern.log4j.Log4j2; | ||
| import org.apache.lucene.search.IndexSearcher; | ||
| import org.apache.lucene.search.Query; | ||
| import org.apache.lucene.search.QueryVisitor; | ||
| import org.apache.lucene.search.ScoreMode; | ||
| import org.apache.lucene.search.Weight; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| /** | ||
| * LuceneEngineKnnVectorQuery is a wrapper around a vector queries for the Lucene engine. | ||
| * This enables us to defer rewrites until weight creation to optimize repeated execution | ||
| * of Lucene based k-NN queries. | ||
| */ | ||
| @AllArgsConstructor | ||
| @Log4j2 | ||
| public class LuceneEngineKnnVectorQuery extends Query { | ||
| @Getter | ||
| private final Query luceneQuery; | ||
|
|
||
| /* | ||
| Prevents repeated rewrites of the query for the Lucene engine. | ||
| */ | ||
| @Override | ||
| public Query rewrite(IndexSearcher indexSearcher) { | ||
| return this; | ||
| } | ||
|
|
||
| /* | ||
| Rewrites the query just before weight creation. | ||
| */ | ||
| @Override | ||
| public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { | ||
| Query rewrittenQuery = luceneQuery.rewrite(searcher); | ||
| return rewrittenQuery.createWeight(searcher, scoreMode, boost); | ||
| } | ||
|
|
||
| @Override | ||
| public String toString(String s) { | ||
| return luceneQuery.toString(); | ||
| } | ||
|
|
||
| @Override | ||
| public void visit(QueryVisitor queryVisitor) { | ||
| queryVisitor.visitLeaf(this); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean equals(Object o) { | ||
| if (this == o) return true; | ||
| if (o == null || getClass() != o.getClass()) return false; | ||
| LuceneEngineKnnVectorQuery otherQuery = (LuceneEngineKnnVectorQuery) o; | ||
| return luceneQuery.equals(otherQuery.luceneQuery); | ||
| } | ||
|
|
||
| @Override | ||
| public int hashCode() { | ||
| return luceneQuery.hashCode(); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
95 changes: 95 additions & 0 deletions
95
src/test/java/org/opensearch/knn/index/query/lucene/LuceneEngineKnnVectorQueryTests.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,95 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.index.query.lucene; | ||
|
|
||
| import org.apache.lucene.search.IndexSearcher; | ||
| import org.apache.lucene.search.Query; | ||
| import org.apache.lucene.search.QueryVisitor; | ||
| import org.apache.lucene.search.ScoreMode; | ||
| import org.apache.lucene.search.Weight; | ||
| import org.mockito.InjectMocks; | ||
| import org.mockito.Mock; | ||
| import org.mockito.Spy; | ||
| import org.opensearch.test.OpenSearchTestCase; | ||
|
|
||
| import static org.mockito.ArgumentMatchers.*; | ||
| import static org.mockito.Mockito.verifyNoInteractions; | ||
| import static org.mockito.Mockito.when; | ||
| import static org.mockito.Mockito.times; | ||
| import static org.mockito.Mockito.verify; | ||
| import static org.mockito.MockitoAnnotations.openMocks; | ||
|
|
||
| public class LuceneEngineKnnVectorQueryTests extends OpenSearchTestCase { | ||
|
|
||
| @Mock | ||
| IndexSearcher indexSearcher; | ||
|
|
||
| @Mock | ||
| Query luceneQuery; | ||
|
|
||
| @Mock | ||
| Weight weight; | ||
|
|
||
| @Mock | ||
| QueryVisitor queryVisitor; | ||
|
|
||
| @Spy | ||
| @InjectMocks | ||
| LuceneEngineKnnVectorQuery objectUnderTest; | ||
|
|
||
| @Override | ||
| public void setUp() throws Exception { | ||
| super.setUp(); | ||
| openMocks(this); | ||
| when(luceneQuery.rewrite(any(IndexSearcher.class))).thenReturn(luceneQuery); | ||
| when(luceneQuery.createWeight(any(IndexSearcher.class), any(ScoreMode.class), anyFloat())).thenReturn(weight); | ||
| } | ||
|
|
||
| public void testRewrite() { | ||
| objectUnderTest.rewrite(indexSearcher); | ||
| objectUnderTest.rewrite(indexSearcher); | ||
| objectUnderTest.rewrite(indexSearcher); | ||
| verifyNoInteractions(luceneQuery); | ||
| verify(objectUnderTest, times(3)).rewrite(indexSearcher); | ||
| } | ||
|
|
||
| public void testCreateWeight() throws Exception { | ||
| objectUnderTest.rewrite(indexSearcher); | ||
| objectUnderTest.rewrite(indexSearcher); | ||
| objectUnderTest.rewrite(indexSearcher); | ||
| verifyNoInteractions(luceneQuery); | ||
| Weight actualWeight = objectUnderTest.createWeight(indexSearcher, ScoreMode.TOP_DOCS, 1.0f); | ||
| verify(luceneQuery, times(1)).rewrite(indexSearcher); | ||
| verify(objectUnderTest, times(3)).rewrite(indexSearcher); | ||
| assertEquals(weight, actualWeight); | ||
| } | ||
|
|
||
| public void testVisit() { | ||
| objectUnderTest.visit(queryVisitor); | ||
| verify(queryVisitor).visitLeaf(objectUnderTest); | ||
| } | ||
|
|
||
| public void testEquals() { | ||
| LuceneEngineKnnVectorQuery mainQuery = new LuceneEngineKnnVectorQuery(luceneQuery); | ||
| LuceneEngineKnnVectorQuery otherQuery = new LuceneEngineKnnVectorQuery(luceneQuery); | ||
| assertEquals(mainQuery, otherQuery); | ||
| assertEquals(mainQuery, mainQuery); | ||
| assertNotEquals(mainQuery, null); | ||
| assertNotEquals(mainQuery, new Object()); | ||
| LuceneEngineKnnVectorQuery otherQuery2 = new LuceneEngineKnnVectorQuery(null); | ||
| assertNotEquals(mainQuery, otherQuery2); | ||
| } | ||
|
|
||
| public void testHashCode() { | ||
| LuceneEngineKnnVectorQuery mainQuery = new LuceneEngineKnnVectorQuery(luceneQuery); | ||
| assertEquals(mainQuery.hashCode(), luceneQuery.hashCode()); | ||
| } | ||
|
|
||
| public void testToString() { | ||
| LuceneEngineKnnVectorQuery mainQuery = new LuceneEngineKnnVectorQuery(luceneQuery); | ||
| assertEquals(mainQuery.toString(), luceneQuery.toString()); | ||
| } | ||
| } |