Fix more tests

Signed-off-by: Harsha Vamsi Kalluri <harshavamsi096@gmail.com>

Author: harshavamsi

server/src/main/java/org/opensearch/index/search/NestedHelper.java

@@ -46,6 +46,7 @@
 import org.apache.lucene.search.TermQuery;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.mapper.ObjectMapper;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 
 /** Utility class to filter parent and children clauses when building nested
  * queries.
@@ -85,6 +86,8 @@ public boolean mightMatchNestedDocs(Query query) {
             return mightMatchNestedDocs(((PointRangeQuery) query).getField());
         } else if (query instanceof IndexOrDocValuesQuery) {
             return mightMatchNestedDocs(((IndexOrDocValuesQuery) query).getIndexQuery());
+        } else if (query instanceof ApproximateScoreQuery) {
+            return mightMatchNestedDocs(((ApproximateScoreQuery) query).getOriginalQuery());
         } else if (query instanceof BooleanQuery) {
             final BooleanQuery bq = (BooleanQuery) query;
             final boolean hasRequiredClauses = bq.clauses().stream().anyMatch(BooleanClause::isRequired);
@@ -155,6 +158,8 @@ public boolean mightMatchNonNestedDocs(Query query, String nestedPath) {
             return mightMatchNonNestedDocs(((PointRangeQuery) query).getField(), nestedPath);
         } else if (query instanceof IndexOrDocValuesQuery) {
             return mightMatchNonNestedDocs(((IndexOrDocValuesQuery) query).getIndexQuery(), nestedPath);
+        } else if (query instanceof ApproximateScoreQuery) {
+            return mightMatchNonNestedDocs(((ApproximateScoreQuery) query).getOriginalQuery(), nestedPath);
         } else if (query instanceof BooleanQuery) {
             final BooleanQuery bq = (BooleanQuery) query;
             final boolean hasRequiredClauses = bq.clauses().stream().anyMatch(BooleanClause::isRequired);

server/src/test/java/org/opensearch/index/mapper/NumberFieldTypeTests.java

@@ -75,6 +75,8 @@
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.search.DocValueFormat;
 import org.opensearch.search.MultiValueMode;
+import org.opensearch.search.approximate.ApproximatePointRangeQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.search.query.BitmapDocValuesQuery;
 import org.opensearch.search.query.BitmapIndexQuery;
 import org.junit.Before;
@@ -96,6 +98,7 @@
 import static org.hamcrest.Matchers.either;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
+import static org.apache.lucene.document.LongPoint.pack;
 
 public class NumberFieldTypeTests extends FieldTypeTestCase {
 
@@ -392,9 +395,9 @@ public void testUnsignedLongRangeQueryWithDecimalParts() {
 
     public void testLongRangeQuery() {
         MappedFieldType ft = new NumberFieldMapper.NumberFieldType("field", NumberFieldMapper.NumberType.LONG);
-        Query expected = new IndexOrDocValuesQuery(
-            LongPoint.newRangeQuery("field", 1, 3),
-            SortedNumericDocValuesField.newSlowRangeQuery("field", 1, 3)
+        Query expected = new ApproximateScoreQuery(
+            new IndexOrDocValuesQuery(LongPoint.newRangeQuery("field", 1, 3), SortedNumericDocValuesField.newSlowRangeQuery("field", 1, 3)),
+            new ApproximatePointRangeQuery("field", pack(1).bytes, pack(3).bytes, 1, ApproximatePointRangeQuery.LONG_FORMAT)
         );
         assertEquals(expected, ft.rangeQuery("1", "3", true, true, null, null, null, MOCK_QSC));
 
@@ -681,7 +684,11 @@ public void doTestDocValueRangeQueries(NumberType type, Supplier<Number> valueSu
                 true,
                 MOCK_QSC
             );
-            assertThat(query, either(instanceOf(IndexOrDocValuesQuery.class)).or(instanceOf(MatchNoDocsQuery.class)));
+            assertThat(
+                query,
+                either(instanceOf(IndexOrDocValuesQuery.class)).or(instanceOf(MatchNoDocsQuery.class))
+                    .or(instanceOf(ApproximateScoreQuery.class))
+            );
             if (query instanceof IndexOrDocValuesQuery) {
                 IndexOrDocValuesQuery indexOrDvQuery = (IndexOrDocValuesQuery) query;
                 assertEquals(searcher.count(indexOrDvQuery.getIndexQuery()), searcher.count(indexOrDvQuery.getRandomAccessQuery()));
@@ -764,10 +771,20 @@ public void doTestIndexSortRangeQueries(NumberType type, Supplier<Number> valueS
                 true,
                 context
             );
-            assertThat(query, instanceOf(IndexSortSortedNumericDocValuesRangeQuery.class));
+            assertThat(
+                query,
+                either(instanceOf(IndexSortSortedNumericDocValuesRangeQuery.class)).or(instanceOf(ApproximateScoreQuery.class))
+            );
 
-            Query fallbackQuery = ((IndexSortSortedNumericDocValuesRangeQuery) query).getFallbackQuery();
-            assertThat(fallbackQuery, instanceOf(IndexOrDocValuesQuery.class));
+            Query fallbackQuery;
+            if (query instanceof IndexSortSortedNumericDocValuesRangeQuery) {
+                fallbackQuery = ((IndexSortSortedNumericDocValuesRangeQuery) query).getFallbackQuery();
+                assertThat(fallbackQuery, instanceOf(IndexOrDocValuesQuery.class));
+            } else {
+                fallbackQuery = ((IndexSortSortedNumericDocValuesRangeQuery) ((ApproximateScoreQuery) query).getOriginalQuery())
+                    .getFallbackQuery();
+                assertThat(fallbackQuery, instanceOf(IndexOrDocValuesQuery.class));
+            }
 
             IndexOrDocValuesQuery indexOrDvQuery = (IndexOrDocValuesQuery) fallbackQuery;
             assertEquals(searcher.count(query), searcher.count(indexOrDvQuery.getIndexQuery()));