Addressing more PR comments

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
harshavamsi · Aug 30, 2024 · 6a84549 · 6a84549
1 parent 3d63ea0
commit 6a84549
Show file tree

Hide file tree

Showing 9 changed files with 79 additions and 1 deletion.
diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java
@@ -50,6 +50,14 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo
         return resolvedQuery.createWeight(searcher, scoreMode, boost);
     }
 
+    @Override
+    public final Query rewrite(IndexSearcher indexSearcher) throws IOException {
+        if (resolvedQuery == null) {
+            throw new IllegalStateException("Cannot create weight without setContext being called");
+        }
+        return resolvedQuery.rewrite(indexSearcher);
+    }
+
     public void setContext(SearchContext context) {
         if (resolvedQuery != null) {
             throw new IllegalStateException("Query already resolved, duplicate call to setContext");

diff --git a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java
@@ -221,7 +221,7 @@ public Weight createWeight(Query query, ScoreMode scoreMode, float boost) throws
             return new ProfileWeight(query, weight, profile);
         } else if (query instanceof ApproximateScoreQuery) {
             ((ApproximateScoreQuery) query).setContext(searchContext);
-            return query.createWeight(this, scoreMode, boost);
+            return super.createWeight(query, scoreMode, boost);
         } else {
             return super.createWeight(query, scoreMode, boost);
         }

diff --git a/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java b/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java
@@ -52,6 +52,7 @@
 import org.opensearch.common.time.DateFormatters;
 import org.opensearch.common.time.DateMathParser;
 import org.opensearch.common.util.BigArrays;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.common.util.io.IOUtils;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.index.fielddata.IndexNumericFieldData;
@@ -72,7 +73,9 @@
 import java.time.ZoneOffset;
 import java.util.Collections;
 
+import static org.hamcrest.CoreMatchers.is;
 import static org.apache.lucene.document.LongPoint.pack;
+import static org.junit.Assume.assumeThat;
 
 public class DateFieldTypeTests extends FieldTypeTestCase {
 
@@ -224,6 +227,11 @@ protected String toString(int dimension, byte[] value) {
             },
             SortedNumericDocValuesField.newSlowRangeQuery("field", instant, instant + 999)
         );
+        assumeThat(
+            "Using Approximate Range Query as default",
+            FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+            is(true)
+        );
         assertEquals(expected, ft.termQuery(date, context));
 
         MappedFieldType unsearchable = new DateFieldType(
@@ -285,6 +293,11 @@ protected String toString(int dimension, byte[] value) {
             },
             SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
         );
+        assumeThat(
+            "Using Approximate Range Query as default",
+            FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+            is(true)
+        );
         assertEquals(
             expected,
             ft.rangeQuery(date1, date2, true, true, null, null, null, context).rewrite(new IndexSearcher(new MultiReader()))
@@ -309,6 +322,11 @@ protected String toString(int dimension, byte[] value) {
                 SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
             )
         );
+        assumeThat(
+            "Using Approximate Range Query as default",
+            FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+            is(true)
+        );
         assertEquals(expected, ft.rangeQuery("now", instant2, true, true, null, null, null, context));
 
         MappedFieldType unsearchable = new DateFieldType(
@@ -386,6 +404,11 @@ protected String toString(int dimension, byte[] value) {
                 dvQuery
             )
         );
+        assumeThat(
+            "Using Approximate Range Query as default",
+            FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+            is(true)
+        );
         assertEquals(expected, ft.rangeQuery(date1, date2, true, true, null, null, null, context));
     }
 

diff --git a/server/src/test/java/org/opensearch/index/mapper/RangeFieldQueryStringQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/mapper/RangeFieldQueryStringQueryBuilderTests.java
@@ -47,6 +47,7 @@
 import org.opensearch.common.compress.CompressedXContent;
 import org.opensearch.common.network.InetAddresses;
 import org.opensearch.common.time.DateMathParser;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.QueryStringQueryBuilder;
 import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
@@ -56,9 +57,11 @@
 import java.io.IOException;
 import java.net.InetAddress;
 
+import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.Matchers.either;
 import static org.hamcrest.core.IsInstanceOf.instanceOf;
 import static org.apache.lucene.document.LongPoint.pack;
+import static org.junit.Assume.assumeThat;
 
 public class RangeFieldQueryStringQueryBuilderTests extends AbstractQueryTestCase<QueryStringQueryBuilder> {
 
@@ -182,6 +185,11 @@ public void testDateRangeQuery() throws Exception {
             parser.parse(lowerBoundExact, () -> 0).toEpochMilli(),
             parser.parse(upperBoundExact, () -> 0).toEpochMilli()
         );
+        assumeThat(
+            "Using Approximate Range Query as default",
+            FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+            is(true)
+        );
         assertEquals(
             new ApproximateIndexOrDocValuesQuery(
                 LongPoint.newRangeQuery(

diff --git a/server/src/test/java/org/opensearch/index/mapper/RangeFieldTypeTests.java b/server/src/test/java/org/opensearch/index/mapper/RangeFieldTypeTests.java
@@ -286,6 +286,11 @@ public void testDateRangeQueryUsingMappingFormatLegacy() {
         // compare lower and upper bounds with what we would get on a `date` field
         DateFieldType dateFieldType = new DateFieldType("field", DateFieldMapper.Resolution.MILLISECONDS, formatter);
         final Query queryOnDateField = dateFieldType.rangeQuery(from, to, true, true, relation, null, fieldType.dateMathParser(), context);
+        assumeThat(
+            "Using Approximate Range Query as default",
+            FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+            is(true)
+        );
         assertEquals(
             "field:[1465975790000 TO 1466062190999]",
             ((IndexOrDocValuesQuery) ((ApproximateIndexOrDocValuesQuery) queryOnDateField).getOriginalQuery()).getIndexQuery().toString()

diff --git a/server/src/test/java/org/opensearch/index/query/QueryStringQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/QueryStringQueryBuilderTests.java
@@ -70,6 +70,7 @@
 import org.opensearch.common.compress.CompressedXContent;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.Fuzziness;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.mapper.FieldNamesFieldMapper;
@@ -99,7 +100,9 @@
 import static org.hamcrest.CoreMatchers.hasItems;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.is;
 import static org.apache.lucene.document.LongPoint.pack;
+import static org.junit.Assume.assumeThat;
 
 public class QueryStringQueryBuilderTests extends AbstractQueryTestCase<QueryStringQueryBuilder> {
 
@@ -855,6 +858,11 @@ public void testToQueryDateWithTimeZone() throws Exception {
         QueryStringQueryBuilder qsq = queryStringQuery(DATE_FIELD_NAME + ":1970-01-01");
         QueryShardContext context = createShardContext();
         Query query = qsq.toQuery(context);
+        assumeThat(
+            "Using Approximate Range Query as default",
+            FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+            is(true)
+        );
         assertThat(query, instanceOf(ApproximateIndexOrDocValuesQuery.class));
         long lower = 0; // 1970-01-01T00:00:00.999 UTC
         long upper = 86399999;  // 1970-01-01T23:59:59.999 UTC

diff --git a/server/src/test/java/org/opensearch/index/query/RangeQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/RangeQueryBuilderTests.java
@@ -48,6 +48,7 @@
 import org.opensearch.OpenSearchParseException;
 import org.opensearch.common.geo.ShapeRelation;
 import org.opensearch.common.lucene.BytesRefs;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.common.ParsingException;
 import org.opensearch.index.mapper.DateFieldMapper;
 import org.opensearch.index.mapper.FieldNamesFieldMapper;
@@ -69,10 +70,12 @@
 import java.util.Map;
 
 import static org.opensearch.index.query.QueryBuilders.rangeQuery;
+import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.sameInstance;
 import static org.apache.lucene.document.LongPoint.pack;
+import static org.junit.Assume.assumeThat;
 
 public class RangeQueryBuilderTests extends AbstractQueryTestCase<RangeQueryBuilder> {
     @Override
@@ -188,6 +191,11 @@ protected void doAssertLuceneQuery(RangeQueryBuilder queryBuilder, Query query,
                 assertThat(termRangeQuery.includesLower(), equalTo(queryBuilder.includeLower()));
                 assertThat(termRangeQuery.includesUpper(), equalTo(queryBuilder.includeUpper()));
             } else if (expectedFieldName.equals(DATE_FIELD_NAME)) {
+                assumeThat(
+                    "Using Approximate Range Query as default",
+                    FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+                    is(true)
+                );
                 assertThat(query, instanceOf(ApproximateIndexOrDocValuesQuery.class));
                 Query approximationQuery = ((ApproximateIndexOrDocValuesQuery) query).getApproximationQuery();
                 assertThat(approximationQuery, instanceOf(ApproximateQuery.class));
@@ -323,6 +331,11 @@ public void testDateRangeQueryFormat() throws IOException {
             + "    }\n"
             + "}";
         Query parsedQuery = parseQuery(query).toQuery(createShardContext());
+        assumeThat(
+            "Using Approximate Range Query as default",
+            FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+            is(true)
+        );
         assertThat(parsedQuery, instanceOf(ApproximateIndexOrDocValuesQuery.class));
         Query approximationQuery = ((ApproximateIndexOrDocValuesQuery) parsedQuery).getApproximationQuery();
         assertThat(approximationQuery, instanceOf(ApproximateQuery.class));
@@ -376,6 +389,11 @@ public void testDateRangeBoundaries() throws IOException {
             + "    }\n"
             + "}\n";
         Query parsedQuery = parseQuery(query).toQuery(createShardContext());
+        assumeThat(
+            "Using Approximate Range Query as default",
+            FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+            is(true)
+        );
         assertThat(parsedQuery, instanceOf(ApproximateIndexOrDocValuesQuery.class));
 
         long lower = DateTime.parse("2014-11-01T00:00:00.000+00").getMillis();
@@ -453,6 +471,11 @@ public void testDateRangeQueryTimezone() throws IOException {
         Query parsedQuery = parseQuery(query).toQuery(context);
         assertThat(parsedQuery, instanceOf(DateRangeIncludingNowQuery.class));
         parsedQuery = ((DateRangeIncludingNowQuery) parsedQuery).getQuery();
+        assumeThat(
+            "Using Approximate Range Query as default",
+            FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY),
+            is(true)
+        );
         assertThat(parsedQuery, instanceOf(ApproximateIndexOrDocValuesQuery.class));
         parsedQuery = ((ApproximateIndexOrDocValuesQuery) parsedQuery).getApproximationQuery();
         assertThat(parsedQuery, instanceOf(ApproximateQuery.class));

diff --git a/...rc/test/java/org/opensearch/search/approximate/ApproximateIndexOrDocValuesQueryTests.java b/...rc/test/java/org/opensearch/search/approximate/ApproximateIndexOrDocValuesQueryTests.java
@@ -100,6 +100,8 @@ public boolean canApproximate(SearchContext context) {
             dvQuery
         );
 
+        approximateIndexOrDocValuesQueryCanApproximate.resolvedQuery = approximateIndexQueryCanApproximate;
+
         Weight approximateIndexOrDocValuesQueryCanApproximateWeight = approximateIndexOrDocValuesQueryCanApproximate.createWeight(
             searcher,
             ScoreMode.COMPLETE,

diff --git a/server/src/test/java/org/opensearch/search/approximate/ApproximateScoreQueryTests.java b/server/src/test/java/org/opensearch/search/approximate/ApproximateScoreQueryTests.java
@@ -54,6 +54,7 @@ protected String toString(int dimension, byte[] value) {
         };
 
         ApproximateScoreQuery query = new ApproximateScoreQuery(originalQuery, approximateQuery);
+        query.resolvedQuery = approximateQuery;
 
         try (Directory directory = newDirectory()) {
             try (RandomIndexWriter iw = new RandomIndexWriter(random(), directory, new WhitespaceAnalyzer())) {