Adding test and refactor

Signed-off-by: bowenlan-amzn <[email protected]>

modules/mapper-extras/src/main/java/org/opensearch/index/mapper/ScaledFloatFieldMapper.java

@@ -188,6 +188,7 @@ public ScaledFloatFieldType(String name, double scalingFactor) {
             this(name, true, false, true, Collections.emptyMap(), scalingFactor, null);
         }
 
+        @Override
         public double getScalingFactor() {
             return scalingFactor;
         }

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

@@ -536,4 +536,8 @@ public Map<String, String> meta() {
     public TextSearchInfo getTextSearchInfo() {
         return textSearchInfo;
     }
+
+    public double getScalingFactor() {
+        throw new IllegalCallerException("Field [" + name() + "] of type [" + typeName() + "] does not support scaling factor");
+    }
 }

server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java

@@ -45,10 +45,10 @@
 import org.opensearch.search.aggregations.bucket.histogram.LongBounds;
 import org.opensearch.search.aggregations.bucket.range.RangeAggregator.Range;
 import org.opensearch.search.aggregations.support.ValuesSource;
+import org.opensearch.search.aggregations.support.ValuesSourceConfig;
 import org.opensearch.search.internal.SearchContext;
 
 import java.io.IOException;
-import java.lang.reflect.InvocationTargetException;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.util.HashMap;
@@ -456,20 +456,32 @@ public DebugInfo tryFastFilterAggregation(
     /**
      * For range aggregation
      */
-    public static class RangeAggregationType implements FastFilterRewriteHelper.AggregationType {
+    public static class RangeAggregationType implements AggregationType {
 
         private final ValuesSource.Numeric source;
+        private final ValuesSourceConfig config;
         private final Range[] ranges;
+        private FieldTypeEnum fieldTypeEnum;
 
-        public RangeAggregationType(ValuesSource.Numeric source, Range[] ranges) {
-            this.source = source;
+        public RangeAggregationType(ValuesSourceConfig config, Range[] ranges) {
+            this.source = (ValuesSource.Numeric) config.getValuesSource();
+            this.config = config;
             this.ranges = ranges;
         }
 
         @Override
         public boolean isRewriteable(Object parent, int subAggLength) {
-            if (parent == null && subAggLength == 0) {
-                // don't accept values source with scripts
+            if (parent == null && subAggLength == 0 && config.script() == null && config.missing() == null) {
+                try {
+                    fieldTypeEnum = FieldTypeEnum.fromTypeName(config.fieldType().typeName());
+                    if (fieldTypeEnum.isScaled()) {
+                        // make sure we can safely get scaling factor later on
+                        config.fieldType().getScalingFactor();
+                    }
+                } catch (Exception e) {
+                    return false;
+                }
+
                 if (source instanceof ValuesSource.Numeric.FieldData) {
                     // ranges are already sorted by from and then to
                     // we want ranges not overlapping with each other
@@ -488,57 +500,16 @@ public boolean isRewriteable(Object parent, int subAggLength) {
 
         @Override
         public Ranges buildRanges(SearchContext ctx, MappedFieldType fieldType) throws IOException {
-            int byteLen = -1;
-            boolean scaled = false;
-            String pointType = "";
-            switch (fieldType.typeName()) {
-                case "half_float":
-                    byteLen = 2;
-                    pointType = "half_float";
-                    break;
-                case "integer":
-                case "short":
-                case "byte":
-                    byteLen = 4;
-                    pointType = "int";
-                    break;
-                case "float":
-                    byteLen = 4;
-                    pointType = "float";
-                    break;
-                case "long":
-                case "date":
-                case "date_nanos":
-                    byteLen = 8;
-                    pointType = "long";
-                    break;
-                case "double":
-                    byteLen = 8;
-                    pointType = "double";
-                    break;
-                case "unsigned_long":
-                    byteLen = 16;
-                    pointType = "big_integer";
-                    break;
-                case "scaled_float":
-                    byteLen = 8;
-                    scaled = true;
-                    pointType = "long";
-                    break;
-            }
-            if (byteLen == -1) {
-                throw new IllegalArgumentException("Field type " + fieldType.name() + " is not supported");
-            }
+            int byteLen = this.fieldTypeEnum.getByteLen();
+            String pointType = this.fieldTypeEnum.getPointType();
 
             byte[][] mins = new byte[ranges.length][];
             byte[][] maxs = new byte[ranges.length][];
             for (int i = 0; i < ranges.length; i++) {
                 double rangeMin = ranges[i].getFrom();
                 double rangeMax = ranges[i].getTo();
-
                 byte[] min = new byte[byteLen];
                 byte[] max = new byte[byteLen];
-                // TODO any heavy operation that shouldn't be in the loop?
                 switch (pointType) {
                     case "half_float":
                         HalfFloatPoint.encodeDimension((float) rangeMin, min, 0);
@@ -557,18 +528,8 @@ public Ranges buildRanges(SearchContext ctx, MappedFieldType fieldType) throws I
                         IntPoint.encodeDimension((int) rangeMax, max, 0);
                         break;
                     case "long":
-                        if (scaled) {
-                            // use reflection to see if fieldType has a method called getScalingFactor
-                            // if it does, use that to scale the rangeMin and rangeMax
-                            // if it doesn't, use the default scaling factor of 1000
-                            double scalingFactor = 100;
-                            try {
-                                // Method scalingFactorMethod = fieldType.getClass().getMethod("getScalingFactor");
-                                scalingFactor = (double) fieldType.getClass().getMethod("getScalingFactor").invoke(fieldType);
-                            } catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
-                                logger.debug("Failed to get scaling factor from fieldType", e);
-                                // TODO fall back to default aggregation method
-                            }
+                        if (this.fieldTypeEnum.isScaled()) {
+                            double scalingFactor = fieldType.getScalingFactor();
                             rangeMin = scalingFactor * rangeMin;
                             rangeMax = scalingFactor * rangeMax;
                         }
@@ -588,6 +549,61 @@ public Ranges buildRanges(SearchContext ctx, MappedFieldType fieldType) throws I
             return new Ranges(mins, maxs, byteLen);
         }
 
+        private enum FieldTypeEnum {
+            HALF_FLOAT("half_float", 2, "half_float"),
+            INTEGER("integer", 4, "int"),
+            SHORT("short", 4, "int"),
+            BYTE("byte", 4, "int"),
+            FLOAT("float", 4, "float"),
+            LONG("long", 8, "long"),
+            DATE("date", 8, "long"),
+            DATE_NANOS("date_nanos", 8, "long"),
+            DOUBLE("double", 8, "double"),
+            UNSIGNED_LONG("unsigned_long", 16, "big_integer"),
+            SCALED_FLOAT("scaled_float", 8, "long", true);
+
+            private final String typeName;
+            private final int byteLen;
+            private final String pointType;
+            private final boolean scaled;
+
+            FieldTypeEnum(String typeName, int byteLen, String pointType) {
+                this(typeName, byteLen, pointType, false);
+            }
+
+            FieldTypeEnum(String typeName, int byteLen, String pointType, boolean scaled) {
+                this.typeName = typeName;
+                this.byteLen = byteLen;
+                this.pointType = pointType;
+                this.scaled = scaled;
+            }
+
+            String getTypeName() {
+                return typeName;
+            }
+
+            int getByteLen() {
+                return byteLen;
+            }
+
+            String getPointType() {
+                return pointType;
+            }
+
+            boolean isScaled() {
+                return scaled;
+            }
+
+            static FieldTypeEnum fromTypeName(String typeName) {
+                for (FieldTypeEnum fieldTypeEnum : values()) {
+                    if (fieldTypeEnum.getTypeName().equals(typeName)) {
+                        return fieldTypeEnum;
+                    }
+                }
+                throw new IllegalArgumentException("Unknown field type: " + typeName);
+            }
+        }
+
         public static BigInteger convertDoubleToBigInteger(double value) {
             // we use big integer to represent unsigned long
             BigInteger maxUnsignedLong = BigInteger.valueOf(2).pow(64).subtract(BigInteger.ONE);
@@ -608,7 +624,7 @@ public static BigInteger convertDoubleToBigInteger(double value) {
 
         @Override
         public Ranges buildRanges(LeafReaderContext leaf, SearchContext ctx, MappedFieldType fieldType) throws IOException {
-            return buildRanges(ctx, fieldType);
+            throw new UnsupportedOperationException("Range aggregation should not build ranges at segment level");
         }
 
         @Override
@@ -833,7 +849,7 @@ public void visit(int docID, byte[] packedValue) throws IOException {
             @Override
             public void visit(DocIdSetIterator iterator, byte[] packedValue) throws IOException {
                 visitPoints(packedValue, () -> {
-                    for (int doc = iterator.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = iterator.nextDoc()) {
+                    for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) {
                         collector.count();
                     }
                 });

server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java

@@ -271,11 +271,10 @@ public RangeAggregator(
         this.format = format;
         this.keyed = keyed;
         this.rangeFactory = rangeFactory;
-
         this.ranges = ranges; // already sorted by the range.from and range.to
 
         fastFilterContext = new FastFilterRewriteHelper.FastFilterContext(context);
-        fastFilterContext.setAggregationType(new FastFilterRewriteHelper.RangeAggregationType(valuesSource, ranges));
+        fastFilterContext.setAggregationType(new FastFilterRewriteHelper.RangeAggregationType(config, ranges));
         if (fastFilterContext.isRewriteable(parent, subAggregators.length)) {
             fastFilterContext.buildRanges(Objects.requireNonNull(config.fieldType()));
         }

server/src/test/java/org/opensearch/search/aggregations/bucket/range/RangeAggregatorTests.java

@@ -32,16 +32,22 @@
 
 package org.opensearch.search.aggregations.bucket.range;
 
+import org.apache.lucene.document.DoubleField;
+import org.apache.lucene.document.DoublePoint;
+import org.apache.lucene.document.Field;
 import org.apache.lucene.document.NumericDocValuesField;
 import org.apache.lucene.document.SortedNumericDocValuesField;
 import org.apache.lucene.document.SortedSetDocValuesField;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
+import org.apache.lucene.tests.util.TestUtil;
 import org.apache.lucene.util.BytesRef;
 import org.opensearch.common.CheckedConsumer;
 import org.opensearch.index.mapper.DateFieldMapper;
@@ -67,6 +73,8 @@ public class RangeAggregatorTests extends AggregatorTestCase {
     private static final String NUMBER_FIELD_NAME = "number";
     private static final String DATE_FIELD_NAME = "date";
 
+    private static final String DOUBLE_FIELD_NAME = "double";
+
     public void testNoMatchingField() throws IOException {
         testCase(new MatchAllDocsQuery(), iw -> {
             iw.addDocument(singleton(new SortedNumericDocValuesField("bogus_field_name", 7)));
@@ -305,6 +313,26 @@ public void testSubAggCollectsFromManyBucketsIfManyRanges() throws IOException {
         });
     }
 
+    public void testDoubleType() throws IOException {
+        RangeAggregationBuilder aggregationBuilder = new RangeAggregationBuilder("range").field(DOUBLE_FIELD_NAME)
+            .addRange(1, 2)
+            .addRange(2, 3);
+
+        testRewriteOptimizationCase(aggregationBuilder, DoublePoint.newRangeQuery(DOUBLE_FIELD_NAME, 0, 5), indexWriter -> {
+            indexWriter.addDocument(singleton(new DoubleField(DOUBLE_FIELD_NAME, 0.1, Field.Store.NO)));
+            indexWriter.addDocument(singleton(new DoubleField(DOUBLE_FIELD_NAME, 1.1, Field.Store.NO)));
+            indexWriter.addDocument(singleton(new DoubleField(DOUBLE_FIELD_NAME, 2.1, Field.Store.NO)));
+        }, range -> {
+            List<? extends InternalRange.Bucket> ranges = range.getBuckets();
+            assertEquals(2, ranges.size());
+            assertEquals("1.0-2.0", ranges.get(0).getKeyAsString());
+            assertEquals(1, ranges.get(0).getDocCount());
+            assertEquals("2.0-3.0", ranges.get(1).getKeyAsString());
+            assertEquals(1, ranges.get(1).getDocCount());
+            assertTrue(AggregationInspectionHelper.hasValue(range));
+        }, new NumberFieldMapper.NumberFieldType(DOUBLE_FIELD_NAME, NumberFieldMapper.NumberType.DOUBLE));
+    }
+
     private void testCase(
         Query query,
         CheckedConsumer<RandomIndexWriter, IOException> buildIndex,
@@ -325,7 +353,7 @@ private void simpleTestCase(
     ) throws IOException {
         MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(NUMBER_FIELD_NAME, NumberFieldMapper.NumberType.INTEGER);
 
-        testCase(aggregationBuilder, new MatchAllDocsQuery(), iw -> {
+        testCase(aggregationBuilder, query, iw -> {
             iw.addDocument(singleton(new SortedNumericDocValuesField(NUMBER_FIELD_NAME, 7)));
             iw.addDocument(singleton(new SortedNumericDocValuesField(NUMBER_FIELD_NAME, 2)));
             iw.addDocument(singleton(new SortedNumericDocValuesField(NUMBER_FIELD_NAME, 3)));
@@ -354,7 +382,32 @@ private void testCase(
                     fieldType
                 );
                 verify.accept(agg);
+            }
+        }
+    }
 
+    private void testRewriteOptimizationCase(
+        RangeAggregationBuilder aggregationBuilder,
+        Query query,
+        CheckedConsumer<IndexWriter, IOException> buildIndex,
+        Consumer<InternalRange<? extends InternalRange.Bucket, ? extends InternalRange>> verify,
+        MappedFieldType fieldType
+    ) throws IOException {
+        try (Directory directory = newDirectory()) {
+            try (IndexWriter indexWriter = new IndexWriter(directory, new IndexWriterConfig().setCodec(TestUtil.getDefaultCodec()))) {
+                buildIndex.accept(indexWriter);
+            }
+
+            try (IndexReader indexReader = DirectoryReader.open(directory)) {
+                IndexSearcher indexSearcher = newSearcher(indexReader, true, true);
+
+                InternalRange<? extends InternalRange.Bucket, ? extends InternalRange> agg = searchAndReduce(
+                    indexSearcher,
+                    query,
+                    aggregationBuilder,
+                    fieldType
+                );
+                verify.accept(agg);
             }
         }
     }