Update pagination_depth datatype from int to Integer (#1094)

* Update pagination_depth datatype from int to Integer Signed-off-by: Varun Jain <[email protected]>
opensearch-project · Jan 14, 2025 · bff857d · bff857d
1 parent f6d8a12
commit bff857d
Show file tree

Hide file tree

Showing 6 changed files with 15 additions and 11 deletions.
diff --git 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<Query> subQueries, final List<Query> 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
@@ -55,7 +55,7 @@ public final class HybridQueryBuilder extends AbstractQueryBuilder<HybridQueryBu
 
     private final List<QueryBuilder> 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<Query> toQueries(Collection<QueryBuilder> 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
@@ -6,14 +6,12 @@
 
 import lombok.Builder;
 import lombok.Getter;
-import lombok.NonNull;
 
 /**
  * Class that holds the low level information of hybrid query in the form of context
  */
 @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
@@ -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
@@ -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
@@ -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