Skip to content

Commit

Permalink
Adding javadoc in hybridquerycontext and addressing few comments from…
Browse files Browse the repository at this point in the history 
… review

Signed-off-by: Varun Jain <[email protected]>
  • Loading branch information
@vibrantvarun
vibrantvarun committed Jan 6, 2025
1 parent a07b714 commit 1474905
Show file tree
Hide file tree
Showing 5 changed files with 46 additions and 21 deletions.
7 changes: 5 additions & 2 deletions src/main/java/org/opensearch/neuralsearch/query/HybridQueryContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,16 @@
*/
package org.opensearch.neuralsearch.query;

import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Getter;
import lombok.NonNull;

@AllArgsConstructor
/**
* Class that holds the low level information of hybrid query in the form of context
*/
@Builder
@Getter
public class HybridQueryContext {
@NonNull
private int paginationDepth;
}
2 changes: 1 addition & 1 deletion src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -494,7 +494,7 @@ private static int getSubqueryResultsRetrievalSize(final SearchContext searchCon
}

/**
* Extract hybrid query from generic query object retrieved from search context
* Unwraps a HybridQuery from either a direct query or a nested BooleanQuery
*/
private static HybridQuery extractHybridQueryFromAbstractQuery(Query query) {
HybridQuery hybridQuery;
Expand Down
6 changes: 4 additions & 2 deletions src/test/java/org/opensearch/neuralsearch/search/query/HybridAggregationProcessorTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,8 @@ public void testCollectorManager_whenHybridQueryAndNotConcurrentSearch_thenSucce
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
TermQueryBuilder termSubQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, TERM_QUERY_TEXT);
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), new HybridQueryContext(10));
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), hybridQueryContext);

when(searchContext.query()).thenReturn(hybridQuery);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -130,7 +131,8 @@ public void testCollectorManager_whenHybridQueryAndConcurrentSearch_thenSuccessf
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
TermQueryBuilder termSubQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, TERM_QUERY_TEXT);
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), new HybridQueryContext(10));
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), hybridQueryContext);

when(searchContext.query()).thenReturn(hybridQuery);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down
49 changes: 34 additions & 15 deletions src/test/java/org/opensearch/neuralsearch/search/query/HybridCollectorManagerTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,8 @@ public void testNewCollector_whenNotConcurrentSearch_thenSuccessful() {
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
TermQueryBuilder termSubQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1);
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), new HybridQueryContext(10));
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), hybridQueryContext);

when(searchContext.query()).thenReturn(hybridQuery);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -125,7 +126,8 @@ public void testNewCollector_whenConcurrentSearch_thenSuccessful() {
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
TermQueryBuilder termSubQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1);
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), new HybridQueryContext(10));
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), hybridQueryContext);

when(searchContext.query()).thenReturn(hybridQuery);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -156,7 +158,8 @@ public void testPostFilter_whenNotConcurrentSearch_thenSuccessful() {
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
TermQueryBuilder termSubQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1);
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), new HybridQueryContext(10));
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), hybridQueryContext);

QueryBuilder postFilterQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, "world");
ParsedQuery parsedQuery = new ParsedQuery(postFilterQuery.toQuery(mockQueryShardContext));
Expand Down Expand Up @@ -200,7 +203,8 @@ public void testPostFilter_whenConcurrentSearch_thenSuccessful() {
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
TermQueryBuilder termSubQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1);
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), new HybridQueryContext(10));
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), hybridQueryContext);

QueryBuilder postFilterQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, "world");
Query pfQuery = postFilterQuery.toQuery(mockQueryShardContext);
Expand Down Expand Up @@ -243,10 +247,11 @@ public void testReduce_whenMatchedDocs_thenSuccessful() {
QueryShardContext mockQueryShardContext = mock(QueryShardContext.class);
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery hybridQueryWithTerm = new HybridQuery(
List.of(QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1).toQuery(mockQueryShardContext)),
new HybridQueryContext(10)
hybridQueryContext
);
when(searchContext.query()).thenReturn(hybridQueryWithTerm);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -347,7 +352,9 @@ public void testNewCollector_whenNotConcurrentSearchAndSortingIsApplied_thenSucc
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
TermQueryBuilder termSubQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1);
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), new HybridQueryContext(10));
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), hybridQueryContext);

when(searchContext.query()).thenReturn(hybridQuery);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -384,7 +391,9 @@ public void testNewCollector_whenNotConcurrentSearchAndSortingAndSearchAfterAreA
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
TermQueryBuilder termSubQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1);
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), new HybridQueryContext(10));
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), hybridQueryContext);

when(searchContext.query()).thenReturn(hybridQuery);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -414,10 +423,11 @@ public void testReduce_whenMatchedDocsAndSortingIsApplied_thenSuccessful() {
QueryShardContext mockQueryShardContext = mock(QueryShardContext.class);
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery hybridQueryWithMatchAll = new HybridQuery(
List.of(QueryBuilders.matchAllQuery().toQuery(mockQueryShardContext)),
new HybridQueryContext(10)
hybridQueryContext
);
when(searchContext.query()).thenReturn(hybridQueryWithMatchAll);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -510,13 +520,14 @@ public void testReduceWithConcurrentSegmentSearch_whenMultipleCollectorsMatchedD
QueryShardContext mockQueryShardContext = mock(QueryShardContext.class);
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery hybridQueryWithTerm = new HybridQuery(
List.of(
QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1).toQuery(mockQueryShardContext),
QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY2).toQuery(mockQueryShardContext)
),
new HybridQueryContext(10)
hybridQueryContext
);
when(searchContext.query()).thenReturn(hybridQueryWithTerm);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -601,10 +612,11 @@ public void testReduceWithConcurrentSegmentSearch_whenMultipleCollectorsMatchedD
QueryShardContext mockQueryShardContext = mock(QueryShardContext.class);
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery hybridQueryWithTerm = new HybridQuery(
List.of(QueryBuilders.matchAllQuery().toQuery(mockQueryShardContext)),
new HybridQueryContext(10)
hybridQueryContext
);
when(searchContext.query()).thenReturn(hybridQueryWithTerm);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -729,13 +741,14 @@ public void testReduceAndRescore_whenMatchedDocsAndRescoreContextPresent_thenSuc
QueryShardContext mockQueryShardContext = mock(QueryShardContext.class);
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery hybridQueryWithTerm = new HybridQuery(
List.of(
QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1).toQuery(mockQueryShardContext),
QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY2).toQuery(mockQueryShardContext)
),
new HybridQueryContext(10)
hybridQueryContext
);
when(searchContext.query()).thenReturn(hybridQueryWithTerm);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -847,14 +860,15 @@ public void testRescoreWithConcurrentSegmentSearch_whenMatchedDocsAndRescore_the
QueryShardContext mockQueryShardContext = mock(QueryShardContext.class);
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery hybridQueryWithTerm = new HybridQuery(
List.of(
QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1).toQuery(mockQueryShardContext),
QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY2).toQuery(mockQueryShardContext),
QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY3).toQuery(mockQueryShardContext)
),
new HybridQueryContext(10)
hybridQueryContext
);
when(searchContext.query()).thenReturn(hybridQueryWithTerm);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -992,13 +1006,14 @@ public void testReduceAndRescore_whenRescorerThrowsException_thenFail() {
QueryShardContext mockQueryShardContext = mock(QueryShardContext.class);
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery hybridQueryWithTerm = new HybridQuery(
List.of(
QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1).toQuery(mockQueryShardContext),
QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY2).toQuery(mockQueryShardContext)
),
new HybridQueryContext(10)
hybridQueryContext
);
when(searchContext.query()).thenReturn(hybridQueryWithTerm);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -1064,8 +1079,10 @@ public void testCreateCollectorManager_whenFromAreEqualToZeroAndPaginationDepthI
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
TermQueryBuilder termSubQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1);
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

// pagination_depth=10
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), new HybridQueryContext(10));
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), hybridQueryContext);

when(searchContext.query()).thenReturn(hybridQuery);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down Expand Up @@ -1098,7 +1115,9 @@ public void testScrollWithHybridQuery_thenFail() {
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
TermQueryBuilder termSubQuery = QueryBuilders.termQuery(TEXT_FIELD_NAME, QUERY1);
HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), new HybridQueryContext(10));
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery hybridQuery = new HybridQuery(List.of(termSubQuery.toQuery(mockQueryShardContext)), hybridQueryContext);

when(searchContext.query()).thenReturn(hybridQuery);
ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class);
Expand Down
3 changes: 2 additions & 1 deletion src/test/java/org/opensearch/neuralsearch/util/HybridQueryUtilTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ public void testIsHybridQueryCheck_whenQueryIsHybridQueryInstance_thenSuccess()
QueryShardContext mockQueryShardContext = mock(QueryShardContext.class);
TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME);
when(mockQueryShardContext.fieldMapper(eq(TEXT_FIELD_NAME))).thenReturn(fieldType);
HybridQueryContext hybridQueryContext = HybridQueryContext.builder().paginationDepth(10).build();

HybridQuery query = new HybridQuery(
List.of(
Expand All @@ -55,7 +56,7 @@ public void testIsHybridQueryCheck_whenQueryIsHybridQueryInstance_thenSuccess()
.toQuery(mockQueryShardContext),
QueryBuilders.termQuery(TEXT_FIELD_NAME, TERM_QUERY_TEXT).toQuery(mockQueryShardContext)
),
new HybridQueryContext(10)
hybridQueryContext
);
SearchContext searchContext = mock(SearchContext.class);

Expand Down

0 comments on commit 1474905

Please sign in to comment.