Skip to content

Commit

Permalink
Fixing Integ Tests
Browse files Browse the repository at this point in the history 
Signed-off-by: Varun Jain <[email protected]>
  • Loading branch information
@vibrantvarun
vibrantvarun committed Nov 4, 2024
1 parent e0ed9cb commit c5d3c46
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 33 deletions.
8 changes: 6 additions & 2 deletions src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,9 @@ public final class HybridQueryBuilder extends AbstractQueryBuilder<HybridQueryBu
private final List<QueryBuilder> queries = new ArrayList<>();

private String fieldName;
private Integer paginationDepth;
private Integer paginationDepth = null;
static final int MAX_NUMBER_OF_SUB_QUERIES = 5;
private final static int DEFAULT_PAGINATION_DEPTH = 10;
private static final int LOWER_BOUND_OF_PAGINATION_DEPTH = 1;
private static final int UPPER_BOUND_OF_PAGINATION_DEPTH = 10000;

Expand Down Expand Up @@ -108,7 +109,9 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
queryBuilder.toXContent(builder, params);
}
builder.endArray();
builder.field(DEPTH_FIELD.getPreferredName(), paginationDepth);
if (isClusterOnOrAfterMinReqVersionForPaginationInHybridQuery()) {
builder.field(DEPTH_FIELD.getPreferredName(), paginationDepth == null ? DEFAULT_PAGINATION_DEPTH : paginationDepth);
}
printBoostAndQueryName(builder);
builder.endObject();
}
Expand Down Expand Up @@ -159,6 +162,7 @@ protected Query doToQuery(QueryShardContext queryShardContext) throws IOExceptio
* @throws IOException
*/
public static HybridQueryBuilder fromXContent(XContentParser parser) throws IOException {
log.info("fromXContent called");
float boost = AbstractQueryBuilder.DEFAULT_BOOST;

Integer paginationDepth = null;
Expand Down
14 changes: 7 additions & 7 deletions src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -485,16 +485,16 @@ private ReduceableSearchResult reduceSearchResults(final List<ReduceableSearchRe
private static int getSubqueryResultsRetrievalSize(final SearchContext searchContext) {
HybridQuery hybridQuery = getHybridQueryFromAbstractQuery(searchContext.query());
Integer paginationDepth = hybridQuery.getPaginationDepth();

// Switch to from+size retrieval size when pagination_depth is null.
if (searchContext.from() == 0) {
return searchContext.from() + searchContext.size();
}
// Pagination is expected to work only pagination_depth is provided to hold the reference of search result.
if (searchContext.from() > 0 && (Objects.isNull(paginationDepth))) {
if (Objects.isNull(paginationDepth)) {
throw new IllegalArgumentException(String.format(Locale.ROOT, "pagination_depth is missing in the search request"));
}
if (paginationDepth != null) {
return paginationDepth;
} else {
// Switch to from+size retrieval size when pagination_depth is null.
return searchContext.from() + searchContext.size();
}
return paginationDepth;
}

private static HybridQuery getHybridQueryFromAbstractQuery(Query query) {
Expand Down
48 changes: 24 additions & 24 deletions src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -639,6 +639,7 @@ public void testWrappedQueryWithFilter_whenIndexAliasHasFilterAndIndexWithNested

HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1));
// hybridQueryBuilder.paginationDepth(10);

Map<String, Object> searchResponseAsMap = search(
alias,
Expand Down Expand Up @@ -845,7 +846,7 @@ public void testPaginationOnSingleShard_whenConcurrentSearchEnabled_thenSuccessf
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
testHybridQuery_whenFromIsGreaterThanZeroAndPaginationDepthIsNotSent_thenFail(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
testHybridQuery_whenFromIsGreaterThanZeroAndPaginationDepthIsNotSent_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
testHybridQuery_whenFromIsGreaterThanTotalResultCount_thenFail(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
testHybridQuery_whenPaginationDepthIsOutOfRange_thenFail(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
} finally {
Expand All @@ -860,7 +861,7 @@ public void testPaginationOnSingleShard_whenConcurrentSearchDisabled_thenSuccess
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
testHybridQuery_whenFromIsGreaterThanZeroAndPaginationDepthIsNotSent_thenFail(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
testHybridQuery_whenFromIsGreaterThanZeroAndPaginationDepthIsNotSent_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
testHybridQuery_whenFromIsGreaterThanTotalResultCount_thenFail(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
testHybridQuery_whenPaginationDepthIsOutOfRange_thenFail(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD);
} finally {
Expand All @@ -875,7 +876,7 @@ public void testPaginationOnMultipleShard_whenConcurrentSearchEnabled_thenSucces
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME);
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME);
testHybridQuery_whenFromIsGreaterThanZeroAndPaginationDepthIsNotSent_thenFail(TEST_MULTI_DOC_INDEX_NAME);
testHybridQuery_whenFromIsGreaterThanZeroAndPaginationDepthIsNotSent_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME);
testHybridQuery_whenFromIsGreaterThanTotalResultCount_thenFail(TEST_MULTI_DOC_INDEX_NAME);
testHybridQuery_whenPaginationDepthIsOutOfRange_thenFail(TEST_MULTI_DOC_INDEX_NAME);
} finally {
Expand All @@ -890,7 +891,7 @@ public void testPaginationOnMultipleShard_whenConcurrentSearchDisabled_thenSucce
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME);
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME);
testHybridQuery_whenFromIsGreaterThanZeroAndPaginationDepthIsNotSent_thenFail(TEST_MULTI_DOC_INDEX_NAME);
testHybridQuery_whenFromIsGreaterThanZeroAndPaginationDepthIsNotSent_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME);
testHybridQuery_whenFromIsGreaterThanTotalResultCount_thenFail(TEST_MULTI_DOC_INDEX_NAME);
testHybridQuery_whenPaginationDepthIsOutOfRange_thenFail(TEST_MULTI_DOC_INDEX_NAME);
} finally {
Expand Down Expand Up @@ -927,31 +928,30 @@ public void testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSucc
}

@SneakyThrows
public void testHybridQuery_whenFromIsGreaterThanZeroAndPaginationDepthIsNotSent_thenFail(String indexName) {
public void testHybridQuery_whenFromIsGreaterThanZeroAndPaginationDepthIsNotSent_thenSuccessful(String indexName) {
HybridQueryBuilder hybridQueryBuilderOnlyMatchAll = new HybridQueryBuilder();
hybridQueryBuilderOnlyMatchAll.add(new MatchAllQueryBuilder());

ResponseException responseException = assertThrows(
ResponseException.class,
() -> search(
indexName,
hybridQueryBuilderOnlyMatchAll,
null,
10,
Map.of("search_pipeline", SEARCH_PIPELINE),
null,
null,
null,
false,
null,
2
)
Map<String, Object> searchResponseAsMap = search(
indexName,
hybridQueryBuilderOnlyMatchAll,
null,
10,
Map.of("search_pipeline", SEARCH_PIPELINE),
null,
null,
null,
false,
null,
2
);

org.hamcrest.MatcherAssert.assertThat(
responseException.getMessage(),
allOf(containsString("pagination_depth is missing in the search request"))
);
assertEquals(2, getHitCount(searchResponseAsMap));
Map<String, Object> total = getTotalHits(searchResponseAsMap);
assertNotNull(total.get("value"));
assertEquals(4, total.get("value"));
assertNotNull(total.get("relation"));
assertEquals(RELATION_EQUAL_TO, total.get("relation"));
}

@SneakyThrows
Expand Down

0 comments on commit c5d3c46

Please sign in to comment.