Address PR comments

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
harshavamsi · Aug 20, 2024 · 7056c52 · 7056c52
1 parent 5abb7ab
commit 7056c52
Show file tree

Hide file tree

Showing 12 changed files with 351 additions and 220 deletions.
diff --git a/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java
@@ -63,7 +63,7 @@
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.search.DocValueFormat;
 import org.opensearch.search.approximate.ApproximatePointRangeQuery;
-import org.opensearch.search.approximate.ApproximateableQuery;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
@@ -467,28 +467,10 @@ public Query rangeQuery(
             }
             DateMathParser parser = forcedDateParser == null ? dateMathParser : forcedDateParser;
             return dateRangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, timeZone, parser, context, resolution, (l, u) -> {
+                Query pointRangeQuery = createPointRangeQuery(l, u);
+                Query dvQuery = hasDocValues() ? SortedNumericDocValuesField.newSlowRangeQuery(name(), l, u) : null;
                 if (isSearchable() && hasDocValues()) {
-                    Query query = new ApproximateableQuery(
-                        new PointRangeQuery(name(), pack(new long[] { l }).bytes, pack(new long[] { u }).bytes, new long[] { l }.length) {
-                            protected String toString(int dimension, byte[] value) {
-                                return Long.toString(LongPoint.decodeDimension(value, 0));
-                            }
-                        },
-                        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));
-                            }
-                        }
-                    );
-
-                    Query dvQuery = SortedNumericDocValuesField.newSlowRangeQuery(name(), l, u);
-                    query = new IndexOrDocValuesQuery(query, dvQuery);
+                    Query query = new IndexOrDocValuesQuery(pointRangeQuery, dvQuery);
 
                     if (context.indexSortedOnField(name())) {
                         query = new IndexSortSortedNumericDocValuesRangeQuery(name(), l, u, query);
@@ -502,25 +484,29 @@ protected String toString(int dimension, byte[] value) {
                     }
                     return query;
                 }
-                return new ApproximateableQuery(
-                    new PointRangeQuery(name(), pack(new long[] { l }).bytes, pack(new long[] { u }).bytes, new long[] { l }.length) {
+                return pointRangeQuery;
+            });
+        }
+
+        private Query createPointRangeQuery(long l, long u) {
+            return new ApproximateScoreQuery(
+                    new PointRangeQuery(name(), pack(new long[]{l}).bytes, pack(new long[]{u}).bytes, new long[]{l}.length) {
                         protected String toString(int dimension, byte[] value) {
                             return Long.toString(LongPoint.decodeDimension(value, 0));
                         }
                     },
                     new ApproximatePointRangeQuery(
-                        name(),
-                        pack(new long[] { l }).bytes,
-                        pack(new long[] { u }).bytes,
-                        new long[] { l }.length
+                            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));
                         }
                     }
-                );
-            });
+            );
         }
 
         public static Query dateRangeQuery(

diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateConstantScoreWeight.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateConstantScoreWeight.java
@@ -0,0 +1,30 @@
+/*
+ * 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.index.PointValues;
+import org.apache.lucene.search.ConstantScoreWeight;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.util.DocIdSetBuilder;
+
+import java.io.IOException;
+
+public abstract class ApproximateConstantScoreWeight extends ConstantScoreWeight {
+
+    protected ApproximateConstantScoreWeight(Query query, float score) {
+        super(query, score);
+    }
+
+    protected abstract long intersectLeft(PointValues.IntersectVisitor visitor, PointValues.PointTree pointTree) throws  IOException;
+
+    protected abstract long intersectRight(PointValues.IntersectVisitor visitor, PointValues.PointTree pointTree) throws IOException;
+
+    protected abstract PointValues.IntersectVisitor getIntersectVisitor(DocIdSetBuilder result) throws IOException;
+
+}