Adding new tests

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
harshavamsi · Aug 28, 2024 · 900a685 · 900a685
1 parent b42fa37
commit 900a685
Show file tree

Hide file tree

Showing 2 changed files with 172 additions and 0 deletions.
diff --git a/...rc/test/java/org/opensearch/search/approximate/ApproximateIndexOrDocValuesQueryTests.java b/...rc/test/java/org/opensearch/search/approximate/ApproximateIndexOrDocValuesQueryTests.java
@@ -0,0 +1,123 @@
+/*
+ * 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.document.LongPoint;
+import org.apache.lucene.document.SortedNumericDocValuesField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.search.ConstantScoreWeight;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.PointRangeQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Weight;
+import org.apache.lucene.store.Directory;
+import org.opensearch.search.internal.SearchContext;
+import org.opensearch.test.OpenSearchTestCase;
+import org.junit.After;
+import org.junit.Before;
+
+import java.io.IOException;
+
+import static org.apache.lucene.document.LongPoint.pack;
+
+public class ApproximateIndexOrDocValuesQueryTests extends OpenSearchTestCase {
+    private Directory dir;
+    private IndexWriter w;
+    private DirectoryReader reader;
+    private IndexSearcher searcher;
+
+    @Before
+    public void initSearcher() throws IOException {
+        dir = newDirectory();
+        w = new IndexWriter(dir, newIndexWriterConfig());
+    }
+
+    @After
+    public void closeAllTheReaders() throws IOException {
+        reader.close();
+        w.close();
+        dir.close();
+    }
+
+    public void testApproximateIndexOrDocValuesQueryWeight() throws IOException {
+
+        long l = Long.MIN_VALUE;
+        long u = Long.MAX_VALUE;
+        Query indexQuery = new PointRangeQuery(
+            "test-index",
+            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));
+            }
+        };
+
+        ApproximateableQuery approximateIndexQuery = new ApproximatePointRangeQuery(
+            "test-index",
+            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));
+            }
+        };
+
+        Query dvQuery = SortedNumericDocValuesField.newSlowRangeQuery("test-index", l, u);
+
+        ApproximateIndexOrDocValuesQuery approximateIndexOrDocValuesQuery = new ApproximateIndexOrDocValuesQuery(
+            indexQuery,
+            approximateIndexQuery,
+            dvQuery
+        );
+
+        reader = DirectoryReader.open(w);
+        searcher = newSearcher(reader);
+
+        Weight weight = approximateIndexOrDocValuesQuery.createWeight(searcher, ScoreMode.COMPLETE, 1f);
+
+        // we only get weight since we're expecting to call IODVQ
+        assertFalse(weight instanceof ConstantScoreWeight);
+
+        ApproximateableQuery approximateIndexQueryCanApproximate = new ApproximatePointRangeQuery(
+            "test-index",
+            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));
+            }
+
+            public boolean canApproximate(SearchContext context) {
+                return true;
+            }
+        };
+
+        ApproximateIndexOrDocValuesQuery approximateIndexOrDocValuesQueryCanApproximate = new ApproximateIndexOrDocValuesQuery(
+            indexQuery,
+            approximateIndexQueryCanApproximate,
+            dvQuery
+        );
+
+        Weight approximateIndexOrDocValuesQueryCanApproximateWeight = approximateIndexOrDocValuesQueryCanApproximate.createWeight(
+            searcher,
+            ScoreMode.COMPLETE,
+            1f
+        );
+
+        // we get ConstantScoreWeight since we're expecting to call ApproximatePointRangeQuery
+        assertTrue(approximateIndexOrDocValuesQueryCanApproximateWeight instanceof ConstantScoreWeight);
+
+    }
+}
diff --git a/server/src/test/java/org/opensearch/search/approximate/ApproximatePointRangeQueryTests.java b/server/src/test/java/org/opensearch/search/approximate/ApproximatePointRangeQueryTests.java
@@ -21,12 +21,14 @@
 import org.apache.lucene.search.TotalHits;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
+import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.sort.SortOrder;
 import org.opensearch.test.OpenSearchTestCase;
 
 import java.io.IOException;
 
 import static org.apache.lucene.document.LongPoint.pack;
+import static org.mockito.Mockito.mock;
 
 public class ApproximatePointRangeQueryTests extends OpenSearchTestCase {
 
@@ -305,4 +307,51 @@ protected String toString(int dimension, byte[] value) {
             }
         }
     }
+
+    public void testSize() {
+        ApproximatePointRangeQuery query = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1) {
+            protected String toString(int dimension, byte[] value) {
+                return Long.toString(LongPoint.decodeDimension(value, 0));
+            }
+        };
+        assertEquals(query.getSize(), 10_000);
+
+        query.setSize(100);
+        assertEquals(query.getSize(), 100);
+
+    }
+
+    public void testSortOrder() {
+        ApproximatePointRangeQuery query = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1) {
+            protected String toString(int dimension, byte[] value) {
+                return Long.toString(LongPoint.decodeDimension(value, 0));
+            }
+        };
+        assertNull(query.getSortOrder());
+
+        query.setSortOrder(SortOrder.ASC);
+        assertEquals(query.getSortOrder(), SortOrder.ASC);
+    }
+
+    public void testCanApproximate() {
+        ApproximatePointRangeQuery query = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1) {
+            protected String toString(int dimension, byte[] value) {
+                return Long.toString(LongPoint.decodeDimension(value, 0));
+            }
+        };
+
+        assertFalse(query.canApproximate(null));
+
+        ApproximatePointRangeQuery queryCanApproximate = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1) {
+            protected String toString(int dimension, byte[] value) {
+                return Long.toString(LongPoint.decodeDimension(value, 0));
+            }
+
+            public boolean canApproximate(SearchContext context) {
+                return true;
+            }
+        };
+        SearchContext searchContext = mock(SearchContext.class);
+        assertTrue(queryCanApproximate.canApproximate(searchContext));
+    }
 }