Approximate match all query with a sort (opensearch-project#17772)

* Adding approximate match all query

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>

* Minor refactor

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>

* Fix doc field sorting

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>

* Fix tests

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>

* Fix more tests

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>

* Fix more tests

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>

* Fix backward tests

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>

---------

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
Signed-off-by: Michael Froh <[email protected]>
Co-authored-by: Michael Froh <[email protected]>
Signed-off-by: Sriram Ganesh <[email protected]>

Author: 2 people

CHANGELOG.md

@@ -36,6 +36,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add update and delete support in pull-based ingestion ([#17822](https://github.com/opensearch-project/OpenSearch/pull/17822))
 - Allow maxPollSize and pollTimeout in IngestionSource to be configurable ([#17863](https://github.com/opensearch-project/OpenSearch/pull/17863))
 - [Star Tree] [Search] Add query changes to support unsigned-long in star tree ([#17275](https://github.com/opensearch-project/OpenSearch/pull/17275))
+- Add `ApproximateMatchAllQuery` that targets match_all queries and approximates sorts ([#17772](https://github.com/opensearch-project/OpenSearch/pull/17772))
 - Add TermsQuery support to Search GRPC endpoint ([#17888](https://github.com/opensearch-project/OpenSearch/pull/17888))
 
 ### Changed

rest-api-spec/src/main/resources/rest-api-spec/test/indices.validate_query/10_basic.yml

@@ -12,7 +12,9 @@ setup:
 
 ---
 "Validate query api":
-
+  - skip:
+      version: ' - 2.99.99'
+      reason: message changed in 3.0.0
 
   - do:
       indices.validate_query:
@@ -64,7 +66,7 @@ setup:
   - is_true: valid
   - match: {_shards.failed: 0}
   - match: {explanations.0.index: 'testing'}
-  - match: {explanations.0.explanation: '*:*'}
+  - match: {explanations.0.explanation: 'ApproximateScoreQuery(originalQuery=*:*, approximationQuery=Approximate(*:*))'}
 
 ---
 "Validate body without query element":

server/src/internalClusterTest/java/org/opensearch/search/sort/FieldSortIT.java

@@ -1145,11 +1145,10 @@ public void testSortMissingNumbersMinMax() throws Exception {
             .get();
         assertNoFailures(searchResponse);
 
-        assertThat(searchResponse.getHits().getTotalHits().value(), equalTo(3L));
+        assertThat(searchResponse.getHits().getTotalHits().value(), equalTo(2L));
         assertThat(searchResponse.getHits().getAt(0).getId(), equalTo("1"));
         // The order here could be unstable (depends on document order) since missing == field value
         assertThat(searchResponse.getHits().getAt(1).getId(), is(oneOf("3", "2")));
-        assertThat(searchResponse.getHits().getAt(2).getId(), is(oneOf("2", "3")));
 
         logger.info("--> sort with missing _last");
         searchResponse = client().prepareSearch()

server/src/internalClusterTest/java/org/opensearch/validate/SimpleValidateQueryIT.java

@@ -294,7 +294,10 @@ public void testExplainNoQuery() {
         assertThat(validateQueryResponse.isValid(), equalTo(true));
         assertThat(validateQueryResponse.getQueryExplanation().size(), equalTo(1));
         assertThat(validateQueryResponse.getQueryExplanation().get(0).getIndex(), equalTo("test"));
-        assertThat(validateQueryResponse.getQueryExplanation().get(0).getExplanation(), equalTo("*:*"));
+        assertThat(
+            validateQueryResponse.getQueryExplanation().get(0).getExplanation(),
+            equalTo("ApproximateScoreQuery(originalQuery=*:*, approximationQuery=Approximate(*:*))")
+        );
     }
 
     public void testExplainFilteredAlias() {

server/src/main/java/org/opensearch/index/query/MatchAllQueryBuilder.java

@@ -40,6 +40,8 @@
 import org.opensearch.core.xcontent.ObjectParser;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.search.approximate.ApproximateMatchAllQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 
 import java.io.IOException;
 
@@ -88,7 +90,7 @@ public static MatchAllQueryBuilder fromXContent(XContentParser parser) {
 
     @Override
     protected Query doToQuery(QueryShardContext context) {
-        return Queries.newMatchAllQuery();
+        return new ApproximateScoreQuery(Queries.newMatchAllQuery(), new ApproximateMatchAllQuery());
     }
 
     @Override

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

@@ -0,0 +1,95 @@
+/*
+ * 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.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.QueryVisitor;
+import org.opensearch.index.mapper.MappedFieldType;
+import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.sort.FieldSortBuilder;
+
+import java.io.IOException;
+import java.util.Objects;
+
+/**
+ * Replaces match-all query with a less expensive query if possible.
+ * <p>
+ * Currently, will rewrite to a bounded range query over the high/low end of a field if a primary sort is specified
+ * on that field.
+ */
+public class ApproximateMatchAllQuery extends ApproximateQuery {
+    private ApproximateQuery approximation = null;
+
+    @Override
+    protected boolean canApproximate(SearchContext context) {
+        approximation = null;
+        if (context == null) {
+            return false;
+        }
+        if (context.aggregations() != null) {
+            return false;
+        }
+
+        if (context.request() != null && context.request().source() != null && context.innerHits().getInnerHits().isEmpty()) {
+            FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source());
+            if (primarySortField != null
+                && primarySortField.missing() == null
+                && !primarySortField.fieldName().equals(FieldSortBuilder.DOC_FIELD_NAME)
+                && !primarySortField.fieldName().equals(FieldSortBuilder.ID_FIELD_NAME)) {
+                MappedFieldType mappedFieldType = context.getQueryShardContext().fieldMapper(primarySortField.fieldName());
+                if (mappedFieldType == null) {
+                    return false;
+                }
+                Query rangeQuery = mappedFieldType.rangeQuery(null, null, false, false, null, null, null, context.getQueryShardContext());
+                if (rangeQuery instanceof ApproximateScoreQuery approximateScoreQuery) {
+                    approximateScoreQuery.setContext(context);
+                    if (approximateScoreQuery.resolvedQuery instanceof ApproximateQuery) {
+                        approximation = (ApproximateQuery) approximateScoreQuery.resolvedQuery;
+                        return true;
+                    }
+                }
+            }
+        }
+        return false;
+    }
+
+    @Override
+    public String toString(String field) {
+        return "Approximate(*:*)";
+    }
+
+    @Override
+    public void visit(QueryVisitor visitor) {
+        visitor.visitLeaf(this);
+
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (sameClassAs(o)) {
+            ApproximateMatchAllQuery other = (ApproximateMatchAllQuery) o;
+            return Objects.equals(approximation, other.approximation);
+        }
+        return false;
+    }
+
+    @Override
+    public int hashCode() {
+        return classHash();
+    }
+
+    @Override
+    public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+        if (approximation == null) {
+            throw new IllegalStateException("rewrite called without setting context or query could not be approximated");
+        }
+        return approximation.rewrite(indexSearcher);
+    }
+}

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

@@ -51,9 +51,6 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
     }
 
     public void setContext(SearchContext context) {
-        if (resolvedQuery != null) {
-            throw new IllegalStateException("Query already resolved, duplicate call to setContext");
-        }
         resolvedQuery = approximationQuery.canApproximate(context) ? approximationQuery : originalQuery;
     };
 

server/src/main/java/org/opensearch/search/sort/FieldSortBuilder.java

@@ -100,6 +100,7 @@ public class FieldSortBuilder extends SortBuilder<FieldSortBuilder> implements W
      * special field name to sort by index order
      */
     public static final String DOC_FIELD_NAME = "_doc";
+    public static final String ID_FIELD_NAME = "_id";
     private static final SortFieldAndFormat SORT_DOC = new SortFieldAndFormat(new SortField(null, SortField.Type.DOC), DocValueFormat.RAW);
     private static final SortFieldAndFormat SORT_DOC_REVERSE = new SortFieldAndFormat(
         new SortField(null, SortField.Type.DOC, true),

server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java

@@ -42,6 +42,8 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParseException;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.search.approximate.ApproximateMatchAllQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.test.AbstractQueryTestCase;
 import org.hamcrest.Matchers;
 
@@ -208,7 +210,10 @@ public void testMinShouldMatchFilterWithoutShouldClauses() throws Exception {
         assertThat(innerBooleanQuery.clauses().size(), equalTo(1));
         BooleanClause innerBooleanClause = innerBooleanQuery.clauses().get(0);
         assertThat(innerBooleanClause.occur(), equalTo(BooleanClause.Occur.MUST));
-        assertThat(innerBooleanClause.query(), instanceOf(MatchAllDocsQuery.class));
+        assertThat(innerBooleanClause.query(), instanceOf(ApproximateScoreQuery.class));
+        ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) innerBooleanClause.query();
+        assertThat(approxQuery.getOriginalQuery(), instanceOf(MatchAllDocsQuery.class));
+        assertThat(approxQuery.getApproximationQuery(), instanceOf(ApproximateMatchAllQuery.class));
     }
 
     public void testMinShouldMatchBiggerThanNumberOfShouldClauses() throws Exception {

server/src/test/java/org/opensearch/index/query/MatchAllQueryBuilderTests.java

@@ -34,6 +34,8 @@
 
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
+import org.opensearch.search.approximate.ApproximateMatchAllQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.test.AbstractQueryTestCase;
 
 import java.io.IOException;
@@ -49,7 +51,10 @@ protected MatchAllQueryBuilder doCreateTestQueryBuilder() {
 
     @Override
     protected void doAssertLuceneQuery(MatchAllQueryBuilder queryBuilder, Query query, QueryShardContext context) throws IOException {
-        assertThat(query, instanceOf(MatchAllDocsQuery.class));
+        assertThat(query, instanceOf(ApproximateScoreQuery.class));
+        ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) query;
+        assertThat(approxQuery.getOriginalQuery(), instanceOf(MatchAllDocsQuery.class));
+        assertThat(approxQuery.getApproximationQuery(), instanceOf(ApproximateMatchAllQuery.class));
     }
 
     public void testFromJson() throws IOException {

server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java

@@ -50,6 +50,8 @@
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder;
 import org.opensearch.index.search.OpenSearchToParentBlockJoinQuery;
+import org.opensearch.search.approximate.ApproximateMatchAllQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.search.fetch.subphase.InnerHitsContext;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.sort.FieldSortBuilder;
@@ -493,7 +495,10 @@ public void testNestedDepthAllowed() throws Exception {
                 .filter(c -> c.occur() == BooleanClause.Occur.MUST)
                 .findFirst();
             assertTrue(childLeg.isPresent());
-            assertEquals(new MatchAllDocsQuery(), childLeg.get().query());
+            assertThat(childLeg.get().query(), instanceOf(ApproximateScoreQuery.class));
+            ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) childLeg.get().query();
+            assertThat(approxQuery.getOriginalQuery(), instanceOf(MatchAllDocsQuery.class));
+            assertThat(approxQuery.getApproximationQuery(), instanceOf(ApproximateMatchAllQuery.class));
         };
         check.accept(createShardContext());
         doWithDepth(randomIntBetween(1, 20), check);

server/src/test/java/org/opensearch/index/query/WrapperQueryBuilderTests.java

@@ -41,12 +41,16 @@
 import org.opensearch.core.common.bytes.BytesArray;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
+import org.opensearch.search.approximate.ApproximateMatchAllQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.test.AbstractQueryTestCase;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.io.UnsupportedEncodingException;
 
+import static org.hamcrest.CoreMatchers.instanceOf;
+
 public class WrapperQueryBuilderTests extends AbstractQueryTestCase<WrapperQueryBuilder> {
 
     @Override
@@ -171,7 +175,10 @@ public void testRewriteInnerQueryToo() throws IOException {
         assertEquals(new TermQuery(new Term(TEXT_FIELD_NAME, "bar")), qb.rewrite(shardContext).toQuery(shardContext));
 
         qb = new WrapperQueryBuilder(new BoolQueryBuilder().toString());
-        assertEquals(new MatchAllDocsQuery(), qb.rewrite(shardContext).toQuery(shardContext));
+        assertThat(qb.rewrite(shardContext).toQuery(shardContext), instanceOf(ApproximateScoreQuery.class));
+        ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) qb.rewrite(shardContext).toQuery(shardContext);
+        assertThat(approxQuery.getOriginalQuery(), instanceOf(MatchAllDocsQuery.class));
+        assertThat(approxQuery.getApproximationQuery(), instanceOf(ApproximateMatchAllQuery.class));
     }
 
     @Override

server/src/test/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java

@@ -67,6 +67,8 @@
 import org.opensearch.script.Script;
 import org.opensearch.script.ScriptType;
 import org.opensearch.search.MultiValueMode;
+import org.opensearch.search.approximate.ApproximateMatchAllQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.test.AbstractQueryTestCase;
 import org.opensearch.test.TestGeoShapeFieldMapperPlugin;
 import org.joda.time.DateTime;
@@ -630,7 +632,10 @@ public void testCustomWeightFactorQueryBuilderWithFunctionScoreWithoutQueryGiven
         Query parsedQuery = parseQuery(functionScoreQuery(weightFactorFunction(1.3f))).toQuery(createShardContext());
         assertThat(parsedQuery, instanceOf(FunctionScoreQuery.class));
         FunctionScoreQuery functionScoreQuery = (FunctionScoreQuery) parsedQuery;
-        assertThat(functionScoreQuery.getSubQuery() instanceof MatchAllDocsQuery, equalTo(true));
+        assertThat(functionScoreQuery.getSubQuery(), CoreMatchers.instanceOf(ApproximateScoreQuery.class));
+        ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) functionScoreQuery.getSubQuery();
+        assertThat(approxQuery.getOriginalQuery(), CoreMatchers.instanceOf(MatchAllDocsQuery.class));
+        assertThat(approxQuery.getApproximationQuery(), CoreMatchers.instanceOf(ApproximateMatchAllQuery.class));
         assertThat((double) (functionScoreQuery.getFunctions()[0]).getWeight(), closeTo(1.3, 0.001));
     }
 

server/src/test/java/org/opensearch/index/search/NestedHelperTests.java

@@ -42,6 +42,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.search.join.ScoreMode;
+import org.opensearch.common.lucene.search.Queries;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -52,6 +53,8 @@
 import org.opensearch.index.query.NestedQueryBuilder;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.TermQueryBuilder;
+import org.opensearch.search.approximate.ApproximateMatchAllQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.test.OpenSearchSingleNodeTestCase;
 
 import java.io.IOException;
@@ -325,7 +328,10 @@ public void testNested() throws IOException {
         NestedQueryBuilder queryBuilder = new NestedQueryBuilder("nested1", new MatchAllQueryBuilder(), ScoreMode.Avg);
         OpenSearchToParentBlockJoinQuery query = (OpenSearchToParentBlockJoinQuery) queryBuilder.toQuery(context);
 
-        Query expectedChildQuery = new BooleanQuery.Builder().add(new MatchAllDocsQuery(), Occur.MUST)
+        Query expectedChildQuery = new BooleanQuery.Builder().add(
+            new ApproximateScoreQuery(Queries.newMatchAllQuery(), new ApproximateMatchAllQuery()),
+            Occur.MUST
+        )
             // we automatically add a filter since the inner query might match non-nested docs
             .add(new TermQuery(new Term(NestedPathFieldMapper.NAME, "nested1")), Occur.FILTER)
             .build();