Expand streaming aggregations to numeric terms

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>

Author: harshavamsi

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

@@ -53,7 +53,6 @@
 import java.io.IOException;
 import java.util.AbstractList;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -260,23 +259,6 @@ public int size() {
         return result;
     }
 
-    /**
-     * Build the sub aggregation results for a list of buckets and set them on
-     * the buckets. This is usually used by aggregations that are selective
-     * in which bucket they build. They use some mechanism of selecting a list
-     * of buckets to build use this method to "finish" building the results.
-     * @param buckets the buckets to finish building
-     * @param bucketToOrd how to convert a bucket into an ordinal
-     * @param setAggs how to set the sub-aggregation results on a bucket
-     */
-    protected final <B> void buildSubAggsForBuckets(B[] buckets, ToLongFunction<B> bucketToOrd, BiConsumer<B, InternalAggregations> setAggs)
-        throws IOException {
-        InternalAggregations[] results = buildSubAggsForBuckets(Arrays.stream(buckets).mapToLong(bucketToOrd).toArray());
-        for (int i = 0; i < buckets.length; i++) {
-            setAggs.accept(buckets[i], results[i]);
-        }
-    }
-
     /**
      * Build the sub aggregation results for a list of buckets and set them on
      * the buckets. This is usually used by aggregations that are selective

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/AbstractStringTermsAggregator.java

@@ -51,11 +51,11 @@
  *
  * @opensearch.internal
  */
-abstract class AbstractStringTermsAggregator extends TermsAggregator {
+public abstract class AbstractStringTermsAggregator extends TermsAggregator {
 
     protected final boolean showTermDocCountError;
 
-    AbstractStringTermsAggregator(
+    protected AbstractStringTermsAggregator(
         String name,
         AggregatorFactories factories,
         SearchContext context,

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/DoubleTerms.java

@@ -58,10 +58,10 @@ public class DoubleTerms extends InternalMappedTerms<DoubleTerms, DoubleTerms.Bu
      *
      * @opensearch.internal
      */
-    static class Bucket extends InternalTerms.Bucket<Bucket> {
-        double term;
+    public static class Bucket extends InternalTerms.Bucket<Bucket> {
+        public double term;
 
-        Bucket(
+        public Bucket(
             double term,
             long docCount,
             InternalAggregations aggregations,

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java

@@ -409,7 +409,7 @@ public void collectDebugInfo(BiConsumer<String, Object> add) {
      *
      * @opensearch.internal
      */
-    static class OrdBucket extends InternalTerms.Bucket<OrdBucket> {
+    public static class OrdBucket extends InternalTerms.Bucket<OrdBucket> {
         long globalOrd;
 
         OrdBucket(boolean showDocCountError, DocValueFormat format) {

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/InternalSignificantTerms.java

@@ -82,13 +82,13 @@ public interface Reader<B extends Bucket<B>> {
             B read(StreamInput in, long subsetSize, long supersetSize, DocValueFormat format) throws IOException;
         }
 
-        long subsetDf;
+        public long subsetDf;
         long subsetSize;
-        long supersetDf;
+        public long supersetDf;
         long supersetSize;
-        long bucketOrd;
+        public long bucketOrd;
         double score;
-        protected InternalAggregations aggregations;
+        public InternalAggregations aggregations;
         final transient DocValueFormat format;
 
         protected Bucket(
@@ -139,7 +139,7 @@ public long getSubsetSize() {
         // TODO we should refactor to remove this, since buckets should be immutable after they are generated.
         // This can lead to confusing bugs if the bucket is re-created (via createBucket() or similar) without
         // the score
-        void updateScore(SignificanceHeuristic significanceHeuristic) {
+        public void updateScore(SignificanceHeuristic significanceHeuristic) {
             score = significanceHeuristic.getScore(subsetDf, subsetSize, supersetDf, supersetSize);
         }
 

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/InternalTerms.java

@@ -104,11 +104,11 @@ public interface Reader<B extends Bucket<B>> {
             B read(StreamInput in, DocValueFormat format, boolean showDocCountError) throws IOException;
         }
 
-        long bucketOrd;
+        public long bucketOrd;
 
-        protected long docCount;
+        public long docCount;
         protected long docCountError;
-        protected InternalAggregations aggregations;
+        public InternalAggregations aggregations;
         protected final boolean showDocCountError;
         protected final DocValueFormat format;
 

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongTerms.java

@@ -59,7 +59,7 @@ public class LongTerms extends InternalMappedTerms<LongTerms, LongTerms.Bucket>
      * @opensearch.internal
      */
     public static class Bucket extends InternalTerms.Bucket<Bucket> {
-        long term;
+        public long term;
 
         public Bucket(
             long term,

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/SignificanceLookup.java

@@ -63,18 +63,18 @@
  *
  * @opensearch.internal
  */
-class SignificanceLookup {
+public class SignificanceLookup {
     /**
      * Lookup frequencies for {@link BytesRef} terms.
      */
-    interface BackgroundFrequencyForBytes extends Releasable {
+    public interface BackgroundFrequencyForBytes extends Releasable {
         long freq(BytesRef term) throws IOException;
     }
 
     /**
      * Lookup frequencies for {@code long} terms.
      */
-    interface BackgroundFrequencyForLong extends Releasable {
+    public interface BackgroundFrequencyForLong extends Releasable {
         long freq(long term) throws IOException;
     }
 
@@ -103,7 +103,7 @@ interface BackgroundFrequencyForLong extends Releasable {
     /**
      * Get the number of docs in the superset.
      */
-    long supersetSize() {
+    public long supersetSize() {
         return supersetNumDocs;
     }
 
@@ -155,7 +155,7 @@ private long getBackgroundFrequency(BytesRef term) throws IOException {
     /**
      * Get the background frequency of a {@code long} term.
      */
-    BackgroundFrequencyForLong longLookup(BigArrays bigArrays, CardinalityUpperBound cardinality) {
+    public BackgroundFrequencyForLong longLookup(BigArrays bigArrays, CardinalityUpperBound cardinality) {
         if (cardinality == CardinalityUpperBound.ONE) {
             return new BackgroundFrequencyForLong() {
                 @Override

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/SignificantLongTerms.java

@@ -56,11 +56,11 @@ public class SignificantLongTerms extends InternalMappedSignificantTerms<Signifi
      *
      * @opensearch.internal
      */
-    static class Bucket extends InternalSignificantTerms.Bucket<Bucket> {
+    public static class Bucket extends InternalSignificantTerms.Bucket<Bucket> {
 
-        long term;
+        public long term;
 
-        Bucket(
+        public Bucket(
             long subsetDf,
             long subsetSize,
             long supersetDf,

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/TermsAggregatorFactory.java

@@ -50,6 +50,8 @@
 import org.opensearch.search.aggregations.bucket.BucketUtils;
 import org.opensearch.search.aggregations.bucket.terms.NumericTermsAggregator.ResultStrategy;
 import org.opensearch.search.aggregations.bucket.terms.TermsAggregator.BucketCountThresholds;
+import org.opensearch.search.aggregations.bucket.terms.stream.StreamNumericTermsAggregator;
+import org.opensearch.search.aggregations.bucket.terms.stream.StreamStringTermsAggregator;
 import org.opensearch.search.aggregations.support.CoreValuesSourceType;
 import org.opensearch.search.aggregations.support.ValuesSource;
 import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
@@ -121,7 +123,7 @@ public Aggregator build(
                     // if user doesn't provide execution mode, and using stream search
                     // we use stream aggregation
                     if (context.isStreamSearch()) {
-                        return createStreamAggregator(
+                        return createStreamStringTermsAggregator(
                             name,
                             factories,
                             valuesSource,
@@ -207,7 +209,6 @@ public Aggregator build(
                             + "include/exclude clauses used to filter numeric fields"
                     );
                 }
-
                 if (subAggCollectMode == null) {
                     subAggCollectMode = pickSubAggCollectMode(factories, bucketCountThresholds.getShardSize(), -1, context);
                 }
@@ -231,6 +232,23 @@ public Aggregator build(
                     }
                     resultStrategy = agg -> agg.new LongTermsResults(showTermDocCountError);
                 }
+                if (context.isStreamSearch()) {
+                    return createStreamNumericTermsAggregator(
+                        name,
+                        factories,
+                        numericValuesSource,
+                        format,
+                        order,
+                        bucketCountThresholds,
+                        context,
+                        parent,
+                        longFilter,
+                        includeExclude,
+                        showTermDocCountError,
+                        cardinality,
+                        metadata
+                    );
+                }
                 return new NumericTermsAggregator(
                     name,
                     factories,
@@ -578,7 +596,7 @@ public String toString() {
         }
     }
 
-    static Aggregator createStreamAggregator(
+    static Aggregator createStreamStringTermsAggregator(
         String name,
         AggregatorFactories factories,
         ValuesSource valuesSource,
@@ -610,6 +628,55 @@ static Aggregator createStreamAggregator(
         }
     }
 
+    static Aggregator createStreamNumericTermsAggregator(
+        String name,
+        AggregatorFactories factories,
+        ValuesSource.Numeric valuesSource,
+        DocValueFormat format,
+        BucketOrder order,
+        BucketCountThresholds bucketCountThresholds,
+        SearchContext aggregationContext,
+        Aggregator parent,
+        IncludeExclude.LongFilter longFilter,
+        IncludeExclude includeExclude,
+        boolean showTermDocCountError,
+        CardinalityUpperBound cardinality,
+        Map<String, Object> metadata
+    ) throws IOException {
+        Function<StreamNumericTermsAggregator, StreamNumericTermsAggregator.ResultStrategy<?, ?>> resultStrategy;
+        if (valuesSource.isFloatingPoint()) {
+            if (includeExclude != null) {
+                longFilter = includeExclude.convertToDoubleFilter();
+            }
+            resultStrategy = agg -> agg.new DoubleTermsResults(showTermDocCountError);
+        } else if (valuesSource.isBigInteger()) {
+            if (includeExclude != null) {
+                longFilter = includeExclude.convertToDoubleFilter();
+            }
+            resultStrategy = agg -> agg.new UnsignedLongTermsResults(showTermDocCountError);
+        } else {
+            if (includeExclude != null) {
+                longFilter = includeExclude.convertToLongFilter(format);
+            }
+            resultStrategy = agg -> agg.new LongTermsResults(showTermDocCountError);
+        }
+        return new StreamNumericTermsAggregator(
+            name,
+            factories,
+            resultStrategy,
+            valuesSource,
+            format,
+            order,
+            bucketCountThresholds,
+            aggregationContext,
+            parent,
+            SubAggCollectionMode.DEPTH_FIRST,
+            longFilter,
+            cardinality,
+            metadata
+        );
+    }
+
     @Override
     protected boolean supportsConcurrentSegmentSearch() {
         return true;