Resolve conflicts with star-tree changes

Signed-off-by: Michael Froh <froh@amazon.com>

server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java

@@ -113,7 +113,8 @@ protected boolean tryPrecomputeAggregationForLeaf(LeafReaderContext ctx) throws
                 // Returning NO_OP_COLLECTOR explicitly because the getLeafCollector() are invoked starting from innermost aggregators
                 return true;
             }
-            getStarTreeCollector(ctx, sub, supportedStarTree);
+            precomputeLeafUsingStarTree(ctx, supportedStarTree);
+            return true;
         }
         return false;
     }
@@ -144,12 +145,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
             }
         }
 
-
-        return getDefaultLeafCollector(ctx, sub);
-    }
-
-    private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
-
         final BigArrays bigArrays = context.bigArrays();
         final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
         final NumericDoubleValues values = MultiValueMode.MAX.select(allValues);
@@ -173,9 +168,9 @@ public void collect(int doc, long bucket) throws IOException {
         };
     }
 
-    public void getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) throws IOException {
+    private void precomputeLeafUsingStarTree(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
         AtomicReference<Double> max = new AtomicReference<>(maxes.get(0));
-        StarTreeQueryHelper.getStarTreeLeafCollector(context, valuesSource, ctx, sub, starTree, MetricStat.MAX.getTypeName(), value -> {
+        StarTreeQueryHelper.precomputeLeafUsingStarTree(context, valuesSource, ctx, starTree, MetricStat.MAX.getTypeName(), value -> {
             max.set(Math.max(max.get(), (NumericUtils.sortableLongToDouble(value))));
         }, () -> maxes.set(0, max.get()));
     }

server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java

@@ -113,7 +113,7 @@ protected boolean tryPrecomputeAggregationForLeaf(LeafReaderContext ctx) throws
                 // Returning NO_OP_COLLECTOR explicitly because the getLeafCollector() are invoked starting from innermost aggregators
                 return true;
             }
-            getStarTreeCollector(ctx, sub, supportedStarTree);
+            precomputeLeafUsingStarTree(ctx, supportedStarTree);
             return true;
         }
 
@@ -145,11 +145,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
             }
         }
 
-
-        return getDefaultLeafCollector(ctx, sub);
-    }
-
-    private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
         final BigArrays bigArrays = context.bigArrays();
         final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
         final NumericDoubleValues values = MultiValueMode.MIN.select(allValues);
@@ -172,9 +167,9 @@ public void collect(int doc, long bucket) throws IOException {
         };
     }
 
-    public void getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) throws IOException {
+    private void precomputeLeafUsingStarTree(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
         AtomicReference<Double> min = new AtomicReference<>(mins.get(0));
-        StarTreeQueryHelper.getStarTreeLeafCollector(context, valuesSource, ctx, sub, starTree, MetricStat.MIN.getTypeName(), value -> {
+        StarTreeQueryHelper.precomputeLeafUsingStarTree(context, valuesSource, ctx, starTree, MetricStat.MIN.getTypeName(), value -> {
             min.set(Math.min(min.get(), (NumericUtils.sortableLongToDouble(value))));
         }, () -> mins.set(0, min.get()));
     }

server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java

@@ -101,7 +101,8 @@ protected boolean tryPrecomputeAggregationForLeaf(LeafReaderContext ctx) throws
                 // Returning NO_OP_COLLECTOR explicitly because the getLeafCollector() are invoked starting from innermost aggregators
                 return true;
             }
-            getStarTreeCollector(ctx, sub, supportedStarTree);
+            precomputeLeafUsingStarTree(ctx, supportedStarTree);
+            return true;
         }
         return false;
     }
@@ -112,12 +113,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
         if (valuesSource == null) {
             return LeafBucketCollector.NO_OP_COLLECTOR;
         }
-
-
-        return getDefaultLeafCollector(ctx, sub);
-    }
-
-    private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
         final BigArrays bigArrays = context.bigArrays();
         final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);
         final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
@@ -147,14 +142,13 @@ public void collect(int doc, long bucket) throws IOException {
         };
     }
 
-    public void getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) throws IOException {
+    private void precomputeLeafUsingStarTree(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
         final CompensatedSum kahanSummation = new CompensatedSum(sums.get(0), compensations.get(0));
 
-        StarTreeQueryHelper.getStarTreeLeafCollector(
+        StarTreeQueryHelper.precomputeLeafUsingStarTree(
             context,
             valuesSource,
             ctx,
-            sub,
             starTree,
             MetricStat.SUM.getTypeName(),
             value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)),

server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java

@@ -97,7 +97,7 @@ protected boolean tryPrecomputeAggregationForLeaf(LeafReaderContext ctx) throws
                     // Returning NO_OP_COLLECTOR explicitly because the getLeafCollector() are invoked starting from innermost aggregators
                     return true;
                 }
-                getStarTreeCollector(ctx, sub, supportedStarTree);
+                precomputeLeafUsingStarTree(ctx, supportedStarTree);
                 return true;
             }
         }
@@ -151,12 +151,11 @@ public void collect(int doc, long bucket) throws IOException {
         };
     }
 
-    public void getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) throws IOException {
-        StarTreeQueryHelper.getStarTreeLeafCollector(
+    private void precomputeLeafUsingStarTree(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
+        StarTreeQueryHelper.precomputeLeafUsingStarTree(
             context,
             (ValuesSource.Numeric) valuesSource,
             ctx,
-            sub,
             starTree,
             MetricStat.VALUE_COUNT.getTypeName(),
             value -> counts.increment(0, value),

server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java

@@ -10,7 +10,6 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SegmentReader;
-import org.apache.lucene.search.CollectionTerminatedException;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.util.FixedBitSet;
 import org.opensearch.common.lucene.Lucene;
@@ -21,7 +20,6 @@
 import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
 import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
 import org.opensearch.index.query.QueryShardContext;
-import org.opensearch.search.aggregations.LeafBucketCollector;
 import org.opensearch.search.aggregations.StarTreeBucketCollector;
 import org.opensearch.search.aggregations.support.ValuesSource;
 import org.opensearch.search.internal.SearchContext;
@@ -70,11 +68,10 @@ public static StarTreeValues getStarTreeValues(LeafReaderContext context, Compos
      * Get the star-tree leaf collector
      * This collector computes the aggregation prematurely and invokes an early termination collector
      */
-    public static void getStarTreeLeafCollector(
+    public static void precomputeLeafUsingStarTree(
         SearchContext context,
         ValuesSource.Numeric valuesSource,
         LeafReaderContext ctx,
-        LeafBucketCollector sub,
         CompositeIndexFieldInfo starTree,
         String metric,
         Consumer<Long> valueConsumer,
@@ -112,10 +109,6 @@ public static void getStarTreeLeafCollector(
 
         // Call the final consumer after processing all entries
         finalConsumer.run();
-
-        // FIXME : Remove after @msfroh PR for precompute
-        // Terminate after pre-computing aggregation
-        throw new CollectionTerminatedException();
     }
 
     /**