Desc debug

Signed-off-by: Prudhvi Godithi <[email protected]>

Author: prudhvigodithi

server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java

@@ -38,7 +38,7 @@
  * An approximate-able version of {@link PointRangeQuery}. It creates an instance of {@link PointRangeQuery} but short-circuits the intersect logic
  * after {@code size} is hit
  */
-public class ApproximatePointRangeQuery extends ApproximateQuery {
+    public class ApproximatePointRangeQuery extends ApproximateQuery {
     public static final Function<byte[], String> LONG_FORMAT = bytes -> Long.toString(LongPoint.decodeDimension(bytes, 0));
     private int size;
 
@@ -449,7 +449,7 @@ public boolean canApproximate(SearchContext context) {
         if (context.from() + context.size() == 0) {
             this.setSize(SearchContext.DEFAULT_TRACK_TOTAL_HITS_UP_TO);
         } else {
-            this.setSize(Math.max(context.from() + context.size(), context.trackTotalHitsUpTo() + 1));
+            this.setSize(Math.max(context.from() + context.size(), context.trackTotalHitsUpTo()));
         }
         if (context.request() != null && context.request().source() != null) {
             FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source());

server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java

@@ -12,6 +12,7 @@
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Weight;
 import org.opensearch.search.internal.SearchContext;
 
@@ -47,7 +48,11 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
             // Default to the original query. This suggests that we were not called from ContextIndexSearcher.
             return originalQuery.rewrite(indexSearcher);
         }
-        return resolvedQuery.rewrite(indexSearcher);
+        Query rewritten = resolvedQuery.rewrite(indexSearcher);
+        if (rewritten != resolvedQuery) {
+            resolvedQuery = rewritten;
+        }
+        return this;
     }
 
     public void setContext(SearchContext context) {
@@ -72,16 +77,23 @@ public void visit(QueryVisitor queryVisitor) {
 
     @Override
     public boolean equals(Object o) {
-        if (!sameClassAs(o)) {
-            return false;
-        }
-        return true;
+        if (!sameClassAs(o)) return false;
+        ApproximateScoreQuery other = (ApproximateScoreQuery) o;
+        return originalQuery.equals(other.originalQuery)
+            && resolvedQuery.equals(other.resolvedQuery)
+            && approximationQuery.equals(other.approximationQuery);
+    }
+
+    @Override
+    public Weight createWeight(IndexSearcher s, ScoreMode sm, float boost) throws IOException {
+        return resolvedQuery.createWeight(s, sm, boost);
     }
 
     @Override
     public int hashCode() {
         int h = classHash();
         h = 31 * h + originalQuery.hashCode();
+        h = 31 * h + resolvedQuery.hashCode();
         h = 31 * h + approximationQuery.hashCode();
         return h;
     }

server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java

@@ -74,6 +74,7 @@
 import org.opensearch.common.util.CachedSupplier;
 import org.opensearch.index.search.OpenSearchToParentBlockJoinQuery;
 import org.opensearch.search.DocValueFormat;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.search.collapse.CollapseContext;
 import org.opensearch.search.internal.ScrollContext;
 import org.opensearch.search.internal.SearchContext;
@@ -724,6 +725,8 @@ static int shortcutTotalHitCount(IndexReader reader, Query query) throws IOExcep
                 query = ((ConstantScoreQuery) query).getQuery();
             } else if (query instanceof BoostQuery) {
                 query = ((BoostQuery) query).getQuery();
+            } else if (query instanceof ApproximateScoreQuery) {
+                query = ((ApproximateScoreQuery) query).getOriginalQuery();
             } else {
                 break;
             }