Top level range query unit test

Signed-off-by: Finn Carroll <[email protected]>
finnegancarroll · Aug 13, 2024 · 7aa2415 · 7aa2415
1 parent b9bd788
commit 7aa2415
Showing 1 changed file with 49 additions and 0 deletions.
diff --git a/...r/src/test/java/org/opensearch/search/aggregations/bucket/range/RangeAggregatorTests.java b/...r/src/test/java/org/opensearch/search/aggregations/bucket/range/RangeAggregatorTests.java
@@ -32,6 +32,7 @@
 
 package org.opensearch.search.aggregations.bucket.range;
 
+import org.apache.lucene.document.IntPoint;
 import org.apache.lucene.document.NumericDocValuesField;
 import org.apache.lucene.document.SortedNumericDocValuesField;
 import org.apache.lucene.document.SortedSetDocValuesField;
@@ -41,6 +42,7 @@
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.PointRangeQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
@@ -54,6 +56,7 @@
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.mapper.NumberFieldMapper.NumberFieldType;
 import org.opensearch.index.mapper.NumberFieldMapper.NumberType;
+import org.opensearch.index.query.TermsQueryBuilder;
 import org.opensearch.search.aggregations.AggregationBuilder;
 import org.opensearch.search.aggregations.AggregatorTestCase;
 import org.opensearch.search.aggregations.CardinalityUpperBound;
@@ -457,6 +460,52 @@ public void testFloatType() throws IOException {
         );
     }
 
+    public void testTopLevelFilterQuery() throws IOException {
+        NumberFieldType fieldType = new NumberFieldType(NumberType.INTEGER.typeName(), NumberType.INTEGER);
+        String fieldName = fieldType.numberType().typeName();
+        Query query = IntPoint.newRangeQuery(fieldName, 5, 20);
+
+        testRewriteOptimizationCase(
+            fieldType,
+            new double[][] { { 0.0, 10.0 }, { 10.0, 20.0 } },
+            query,
+            new Number[] { 0.1, 4.0, 9, 11, 12, 19 },
+            range -> {
+                List<? extends InternalRange.Bucket> ranges = range.getBuckets();
+                assertEquals(2, ranges.size());
+                assertEquals("0.0-10.0", ranges.get(0).getKeyAsString());
+                assertEquals(1, ranges.get(0).getDocCount());
+                assertEquals("10.0-20.0", ranges.get(1).getKeyAsString());
+                assertEquals(3, ranges.get(1).getDocCount());
+                assertTrue(AggregationInspectionHelper.hasValue(range));
+            },
+            false
+        );
+    }
+
+    public void testTopLevelTermQuery() throws IOException {
+        NumberFieldType fieldType = new NumberFieldType(NumberType.INTEGER.typeName(), NumberType.INTEGER);
+        String fieldName = fieldType.numberType().typeName();
+        Query query = IntPoint.newRangeQuery(fieldName, 5, 20);
+
+        testRewriteOptimizationCase(
+            fieldType,
+            new double[][] { { 0.0, 10.0 }, { 10.0, 20.0 } },
+            query,
+            new Number[] { 0.1, 4.0, 9, 11, 12, 19 },
+            range -> {
+                List<? extends InternalRange.Bucket> ranges = range.getBuckets();
+                assertEquals(2, ranges.size());
+                assertEquals("0.0-10.0", ranges.get(0).getKeyAsString());
+                assertEquals(1, ranges.get(0).getDocCount());
+                assertEquals("10.0-20.0", ranges.get(1).getKeyAsString());
+                assertEquals(3, ranges.get(1).getDocCount());
+                assertTrue(AggregationInspectionHelper.hasValue(range));
+            },
+            false
+        );
+    }
+
     public void testUnsignedLongType() throws IOException {
         testRewriteOptimizationCase(
             new NumberFieldType(NumberType.UNSIGNED_LONG.typeName(), NumberType.UNSIGNED_LONG),