diff --git a/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java index 234b40d4e40e3..91d96b3e2c2f7 100644 --- a/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java @@ -467,7 +467,7 @@ public Query rangeQuery( } DateMathParser parser = forcedDateParser == null ? dateMathParser : forcedDateParser; return dateRangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, timeZone, parser, context, resolution, (l, u) -> { - Query pointRangeQuery = createPointRangeQuery(l, u); + Query pointRangeQuery = isSearchable() ? createPointRangeQuery(l, u) : null; Query dvQuery = hasDocValues() ? SortedNumericDocValuesField.newSlowRangeQuery(name(), l, u) : null; if (isSearchable() && hasDocValues()) { Query query = new IndexOrDocValuesQuery(pointRangeQuery, dvQuery); 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 55dfa9953293d..19926e96b2aa6 100644 --- a/server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java @@ -14,9 +14,9 @@ import org.apache.lucene.search.ConstantScoreScorer; import org.apache.lucene.search.ConstantScoreWeight; import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.IndexOrDocValuesQuery; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.PointRangeQuery; -import org.apache.lucene.search.Query; import org.apache.lucene.search.QueryVisitor; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.search.Scorer; @@ -25,6 +25,9 @@ import org.apache.lucene.util.ArrayUtil; import org.apache.lucene.util.DocIdSetBuilder; import org.apache.lucene.util.IntsRef; +import org.opensearch.index.query.RangeQueryBuilder; +import org.opensearch.search.internal.SearchContext; +import org.opensearch.search.sort.FieldSortBuilder; import org.opensearch.search.sort.SortOrder; import java.io.IOException; @@ -35,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 abstract class ApproximatePointRangeQuery extends Query { +public abstract class ApproximatePointRangeQuery extends ApproximateableQuery { private int size; private SortOrder sortOrder; @@ -419,6 +422,29 @@ public boolean isCacheable(LeafReaderContext ctx) { }; } + @Override + public boolean canApproximate(SearchContext context) { + if (!(context.query() instanceof IndexOrDocValuesQuery + && ((IndexOrDocValuesQuery) context.query()).getIndexQuery() instanceof ApproximateScoreQuery + && ((ApproximateScoreQuery) ((IndexOrDocValuesQuery) context.query()).getIndexQuery()) + .getOriginalQuery() instanceof PointRangeQuery)) { + return false; + } + ApproximateScoreQuery query = ((ApproximateScoreQuery) ((IndexOrDocValuesQuery) context.query()).getIndexQuery()); + ((ApproximatePointRangeQuery) query.getApproximationQuery()).setSize(Math.max(context.size(), context.trackTotalHitsUpTo())); + if (context.request() != null && context.request().source() != null) { + FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source()); + if (primarySortField != null + && primarySortField.missing() == null + && primarySortField.getFieldName().equals(((RangeQueryBuilder) context.request().source().query()).fieldName())) { + if (primarySortField.order() == SortOrder.DESC) { + ((ApproximatePointRangeQuery) query.getApproximationQuery()).setSortOrder(SortOrder.DESC); + } + } + } + return true; + } + @Override public final int hashCode() { return pointRangeQuery.hashCode(); 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 b102f7e727b4c..7c3206f259df7 100644 --- a/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java @@ -19,6 +19,7 @@ import org.apache.lucene.search.Scorer; import org.apache.lucene.search.ScorerSupplier; import org.apache.lucene.search.Weight; +import org.opensearch.search.internal.SearchContext; import java.io.IOException; @@ -29,17 +30,20 @@ */ public final class ApproximateScoreQuery extends Query { - private final Query originalQuery, approximationQuery; + private final Query originalQuery; + private final ApproximateableQuery approximationQuery; private Weight originalQueryWeight, approximationQueryWeight; - public ApproximateScoreQuery(Query originalQuery, Query approximationQuery) { + private SearchContext context; + + public ApproximateScoreQuery(Query originalQuery, ApproximateableQuery approximationQuery) { this(originalQuery, approximationQuery, null, null); } public ApproximateScoreQuery( Query originalQuery, - Query approximationQuery, + ApproximateableQuery approximationQuery, Weight originalQueryWeight, Weight approximationQueryWeight ) { @@ -49,15 +53,11 @@ public ApproximateScoreQuery( this.approximationQueryWeight = approximationQueryWeight; } - public Weight getApproximationQueryWeight() { - return approximationQueryWeight; - } - public Query getOriginalQuery() { return originalQuery; } - public Query getApproximationQuery() { + public ApproximateableQuery getApproximationQuery() { return approximationQuery; } @@ -88,8 +88,9 @@ public ScorerSupplier scorerSupplier(LeafReaderContext leafReaderContext) throws return new ScorerSupplier() { @Override public Scorer get(long l) throws IOException { - // TODO: we need to figure out how to compute the cost of running two different queries, by default return the - // original query's scoreSupplier + if (approximationQuery.canApproximate(context)) { + return originalQueryScoreSupplier.get(l); + } return originalQueryScoreSupplier.get(l); } @@ -116,6 +117,10 @@ public boolean isCacheable(LeafReaderContext leafReaderContext) { }; } + public void setContext(SearchContext context) { + this.context = context; + }; + @Override public String toString(String s) { return "ApproximateScoreQuery(originalQuery=" diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateableQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateableQuery.java new file mode 100644 index 0000000000000..069b3f49ee817 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximateableQuery.java @@ -0,0 +1,18 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.approximate; + +import org.apache.lucene.search.Query; +import org.opensearch.search.internal.SearchContext; + +public abstract class ApproximateableQuery extends Query { + + protected abstract boolean canApproximate(SearchContext context); + +} diff --git a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java index ddc4961bffe3a..023cdd8320a1f 100644 --- a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java +++ b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java @@ -49,7 +49,6 @@ import org.apache.lucene.search.IndexOrDocValuesQuery; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.LeafCollector; -import org.apache.lucene.search.PointRangeQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.QueryCache; import org.apache.lucene.search.QueryCachingPolicy; @@ -69,10 +68,8 @@ import org.opensearch.common.annotation.PublicApi; import org.opensearch.common.lease.Releasable; import org.opensearch.common.lucene.search.TopDocsAndMaxScore; -import org.opensearch.index.query.RangeQueryBuilder; import org.opensearch.search.DocValueFormat; import org.opensearch.search.SearchService; -import org.opensearch.search.approximate.ApproximatePointRangeQuery; import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.search.dfs.AggregatedDfs; import org.opensearch.search.profile.ContextualProfileBreakdown; @@ -84,7 +81,6 @@ import org.opensearch.search.query.QuerySearchResult; import org.opensearch.search.sort.FieldSortBuilder; import org.opensearch.search.sort.MinAndMax; -import org.opensearch.search.sort.SortOrder; import java.io.IOException; import java.util.ArrayList; @@ -333,22 +329,10 @@ private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collecto // catch early terminated exception and rethrow? Bits liveDocs = ctx.reader().getLiveDocs(); BitSet liveDocsBitSet = getSparseBitSetOrNull(liveDocs); - if (isApproximateableRangeQuery()) { + if (searchContext.query() instanceof IndexOrDocValuesQuery + && ((IndexOrDocValuesQuery) searchContext.query()).getIndexQuery() instanceof ApproximateScoreQuery) { ApproximateScoreQuery query = ((ApproximateScoreQuery) ((IndexOrDocValuesQuery) searchContext.query()).getIndexQuery()); - ((ApproximatePointRangeQuery) query.getApproximationQuery()).setSize( - Math.max(searchContext.size(), searchContext.trackTotalHitsUpTo()) - ); - if (searchContext.request() != null && searchContext.request().source() != null) { - FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(searchContext.request().source()); - if (primarySortField != null - && primarySortField.missing() == null - && primarySortField.getFieldName().equals(((RangeQueryBuilder) searchContext.request().source().query()).fieldName())) { - if (primarySortField.order() == SortOrder.DESC) { - ((ApproximatePointRangeQuery) query.getApproximationQuery()).setSortOrder(SortOrder.DESC); - } - } - } - weight = query.getApproximationQueryWeight(); + query.setContext(searchContext); } if (liveDocsBitSet == null) { BulkScorer bulkScorer = weight.bulkScorer(ctx); @@ -441,13 +425,6 @@ private static BitSet getSparseBitSetOrNull(Bits liveDocs) { } - private boolean isApproximateableRangeQuery() { - return searchContext.query() instanceof IndexOrDocValuesQuery - && ((IndexOrDocValuesQuery) searchContext.query()).getIndexQuery() instanceof ApproximateScoreQuery - && ((ApproximateScoreQuery) ((IndexOrDocValuesQuery) searchContext.query()).getIndexQuery()) - .getOriginalQuery() instanceof PointRangeQuery; - } - static void intersectScorerAndBitSet(Scorer scorer, BitSet acceptDocs, LeafCollector collector, Runnable checkCancelled) throws IOException { collector.setScorer(scorer); diff --git a/server/src/test/java/org/opensearch/search/approximate/ApproximateScoreQueryTests.java b/server/src/test/java/org/opensearch/search/approximate/ApproximateScoreQueryTests.java index d891235b4dec6..83fa4a500de76 100644 --- a/server/src/test/java/org/opensearch/search/approximate/ApproximateScoreQueryTests.java +++ b/server/src/test/java/org/opensearch/search/approximate/ApproximateScoreQueryTests.java @@ -42,7 +42,7 @@ protected String toString(int dimension, byte[] value) { } }; - Query approximateQuery = new ApproximatePointRangeQuery( + ApproximateableQuery approximateQuery = new ApproximatePointRangeQuery( "test-index", pack(new long[] { l }).bytes, pack(new long[] { u }).bytes,