Remove ApproximateIndexOrDocValuesQuery (opensearch-project#16273)

Looking into the query approximation framework, I realized that we don't actually need ApproximateIndexOrDocValuesQuery. We already have ApproximateScoreQuery that is able to choose between the approximate query and the original query. Assuming the range query is over an indexed field, it is (potentially) possible to approximate. If there are doc values, then the original query will be an IndexOrDocValuesQuery. Otherwise, it will just be a PointRangeQuery. Either way, we can wrap the original query into a generic ApproximateScoreQuery along with the approximate query. The approximation logic doesn't need to be specialized to the IndexOrDocValue case. --------- Signed-off-by: Michael Froh <[email protected]>
amberzsy · Oct 11, 2024 · 53c9ddf · 53c9ddf
1 parent b3459fd
commit 53c9ddf
Show file tree

Hide file tree

Showing 13 changed files with 119 additions and 272 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -54,6 +54,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Remove identity-related feature flagged code from the RestController ([#15430](https://github.com/opensearch-project/OpenSearch/pull/15430))
 - Remove Identity FeatureFlag ([#16024](https://github.com/opensearch-project/OpenSearch/pull/16024))
 - Ensure RestHandler.Wrapper delegates all implementations to the wrapped handler ([#16154](https://github.com/opensearch-project/OpenSearch/pull/16154))
+- Code cleanup: Remove ApproximateIndexOrDocValuesQuery ([#16273](https://github.com/opensearch-project/OpenSearch/pull/16273))
 
 
 ### Deprecated

diff --git a/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java
@@ -62,8 +62,8 @@
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.search.DocValueFormat;
-import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
 import org.opensearch.search.approximate.ApproximatePointRangeQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
@@ -113,21 +113,6 @@ public static DateFormatter getDefaultDateTimeFormatter() {
             : LEGACY_DEFAULT_DATE_TIME_FORMATTER;
     }
 
-    public static Query getDefaultQuery(Query pointRangeQuery, Query dvQuery, String name, long l, long u) {
-        return FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY_SETTING)
-            ? new ApproximateIndexOrDocValuesQuery(
-                pointRangeQuery,
-                new ApproximatePointRangeQuery(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));
-                    }
-                },
-                dvQuery
-            )
-            : new IndexOrDocValuesQuery(pointRangeQuery, dvQuery);
-    }
-
     /**
      * Resolution of the date time
      *
@@ -488,22 +473,42 @@ public Query rangeQuery(
             }
             DateMathParser parser = forcedDateParser == null ? dateMathParser : forcedDateParser;
             return dateRangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, timeZone, parser, context, resolution, (l, u) -> {
-                Query pointRangeQuery = isSearchable() ? LongPoint.newRangeQuery(name(), l, u) : null;
                 Query dvQuery = hasDocValues() ? SortedNumericDocValuesField.newSlowRangeQuery(name(), l, u) : null;
-                if (isSearchable() && hasDocValues()) {
-                    Query query = getDefaultQuery(pointRangeQuery, dvQuery, name(), l, u);
-                    if (context.indexSortedOnField(name())) {
-                        query = new IndexSortSortedNumericDocValuesRangeQuery(name(), l, u, query);
+                if (isSearchable()) {
+                    Query pointRangeQuery = LongPoint.newRangeQuery(name(), l, u);
+                    Query query;
+                    if (dvQuery != null) {
+                        query = new IndexOrDocValuesQuery(pointRangeQuery, dvQuery);
+                        if (context.indexSortedOnField(name())) {
+                            query = new IndexSortSortedNumericDocValuesRangeQuery(name(), l, u, query);
+                        }
+                    } else {
+                        query = pointRangeQuery;
+                    }
+                    if (FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY_SETTING)) {
+                        return new ApproximateScoreQuery(
+                            query,
+                            new ApproximatePointRangeQuery(
+                                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));
+                                }
+                            }
+                        );
                     }
                     return query;
                 }
-                if (hasDocValues()) {
-                    if (context.indexSortedOnField(name())) {
-                        dvQuery = new IndexSortSortedNumericDocValuesRangeQuery(name(), l, u, dvQuery);
-                    }
-                    return dvQuery;
+
+                // Not searchable. Must have doc values.
+                if (context.indexSortedOnField(name())) {
+                    dvQuery = new IndexSortSortedNumericDocValuesRangeQuery(name(), l, u, dvQuery);
                 }
-                return pointRangeQuery;
+                return dvQuery;
             });
         }
 

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
@@ -23,7 +23,7 @@
 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.ApproximateIndexOrDocValuesQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -55,7 +55,7 @@ private Helper() {}
         queryWrappers.put(FunctionScoreQuery.class, q -> ((FunctionScoreQuery) q).getSubQuery());
         queryWrappers.put(DateRangeIncludingNowQuery.class, q -> ((DateRangeIncludingNowQuery) q).getQuery());
         queryWrappers.put(IndexOrDocValuesQuery.class, q -> ((IndexOrDocValuesQuery) q).getIndexQuery());
-        queryWrappers.put(ApproximateIndexOrDocValuesQuery.class, q -> ((ApproximateIndexOrDocValuesQuery) q).getOriginalQuery());
+        queryWrappers.put(ApproximateScoreQuery.class, q -> ((ApproximateScoreQuery) q).getOriginalQuery());
     }
 
     /**

diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateIndexOrDocValuesQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateIndexOrDocValuesQuery.java
diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java
@@ -433,9 +433,6 @@ public boolean canApproximate(SearchContext context) {
         if (context.aggregations() != null) {
             return false;
         }
-        if (!(context.query() instanceof ApproximateIndexOrDocValuesQuery)) {
-            return false;
-        }
         // size 0 could be set for caching
         if (context.from() + context.size() == 0) {
             this.setSize(10_000);

diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java
@@ -21,12 +21,12 @@
  * Entry-point for the approximation framework.
  * This class is heavily inspired by {@link org.apache.lucene.search.IndexOrDocValuesQuery}. It acts as a wrapper that consumer two queries, a regular query and an approximate version of the same. By default, it executes the regular query and returns {@link Weight#scorer} for the original query. At run-time, depending on certain constraints, we can re-write the {@code Weight} to use the approximate weight instead.
  */
-public class ApproximateScoreQuery extends Query {
+public final class ApproximateScoreQuery extends Query {
 
     private final Query originalQuery;
     private final ApproximateQuery approximationQuery;
 
-    protected Query resolvedQuery;
+    Query resolvedQuery;
 
     public ApproximateScoreQuery(Query originalQuery, ApproximateQuery approximationQuery) {
         this.originalQuery = originalQuery;

diff --git a/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java b/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java
@@ -41,6 +41,7 @@
 import org.apache.lucene.index.MultiReader;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.IndexSortSortedNumericDocValuesRangeQuery;
 import org.apache.lucene.search.Query;
@@ -65,8 +66,8 @@
 import org.opensearch.index.query.DateRangeIncludingNowQuery;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
-import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
 import org.opensearch.search.approximate.ApproximatePointRangeQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.joda.time.DateTimeZone;
 
 import java.io.IOException;
@@ -212,8 +213,11 @@ public void testTermQuery() {
         MappedFieldType ft = new DateFieldType("field");
         String date = "2015-10-12T14:10:55";
         long instant = DateFormatters.from(DateFieldMapper.getDefaultDateTimeFormatter().parse(date)).toInstant().toEpochMilli();
-        Query expected = new ApproximateIndexOrDocValuesQuery(
-            LongPoint.newRangeQuery("field", instant, instant + 999),
+        Query expected = new ApproximateScoreQuery(
+            new IndexOrDocValuesQuery(
+                LongPoint.newRangeQuery("field", instant, instant + 999),
+                SortedNumericDocValuesField.newSlowRangeQuery("field", instant, instant + 999)
+            ),
             new ApproximatePointRangeQuery(
                 "field",
                 pack(new long[] { instant }).bytes,
@@ -224,8 +228,7 @@ public void testTermQuery() {
                 protected String toString(int dimension, byte[] value) {
                     return Long.toString(LongPoint.decodeDimension(value, 0));
                 }
-            },
-            SortedNumericDocValuesField.newSlowRangeQuery("field", instant, instant + 999)
+            }
         );
         assumeThat(
             "Using Approximate Range Query as default",
@@ -278,8 +281,11 @@ public void testRangeQuery() throws IOException {
         String date2 = "2016-04-28T11:33:52";
         long instant1 = DateFormatters.from(DateFieldMapper.getDefaultDateTimeFormatter().parse(date1)).toInstant().toEpochMilli();
         long instant2 = DateFormatters.from(DateFieldMapper.getDefaultDateTimeFormatter().parse(date2)).toInstant().toEpochMilli() + 999;
-        Query expected = new ApproximateIndexOrDocValuesQuery(
-            LongPoint.newRangeQuery("field", instant1, instant2),
+        Query expected = new ApproximateScoreQuery(
+            new IndexOrDocValuesQuery(
+                LongPoint.newRangeQuery("field", instant1, instant2),
+                SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
+            ),
             new ApproximatePointRangeQuery(
                 "field",
                 pack(new long[] { instant1 }).bytes,
@@ -290,8 +296,7 @@ public void testRangeQuery() throws IOException {
                 protected String toString(int dimension, byte[] value) {
                     return Long.toString(LongPoint.decodeDimension(value, 0));
                 }
-            },
-            SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
+            }
         );
         assumeThat(
             "Using Approximate Range Query as default",
@@ -306,8 +311,11 @@ protected String toString(int dimension, byte[] value) {
         instant1 = nowInMillis;
         instant2 = instant1 + 100;
         expected = new DateRangeIncludingNowQuery(
-            new ApproximateIndexOrDocValuesQuery(
-                LongPoint.newRangeQuery("field", instant1, instant2),
+            new ApproximateScoreQuery(
+                new IndexOrDocValuesQuery(
+                    LongPoint.newRangeQuery("field", instant1, instant2),
+                    SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
+                ),
                 new ApproximatePointRangeQuery(
                     "field",
                     pack(new long[] { instant1 }).bytes,
@@ -318,8 +326,7 @@ protected String toString(int dimension, byte[] value) {
                     protected String toString(int dimension, byte[] value) {
                         return Long.toString(LongPoint.decodeDimension(value, 0));
                     }
-                },
-                SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
+                }
             )
         );
         assumeThat(
@@ -388,8 +395,8 @@ public void testRangeQueryWithIndexSort() {
             "field",
             instant1,
             instant2,
-            new ApproximateIndexOrDocValuesQuery(
-                LongPoint.newRangeQuery("field", instant1, instant2),
+            new ApproximateScoreQuery(
+                new IndexOrDocValuesQuery(LongPoint.newRangeQuery("field", instant1, instant2), dvQuery),
                 new ApproximatePointRangeQuery(
                     "field",
                     pack(new long[] { instant1 }).bytes,
@@ -400,8 +407,7 @@ public void testRangeQueryWithIndexSort() {
                     protected String toString(int dimension, byte[] value) {
                         return Long.toString(LongPoint.decodeDimension(value, 0));
                     }
-                },
-                dvQuery
+                }
             )
         );
         assumeThat(

diff --git a/server/src/test/java/org/opensearch/index/mapper/RangeFieldQueryStringQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/mapper/RangeFieldQueryStringQueryBuilderTests.java
@@ -50,8 +50,8 @@
 import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.QueryStringQueryBuilder;
-import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
 import org.opensearch.search.approximate.ApproximatePointRangeQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.test.AbstractQueryTestCase;
 
 import java.io.IOException;
@@ -191,11 +191,14 @@ public void testDateRangeQuery() throws Exception {
             is(true)
         );
         assertEquals(
-            new ApproximateIndexOrDocValuesQuery(
-                LongPoint.newRangeQuery(
-                    DATE_FIELD_NAME,
-                    parser.parse(lowerBoundExact, () -> 0).toEpochMilli(),
-                    parser.parse(upperBoundExact, () -> 0).toEpochMilli()
+            new ApproximateScoreQuery(
+                new IndexOrDocValuesQuery(
+                    LongPoint.newRangeQuery(
+                        DATE_FIELD_NAME,
+                        parser.parse(lowerBoundExact, () -> 0).toEpochMilli(),
+                        parser.parse(upperBoundExact, () -> 0).toEpochMilli()
+                    ),
+                    controlDv
                 ),
                 new ApproximatePointRangeQuery(
                     DATE_FIELD_NAME,
@@ -207,8 +210,7 @@ public void testDateRangeQuery() throws Exception {
                     protected String toString(int dimension, byte[] value) {
                         return Long.toString(LongPoint.decodeDimension(value, 0));
                     }
-                },
-                controlDv
+                }
             ),
             queryOnDateField
         );

diff --git a/server/src/test/java/org/opensearch/index/mapper/RangeFieldTypeTests.java b/server/src/test/java/org/opensearch/index/mapper/RangeFieldTypeTests.java
@@ -57,7 +57,7 @@
 import org.opensearch.index.mapper.RangeFieldMapper.RangeFieldType;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.QueryShardException;
-import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.test.IndexSettingsModule;
 import org.joda.time.DateTime;
 import org.junit.Before;
@@ -293,7 +293,7 @@ public void testDateRangeQueryUsingMappingFormatLegacy() {
         );
         assertEquals(
             "field:[1465975790000 TO 1466062190999]",
-            ((IndexOrDocValuesQuery) ((ApproximateIndexOrDocValuesQuery) queryOnDateField).getOriginalQuery()).getIndexQuery().toString()
+            ((IndexOrDocValuesQuery) ((ApproximateScoreQuery) queryOnDateField).getOriginalQuery()).getIndexQuery().toString()
         );
     }
 

diff --git a/server/src/test/java/org/opensearch/index/query/MatchPhraseQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/MatchPhraseQueryBuilderTests.java
@@ -42,7 +42,7 @@
 import org.apache.lucene.search.TermQuery;
 import org.opensearch.core.common.ParsingException;
 import org.opensearch.index.search.MatchQuery.ZeroTermsQuery;
-import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.test.AbstractQueryTestCase;
 
 import java.io.IOException;
@@ -131,7 +131,7 @@ protected void doAssertLuceneQuery(MatchPhraseQueryBuilder queryBuilder, Query q
                 .or(instanceOf(PointRangeQuery.class))
                 .or(instanceOf(IndexOrDocValuesQuery.class))
                 .or(instanceOf(MatchNoDocsQuery.class))
-                .or(instanceOf(ApproximateIndexOrDocValuesQuery.class))
+                .or(instanceOf(ApproximateScoreQuery.class))
         );
     }