diff --git a/src/main/java/org/opensearch/neuralsearch/query/HybridQuery.java b/src/main/java/org/opensearch/neuralsearch/query/HybridQuery.java index f08120ce9..d1e339bd5 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/HybridQuery.java +++ b/src/main/java/org/opensearch/neuralsearch/query/HybridQuery.java @@ -46,7 +46,8 @@ public HybridQuery(final Collection subQueries, final List filterQ if (subQueries.isEmpty()) { throw new IllegalArgumentException("collection of queries must not be empty"); } - if (hybridQueryContext.getPaginationDepth() == 0) { + Integer paginationDepth = hybridQueryContext.getPaginationDepth(); + if (Objects.nonNull(paginationDepth) && paginationDepth == 0) { throw new IllegalArgumentException("pagination_depth must not be zero"); } if (Objects.isNull(filterQueries) || filterQueries.isEmpty()) { diff --git a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java index 55f6b4a96..bea94e603 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java @@ -55,7 +55,7 @@ public final class HybridQueryBuilder extends AbstractQueryBuilder queries = new ArrayList<>(); - private int paginationDepth; + private Integer paginationDepth; static final int MAX_NUMBER_OF_SUB_QUERIES = 5; private final static int DEFAULT_PAGINATION_DEPTH = 10; @@ -65,7 +65,7 @@ public HybridQueryBuilder(StreamInput in) throws IOException { super(in); queries.addAll(readQueries(in)); if (isClusterOnOrAfterMinReqVersionForPaginationInHybridQuery()) { - paginationDepth = in.readInt(); + paginationDepth = in.readOptionalInt(); } } @@ -78,7 +78,7 @@ public HybridQueryBuilder(StreamInput in) throws IOException { protected void doWriteTo(StreamOutput out) throws IOException { writeQueries(out, queries); if (isClusterOnOrAfterMinReqVersionForPaginationInHybridQuery()) { - out.writeInt(paginationDepth); + out.writeOptionalInt(paginationDepth); } } @@ -109,8 +109,9 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep queryBuilder.toXContent(builder, params); } builder.endArray(); - if (isClusterOnOrAfterMinReqVersionForPaginationInHybridQuery()) { - builder.field(PAGINATION_DEPTH_FIELD.getPreferredName(), paginationDepth == 0 ? DEFAULT_PAGINATION_DEPTH : paginationDepth); + // TODO https://github.com/opensearch-project/neural-search/issues/1097 + if (Objects.nonNull(paginationDepth)) { + builder.field(PAGINATION_DEPTH_FIELD.getPreferredName(), paginationDepth); } printBoostAndQueryName(builder); builder.endObject(); @@ -324,6 +325,9 @@ private Collection toQueries(Collection queryBuilders, Quer } private static void validatePaginationDepth(final int paginationDepth, final QueryShardContext queryShardContext) { + if (Objects.isNull(paginationDepth)) { + return; + } if (paginationDepth < LOWER_BOUND_OF_PAGINATION_DEPTH) { throw new IllegalArgumentException( String.format(Locale.ROOT, "pagination_depth should be greater than %s", LOWER_BOUND_OF_PAGINATION_DEPTH) diff --git a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryContext.java b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryContext.java index 3be741bcd..34706e6e7 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryContext.java +++ b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryContext.java @@ -6,7 +6,6 @@ import lombok.Builder; import lombok.Getter; -import lombok.NonNull; /** * Class that holds the low level information of hybrid query in the form of context @@ -14,6 +13,5 @@ @Builder @Getter public class HybridQueryContext { - @NonNull - private int paginationDepth; + private Integer paginationDepth; } diff --git a/src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java b/src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java index 03290f421..3c6a7271f 100644 --- a/src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java +++ b/src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java @@ -491,6 +491,7 @@ private static int getSubqueryResultsRetrievalSize(final SearchContext searchCon if (searchContext.from() == 0) { return searchContext.size(); } + log.info("pagination_depth is {}", paginationDepth); return paginationDepth; } diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java index c22d174ec..a6cf4d29e 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java @@ -426,7 +426,7 @@ public void testFromXContent_whenMultipleSubQueries_thenBuildSuccessfully() { assertEquals(2, queryTwoSubQueries.queries().size()); assertTrue(queryTwoSubQueries.queries().get(0) instanceof NeuralQueryBuilder); assertTrue(queryTwoSubQueries.queries().get(1) instanceof TermQueryBuilder); - assertEquals(10, queryTwoSubQueries.paginationDepth()); + assertEquals(10, queryTwoSubQueries.paginationDepth().intValue()); // verify knn vector query NeuralQueryBuilder neuralQueryBuilder = (NeuralQueryBuilder) queryTwoSubQueries.queries().get(0); assertEquals(VECTOR_FIELD_NAME, neuralQueryBuilder.fieldName()); diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryTests.java index 8e429e270..26babdbce 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryTests.java @@ -99,7 +99,7 @@ public void testQueryBasics_whenMultipleDifferentQueries_thenSuccessful() { countOfQueries++; } assertEquals(2, countOfQueries); - assertEquals(10, query3.getQueryContext().getPaginationDepth()); + assertEquals(10, query3.getQueryContext().getPaginationDepth().intValue()); } @SneakyThrows