Approximate match all with sort with a bounded range query

CHANGELOG.md

@@ -26,6 +26,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add _list/shards API as paginated alternate to _cat/shards ([#14641](https://github.com/opensearch-project/OpenSearch/pull/14641))
 - Latency and Memory allocation improvements to Multi Term Aggregation queries ([#14993](https://github.com/opensearch-project/OpenSearch/pull/14993))
 - Flat object field use IndexOrDocValuesQuery to optimize query ([#14383](https://github.com/opensearch-project/OpenSearch/issues/14383))
+- [Experimental] Approximate match_all query with sort by range query ([#16321](https://github.com/opensearch-project/OpenSearch/pull/16321))
 
 ### Dependencies
 - Bump `com.azure:azure-identity` from 1.13.0 to 1.13.2 ([#15578](https://github.com/opensearch-project/OpenSearch/pull/15578))
@@ -85,7 +86,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Streaming Indexing] Fix intermittent 'The bulk request must be terminated by a newline [\n]' failures [#16337](https://github.com/opensearch-project/OpenSearch/pull/16337))
 - Fix wrong default value when setting `index.number_of_routing_shards` to null on index creation ([#16331](https://github.com/opensearch-project/OpenSearch/pull/16331))
 - Fix disk usage exceeds threshold cluster can't spin up issue ([#15258](https://github.com/opensearch-project/OpenSearch/pull/15258)))
-
 - Fix inefficient Stream API call chains ending with count() ([#15386](https://github.com/opensearch-project/OpenSearch/pull/15386))
 
 ### Security

server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java

@@ -39,6 +39,7 @@ protected FeatureFlagSettings(
         FeatureFlags.STAR_TREE_INDEX_SETTING,
         FeatureFlags.APPLICATION_BASED_CONFIGURATION_TEMPLATES_SETTING,
         FeatureFlags.READER_WRITER_SPLIT_EXPERIMENTAL_SETTING,
-        FeatureFlags.TERM_VERSION_PRECOMMIT_ENABLE_SETTING
+        FeatureFlags.TERM_VERSION_PRECOMMIT_ENABLE_SETTING,
+        FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY_SETTING // TODO: Copy set from FeatureFlags.ALL_FEATURE_FLAG_SETTINGS
     );
 }

server/src/main/java/org/opensearch/common/util/FeatureFlags.java

@@ -138,7 +138,8 @@ public class FeatureFlags {
         STAR_TREE_INDEX_SETTING,
         APPLICATION_BASED_CONFIGURATION_TEMPLATES_SETTING,
         READER_WRITER_SPLIT_EXPERIMENTAL_SETTING,
-        TERM_VERSION_PRECOMMIT_ENABLE_SETTING
+        TERM_VERSION_PRECOMMIT_ENABLE_SETTING,
+        APPROXIMATE_POINT_RANGE_QUERY_SETTING
     );
 
     /**

server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java

@@ -492,13 +492,9 @@ public Query rangeQuery(
                                 name(),
                                 pack(new long[] { l }).bytes,
                                 pack(new long[] { u }).bytes,
-                                new long[] { l }.length
-                            ) {
-                                @Override
-                                protected String toString(int dimension, byte[] value) {
-                                    return Long.toString(LongPoint.decodeDimension(value, 0));
-                                }
-                            }
+                                new long[] { l }.length,
+                                ApproximatePointRangeQuery.LONG_FORMAT
+                            )
                         );
                     }
                     return query;

server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java

@@ -60,6 +60,7 @@
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Setting.Property;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.common.bytes.BytesArray;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.core.xcontent.XContentParser.Token;
@@ -71,6 +72,8 @@
 import org.opensearch.index.fielddata.plain.SortedNumericIndexFieldData;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.search.DocValueFormat;
+import org.opensearch.search.approximate.ApproximatePointRangeQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.search.lookup.SearchLookup;
 import org.opensearch.search.query.BitmapDocValuesQuery;
 
@@ -1064,24 +1067,36 @@
                 QueryShardContext context
             ) {
                 return longRangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, (l, u) -> {
-                    if (isSearchable && hasDocValues) {
-                        Query query = LongPoint.newRangeQuery(field, l, u);
-                        Query dvQuery = SortedNumericDocValuesField.newSlowRangeQuery(field, l, u);
-                        query = new IndexOrDocValuesQuery(query, dvQuery);
-                        if (context.indexSortedOnField(field)) {
-                            query = new IndexSortSortedNumericDocValuesRangeQuery(field, l, u, query);
+                    Query dvQuery = hasDocValues ? SortedNumericDocValuesField.newSlowRangeQuery(field, l, u) : null;
+                    if (isSearchable) {
+                        Query pointRangeQuery = LongPoint.newRangeQuery(field, l, u);
+                        Query query;
+                        if (dvQuery != null) {
+                            query = new IndexOrDocValuesQuery(pointRangeQuery, dvQuery);
+                            if (context.indexSortedOnField(field)) {
+                                query = new IndexSortSortedNumericDocValuesRangeQuery(field, l, u, query);
+                            }
+                        } else {
+                            query = pointRangeQuery;
                         }
-                        return query;
-                    }
-                    if (hasDocValues) {
-                        Query query = SortedNumericDocValuesField.newSlowRangeQuery(field, l, u);
-                        if (context.indexSortedOnField(field)) {
-                            query = new IndexSortSortedNumericDocValuesRangeQuery(field, l, u, query);
+                        if (FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY_SETTING)) {
+                            return new ApproximateScoreQuery(
+                                query,
+                                new ApproximatePointRangeQuery(
+                                    field,
+                                    LongPoint.pack(new long[] { l }).bytes,
+                                    LongPoint.pack(new long[] { u }).bytes,
+                                    new long[] { l }.length,
+                                    ApproximatePointRangeQuery.LONG_FORMAT
+                                )
+                            );
                         }
                         return query;
                     }
-                    return LongPoint.newRangeQuery(field, l, u);
-
+                    if (context.indexSortedOnField(field)) {
+                        dvQuery = new IndexSortSortedNumericDocValuesRangeQuery(field, l, u, dvQuery);
+                    }
+                    return dvQuery;
                 });
             }
 

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

@@ -34,12 +34,15 @@
 
 import org.apache.lucene.search.Query;
 import org.opensearch.common.lucene.search.Queries;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.common.ParsingException;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 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 +91,11 @@
 
     @Override
     protected Query doToQuery(QueryShardContext context) {
-        return Queries.newMatchAllQuery();
+        Query query = Queries.newMatchAllQuery();
+        if (FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY_SETTING)) {
+            return new ApproximateScoreQuery(query, new ApproximateMatchAllQuery());
+        }
+        return query;
     }
 
     @Override

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

@@ -0,0 +1,90 @@
+/*
+ * 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) {
+            FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source());
+            if (primarySortField != null && primarySortField.missing() == null) {
+                MappedFieldType mappedFieldType = context.getQueryShardContext().fieldMapper(primarySortField.fieldName());
+                Query rangeQuery = mappedFieldType.rangeQuery(null, null, false, false, null, null, null, context.getQueryShardContext());
+                if (rangeQuery instanceof ApproximateScoreQuery) {
+                    ApproximateScoreQuery approximateScoreQuery = (ApproximateScoreQuery) rangeQuery;
+                    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);
+    }
+}