Adding yaml rest tests to test sorting

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>

rest-api-spec/src/main/resources/rest-api-spec/test/search/370_approximate_range.yml

@@ -0,0 +1,72 @@
+---
+"search with approximate range":
+  - do:
+      indices.create:
+        index: test
+        body:
+          mappings:
+            properties:
+              date:
+                type: date
+                index: true
+                doc_values: true
+
+  - do:
+      bulk:
+        index: test
+        refresh: true
+        body:
+          - '{"index": {"_index": "test", "_id": "1" }}'
+          - '{ "date": "2018-10-29T12:12:12.987Z" }'
+          - '{ "index": { "_index": "test", "_id": "2" }}'
+          - '{ "date": "2020-10-29T12:12:12.987Z" }'
+          - '{ "index": { "_index": "test", "_id": "3" } }'
+          - '{ "date": "2024-10-29T12:12:12.987Z" }'
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        index: test
+        body:
+          query:
+            range: {
+              date: {
+                gte: "2018-10-29T12:12:12.987Z"
+              },
+            }
+
+  - match: { hits.total: 3 }
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        index: test
+        body:
+          sort: [{ date: asc }]
+          query:
+            range: {
+              date: {
+                gte: "2018-10-29T12:12:12.987Z"
+              },
+            }
+
+
+  - match: { hits.total: 3 }
+  - match: { hits.hits.0._id: "1" }
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        index: test
+        body:
+          sort: [{ date: desc }]
+          query:
+            range: {
+              date: {
+                gte: "2018-10-29T12:12:12.987Z",
+                lte: "2020-10-29T12:12:12.987Z"
+              },
+            }
+
+  - match: { hits.total: 2 }
+  - match: { hits.hits.0._id: "2" }

server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java

@@ -56,6 +56,10 @@ protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upp
         this(field, lowerPoint, upperPoint, numDims, 10_000, SortOrder.ASC);
     }
 
+    protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upperPoint, int numDims, int size) {
+        this(field, lowerPoint, upperPoint, numDims, size, SortOrder.ASC);
+    }
+
     protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upperPoint, int numDims, int size, SortOrder sortOrder) {
         checkArgs(field, lowerPoint, upperPoint);
         this.field = field;
@@ -234,6 +238,7 @@ private void intersectRight(PointValues.PointTree pointTree, PointValues.Interse
                 assert pointTree.moveToParent() == false;
             }
 
+            // custom intersect visitor to walk the left of the tree
             private long intersectLeft(PointValues.IntersectVisitor visitor, PointValues.PointTree pointTree, int count)
                 throws IOException {
                 PointValues.Relation r = visitor.compare(pointTree.getMinPackedValue(), pointTree.getMaxPackedValue());
@@ -246,16 +251,14 @@ private long intersectLeft(PointValues.IntersectVisitor visitor, PointValues.Poi
                         break;
                     case CELL_INSIDE_QUERY:
                         // If the cell is fully inside, we keep moving to child until we reach a point where we can no longer move or when
-                        // we have sufficient doc count
+                        // we have sufficient doc count. We first move down and then move to the left child
                         if (pointTree.moveToChild()) {
                             do {
                                 docCount[0] += intersectLeft(visitor, pointTree, count);
                             } while (pointTree.moveToSibling() && docCount[0] <= count);
                             pointTree.moveToParent();
                         } else {
-                            // TODO: we can assert that the first value here in fact matches what the pointTree
-                            // claimed?
-                            // Leaf node; scan and filter all points in this block:
+                            // we're at the leaf node, if we're under the count, visit all the docIds in this node.
                             if (docCount[0] <= count) {
                                 pointTree.visitDocIDs(visitor);
                                 docCount[0] += pointTree.size();
@@ -287,6 +290,7 @@ private long intersectLeft(PointValues.IntersectVisitor visitor, PointValues.Poi
                 return docCount[0] > 0 ? docCount[0] : 0;
             }
 
+            // custom intersect visitor to walk the right of tree
             private long intersectRight(PointValues.IntersectVisitor visitor, PointValues.PointTree pointTree, int count)
                 throws IOException {
                 PointValues.Relation r = visitor.compare(pointTree.getMinPackedValue(), pointTree.getMaxPackedValue());
@@ -299,16 +303,14 @@ private long intersectRight(PointValues.IntersectVisitor visitor, PointValues.Po
                         break;
                     case CELL_INSIDE_QUERY:
                         // If the cell is fully inside, we keep moving to child until we reach a point where we can no longer move or when
-                        // we have sufficient doc count
+                        // we have sufficient doc count. We first move down and then move right
                         if (pointTree.moveToChild()) {
                             while (pointTree.moveToSibling() && docCount[0] <= count) {
                                 docCount[0] += intersectRight(visitor, pointTree, count);
                             }
                             pointTree.moveToParent();
                         } else {
-                            // TODO: we can assert that the first value here in fact matches what the pointTree
-                            // claimed?
-                            // Leaf node; scan and filter all points in this block:
+                            // we're at the leaf node, if we're under the count, visit all the docIds in this node.
                             if (docCount[0] <= count) {
                                 pointTree.visitDocIDs(visitor);
                                 docCount[0] += pointTree.size();

server/src/test/java/org/opensearch/search/approximate/ApproximatePointRangeQueryTests.java

@@ -101,8 +101,7 @@ public void testApproximateRangeWithSize() throws IOException {
                             pack(lower).bytes,
                             pack(upper).bytes,
                             dims,
-                            10,
-                            SortOrder.ASC
+                            10
                         ) {
                             protected String toString(int dimension, byte[] value) {
                                 return Long.toString(LongPoint.decodeDimension(value, 0));
@@ -113,8 +112,7 @@ protected String toString(int dimension, byte[] value) {
                             pack(lower).bytes,
                             pack(upper).bytes,
                             dims,
-                            100,
-                            SortOrder.ASC
+                            100
                         ) {
                             protected String toString(int dimension, byte[] value) {
                                 return Long.toString(LongPoint.decodeDimension(value, 0));
@@ -159,6 +157,55 @@ public void testApproximateRangeShortCircuit() throws IOException {
                     try {
                         long lower = 0;
                         long upper = 100;
+                        Query approximateQuery = new ApproximatePointRangeQuery("point", pack(lower).bytes, pack(upper).bytes, dims, 10) {
+                            protected String toString(int dimension, byte[] value) {
+                                return Long.toString(LongPoint.decodeDimension(value, 0));
+                            }
+                        };
+                        Query query = new PointRangeQuery("point", pack(lower).bytes, pack(upper).bytes, dims) {
+                            protected String toString(int dimension, byte[] value) {
+                                return Long.toString(LongPoint.decodeDimension(value, 0));
+                            }
+                        };
+                        IndexSearcher searcher = new IndexSearcher(reader);
+                        TopDocs topDocs = searcher.search(approximateQuery, 10);
+                        TopDocs topDocs1 = searcher.search(query, 10);
+
+                        // since we short-circuit from the approx range at the end of size these will not be equal
+                        assertNotEquals(topDocs.totalHits, topDocs1.totalHits);
+                        assertEquals(topDocs.totalHits, new TotalHits(11, TotalHits.Relation.EQUAL_TO));
+                        assertEquals(topDocs1.totalHits, new TotalHits(101, TotalHits.Relation.EQUAL_TO));
+
+                    } catch (IOException e) {
+                        throw new RuntimeException(e);
+                    }
+
+                }
+            }
+        }
+    }
+
+    public void testApproximateRangeShortCircuitAscSort() throws IOException {
+        try (Directory directory = newDirectory()) {
+            try (RandomIndexWriter iw = new RandomIndexWriter(random(), directory, new WhitespaceAnalyzer())) {
+                int dims = 1;
+
+                long[] scratch = new long[dims];
+                int numPoints = 1000;
+                for (int i = 0; i < numPoints; i++) {
+                    Document doc = new Document();
+                    for (int v = 0; v < dims; v++) {
+                        scratch[v] = i;
+                    }
+                    doc.add(new LongPoint("point", scratch));
+                    iw.addDocument(doc);
+                    // if (i % 10 == 0) iw.flush();
+                }
+                iw.flush();
+                try (IndexReader reader = iw.getReader()) {
+                    try {
+                        long lower = 0;
+                        long upper = 20;
                         Query approximateQuery = new ApproximatePointRangeQuery(
                             "point",
                             pack(lower).bytes,
@@ -183,7 +230,13 @@ protected String toString(int dimension, byte[] value) {
                         // since we short-circuit from the approx range at the end of size these will not be equal
                         assertNotEquals(topDocs.totalHits, topDocs1.totalHits);
                         assertEquals(topDocs.totalHits, new TotalHits(11, TotalHits.Relation.EQUAL_TO));
-                        assertEquals(topDocs1.totalHits, new TotalHits(101, TotalHits.Relation.EQUAL_TO));
+                        assertEquals(topDocs1.totalHits, new TotalHits(21, TotalHits.Relation.EQUAL_TO));
+                        assertEquals(topDocs.scoreDocs[0].doc, 0);
+                        assertEquals(topDocs.scoreDocs[1].doc, 1);
+                        assertEquals(topDocs.scoreDocs[2].doc, 2);
+                        assertEquals(topDocs.scoreDocs[3].doc, 3);
+                        assertEquals(topDocs.scoreDocs[4].doc, 4);
+                        assertEquals(topDocs.scoreDocs[5].doc, 5);
 
                     } catch (IOException e) {
                         throw new RuntimeException(e);
@@ -193,5 +246,4 @@ protected String toString(int dimension, byte[] value) {
             }
         }
     }
-
 }