From ff8a2fd6ffe504f9350f2e88398a1b68e2376bc3 Mon Sep 17 00:00:00 2001 From: Harsha Vamsi Kalluri Date: Thu, 22 Aug 2024 15:42:36 -0400 Subject: [PATCH] Fix date histogram test Signed-off-by: Harsha Vamsi Kalluri --- .../aggregations/bucket/filterrewrite/Helper.java | 14 ++++++++++++++ .../approximate/ApproximatePointRangeQuery.java | 5 ++++- .../search/approximate/ApproximateScoreQuery.java | 8 ++++++++ 3 files changed, 26 insertions(+), 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 index 7493754d8efa2..d89a0fb0cc1d5 100644 --- 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 index 4870a378c34b1..3932d2d04767b 100644 --- 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 index fad46712136a8..1673d345ac1bf 100644 --- 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); + } }; }