Skip to content

Commit 094d089

Browse files
sarthakraghuvanshisarthakraghuvanshi
authored
Add BWC tests for search_after in hybrid query (#1903)
Add restart-upgrade and rolling-upgrade BWC tests validating sort + search_after pagination with hybrid queries, and gate them behind versions >= 2.16 where the feature was introduced. Signed-off-by: Sarthak Raghuvanshi <sarthakraghuvanshi1005@gmail.com>
1 parent 11b2e13 commit 094d089

4 files changed

Lines changed: 370 additions & 0 deletions

File tree

‎qa/restart-upgrade/build.gradle‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,9 +121,11 @@ task testAgainstOldCluster(type: StandaloneRestIntegTestTask) {
121121
}
122122

123123
// Excluding the batching processor tests because we introduce this feature in 2.16
124+
// Excluding hybrid search with search_after tests because sorting and search_after in hybrid query were introduced in 2.16
124125
if (versionsBelow2_16.any { ext.neural_search_bwc_version.startsWith(it) }){
125126
filter {
126127
excludeTestsMatching "org.opensearch.neuralsearch.bwc.restart.BatchIngestionIT.*"
128+
excludeTestsMatching "org.opensearch.neuralsearch.bwc.restart.HybridSearchWithSearchAfterIT.*"
127129
}
128130
}
129131

@@ -228,9 +230,11 @@ task testAgainstNewCluster(type: StandaloneRestIntegTestTask) {
228230
}
229231

230232
// Excluding the batch processor tests because we introduce this feature in 2.16
233+
// Excluding hybrid search with search_after tests because sorting and search_after in hybrid query were introduced in 2.16
231234
if (versionsBelow2_16.any { ext.neural_search_bwc_version.startsWith(it) }){
232235
filter {
233236
excludeTestsMatching "org.opensearch.neuralsearch.bwc.restart.BatchIngestionIT.*"
237+
excludeTestsMatching "org.opensearch.neuralsearch.bwc.restart.HybridSearchWithSearchAfterIT.*"
234238
}
235239
}
236240

‎qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithSearchAfterIT.java‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package org.opensearch.neuralsearch.bwc.restart;
6+
7+
import java.io.IOException;
8+
import java.nio.file.Files;
9+
import java.nio.file.Path;
10+
import java.util.ArrayList;
11+
import java.util.Arrays;
12+
import java.util.LinkedHashMap;
13+
import java.util.List;
14+
import java.util.Map;
15+
16+
import org.opensearch.client.Request;
17+
import org.opensearch.common.xcontent.XContentFactory;
18+
import org.opensearch.core.xcontent.XContentBuilder;
19+
import org.opensearch.index.query.MatchQueryBuilder;
20+
import org.opensearch.neuralsearch.query.HybridQueryBuilder;
21+
import org.opensearch.neuralsearch.query.NeuralQueryBuilder;
22+
import org.opensearch.search.sort.SortOrder;
23+
24+
import static org.opensearch.neuralsearch.util.TestUtils.DEFAULT_COMBINATION_METHOD;
25+
import static org.opensearch.neuralsearch.util.TestUtils.DEFAULT_NORMALIZATION_METHOD;
26+
import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER;
27+
import static org.opensearch.neuralsearch.util.TestUtils.PARAM_NAME_WEIGHTS;
28+
import static org.opensearch.neuralsearch.util.TestUtils.TEXT_EMBEDDING_PROCESSOR;
29+
import static org.opensearch.neuralsearch.util.TestUtils.getModelId;
30+
31+
public class HybridSearchWithSearchAfterIT extends AbstractRestartUpgradeRestTestCase {
32+
33+
private static final String PIPELINE_NAME = "nlp-hybrid-search-after-pipeline";
34+
private static final String SEARCH_PIPELINE_NAME = "nlp-hybrid-search-after-search-pipeline";
35+
private static final String TEST_FIELD = "passage_text";
36+
private static final String SORT_FIELD = "stock";
37+
private static final String VECTOR_EMBEDDING_FIELD = "passage_embedding";
38+
private static final String QUERY = "Hi world";
39+
private static final int QUERY_SIZE = 10;
40+
private static final List<String> TEXTS = List.of(
41+
"Hello world",
42+
"Hi planet",
43+
"Hi earth",
44+
"Hi amazon",
45+
"Hi mars",
46+
"Hi opensearch",
47+
"Hi neptune"
48+
);
49+
// stock value for doc with id i is (i + 1) * 10
50+
private static String modelId = "";
51+
52+
// Test rolling-upgrade with hybrid query using sort and search_after (deep pagination)
53+
// Create Text Embedding Processor, Ingestion Pipeline, add documents with a numeric sort field,
54+
// and a search pipeline with normalization processor.
55+
// Validate that sort + search_after returns correctly ordered pages in mixed and upgraded clusters.
56+
public void testHybridSearchWithSearchAfter_E2EFlow() throws Exception {
57+
waitForClusterHealthGreen(NODES_BWC_CLUSTER);
58+
if (isRunningAgainstOldCluster()) {
59+
modelId = uploadTextEmbeddingModel();
60+
createPipelineProcessor(modelId, PIPELINE_NAME);
61+
createIndexWithConfiguration(
62+
getIndexNameForTest(),
63+
Files.readString(Path.of(classLoader.getResource("processor/IndexMappingSingleShard.json").toURI())),
64+
PIPELINE_NAME
65+
);
66+
// docs 0..4 with stock values 10, 20, 30, 40, 50
67+
for (int docId = 0; docId < 5; docId++) {
68+
addDocumentWithSortField(getIndexNameForTest(), String.valueOf(docId), TEXTS.get(docId), (docId + 1) * 10);
69+
}
70+
createSearchPipeline(
71+
SEARCH_PIPELINE_NAME,
72+
DEFAULT_NORMALIZATION_METHOD,
73+
DEFAULT_COMBINATION_METHOD,
74+
Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.7f }))
75+
);
76+
} else {
77+
try {
78+
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
79+
loadAndWaitForModelToBeReady(modelId);
80+
// doc 5 with stock value 60 and doc 6 with stock value 70
81+
addDocumentWithSortField(getIndexNameForTest(), "5", TEXTS.get(5), 60);
82+
addDocumentWithSortField(getIndexNameForTest(), "6", TEXTS.get(6), 70);
83+
validateSearchAfterQuery(7, 65, List.of(60, 50, 40, 30, 20, 10));
84+
validateSearchAfterQuery(7, 35, List.of(30, 20, 10));
85+
} finally {
86+
wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, SEARCH_PIPELINE_NAME);
87+
}
88+
}
89+
}
90+
91+
private void validateSearchAfterQuery(final int expectedDocCount, final int searchAfterValue, final List<Integer> expectedStockValues) {
92+
int docCount = getDocCount(getIndexNameForTest());
93+
assertEquals(expectedDocCount, docCount);
94+
95+
Map<String, SortOrder> fieldSortOrderMap = new LinkedHashMap<>();
96+
fieldSortOrderMap.put(SORT_FIELD, SortOrder.DESC);
97+
List<Object> searchAfter = new ArrayList<>();
98+
searchAfter.add(searchAfterValue);
99+
100+
Map<String, Object> searchResponseAsMap = search(
101+
getIndexNameForTest(),
102+
getQueryBuilder(modelId),
103+
null,
104+
QUERY_SIZE,
105+
Map.of("search_pipeline", SEARCH_PIPELINE_NAME),
106+
null,
107+
null,
108+
createSortBuilders(fieldSortOrderMap, false),
109+
false,
110+
searchAfter,
111+
0,
112+
null
113+
);
114+
assertNotNull(searchResponseAsMap);
115+
assertEquals(expectedStockValues.size(), getHitCount(searchResponseAsMap));
116+
List<Integer> actualStockValues = getStockValuesFromSortFields(searchResponseAsMap);
117+
assertEquals(expectedStockValues, actualStockValues);
118+
}
119+
120+
@SuppressWarnings("unchecked")
121+
private List<Integer> getStockValuesFromSortFields(final Map<String, Object> searchResponseAsMap) {
122+
Map<String, Object> hitsMap = (Map<String, Object>) searchResponseAsMap.get("hits");
123+
List<Map<String, Object>> hitsList = (List<Map<String, Object>>) hitsMap.get("hits");
124+
List<Integer> stockValues = new ArrayList<>();
125+
for (Map<String, Object> hit : hitsList) {
126+
List<Object> sortValues = (List<Object>) hit.get("sort");
127+
assertNotNull(sortValues);
128+
assertEquals(1, sortValues.size());
129+
stockValues.add(((Number) sortValues.get(0)).intValue());
130+
}
131+
return stockValues;
132+
}
133+
134+
private void addDocumentWithSortField(final String index, final String docId, final String text, final int sortFieldValue)
135+
throws IOException {
136+
Request request = new Request("PUT", "/" + index + "/_doc/" + docId + "?refresh=true");
137+
XContentBuilder builder = XContentFactory.jsonBuilder().startObject();
138+
builder.field(TEST_FIELD, text);
139+
builder.field(SORT_FIELD, sortFieldValue);
140+
builder.endObject();
141+
request.setJsonEntity(builder.toString());
142+
client().performRequest(request);
143+
}
144+
145+
private HybridQueryBuilder getQueryBuilder(final String modelId) {
146+
NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
147+
.fieldName(VECTOR_EMBEDDING_FIELD)
148+
.modelId(modelId)
149+
.queryText(QUERY)
150+
// k is intentionally larger than the total document count so that every document
151+
// is a candidate of the neural sub-query and pages have deterministic sizes
152+
.k(100)
153+
.build();
154+
155+
MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_FIELD, QUERY);
156+
157+
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
158+
hybridQueryBuilder.add(matchQueryBuilder);
159+
hybridQueryBuilder.add(neuralQueryBuilder);
160+
return hybridQueryBuilder;
161+
}
162+
}

‎qa/rolling-upgrade/build.gradle‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,11 @@ task testAgainstOldCluster(type: StandaloneRestIntegTestTask) {
8888
nonInputProperties.systemProperty('tests.rest.cluster', "${-> testClusters."${baseName}".allHttpSocketURI.join(",")}")
8989
nonInputProperties.systemProperty('tests.clustername', "${-> testClusters."${baseName}".getName()}")
9090
systemProperty 'tests.security.manager', 'false'
91+
if (versionsBelow2_16.any { ext.neural_search_bwc_version.startsWith(it) }){
92+
filter {
93+
excludeTestsMatching "org.opensearch.neuralsearch.bwc.rolling.HybridSearchWithSearchAfterIT.*"
94+
}
95+
}
9196

9297
// Excluding stats tests because we introduce this feature in 3.0
9398
// Excluding semantic highlighting BWC tests because we introduce this feature in 3.0
@@ -164,6 +169,12 @@ task testAgainstOneThirdUpgradedCluster(type: StandaloneRestIntegTestTask) {
164169
nonInputProperties.systemProperty('tests.clustername', "${-> testClusters."${baseName}".getName()}")
165170
systemProperty 'tests.security.manager', 'false'
166171

172+
if (versionsBelow2_16.any { ext.neural_search_bwc_version.startsWith(it) }){
173+
filter {
174+
excludeTestsMatching "org.opensearch.neuralsearch.bwc.rolling.HybridSearchWithSearchAfterIT.*"
175+
}
176+
}
177+
167178
// Excluding stats tests because we introduce this feature in 3.0
168179
// Excluding semantic highlighting BWC tests because we introduce this feature in 3.0
169180
if (versionsBelow3_0.any { ext.neural_search_bwc_version.startsWith(it) }){
@@ -238,6 +249,11 @@ task testAgainstTwoThirdsUpgradedCluster(type: StandaloneRestIntegTestTask) {
238249
nonInputProperties.systemProperty('tests.clustername', "${-> testClusters."${baseName}".getName()}")
239250
systemProperty 'tests.security.manager', 'false'
240251

252+
if (versionsBelow2_16.any { ext.neural_search_bwc_version.startsWith(it) }){
253+
filter {
254+
excludeTestsMatching "org.opensearch.neuralsearch.bwc.rolling.HybridSearchWithSearchAfterIT.*"
255+
}
256+
}
241257
// Excluding stats tests because we introduce this feature in 3.0
242258
// Excluding semantic highlighting BWC tests because we introduce this feature in 3.0
243259
if (versionsBelow3_0.any { ext.neural_search_bwc_version.startsWith(it) }){
@@ -313,6 +329,11 @@ task testRollingUpgrade(type: StandaloneRestIntegTestTask) {
313329
nonInputProperties.systemProperty('tests.clustername', "${-> testClusters."${baseName}".getName()}")
314330
systemProperty 'tests.security.manager', 'false'
315331

332+
if (versionsBelow2_16.any { ext.neural_search_bwc_version.startsWith(it) }){
333+
filter {
334+
excludeTestsMatching "org.opensearch.neuralsearch.bwc.rolling.HybridSearchWithSearchAfterIT.*"
335+
}
336+
}
316337
// Excluding stats tests because we introduce this feature in 3.0
317338
// Excluding semantic highlighting BWC tests because we introduce this feature in 3.0
318339
if (versionsBelow3_0.any { ext.neural_search_bwc_version.startsWith(it) }){

0 commit comments

Comments
 (0)