Skip to content

Commit

Permalink
Some update for comments.
Browse files Browse the repository at this point in the history 
Signed-off-by: conggguan <[email protected]>
  • Loading branch information
conggguan committed Apr 23, 2024
1 parent aec35b6 commit 6857ef2
Show file tree
Hide file tree
Showing 5 changed files with 137 additions and 60 deletions.
2 changes: 1 addition & 1 deletion CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),

## [Unreleased 2.x](https://github.com/opensearch-project/neural-search/compare/2.13...2.x)
### Features
- Enhance neural_sparse query's latency performance with two-phase rescore query([#695](https://github.com/opensearch-project/neural-search/pull/695/files)).
### Enhancements
- BWC tests for text chunking processor ([#661](https://github.com/opensearch-project/neural-search/pull/661))
- Allowing execution of hybrid query on index alias with filters ([#670](https://github.com/opensearch-project/neural-search/pull/670))
- Allowing query by raw tokens in neural_sparse query ([#693](https://github.com/opensearch-project/neural-search/pull/693))
- Enhance neural_sparse query's latency performance with two-phase rescore query([#695](https://github.com/opensearch-project/neural-search/pull/695/files)).
### Bug Fixes
- Add support for request_cache flag in hybrid query ([#663](https://github.com/opensearch-project/neural-search/pull/663))
### Infrastructure
Expand Down
6 changes: 4 additions & 2 deletions src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -345,13 +345,15 @@ protected Query doToQuery(QueryShardContext context) throws IOException {
}
// in the last step we make sure neuralSparseTwoPhaseParameters is not null
float ratio = neuralSparseTwoPhaseParameters.pruning_ratio();
Query allTokenQuery = buildFeatureFieldQueryFromTokens(getAllTokens(), fieldName);
Map<String, Float> highScoreTokens = getHighScoreTokens(ratio);
Map<String, Float> lowScoreTokens = getLowScoreTokens(ratio);
Map<String, Float> allTokens = getAllTokens();
Query allTokenQuery = buildFeatureFieldQueryFromTokens(allTokens, fieldName);
// if all token are valid score that we don't need the two-phase optimize, return allTokenQuery.
if (lowScoreTokens.isEmpty()) {
return allTokenQuery;
}
Query highScoreTokenQuery = buildFeatureFieldQueryFromTokens(getHighScoreTokens(ratio), fieldName);
Query highScoreTokenQuery = buildFeatureFieldQueryFromTokens(highScoreTokens, fieldName);
Query lowScoreTokenQuery = buildFeatureFieldQueryFromTokens(lowScoreTokens, fieldName);
return new NeuralSparseQuery(
allTokenQuery,
Expand Down
18 changes: 10 additions & 8 deletions src/main/java/org/opensearch/neuralsearch/search/util/NeuralSparseTwoPhaseUtil.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,6 @@
import java.util.Map;

import static java.lang.Float.max;
import static java.lang.Integer.min;
import static org.opensearch.index.IndexSettings.MAX_RESCORE_WINDOW_SETTING;

/**
* Util class for do two phase preprocess for the NeuralSparseQuery.
Expand All @@ -44,12 +42,16 @@ public static void addRescoreContextFromNeuralSparseQuery(final Query query, Sea
twoPhaseQuery = getNestedTwoPhaseQuery(query2weight);
}
int curWindowSize = (int) (searchContext.size() * windowSizeExpansion);
if (curWindowSize < 0
|| curWindowSize > min(
NeuralSparseTwoPhaseParameters.MAX_WINDOW_SIZE,
MAX_RESCORE_WINDOW_SETTING.get(searchContext.getQueryShardContext().getIndexSettings().getSettings())
)) {
throw new IllegalArgumentException("Two phase final windowSize out of score with value " + curWindowSize + ".");
if (curWindowSize < 0 || curWindowSize > NeuralSparseTwoPhaseParameters.MAX_WINDOW_SIZE) {
throw new IllegalArgumentException(
"Two phase final windowSize "
+ curWindowSize
+ " out of score with limit "
+ NeuralSparseTwoPhaseParameters.MAX_WINDOW_SIZE
+ "."
+ "You can change the value of cluster setting "
+ "[plugins.neural_search.neural_sparse.two_phase.max_window_size] to a integer at least 50."
);
}
QueryRescorer.QueryRescoreContext rescoreContext = new QueryRescorer.QueryRescoreContext(curWindowSize);
rescoreContext.setQuery(twoPhaseQuery);
Expand Down
100 changes: 84 additions & 16 deletions src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilderTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import org.apache.lucene.document.FeatureField;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.Query;
import org.junit.Before;
import org.opensearch.Version;
import org.opensearch.client.Client;
Expand Down Expand Up @@ -80,6 +81,7 @@ public class NeuralSparseQueryBuilderTests extends OpenSearchTestCase {
private Settings settings;
private ClusterSettings clusterSettings;
private ClusterService clusterService;
private QueryShardContext mockQueryShardContext = mock(QueryShardContext.class);

@Before
public void setUpNeuralSparseTwoPhaseParameters() {
Expand All @@ -97,6 +99,10 @@ public void setUpNeuralSparseTwoPhaseParameters() {
clusterService = mock(ClusterService.class);
when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
NeuralSparseTwoPhaseParameters.initialize(clusterService, settings);
// initialize mockQueryShardContext
MappedFieldType mappedFieldType = mock(MappedFieldType.class);
when(mappedFieldType.typeName()).thenReturn("rank_features");
when(mockQueryShardContext.fieldMapper(anyString())).thenReturn(mappedFieldType);
}

@Before
Expand Down Expand Up @@ -241,7 +247,7 @@ public void testFromXContent_whenBuiltWithIllegalTwoPhaseWindowSizeExpansion_the
"window_size_expansion": 0.4,
"pruning_ratio": 0.4,
"enabled": true
}
}
}
}
*/
Expand Down Expand Up @@ -491,6 +497,36 @@ public void testFromXContent_whenBuildWithEmptyModelId_thenFail() {
expectThrows(IllegalArgumentException.class, () -> NeuralSparseQueryBuilder.fromXContent(contentParser));
}

@SneakyThrows
public void testFromXContent_whenBuiltWithEmptyTwoPhaseParams_thenThrowException() {
/*
{
"VECTOR_FIELD": {
"query_text": "string",
"model_id": "string",
"two_phase_settings":{
"window_size_expansion": 5,
"pruning_ratio": 0.4,
"enabled": false
}
}
}
*/
NeuralSparseTwoPhaseParameters parameters = new NeuralSparseTwoPhaseParameters().enabled(null)
.pruning_ratio(null)
.window_size_expansion(null);
XContentBuilder xContentBuilder = XContentFactory.jsonBuilder()
.startObject()
.startObject(FIELD_NAME)
.field(QUERY_TEXT_FIELD.getPreferredName(), QUERY_TEXT)
.field(MODEL_ID_FIELD.getPreferredName(), MODEL_ID);
parameters.doXContent(xContentBuilder);
xContentBuilder.endObject().endObject();
XContentParser contentParser = createParser(xContentBuilder);
contentParser.nextToken();
expectThrows(ParsingException.class, () -> NeuralSparseQueryBuilder.fromXContent(contentParser));
}

@SuppressWarnings("unchecked")
@SneakyThrows
public void testToXContentWithFullField() {
Expand Down Expand Up @@ -937,15 +973,10 @@ public void testTokenDividedByScores_whenDefaultSettings() {
.neuralSparseTwoPhaseParameters(NeuralSparseTwoPhaseParameters.getDefaultSettings())
.modelId(MODEL_ID)
.queryTokensSupplier(tokenSupplier);
QueryShardContext context = mock(QueryShardContext.class);
MappedFieldType mappedFieldType = mock(MappedFieldType.class);
when(mappedFieldType.typeName()).thenReturn("rank_features");
when(context.fieldMapper(anyString())).thenReturn(mappedFieldType);
NeuralSparseQuery neuralSparseQuery = (NeuralSparseQuery) sparseEncodingQueryBuilder.doToQuery(context);

NeuralSparseQuery neuralSparseQuery = (NeuralSparseQuery) sparseEncodingQueryBuilder.doToQuery(mockQueryShardContext);
BooleanQuery highScoreTokenQuery = (BooleanQuery) neuralSparseQuery.getHighScoreTokenQuery();
BooleanQuery lowScoreTokenQuery = (BooleanQuery) neuralSparseQuery.getLowScoreTokenQuery();
assertNotNull(highScoreTokenQuery.clauses());
assertNotNull(lowScoreTokenQuery.clauses());
assertEquals(highScoreTokenQuery.clauses().size(), 7);
assertEquals(lowScoreTokenQuery.clauses().size(), 3);
sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName("rank_features")
Expand All @@ -955,31 +986,68 @@ public void testTokenDividedByScores_whenDefaultSettings() {
)
.modelId(MODEL_ID)
.queryTokensSupplier(tokenSupplier);
neuralSparseQuery = (NeuralSparseQuery) sparseEncodingQueryBuilder.doToQuery(context);
neuralSparseQuery = (NeuralSparseQuery) sparseEncodingQueryBuilder.doToQuery(mockQueryShardContext);
highScoreTokenQuery = (BooleanQuery) neuralSparseQuery.getHighScoreTokenQuery();
lowScoreTokenQuery = (BooleanQuery) neuralSparseQuery.getLowScoreTokenQuery();
assertNotNull(highScoreTokenQuery.clauses());
assertNotNull(lowScoreTokenQuery.clauses());
assertEquals(highScoreTokenQuery.clauses().size(), 5);
assertEquals(lowScoreTokenQuery.clauses().size(), 5);
}

@SneakyThrows
public void testDoToQuery_whenTwoPhaseParaDisabled_thenDegradeSuccess() {
NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(FIELD_NAME)
.maxTokenScore(MAX_TOKEN_SCORE)
.queryText(QUERY_TEXT)
.modelId(MODEL_ID)
.queryTokensSupplier(QUERY_TOKENS_SUPPLIER)
.neuralSparseTwoPhaseParameters(
new NeuralSparseTwoPhaseParameters().enabled(false).pruning_ratio(0.4f).window_size_expansion(6.0f)
);
Query query = sparseEncodingQueryBuilder.doToQuery(mockQueryShardContext);
assertTrue(query instanceof BooleanQuery);
List<BooleanClause> booleanClauseList = ((BooleanQuery) query).clauses();
assertEquals(2, ((BooleanQuery) query).clauses().size());
BooleanClause firstClause = booleanClauseList.get(0);
BooleanClause secondClause = booleanClauseList.get(1);

Query firstFeatureQuery = firstClause.getQuery();
assertEquals(firstFeatureQuery, FeatureField.newLinearQuery(FIELD_NAME, "world", 2.f));
Query secondFeatureQuery = secondClause.getQuery();
assertEquals(secondFeatureQuery, FeatureField.newLinearQuery(FIELD_NAME, "hello", 1.f));
}

@SneakyThrows
public void testDoToQuery_whenTwoPhaseParaEmpty_thenDegradeSuccess() {
NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(FIELD_NAME)
.maxTokenScore(MAX_TOKEN_SCORE)
.queryText(QUERY_TEXT)
.modelId(MODEL_ID)
.queryTokensSupplier(QUERY_TOKENS_SUPPLIER);
Query query = sparseEncodingQueryBuilder.doToQuery(mockQueryShardContext);
assertTrue(query instanceof BooleanQuery);
List<BooleanClause> booleanClauseList = ((BooleanQuery) query).clauses();
assertEquals(2, ((BooleanQuery) query).clauses().size());
BooleanClause firstClause = booleanClauseList.get(0);
BooleanClause secondClause = booleanClauseList.get(1);

Query firstFeatureQuery = firstClause.getQuery();
assertEquals(firstFeatureQuery, FeatureField.newLinearQuery(FIELD_NAME, "world", 2.f));
Query secondFeatureQuery = secondClause.getQuery();
assertEquals(secondFeatureQuery, FeatureField.newLinearQuery(FIELD_NAME, "hello", 1.f));
}

@SneakyThrows
public void testDoToQuery_successfulDoToQuery() {
NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(FIELD_NAME)
.maxTokenScore(MAX_TOKEN_SCORE)
.queryText(QUERY_TEXT)
.modelId(MODEL_ID)
.queryTokensSupplier(QUERY_TOKENS_SUPPLIER);
QueryShardContext mockedQueryShardContext = mock(QueryShardContext.class);
MappedFieldType mockedMappedFieldType = mock(MappedFieldType.class);
doAnswer(invocation -> "rank_features").when(mockedMappedFieldType).typeName();
doAnswer(invocation -> mockedMappedFieldType).when(mockedQueryShardContext).fieldMapper(any());

BooleanQuery.Builder targetQueryBuilder = new BooleanQuery.Builder();
targetQueryBuilder.add(FeatureField.newLinearQuery(FIELD_NAME, "hello", 1.f), BooleanClause.Occur.SHOULD);
targetQueryBuilder.add(FeatureField.newLinearQuery(FIELD_NAME, "world", 2.f), BooleanClause.Occur.SHOULD);

assertEquals(sparseEncodingQueryBuilder.doToQuery(mockedQueryShardContext), targetQueryBuilder.build());
assertEquals(sparseEncodingQueryBuilder.doToQuery(mockQueryShardContext), targetQueryBuilder.build());
}
}
71 changes: 38 additions & 33 deletions src/test/java/org/opensearch/neuralsearch/search/util/NeuralSparseTwoPhaseUtilTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,10 @@
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Query;
import org.junit.Before;
import org.opensearch.Version;
import org.opensearch.cluster.metadata.IndexMetadata;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.settings.ClusterSettings;
import org.opensearch.common.settings.Setting;
import org.opensearch.common.settings.Settings;
import org.opensearch.index.IndexSettings;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.neuralsearch.query.HybridQuery;
import org.opensearch.neuralsearch.query.NeuralSparseQuery;
Expand Down Expand Up @@ -57,42 +54,15 @@ public class NeuralSparseTwoPhaseUtilTests extends OpenSearchTestCase {
private final Query highScoreTokenQuery = mock(Query.class);
private final Query lowScoreTokenQuery = mock(Query.class);

protected IndexSettings createIndexSettings() {
return new IndexSettings(
IndexMetadata.builder("_index")
.settings(
Settings.builder().put("index.max_rescore_window", 10000).put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT)
)
.numberOfShards(1)
.numberOfReplicas(0)
.creationDate(System.currentTimeMillis())
.build(),
Settings.EMPTY
);
}
private ClusterSettings clusterSettings;

@SneakyThrows
@Before
public void testInitialize() {
normalNeuralSparseQuery = new NeuralSparseQuery(currentQuery, highScoreTokenQuery, lowScoreTokenQuery, 5f);
IndexSettings indexSettings = createIndexSettings();
when(mockSearchContext.getQueryShardContext()).thenReturn(mockQueryShardContext);
when(mockSearchContext.size()).thenReturn(10);
when(mockQueryShardContext.getIndexSettings()).thenReturn(indexSettings);
Settings settings = Settings.builder().build();
final Set<Setting<?>> settingsSet = Stream.concat(
ClusterSettings.BUILT_IN_CLUSTER_SETTINGS.stream(),
Stream.of(
NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_DEFAULT_ENABLED,
NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_DEFAULT_WINDOW_SIZE_EXPANSION,
NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_DEFAULT_PRUNING_RATIO,
NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_MAX_WINDOW_SIZE
)
).collect(Collectors.toSet());
ClusterSettings clusterSettings = new ClusterSettings(settings, settingsSet);
ClusterService clusterService = mock(ClusterService.class);
when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
NeuralSparseTwoPhaseParameters.initialize(clusterService, settings);
updateClusterSettingForTwoPhaseUtilTest();
}

@SneakyThrows
Expand Down Expand Up @@ -189,7 +159,6 @@ public void testWindowSize_whenNormalConditions_thenWindowSizeIsAsSet() {

@SneakyThrows
public void testWindowSize_whenBoundaryConditions_thenThrowException() {

NeuralSparseQuery query = new NeuralSparseQuery(new MatchAllDocsQuery(), new MatchAllDocsQuery(), new MatchAllDocsQuery(), 5000f);
NeuralSparseQuery finalQuery1 = query;
expectThrows(IllegalArgumentException.class, () -> { addRescoreContextFromNeuralSparseQuery(finalQuery1, mockSearchContext); });
Expand Down Expand Up @@ -225,4 +194,40 @@ public void testEmptyRescoreListWeight_whenRescoreListEmpty_thenDefaultWeightUse
assertEquals(context.rescoreQueryWeight(), 1.0f, 0.01f);
}

@SneakyThrows
public void testAddRescoreContext_whenOverEdgeWindowSize_thenThrowException() {
when(mockSearchContext.size()).thenReturn(100);
updateMaxWindowSizeClusterSettingForTwoPhaseUtilTest(50);
expectThrows(
IllegalArgumentException.class,
()->addRescoreContextFromNeuralSparseQuery(normalNeuralSparseQuery,mockSearchContext)
);
when(mockSearchContext.size()).thenReturn(10);
addRescoreContextFromNeuralSparseQuery(normalNeuralSparseQuery,mockSearchContext);
ArgumentCaptor<QueryRescorer.QueryRescoreContext> rescoreContextArgumentCaptor = ArgumentCaptor.forClass(
QueryRescorer.QueryRescoreContext.class
);
verify(mockSearchContext).addRescore(rescoreContextArgumentCaptor.capture());
assertEquals(50, rescoreContextArgumentCaptor.getValue().getWindowSize());
updateClusterSettingForTwoPhaseUtilTest();
}

private void updateClusterSettingForTwoPhaseUtilTest() {
Settings settings = Settings.builder().build();
final Set<Setting<?>> settingsSet = Stream.of(
NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_DEFAULT_ENABLED,
NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_DEFAULT_WINDOW_SIZE_EXPANSION,
NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_DEFAULT_PRUNING_RATIO,
NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_MAX_WINDOW_SIZE
).collect(Collectors.toSet());
clusterSettings = new ClusterSettings(settings, settingsSet);
ClusterService clusterService = mock(ClusterService.class);
when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
NeuralSparseTwoPhaseParameters.initialize(clusterService, settings);
}

private void updateMaxWindowSizeClusterSettingForTwoPhaseUtilTest(int size) {
Settings newSettings = Settings.builder().put("plugins.neural_search.neural_sparse.two_phase.max_window_size", size).build();
clusterSettings.applySettings(newSettings);
}
}

0 comments on commit 6857ef2

Please sign in to comment.