[NOCOMMIT] Approximate match all query benchmark

.github/benchmark-configs.json

@@ -221,5 +221,23 @@
       "data_instance_config": "4vCPU, 32G Mem, 16G Heap"
     },
     "baseline_cluster_config": "x64-r5.xlarge-1-shard-0-replica-snapshot-baseline"
+  },
+  "id_14": {
+    "description": "Search only test-procedure for big5, uses snapshot to restore the data for OS-3.0.0. Enables range query approximation.",
+    "supported_major_versions": ["3"],
+    "cluster-benchmark-configs": {
+      "SINGLE_NODE_CLUSTER": "true",
+      "MIN_DISTRIBUTION": "true",
+      "TEST_WORKLOAD": "big5",
+      "ADDITIONAL_CONFIG": "opensearch.experimental.feature.approximate_point_range_query.enabled:true",
+      "WORKLOAD_PARAMS": "{\"snapshot_repo_name\":\"benchmark-workloads-repo-300\",\"snapshot_bucket_name\":\"benchmark-workload-snapshots\",\"snapshot_region\":\"us-east-1\",\"snapshot_base_path\":\"workload-snapshots-300\",\"snapshot_name\":\"big5_1_shard_ordered\"}",
+      "CAPTURE_NODE_STAT": "true",
+      "TEST_PROCEDURE": "restore-from-snapshot"
+    },
+    "cluster_configuration": {
+      "size": "Single-Node",
+      "data_instance_config": "4vCPU, 32G Mem, 16G Heap"
+    },
+    "baseline_cluster_config": "x64-r5.xlarge-1-shard-0-replica-snapshot-baseline"
   }
 }

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,40 @@ public Query rangeQuery(
                 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, pointRangeQuery);
+                            }
+                        } 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
+                                ) {
+                                    @Override
+                                    protected String toString(int dimension, byte[] value) {
+                                        return Long.toString(LongPoint.decodeDimension(value, 0));
+                                    }
+                                }
+                            );
                         }
                         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 @@ public static MatchAllQueryBuilder fromXContent(XContentParser parser) {
 
     @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,78 @@
+/*
+ * 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;
+
+public class ApproximateMatchAllQuery extends ApproximateQuery {
+    private ApproximateQuery approximation = null;
+
+    @Override
+    protected boolean canApproximate(SearchContext context) {
+        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) {
+        return sameClassAs(o);
+    }
+
+    @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;
+    }
+}

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

@@ -42,7 +42,7 @@ public ApproximateQuery getApproximationQuery() {
     }
 
     @Override
-    public final Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    public Query rewrite(IndexSearcher indexSearcher) throws IOException {
         if (resolvedQuery == null) {
             throw new IllegalStateException("Cannot rewrite resolved query without setContext being called");
         }