Skip to content

Commit

Permalink
Simplify a function name.
Browse files Browse the repository at this point in the history
Signed-off-by: conggguan <congguan@amazon.com>
  • Loading branch information
conggguan committed Apr 22, 2024
1 parent a511b0a commit df0c9fc
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 16 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@

import lombok.extern.log4j.Log4j2;

import static org.opensearch.neuralsearch.search.util.NeuralSparseTwoPhaseUtil.addSecondPhaseRescoreContextFromValidNeuralSparseQuery;
import static org.opensearch.neuralsearch.search.util.NeuralSparseTwoPhaseUtil.addRescoreContextFromNeuralSparseSparseQuery;
import static org.opensearch.neuralsearch.util.HybridQueryUtil.hasAliasFilter;
import static org.opensearch.neuralsearch.util.HybridQueryUtil.hasNestedFieldOrNestedDocs;
import static org.opensearch.neuralsearch.util.HybridQueryUtil.isHybridQuery;
Expand All @@ -45,7 +45,7 @@ public boolean searchWith(
final boolean hasFilterCollector,
final boolean hasTimeout
) throws IOException {
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(query, searchContext);
addRescoreContextFromNeuralSparseSparseQuery(query, searchContext);
if (!isHybridQuery(query, searchContext)) {
validateQuery(searchContext, query);
return super.searchWith(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ private static Query getNestedTwoPhaseQuery(Map<Query, Float> query2weight) {
* @param query The whole query include neuralSparseQuery to executed.
* @param searchContext The searchContext with this query.
*/
public static void addSecondPhaseRescoreContextFromValidNeuralSparseQuery(final Query query, SearchContext searchContext) {
public static void addRescoreContextFromNeuralSparseSparseQuery(final Query query, SearchContext searchContext) {
Map<Query, Float> query2weight = new HashMap<>();
float windowSizeExpansion = populateQueryWeightsMapAndGetWindowSizeExpansion(query, query2weight, 1.0f, 1.0f);
Query twoPhaseQuery;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.opensearch.neuralsearch.search.util.NeuralSparseTwoPhaseUtil.addSecondPhaseRescoreContextFromValidNeuralSparseQuery;
import static org.opensearch.neuralsearch.search.util.NeuralSparseTwoPhaseUtil.addRescoreContextFromNeuralSparseSparseQuery;

public class NeuralSparseTwoPhaseUtilTests extends OpenSearchTestCase {

Expand Down Expand Up @@ -98,30 +98,30 @@ public void testInitialize() {
@SneakyThrows
public void testAddTwoPhaseNeuralSparseQuery_whenQuery2WeightEmpty_thenNoRescoreAdded() {
Query query = mock(Query.class);
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(query, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(query, mockSearchContext);
verify(mockSearchContext, never()).addRescore(any());
}

@SneakyThrows
public void testAddTwoPhaseNeuralSparseQuery_whenUnSupportedQuery_thenNoRescoreAdded() {
FunctionScoreQuery functionScoreQuery = new FunctionScoreQuery(normalNeuralSparseQuery, mock(DoubleValuesSource.class));
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(functionScoreQuery, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(functionScoreQuery, mockSearchContext);
DisjunctionMaxQuery disjunctionMaxQuery = new DisjunctionMaxQuery(Collections.emptyList(), 1.0f);
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(disjunctionMaxQuery, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(disjunctionMaxQuery, mockSearchContext);
List<Query> subQueries = new ArrayList<>();
List<Query> filterQueries = new ArrayList<>();
subQueries.add(normalNeuralSparseQuery);
filterQueries.add(new MatchAllDocsQuery());
HybridQuery hybridQuery = new HybridQuery(subQueries, filterQueries);
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(hybridQuery, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(hybridQuery, mockSearchContext);
assertEquals(normalNeuralSparseQuery.getCurrentQuery(), currentQuery);
verify(mockSearchContext, never()).addRescore(any());
}

@SneakyThrows
public void testAddTwoPhaseNeuralSparseQuery_whenSingleEntryInQuery2Weight_thenRescoreAdded() {
NeuralSparseQuery neuralSparseQuery = new NeuralSparseQuery(mock(Query.class), mock(Query.class), mock(Query.class), 5.0f);
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(neuralSparseQuery, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(neuralSparseQuery, mockSearchContext);
verify(mockSearchContext).addRescore(any(QueryRescorer.QueryRescoreContext.class));
}

Expand All @@ -135,7 +135,7 @@ public void testAddTwoPhaseNeuralSparseQuery_whenCompoundBooleanQuery_thenRescor
queryBuilder.add(boostQuery1, BooleanClause.Occur.SHOULD);
queryBuilder.add(boostQuery2, BooleanClause.Occur.SHOULD);
BooleanQuery booleanQuery = queryBuilder.build();
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(booleanQuery, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(booleanQuery, mockSearchContext);
verify(mockSearchContext).addRescore(any(QueryRescorer.QueryRescoreContext.class));
}

Expand All @@ -155,7 +155,7 @@ public void testAddTwoPhaseNeuralSparseQuery_whenBooleanClauseType_thenVerifyBoo
queryBuilder.add(boostQuery3, BooleanClause.Occur.FILTER);
queryBuilder.add(boostQuery4, BooleanClause.Occur.MUST_NOT);
BooleanQuery booleanQuery = queryBuilder.build();
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(booleanQuery, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(booleanQuery, mockSearchContext);
ArgumentCaptor<RescoreContext> rtxCaptor = ArgumentCaptor.forClass(RescoreContext.class);
verify(mockSearchContext).addRescore(rtxCaptor.capture());
QueryRescorer.QueryRescoreContext context = (QueryRescorer.QueryRescoreContext) rtxCaptor.getValue();
Expand All @@ -179,7 +179,7 @@ public void testAddTwoPhaseNeuralSparseQuery_whenBooleanClauseType_thenVerifyBoo
@SneakyThrows
public void testWindowSize_whenNormalConditions_thenWindowSizeIsAsSet() {
NeuralSparseQuery query = normalNeuralSparseQuery;
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(query, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(query, mockSearchContext);
ArgumentCaptor<QueryRescorer.QueryRescoreContext> rescoreContextArgumentCaptor = ArgumentCaptor.forClass(
QueryRescorer.QueryRescoreContext.class
);
Expand All @@ -193,12 +193,12 @@ public void testWindowSize_whenBoundaryConditions_thenThrowException() {
NeuralSparseQuery query = new NeuralSparseQuery(new MatchAllDocsQuery(), new MatchAllDocsQuery(), new MatchAllDocsQuery(), 5000f);
NeuralSparseQuery finalQuery1 = query;
expectThrows(IllegalArgumentException.class, () -> {
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(finalQuery1, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(finalQuery1, mockSearchContext);
});
query = new NeuralSparseQuery(new MatchAllDocsQuery(), new MatchAllDocsQuery(), new MatchAllDocsQuery(), Float.MAX_VALUE);
NeuralSparseQuery finalQuery = query;
expectThrows(IllegalArgumentException.class, () -> {
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(finalQuery, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(finalQuery, mockSearchContext);
});
}

Expand All @@ -211,7 +211,7 @@ public void testRescoreListWeightCalculation_whenMultipleRescoreContexts_thenCal
List<RescoreContext> rescoreContextList = Arrays.asList(mockContext1, mockContext2);
when(mockSearchContext.rescore()).thenReturn(rescoreContextList);
NeuralSparseQuery query = normalNeuralSparseQuery;
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(query, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(query, mockSearchContext);
ArgumentCaptor<RescoreContext> rtxCaptor = ArgumentCaptor.forClass(RescoreContext.class);
verify(mockSearchContext).addRescore(rtxCaptor.capture());
QueryRescorer.QueryRescoreContext context = (QueryRescorer.QueryRescoreContext) rtxCaptor.getValue();
Expand All @@ -222,7 +222,7 @@ public void testRescoreListWeightCalculation_whenMultipleRescoreContexts_thenCal
public void testEmptyRescoreListWeight_whenRescoreListEmpty_thenDefaultWeightUsed() {
when(mockSearchContext.rescore()).thenReturn(Collections.emptyList());
NeuralSparseQuery query = normalNeuralSparseQuery;
addSecondPhaseRescoreContextFromValidNeuralSparseQuery(query, mockSearchContext);
addRescoreContextFromNeuralSparseSparseQuery(query, mockSearchContext);
ArgumentCaptor<RescoreContext> rtxCaptor = ArgumentCaptor.forClass(RescoreContext.class);
verify(mockSearchContext).addRescore(rtxCaptor.capture());
QueryRescorer.QueryRescoreContext context = (QueryRescorer.QueryRescoreContext) rtxCaptor.getValue();
Expand Down

0 comments on commit df0c9fc

Please sign in to comment.