prepareFromSegment now doesn't return Ranges

Signed-off-by: bowenlan-amzn <bowenlan23@gmail.com>

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

@@ -44,7 +44,7 @@ void setOptimizationContext(OptimizationContext optimizationContext) {
 
     protected abstract void prepare() throws IOException;
 
-    abstract Ranges prepare(LeafReaderContext leaf) throws IOException;
+    abstract void prepareFromSegment(LeafReaderContext leaf) throws IOException;
 
     abstract void tryOptimize(PointValues values, BiConsumer<Long, Long> incrementDocCount, Ranges ranges) throws IOException;
 

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

@@ -65,13 +65,13 @@ protected boolean canOptimize(CompositeValuesSourceConfig[] sourceConfigs) {
 
     protected void buildRanges(SearchContext context) throws IOException {
         long[] bounds = Helper.getDateHistoAggBounds(context, fieldType.name());
-        this.optimizationContext.setRanges(buildRanges(bounds));
+        optimizationContext.setRanges(buildRanges(bounds));
     }
 
     @Override
-    OptimizationContext.Ranges prepare(LeafReaderContext leaf) throws IOException {
+    void prepareFromSegment(LeafReaderContext leaf) throws IOException {
         long[] bounds = Helper.getSegmentBounds(leaf, fieldType.name());
-        return buildRanges(bounds);
+        optimizationContext.setRangesFromSegment(buildRanges(bounds));
     }
 
     private OptimizationContext.Ranges buildRanges(long[] bounds) {
@@ -97,7 +97,7 @@ private OptimizationContext.Ranges buildRanges(long[] bounds) {
             getRoundingPrepared(),
             bounds[0],
             bounds[1],
-            this.optimizationContext.maxAggRewriteFilters
+            optimizationContext.maxAggRewriteFilters
         );
     }
 
@@ -153,7 +153,7 @@ final void tryOptimize(PointValues values, BiConsumer<Long, Long> incrementDocCo
             incrementDocCount.accept(ord, (long) docCount);
         };
 
-        this.optimizationContext.consumeDebugInfo(multiRangesTraverse(values.getPointTree(), ranges, incrementFunc, size));
+        optimizationContext.consumeDebugInfo(multiRangesTraverse(values.getPointTree(), ranges, incrementFunc, size));
     }
 
     private static long getBucketOrd(long bucketOrd) {

server/src/main/java/org/opensearch/search/optimization/ranges/OptimizationContext.java

@@ -47,6 +47,7 @@ public final class OptimizationContext {
     String shardId;
 
     private Ranges ranges;
+    private Ranges rangesFromSegment;
 
     // debug info related fields
     private int leaf;
@@ -75,20 +76,24 @@ public boolean canOptimize(final Object parent, final int subAggLength, SearchCo
 
     public void prepare() throws IOException {
         assert ranges == null : "Ranges should only be built once at shard level, but they are already built";
-        this.aggregatorBridge.prepare();
+        aggregatorBridge.prepare();
         if (ranges != null) {
             preparedAtShardLevel = true;
         }
     }
 
-    private Ranges prepare(LeafReaderContext leaf) throws IOException {
-        return this.aggregatorBridge.prepare(leaf);
+    private void prepareFromSegment(LeafReaderContext leaf) throws IOException {
+        aggregatorBridge.prepareFromSegment(leaf);
     }
 
     void setRanges(Ranges ranges) {
         this.ranges = ranges;
     }
 
+    void setRangesFromSegment(Ranges ranges) {
+        this.rangesFromSegment = ranges;
+    }
+
     /**
      * Try to populate the bucket doc counts for aggregation
      * <p>
@@ -119,7 +124,7 @@ public boolean tryOptimize(final LeafReaderContext leafCtx, final BiConsumer<Lon
             return false;
         }
 
-        Ranges ranges = prepareFromSegment(leafCtx);
+        Ranges ranges = prepare(leafCtx);
         if (ranges == null) return false;
 
         aggregatorBridge.tryOptimize(values, incrementDocCount, ranges);
@@ -134,15 +139,15 @@ public boolean tryOptimize(final LeafReaderContext leafCtx, final BiConsumer<Lon
      * Even when ranges cannot be built at shard level, we can still build ranges
      * at segment level when it's functionally match-all at segment level
      */
-    private Ranges prepareFromSegment(LeafReaderContext leafCtx) throws IOException {
+    private Ranges prepare(LeafReaderContext leafCtx) throws IOException {
         if (!preparedAtShardLevel && !aggregatorBridge.segmentMatchAll(leafCtx)) {
             return null;
         }
 
-        Ranges ranges = this.ranges;
         if (ranges == null) { // not built at shard level but segment match all
             logger.debug("Shard {} segment {} functionally match all documents. Build the fast filter", shardId, leafCtx.ord);
-            ranges = prepare(leafCtx);
+            prepareFromSegment(leafCtx);
+            return rangesFromSegment;
         }
         return ranges;
     }

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

@@ -68,7 +68,7 @@ protected void buildRanges(RangeAggregator.Range[] ranges) {
     }
 
     @Override
-    OptimizationContext.Ranges prepare(LeafReaderContext leaf) {
+    void prepareFromSegment(LeafReaderContext leaf) {
         throw new UnsupportedOperationException("Range aggregation should not build ranges at segment level");
     }