Replace generic object lambda in range agg with long primitive

Signed-off-by: Finn Carroll <[email protected]>

server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java

@@ -90,6 +90,7 @@
 import java.util.Map;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
+import java.util.function.LongFunction;
 import java.util.function.LongUnaryOperator;
 import java.util.stream.Collectors;
 
@@ -219,7 +220,7 @@ protected int getSize() {
             }
 
             @Override
-            protected Function<Long, Long> bucketOrdProducer() {
+            protected LongFunction<Long> bucketOrdProducer() {
                 return (key) -> bucketOrds.add(0, getRoundingPrepared().round((long) key));
             }
         });

server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java

@@ -66,6 +66,7 @@
 import java.util.Map;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
+import java.util.function.LongFunction;
 import java.util.function.LongToIntFunction;
 
 import static org.opensearch.search.optimization.filterrewrite.DateHistogramAggregatorBridge.segmentMatchAll;
@@ -200,7 +201,7 @@ protected Prepared getRoundingPrepared() {
             }
 
             @Override
-            protected Function<Long, Long> bucketOrdProducer() {
+            protected LongFunction<Long> bucketOrdProducer() {
                 return (key) -> getBucketOrds().add(0, preparedRounding.round((long) key));
             }
 

server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java

@@ -60,6 +60,7 @@
 import java.util.Map;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
+import java.util.function.LongFunction;
 
 import static org.opensearch.search.optimization.filterrewrite.DateHistogramAggregatorBridge.segmentMatchAll;
 
@@ -146,7 +147,7 @@ protected long[] processHardBounds(long[] bounds) {
             }
 
             @Override
-            protected Function<Long, Long> bucketOrdProducer() {
+            protected LongFunction<Long> bucketOrdProducer() {
                 return (key) -> bucketOrds.add(0, preparedRounding.round((long) key));
             }
 

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

@@ -68,6 +68,7 @@
 import java.util.Objects;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
+import java.util.function.LongFunction;
 
 import static org.opensearch.core.xcontent.ConstructingObjectParser.optionalConstructorArg;
 
@@ -293,7 +294,7 @@ public void prepare() {
             }
 
             @Override
-            protected Function<Object, Long> bucketOrdProducer() {
+            protected LongFunction<Long> bucketOrdProducer() {
                 return (activeIndex) -> subBucketOrdinal(0, (int) activeIndex);
             }
         });

server/src/main/java/org/opensearch/search/optimization/filterrewrite/AggregatorBridge.java

@@ -15,6 +15,8 @@
 
 import java.io.IOException;
 import java.util.function.BiConsumer;
+import java.util.function.Function;
+import java.util.function.LongFunction;
 
 /**
  * This interface provides a bridge between an aggregator and the optimization context, allowing
@@ -75,4 +77,9 @@ void setOptimizationContext(OptimizationContext context) {
      */
     public abstract void tryOptimize(PointValues values, BiConsumer<Long, Long> incrementDocCount, final LeafBucketCollector sub)
         throws IOException;
+
+    /**
+     * Provides a function to produce bucket ordinals from index of the corresponding range in the range array
+     */
+    protected abstract LongFunction<Long> bucketOrdProducer();
 }

server/src/main/java/org/opensearch/search/optimization/filterrewrite/DateHistogramAggregatorBridge.java

@@ -171,11 +171,6 @@ protected static long getBucketOrd(long bucketOrd) {
         return bucketOrd;
     }
 
-    /**
-    * Provides a function to produce bucket ordinals from the lower bound of the range
-    */
-    protected abstract Function<Long, Long> bucketOrdProducer();
-
     /**
      * Checks whether the top level query matches all documents on the segment
      *

server/src/main/java/org/opensearch/search/optimization/filterrewrite/RangeAggregatorBridge.java

@@ -84,7 +84,7 @@ public final void tryOptimize(PointValues values, BiConsumer<Long, Long> increme
                 optimizationContext.getRanges(),
                 Integer.MAX_VALUE,
                 (activeIndex, docID) -> {
-                    long ord = bucketOrdProducer().apply(activeIndex);
+                    long ord = bucketOrdProducer().apply((long) activeIndex);
 
                     try {
                         incrementDocCount.accept(ord, (long) 1);
@@ -100,17 +100,12 @@ public final void tryOptimize(PointValues values, BiConsumer<Long, Long> increme
                 optimizationContext.getRanges(),
                 Integer.MAX_VALUE,
                 (activeIndex, docCount) -> {
-                    long ord = bucketOrdProducer().apply(activeIndex);
+                    long ord = bucketOrdProducer().apply((long) activeIndex);
                     incrementDocCount.accept(ord, (long) docCount);
                 }
             );
         }
 
         optimizationContext.consumeDebugInfo(multiRangesTraverse(treeVisitor));
     }
-
-    /**
-     * Provides a function to produce bucket ordinals from index of the corresponding range in the range array
-     */
-    protected abstract Function<Object, Long> bucketOrdProducer();
 }

server/src/main/java/org/opensearch/search/optimization/filterrewrite/TreeTraversal.java

@@ -72,7 +72,7 @@ public RangeAwareIntersectVisitor(PointValues.PointTree pointTree, PackedValueRa
             this.activeIndex = packedValueRanges.firstRangeIndex(pointTree.getMinPackedValue(), pointTree.getMaxPackedValue());
         }
 
-        public int getActiveIndex() {
+        public long getActiveIndex() {
             return activeIndex;
         }