Added approximation support for range queries with now in date field

CHANGELOG.md

@@ -36,6 +36,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Star-Tree] Add star-tree search related stats ([#18707](https://github.com/opensearch-project/OpenSearch/pull/18707))
 - Add support for plugins to profile information ([#18656](https://github.com/opensearch-project/OpenSearch/pull/18656))
 - Add support for Combined Fields query ([#18724](https://github.com/opensearch-project/OpenSearch/pull/18724))
+- Added approximation support for range queries with now in date field ([#18511](https://github.com/opensearch-project/OpenSearch/pull/18511))
 
 ### Changed
 - Update Subject interface to use CheckedRunnable ([#18570](https://github.com/opensearch-project/OpenSearch/issues/18570))

modules/percolator/src/main/java/org/opensearch/percolator/QueryAnalyzer.java

@@ -55,6 +55,7 @@
 import org.opensearch.common.lucene.search.function.FunctionScoreQuery;
 import org.opensearch.index.query.DateRangeIncludingNowQuery;
 import org.opensearch.lucene.queries.BlendedTermQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -180,7 +181,8 @@ Result getResult() {
 
         @Override
         public QueryVisitor getSubVisitor(Occur occur, Query parent) {
-            if (parent instanceof DateRangeIncludingNowQuery) {
+            if (parent instanceof ApproximateScoreQuery
+                && ((ApproximateScoreQuery) parent).getOriginalQuery() instanceof DateRangeIncludingNowQuery) {
                 terms.add(Result.UNKNOWN);
                 return QueryVisitor.EMPTY_VISITOR;
             }

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

@@ -46,6 +46,7 @@
 import org.opensearch.Version;
 import org.opensearch.common.Explicit;
 import org.opensearch.common.Nullable;
+import org.opensearch.common.TriFunction;
 import org.opensearch.common.geo.ShapeRelation;
 import org.opensearch.common.logging.DeprecationLogger;
 import org.opensearch.common.lucene.BytesRefs;
@@ -80,7 +81,6 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
-import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
@@ -511,11 +511,30 @@ public Query rangeQuery(
                 throw new IllegalArgumentException("Field [" + name() + "] of type [" + typeName() + "] does not support DISJOINT ranges");
             }
             DateMathParser parser = forcedDateParser == null ? dateMathParser : forcedDateParser;
-            return dateRangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, timeZone, parser, context, resolution, (l, u) -> {
-                Query dvQuery = hasDocValues() ? SortedNumericDocValuesField.newSlowRangeQuery(name(), l, u) : null;
-                if (isSearchable()) {
+            return dateRangeQuery(
+                lowerTerm,
+                upperTerm,
+                includeLower,
+                includeUpper,
+                timeZone,
+                parser,
+                context,
+                resolution,
+                (l, u, nowUsed) -> {
+                    Query dvQuery = hasDocValues() ? SortedNumericDocValuesField.newSlowRangeQuery(name(), l, u) : null;
+
+                    // Not searchable. Must have doc values.
+                    if (!isSearchable()) {
+                        if (context.indexSortedOnField(name())) {
+                            dvQuery = new IndexSortSortedNumericDocValuesRangeQuery(name(), l, u, dvQuery);
+                        }
+                        return nowUsed[0] ? new DateRangeIncludingNowQuery(dvQuery) : dvQuery;
+                    }
+
+                    // Field is searchable
                     Query pointRangeQuery = LongPoint.newRangeQuery(name(), l, u);
                     Query query;
+
                     if (dvQuery != null) {
                         query = new IndexOrDocValuesQuery(pointRangeQuery, dvQuery);
                         if (context.indexSortedOnField(name())) {
@@ -524,24 +543,21 @@ public Query rangeQuery(
                     } else {
                         query = pointRangeQuery;
                     }
-                    return new ApproximateScoreQuery(
-                        query,
-                        new ApproximatePointRangeQuery(
-                            name(),
-                            pack(new long[] { l }).bytes,
-                            pack(new long[] { u }).bytes,
-                            new long[] { l }.length,
-                            ApproximatePointRangeQuery.LONG_FORMAT
-                        )
+
+                    ApproximatePointRangeQuery approxQuery = new ApproximatePointRangeQuery(
+                        name(),
+                        pack(new long[] { l }).bytes,
+                        pack(new long[] { u }).bytes,
+                        new long[] { l }.length,
+                        ApproximatePointRangeQuery.LONG_FORMAT
                     );
-                }
 
-                // Not searchable. Must have doc values.
-                if (context.indexSortedOnField(name())) {
-                    dvQuery = new IndexSortSortedNumericDocValuesRangeQuery(name(), l, u, dvQuery);
+                    return nowUsed[0]
+                        ? new ApproximateScoreQuery(new DateRangeIncludingNowQuery(query), approxQuery)
+                        : new ApproximateScoreQuery(query, approxQuery);
+
                 }
-                return dvQuery;
-            });
+            );
         }
 
         public static Query dateRangeQuery(
@@ -553,44 +569,31 @@ public static Query dateRangeQuery(
             DateMathParser parser,
             QueryShardContext context,
             Resolution resolution,
-            BiFunction<Long, Long, Query> builder
+            TriFunction<Long, Long, boolean[], Query> builder
         ) {
-            return handleNow(context, nowSupplier -> {
-                long l, u;
-                if (lowerTerm == null) {
-                    l = Long.MIN_VALUE;
-                } else {
-                    l = parseToLong(lowerTerm, !includeLower, timeZone, parser, nowSupplier, resolution);
-                    if (includeLower == false) {
-                        ++l;
-                    }
-                }
-                if (upperTerm == null) {
-                    u = Long.MAX_VALUE;
-                } else {
-                    u = parseToLong(upperTerm, includeUpper, timeZone, parser, nowSupplier, resolution);
-                    if (includeUpper == false) {
-                        --u;
-                    }
-                }
-                return builder.apply(l, u);
-            });
-        }
-
-        /**
-         * Handle {@code now} in queries.
-         * @param context context from which to read the current time
-         * @param builder build the query
-         * @return the result of the builder, wrapped in {@link DateRangeIncludingNowQuery} if {@code now} was used.
-         */
-        public static Query handleNow(QueryShardContext context, Function<LongSupplier, Query> builder) {
             boolean[] nowUsed = new boolean[1];
             LongSupplier nowSupplier = () -> {
                 nowUsed[0] = true;
                 return context.nowInMillis();
             };
-            Query query = builder.apply(nowSupplier);
-            return nowUsed[0] ? new DateRangeIncludingNowQuery(query) : query;
+            long l, u;
+            if (lowerTerm == null) {
+                l = Long.MIN_VALUE;
+            } else {
+                l = parseToLong(lowerTerm, !includeLower, timeZone, parser, nowSupplier, resolution);
+                if (includeLower == false) {
+                    ++l;
+                }
+            }
+            if (upperTerm == null) {
+                u = Long.MAX_VALUE;
+            } else {
+                u = parseToLong(upperTerm, includeUpper, timeZone, parser, nowSupplier, resolution);
+                if (includeUpper == false) {
+                    --u;
+                }
+            }
+            return builder.apply(l, u, nowUsed);
         }
 
         public long parseToLong(Object value, boolean roundUp, @Nullable ZoneId zone, DateMathParser dateParser, LongSupplier now) {

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

@@ -308,19 +308,20 @@ public void testRangeQuery() throws IOException {
 
         instant1 = nowInMillis;
         instant2 = instant1 + 100;
-        expected = new DateRangeIncludingNowQuery(
-            new ApproximateScoreQuery(
+        expected = new ApproximateScoreQuery(
+            new DateRangeIncludingNowQuery(
                 new IndexOrDocValuesQuery(
                     LongPoint.newRangeQuery("field", instant1, instant2),
                     SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
-                ),
-                new ApproximatePointRangeQuery(
-                    "field",
-                    pack(new long[] { instant1 }).bytes,
-                    pack(new long[] { instant2 }).bytes,
-                    new long[] { instant1 }.length,
-                    ApproximatePointRangeQuery.LONG_FORMAT
                 )
+            ),
+            new ApproximatePointRangeQuery(
+                "field",
+                pack(new long[] { instant1 }).bytes,
+                pack(new long[] { instant2 }).bytes,
+                new long[] { instant1 }.length,
+                ApproximatePointRangeQuery.LONG_FORMAT
+
             )
         );
         assertEquals(expected, ft.rangeQuery("now", instant2, true, true, null, null, null, context));

server/src/test/java/org/opensearch/index/query/RangeQueryBuilderTests.java

@@ -447,14 +447,22 @@ public void testDateRangeQueryTimezone() throws IOException {
             + "        }\n"
             + "    }\n"
             + "}";
+
+        // TODO what else can we assert
         QueryShardContext context = createShardContext();
         Query parsedQuery = parseQuery(query).toQuery(context);
-        assertThat(parsedQuery, instanceOf(DateRangeIncludingNowQuery.class));
-        parsedQuery = ((DateRangeIncludingNowQuery) parsedQuery).getQuery();
+
         assertThat(parsedQuery, instanceOf(ApproximateScoreQuery.class));
-        parsedQuery = ((ApproximateScoreQuery) parsedQuery).getApproximationQuery();
-        assertThat(parsedQuery, instanceOf(ApproximateQuery.class));
-        // TODO what else can we assert
+
+        // Get the exact query from ApproximateScoreQuery (which should be DateRangeIncludingNowQuery)
+        ApproximateScoreQuery approximateScoreQuery = (ApproximateScoreQuery) parsedQuery;
+        Query exactQuery = approximateScoreQuery.getOriginalQuery();
+
+        // The exact query should be DateRangeIncludingNowQuery
+        assertThat(exactQuery, instanceOf(DateRangeIncludingNowQuery.class));
+
+        ApproximateQuery approximationQuery = approximateScoreQuery.getApproximationQuery();
+        assertThat(approximationQuery, instanceOf(ApproximatePointRangeQuery.class));
 
         query = "{\n"
             + "    \"range\" : {\n"