Fix date histogram test

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
harshavamsi · Aug 22, 2024 · ff8a2fd · ff8a2fd
1 parent 15fd737
commit ff8a2fd
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 1 deletion.
diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/Helper.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/Helper.java
@@ -23,6 +23,8 @@
 import org.opensearch.common.lucene.search.function.FunctionScoreQuery;
 import org.opensearch.index.mapper.DateFieldMapper;
 import org.opensearch.index.query.DateRangeIncludingNowQuery;
+import org.opensearch.search.approximate.ApproximatePointRangeQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -123,6 +125,18 @@ public static long[] getDateHistoAggBounds(final SearchContext context, final St
             final long[] indexBounds = getShardBounds(leaves, fieldName);
             if (indexBounds == null) return null;
             return getBoundsWithRangeQuery(prq, fieldName, indexBounds);
+        } else if (cq instanceof ApproximateScoreQuery) {
+            final ApproximateScoreQuery asq = (ApproximateScoreQuery) cq;
+            final long[] indexBounds = getShardBounds(leaves, fieldName);
+            if (indexBounds == null) return null;
+            if ((asq.getApproximationQuery() instanceof ApproximatePointRangeQuery)) {
+                ApproximatePointRangeQuery aprq = (ApproximatePointRangeQuery) asq.getApproximationQuery();
+                if (aprq.canApproximate(context)) {
+                    return getBoundsWithRangeQuery(aprq.pointRangeQuery, fieldName, indexBounds);
+                }
+                final PointRangeQuery prq = (PointRangeQuery) asq.getOriginalQuery();
+                return getBoundsWithRangeQuery(prq, fieldName, indexBounds);
+            }
         } else if (cq instanceof MatchAllDocsQuery) {
             return getShardBounds(leaves, fieldName);
         } else if (cq instanceof FieldExistsQuery) {

diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java
@@ -45,7 +45,7 @@ public abstract class ApproximatePointRangeQuery extends ApproximateableQuery {
 
     private long[] docCount = { 0 };
 
-    private final PointRangeQuery pointRangeQuery;
+    public final PointRangeQuery pointRangeQuery;
 
     protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upperPoint, int numDims) {
         this(field, lowerPoint, upperPoint, numDims, 10_000, null);
@@ -427,6 +427,9 @@ public boolean canApproximate(SearchContext context) {
         if (context == null) {
             return false;
         }
+        if (context.aggregations() != null) {
+            return false;
+        }
         if (!(context.query() instanceof IndexOrDocValuesQuery
             && ((IndexOrDocValuesQuery) context.query()).getIndexQuery() instanceof ApproximateScoreQuery
             && ((ApproximateScoreQuery) ((IndexOrDocValuesQuery) context.query()).getIndexQuery())

diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java
@@ -114,6 +114,14 @@ public Scorer scorer(LeafReaderContext leafReaderContext) throws IOException {
             public boolean isCacheable(LeafReaderContext leafReaderContext) {
                 return originalQueryWeight.isCacheable(leafReaderContext);
             }
+
+            @Override
+            public int count(LeafReaderContext leafReaderContext) throws IOException {
+                if (approximationQuery.canApproximate(context)) {
+                    return approximationQueryWeight.count(leafReaderContext);
+                }
+                return originalQueryWeight.count(leafReaderContext);
+            }
         };
     }