Remove supersetSize and subsetSize from InternalSignificantTerms.Bucket (#117574)

Those fields are only used to update the score and not serialized in the bucket so they can be removed.

Author: iverase

server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/bucket/SignificantTermsSignificanceScoreIT.java

@@ -495,7 +495,7 @@ public void testScriptScore() throws ExecutionException, InterruptedException, I
                 for (SignificantTerms.Bucket bucket : sigTerms.getBuckets()) {
                     assertThat(
                         bucket.getSignificanceScore(),
-                        is((double) bucket.getSubsetDf() + bucket.getSubsetSize() + bucket.getSupersetDf() + bucket.getSupersetSize())
+                        is((double) bucket.getSubsetDf() + sigTerms.getSubsetSize() + bucket.getSupersetDf() + sigTerms.getSupersetSize())
                     );
                 }
             }

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

@@ -989,7 +989,7 @@ SignificantStringTerms.Bucket[] buildBuckets(int size) {
 
         @Override
         SignificantStringTerms.Bucket buildEmptyTemporaryBucket() {
-            return new SignificantStringTerms.Bucket(new BytesRef(), 0, 0, 0, 0, null, format, 0);
+            return new SignificantStringTerms.Bucket(new BytesRef(), 0, 0, null, format, 0);
         }
 
         private long subsetSize(long owningBucketOrd) {
@@ -998,22 +998,19 @@ private long subsetSize(long owningBucketOrd) {
         }
 
         @Override
-        BucketUpdater<SignificantStringTerms.Bucket> bucketUpdater(long owningBucketOrd, GlobalOrdLookupFunction lookupGlobalOrd)
-            throws IOException {
+        BucketUpdater<SignificantStringTerms.Bucket> bucketUpdater(long owningBucketOrd, GlobalOrdLookupFunction lookupGlobalOrd) {
             long subsetSize = subsetSize(owningBucketOrd);
             return (spare, globalOrd, bucketOrd, docCount) -> {
                 spare.bucketOrd = bucketOrd;
                 oversizedCopy(lookupGlobalOrd.apply(globalOrd), spare.termBytes);
                 spare.subsetDf = docCount;
-                spare.subsetSize = subsetSize;
                 spare.supersetDf = backgroundFrequencies.freq(spare.termBytes);
-                spare.supersetSize = supersetSize;
                 /*
                  * During shard-local down-selection we use subset/superset stats
                  * that are for this shard only. Back at the central reducer these
                  * properties will be updated with global stats.
                  */
-                spare.updateScore(significanceHeuristic);
+                spare.updateScore(significanceHeuristic, subsetSize, supersetSize);
             };
         }
 

server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalMappedSignificantTerms.java

@@ -59,7 +59,7 @@ protected InternalMappedSignificantTerms(StreamInput in, Bucket.Reader<B> bucket
         subsetSize = in.readVLong();
         supersetSize = in.readVLong();
         significanceHeuristic = in.readNamedWriteable(SignificanceHeuristic.class);
-        buckets = in.readCollectionAsList(stream -> bucketReader.read(stream, subsetSize, supersetSize, format));
+        buckets = in.readCollectionAsList(stream -> bucketReader.read(stream, format));
     }
 
     @Override
@@ -91,12 +91,12 @@ public B getBucketByKey(String term) {
     }
 
     @Override
-    protected long getSubsetSize() {
+    public long getSubsetSize() {
         return subsetSize;
     }
 
     @Override
-    protected long getSupersetSize() {
+    public long getSupersetSize() {
         return supersetSize;
     }
 

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

@@ -53,13 +53,11 @@ public abstract static class Bucket<B extends Bucket<B>> extends InternalMultiBu
          */
         @FunctionalInterface
         public interface Reader<B extends Bucket<B>> {
-            B read(StreamInput in, long subsetSize, long supersetSize, DocValueFormat format) throws IOException;
+            B read(StreamInput in, DocValueFormat format) throws IOException;
         }
 
         long subsetDf;
-        long subsetSize;
         long supersetDf;
-        long supersetSize;
         /**
          * Ordinal of the bucket while it is being built. Not used after it is
          * returned from {@link Aggregator#buildAggregations(org.elasticsearch.common.util.LongArray)} and not
@@ -70,16 +68,7 @@ public interface Reader<B extends Bucket<B>> {
         protected InternalAggregations aggregations;
         final transient DocValueFormat format;
 
-        protected Bucket(
-            long subsetDf,
-            long subsetSize,
-            long supersetDf,
-            long supersetSize,
-            InternalAggregations aggregations,
-            DocValueFormat format
-        ) {
-            this.subsetSize = subsetSize;
-            this.supersetSize = supersetSize;
+        protected Bucket(long subsetDf, long supersetDf, InternalAggregations aggregations, DocValueFormat format) {
             this.subsetDf = subsetDf;
             this.supersetDf = supersetDf;
             this.aggregations = aggregations;
@@ -89,9 +78,7 @@ protected Bucket(
         /**
          * Read from a stream.
          */
-        protected Bucket(StreamInput in, long subsetSize, long supersetSize, DocValueFormat format) {
-            this.subsetSize = subsetSize;
-            this.supersetSize = supersetSize;
+        protected Bucket(StreamInput in, DocValueFormat format) {
             this.format = format;
         }
 
@@ -105,20 +92,10 @@ public long getSupersetDf() {
             return supersetDf;
         }
 
-        @Override
-        public long getSupersetSize() {
-            return supersetSize;
-        }
-
-        @Override
-        public long getSubsetSize() {
-            return subsetSize;
-        }
-
         // 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) {
+        void updateScore(SignificanceHeuristic significanceHeuristic, long subsetSize, long supersetSize) {
             score = significanceHeuristic.getScore(subsetDf, subsetSize, supersetDf, supersetSize);
         }
 
@@ -262,13 +239,11 @@ public InternalAggregation get() {
                     buckets.forEach(entry -> {
                         final B b = createBucket(
                             entry.value.subsetDf[0],
-                            globalSubsetSize,
                             entry.value.supersetDf[0],
-                            globalSupersetSize,
                             entry.value.reducer.getAggregations(),
                             entry.value.reducer.getProto()
                         );
-                        b.updateScore(heuristic);
+                        b.updateScore(heuristic, globalSubsetSize, globalSupersetSize);
                         if (((b.score > 0) && (b.subsetDf >= minDocCount)) || reduceContext.isFinalReduce() == false) {
                             final B removed = ordered.insertWithOverflow(b);
                             if (removed == null) {
@@ -317,9 +292,7 @@ public InternalAggregation finalizeSampling(SamplingContext samplingContext) {
                 .map(
                     b -> createBucket(
                         samplingContext.scaleUp(b.subsetDf),
-                        subsetSize,
                         samplingContext.scaleUp(b.supersetDf),
-                        supersetSize,
                         InternalAggregations.finalizeSampling(b.aggregations, samplingContext),
                         b
                     )
@@ -328,14 +301,7 @@ public InternalAggregation finalizeSampling(SamplingContext samplingContext) {
         );
     }
 
-    abstract B createBucket(
-        long subsetDf,
-        long subsetSize,
-        long supersetDf,
-        long supersetSize,
-        InternalAggregations aggregations,
-        B prototype
-    );
+    abstract B createBucket(long subsetDf, long supersetDf, InternalAggregations aggregations, B prototype);
 
     protected abstract A create(long subsetSize, long supersetSize, List<B> buckets);
 
@@ -344,10 +310,6 @@ abstract B createBucket(
      */
     protected abstract B[] createBucketsArray(int size);
 
-    protected abstract long getSubsetSize();
-
-    protected abstract long getSupersetSize();
-
     protected abstract SignificanceHeuristic getSignificanceHeuristic();
 
     @Override

server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/MapStringTermsAggregator.java

@@ -47,7 +47,6 @@
 import java.util.function.BiConsumer;
 import java.util.function.Function;
 import java.util.function.LongConsumer;
-import java.util.function.Supplier;
 
 import static org.elasticsearch.search.aggregations.InternalOrder.isKeyOrder;
 
@@ -296,7 +295,7 @@ private InternalAggregation[] buildAggregations(LongArray owningBucketOrds) thro
                     try (ObjectArrayPriorityQueue<B> ordered = buildPriorityQueue(size)) {
                         B spare = null;
                         BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningOrd);
-                        Supplier<B> emptyBucketBuilder = emptyBucketBuilder(owningOrd);
+                        BucketUpdater<B> bucketUpdater = bucketUpdater(owningOrd);
                         while (ordsEnum.next()) {
                             long docCount = bucketDocCount(ordsEnum.ord());
                             otherDocCounts.increment(ordIdx, docCount);
@@ -305,9 +304,9 @@ private InternalAggregation[] buildAggregations(LongArray owningBucketOrds) thro
                             }
                             if (spare == null) {
                                 checkRealMemoryCBForInternalBucket();
-                                spare = emptyBucketBuilder.get();
+                                spare = buildEmptyBucket();
                             }
-                            updateBucket(spare, ordsEnum, docCount);
+                            bucketUpdater.updateBucket(spare, ordsEnum, docCount);
                             spare = ordered.insertWithOverflow(spare);
                         }
 
@@ -348,9 +347,9 @@ private InternalAggregation[] buildAggregations(LongArray owningBucketOrds) thro
         abstract void collectZeroDocEntriesIfNeeded(long owningBucketOrd, boolean excludeDeletedDocs) throws IOException;
 
         /**
-         * Build an empty temporary bucket.
+         * Build an empty bucket.
          */
-        abstract Supplier<B> emptyBucketBuilder(long owningBucketOrd);
+        abstract B buildEmptyBucket();
 
         /**
          * Build a {@link PriorityQueue} to sort the buckets. After we've
@@ -362,7 +361,7 @@ private InternalAggregation[] buildAggregations(LongArray owningBucketOrds) thro
          * Update fields in {@code spare} to reflect information collected for
          * this bucket ordinal.
          */
-        abstract void updateBucket(B spare, BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum, long docCount) throws IOException;
+        abstract BucketUpdater<B> bucketUpdater(long owningBucketOrd);
 
         /**
          * Build an array to hold the "top" buckets for each ordinal.
@@ -399,6 +398,10 @@ private InternalAggregation[] buildAggregations(LongArray owningBucketOrds) thro
         abstract R buildEmptyResult();
     }
 
+    interface BucketUpdater<B extends InternalMultiBucketAggregation.InternalBucket> {
+        void updateBucket(B spare, BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum, long docCount) throws IOException;
+    }
+
     /**
      * Builds results for the standard {@code terms} aggregation.
      */
@@ -490,8 +493,8 @@ private void collectZeroDocEntries(BinaryDocValues values, Bits liveDocs, int ma
         }
 
         @Override
-        Supplier<StringTerms.Bucket> emptyBucketBuilder(long owningBucketOrd) {
-            return () -> new StringTerms.Bucket(new BytesRef(), 0, null, showTermDocCountError, 0, format);
+        StringTerms.Bucket buildEmptyBucket() {
+            return new StringTerms.Bucket(new BytesRef(), 0, null, showTermDocCountError, 0, format);
         }
 
         @Override
@@ -500,10 +503,12 @@ ObjectArrayPriorityQueue<StringTerms.Bucket> buildPriorityQueue(int size) {
         }
 
         @Override
-        void updateBucket(StringTerms.Bucket spare, BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum, long docCount) throws IOException {
-            ordsEnum.readValue(spare.termBytes);
-            spare.docCount = docCount;
-            spare.bucketOrd = ordsEnum.ord();
+        BucketUpdater<StringTerms.Bucket> bucketUpdater(long owningBucketOrd) {
+            return (spare, ordsEnum, docCount) -> {
+                ordsEnum.readValue(spare.termBytes);
+                spare.docCount = docCount;
+                spare.bucketOrd = ordsEnum.ord();
+            };
         }
 
         @Override
@@ -615,9 +620,8 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
         void collectZeroDocEntriesIfNeeded(long owningBucketOrd, boolean excludeDeletedDocs) throws IOException {}
 
         @Override
-        Supplier<SignificantStringTerms.Bucket> emptyBucketBuilder(long owningBucketOrd) {
-            long subsetSize = subsetSizes.get(owningBucketOrd);
-            return () -> new SignificantStringTerms.Bucket(new BytesRef(), 0, subsetSize, 0, 0, null, format, 0);
+        SignificantStringTerms.Bucket buildEmptyBucket() {
+            return new SignificantStringTerms.Bucket(new BytesRef(), 0, 0, null, format, 0);
         }
 
         @Override
@@ -626,20 +630,20 @@ ObjectArrayPriorityQueue<SignificantStringTerms.Bucket> buildPriorityQueue(int s
         }
 
         @Override
-        void updateBucket(SignificantStringTerms.Bucket spare, BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum, long docCount)
-            throws IOException {
-
-            ordsEnum.readValue(spare.termBytes);
-            spare.bucketOrd = ordsEnum.ord();
-            spare.subsetDf = docCount;
-            spare.supersetDf = backgroundFrequencies.freq(spare.termBytes);
-            spare.supersetSize = supersetSize;
-            /*
-             * During shard-local down-selection we use subset/superset stats
-             * that are for this shard only. Back at the central reducer these
-             * properties will be updated with global stats.
-             */
-            spare.updateScore(significanceHeuristic);
+        BucketUpdater<SignificantStringTerms.Bucket> bucketUpdater(long owningBucketOrd) {
+            long subsetSize = subsetSizes.get(owningBucketOrd);
+            return (spare, ordsEnum, docCount) -> {
+                ordsEnum.readValue(spare.termBytes);
+                spare.bucketOrd = ordsEnum.ord();
+                spare.subsetDf = docCount;
+                spare.supersetDf = backgroundFrequencies.freq(spare.termBytes);
+                /*
+                 * During shard-local down-selection we use subset/superset stats
+                 * that are for this shard only. Back at the central reducer these
+                 * properties will be updated with global stats.
+                 */
+                spare.updateScore(significanceHeuristic, subsetSize, supersetSize);
+            };
         }
 
         @Override