InternalMultiBucketAggregation.InternalBucket does not implement writable anymore (#117310)

iverase · web-flow · commit 78400b8d05c0 · 2024-11-26T07:53:40.000+01:00
This allows to make some Bucket implementations leaner, in particular terms and multi-terms aggregations
diff --git a/modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/adjacency/InternalAdjacencyMatrix.java b/modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/adjacency/InternalAdjacencyMatrix.java
@@ -33,7 +33,7 @@
 public class InternalAdjacencyMatrix extends InternalMultiBucketAggregation<InternalAdjacencyMatrix, InternalAdjacencyMatrix.InternalBucket>
     implements
         AdjacencyMatrix {
-    public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucket implements AdjacencyMatrix.Bucket {
+    public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucketWritable implements AdjacencyMatrix.Bucket {
 
         private final String key;
         private final long docCount;
diff --git a/modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/timeseries/InternalTimeSeries.java b/modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/timeseries/InternalTimeSeries.java
@@ -34,7 +34,7 @@ public class InternalTimeSeries extends InternalMultiBucketAggregation<InternalT
     /**
      * A bucket associated with a specific time series (identified by its key)
      */
-    public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucket {
+    public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucketWritable {
         protected long bucketOrd;
         protected final BytesRef key;
         // TODO: make computing docCount optional
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/InternalMultiBucketAggregation.java b/server/src/main/java/org/elasticsearch/search/aggregations/InternalMultiBucketAggregation.java
@@ -220,7 +220,7 @@ private List<B> reducePipelineBuckets(AggregationReduceContext reduceContext, Pi
         return reducedBuckets;
     }
 
-    public abstract static class InternalBucket implements Bucket, Writeable {
+    public abstract static class InternalBucket implements Bucket {
 
         public Object getProperty(String containingAggName, List<String> path) {
             if (path.isEmpty()) {
@@ -248,4 +248,8 @@ public Object getProperty(String containingAggName, List<String> path) {
             return aggregation.getProperty(path.subList(1, path.size()));
         }
     }
+
+    /** A {@link InternalBucket} that implements the {@link Writeable} interface. Most implementation might want
+     * to use this one except when specific logic is need to write into the stream. */
+    public abstract static class InternalBucketWritable extends InternalBucket implements Writeable {}
 }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/InternalComposite.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/InternalComposite.java
@@ -320,7 +320,9 @@ public int hashCode() {
         return Objects.hash(super.hashCode(), size, buckets, afterKey, Arrays.hashCode(reverseMuls), Arrays.hashCode(missingOrders));
     }
 
-    public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucket implements CompositeAggregation.Bucket {
+    public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucketWritable
+        implements
+            CompositeAggregation.Bucket {
 
         private final CompositeKey key;
         private final long docCount;
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/filter/InternalFilters.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/filter/InternalFilters.java
@@ -30,7 +30,7 @@
 import java.util.Objects;
 
 public class InternalFilters extends InternalMultiBucketAggregation<InternalFilters, InternalFilters.InternalBucket> implements Filters {
-    public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucket implements Filters.Bucket {
+    public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucketWritable implements Filters.Bucket {
 
         private final String key;
         private long docCount;
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/InternalGeoGridBucket.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/InternalGeoGridBucket.java
@@ -19,7 +19,7 @@
 import java.io.IOException;
 import java.util.Objects;
 
-public abstract class InternalGeoGridBucket extends InternalMultiBucketAggregation.InternalBucket
+public abstract class InternalGeoGridBucket extends InternalMultiBucketAggregation.InternalBucketWritable
     implements
         GeoGrid.Bucket,
         Comparable<InternalGeoGridBucket> {
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/AbstractHistogramBucket.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/AbstractHistogramBucket.java
@@ -16,7 +16,7 @@
 /**
  * A bucket in the histogram where documents fall in
  */
-public abstract class AbstractHistogramBucket extends InternalMultiBucketAggregation.InternalBucket {
+public abstract class AbstractHistogramBucket extends InternalMultiBucketAggregation.InternalBucketWritable {
 
     protected final long docCount;
     protected final InternalAggregations aggregations;
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/prefix/InternalIpPrefix.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/prefix/InternalIpPrefix.java
@@ -33,7 +33,7 @@
 
 public class InternalIpPrefix extends InternalMultiBucketAggregation<InternalIpPrefix, InternalIpPrefix.Bucket> {
 
-    public static class Bucket extends InternalMultiBucketAggregation.InternalBucket
+    public static class Bucket extends InternalMultiBucketAggregation.InternalBucketWritable
         implements
             IpPrefix.Bucket,
             KeyComparable<InternalIpPrefix.Bucket> {
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/InternalBinaryRange.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/InternalBinaryRange.java
@@ -36,7 +36,7 @@ public final class InternalBinaryRange extends InternalMultiBucketAggregation<In
     implements
         Range {
 
-    public static class Bucket extends InternalMultiBucketAggregation.InternalBucket implements Range.Bucket {
+    public static class Bucket extends InternalMultiBucketAggregation.InternalBucketWritable implements Range.Bucket {
 
         private final transient DocValueFormat format;
         private final String key;
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/InternalRange.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/InternalRange.java
@@ -36,7 +36,7 @@ public class InternalRange<B extends InternalRange.Bucket, R extends InternalRan
     @SuppressWarnings("rawtypes")
     static final Factory FACTORY = new Factory();
 
-    public static class Bucket extends InternalMultiBucketAggregation.InternalBucket implements Range.Bucket {
+    public static class Bucket extends InternalMultiBucketAggregation.InternalBucketWritable implements Range.Bucket {
 
         protected final transient DocValueFormat format;
         protected final double from;
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/AbstractInternalTerms.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/AbstractInternalTerms.java
@@ -63,8 +63,6 @@ public abstract static class AbstractTermsBucket<B extends AbstractTermsBucket<B
 
         protected abstract void setDocCountError(long docCountError);
 
-        protected abstract boolean getShowDocCountError();
-
         protected abstract long getDocCountError();
 
         protected abstract void bucketToXContent(XContentBuilder builder, Params params, boolean showDocCountError) throws IOException;
@@ -91,6 +89,8 @@ public abstract static class AbstractTermsBucket<B extends AbstractTermsBucket<B
 
     protected abstract int getRequiredSize();
 
+    protected abstract boolean getShowDocCountError();
+
     protected abstract B createBucket(long docCount, InternalAggregations aggs, long docCountError, B prototype);
 
     private B reduceBucket(List<B> buckets, AggregationReduceContext context) {
@@ -104,7 +104,7 @@ private B reduceBucket(List<B> buckets, AggregationReduceContext context) {
         for (B bucket : buckets) {
             docCount += bucket.getDocCount();
             if (docCountError != -1) {
-                if (bucket.getShowDocCountError() == false || bucket.getDocCountError() == -1) {
+                if (getShowDocCountError() == false || bucket.getDocCountError() == -1) {
                     docCountError = -1;
                 } else {
                     docCountError += bucket.getDocCountError();
@@ -257,23 +257,25 @@ public void accept(InternalAggregation aggregation) {
             }
             otherDocCount[0] += terms.getSumOfOtherDocCounts();
             final long thisAggDocCountError = getDocCountError(terms);
+            setDocCountError(thisAggDocCountError);
             if (sumDocCountError != -1) {
                 if (thisAggDocCountError == -1) {
                     sumDocCountError = -1;
                 } else {
                     sumDocCountError += thisAggDocCountError;
                 }
             }
-            setDocCountError(thisAggDocCountError);
-            for (B bucket : terms.getBuckets()) {
-                // If there is already a doc count error for this bucket
-                // subtract this aggs doc count error from it to make the
-                // new value for the bucket. This then means that when the
-                // final error for the bucket is calculated below we account
-                // for the existing error calculated in a previous reduce.
-                // Note that if the error is unbounded (-1) this will be fixed
-                // later in this method.
-                bucket.updateDocCountError(-thisAggDocCountError);
+            if (getShowDocCountError()) {
+                for (B bucket : terms.getBuckets()) {
+                    // If there is already a doc count error for this bucket
+                    // subtract this aggs doc count error from it to make the
+                    // new value for the bucket. This then means that when the
+                    // final error for the bucket is calculated below we account
+                    // for the existing error calculated in a previous reduce.
+                    // Note that if the error is unbounded (-1) this will be fixed
+                    // later in this method.
+                    bucket.updateDocCountError(-thisAggDocCountError);
+                }
             }
             if (terms.getBuckets().isEmpty() == false) {
                 bucketsList.add(terms.getBuckets());
@@ -319,17 +321,17 @@ public InternalAggregation get() {
                     result.add(bucket.reduced(AbstractInternalTerms.this::reduceBucket, reduceContext));
                 });
             }
-            for (B r : result) {
-                if (sumDocCountError == -1) {
-                    r.setDocCountError(-1);
-                } else {
-                    r.updateDocCountError(sumDocCountError);
+            if (getShowDocCountError()) {
+                for (B r : result) {
+                    if (sumDocCountError == -1) {
+                        r.setDocCountError(-1);
+                    } else {
+                        r.updateDocCountError(sumDocCountError);
+                    }
                 }
             }
-            long docCountError;
-            if (sumDocCountError == -1) {
-                docCountError = -1;
-            } else {
+            long docCountError = -1;
+            if (sumDocCountError != -1) {
                 docCountError = size == 1 ? 0 : sumDocCountError;
             }
             return create(name, result, reduceContext.isFinalReduce() ? getOrder() : thisReduceOrder, docCountError, otherDocCount[0]);
@@ -349,7 +351,7 @@ public InternalAggregation finalizeSampling(SamplingContext samplingContext) {
                     b -> createBucket(
                         samplingContext.scaleUp(b.getDocCount()),
                         InternalAggregations.finalizeSampling(b.getAggregations(), samplingContext),
-                        b.getShowDocCountError() ? samplingContext.scaleUp(b.getDocCountError()) : 0,
+                        getShowDocCountError() ? samplingContext.scaleUp(b.getDocCountError()) : 0,
                         b
                     )
                 )
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/DoubleTerms.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/DoubleTerms.java
@@ -164,8 +164,8 @@ public Bucket createBucket(InternalAggregations aggregations, Bucket prototype)
             prototype.term,
             prototype.docCount,
             aggregations,
-            prototype.showDocCountError,
-            prototype.docCountError,
+            showTermDocCountError,
+            prototype.getDocCountError(),
             prototype.format
         );
     }
@@ -216,6 +216,6 @@ public void close() {
 
     @Override
     protected Bucket createBucket(long docCount, InternalAggregations aggs, long docCountError, DoubleTerms.Bucket prototype) {
-        return new Bucket(prototype.term, docCount, aggs, prototype.showDocCountError, docCountError, format);
+        return new Bucket(prototype.term, docCount, aggs, showTermDocCountError, docCountError, format);
     }
 }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java
@@ -880,7 +880,6 @@ StringTerms.Bucket convertTempBucketToRealBucket(OrdBucket temp, GlobalOrdLookup
             BytesRef term = BytesRef.deepCopyOf(lookupGlobalOrd.apply(temp.globalOrd));
             StringTerms.Bucket result = new StringTerms.Bucket(term, temp.docCount, null, showTermDocCountError, 0, format);
             result.bucketOrd = temp.bucketOrd;
-            result.docCountError = 0;
             return result;
         }
 
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalMappedTerms.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalMappedTerms.java
@@ -87,14 +87,22 @@ protected final void writeTermTypeInfoTo(StreamOutput out) throws IOException {
         writeSize(shardSize, out);
         out.writeBoolean(showTermDocCountError);
         out.writeVLong(otherDocCount);
-        out.writeCollection(buckets);
+        out.writeVInt(buckets.size());
+        for (var bucket : buckets) {
+            bucket.writeTo(out, showTermDocCountError);
+        }
     }
 
     @Override
     protected void setDocCountError(long docCountError) {
         this.docCountError = docCountError;
     }
 
+    @Override
+    protected boolean getShowDocCountError() {
+        return showTermDocCountError;
+    }
+
     @Override
     protected int getShardSize() {
         return shardSize;
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalRareTerms.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalRareTerms.java
@@ -10,6 +10,7 @@
 
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.util.SetBackedScalingCuckooFilter;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.BucketOrder;
@@ -29,10 +30,11 @@ public abstract class InternalRareTerms<A extends InternalRareTerms<A, B>, B ext
     implements
         RareTerms {
 
-    public abstract static class Bucket<B extends Bucket<B>> extends InternalMultiBucketAggregation.InternalBucket
+    public abstract static class Bucket<B extends Bucket<B>> extends InternalMultiBucketAggregation.InternalBucketWritable
         implements
             RareTerms.Bucket,
-            KeyComparable<B> {
+            KeyComparable<B>,
+            Writeable {
         /**
          * Reads a bucket. Should be a constructor reference.
          */
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalSignificantTerms.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalSignificantTerms.java
@@ -45,7 +45,7 @@ public abstract class InternalSignificantTerms<A extends InternalSignificantTerm
     public static final String BG_COUNT = "bg_count";
 
     @SuppressWarnings("PMD.ConstructorCallsOverridableMethod")
-    public abstract static class Bucket<B extends Bucket<B>> extends InternalMultiBucketAggregation.InternalBucket
+    public abstract static class Bucket<B extends Bucket<B>> extends InternalMultiBucketAggregation.InternalBucketWritable
         implements
             SignificantTerms.Bucket {
         /**
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalTerms.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalTerms.java
@@ -41,9 +41,8 @@ public interface Reader<B extends Bucket<B>> {
         long bucketOrd;
 
         protected long docCount;
-        protected long docCountError;
+        private long docCountError;
         protected InternalAggregations aggregations;
-        protected final boolean showDocCountError;
         protected final DocValueFormat format;
 
         protected Bucket(
@@ -53,29 +52,23 @@ protected Bucket(
             long docCountError,
             DocValueFormat formatter
         ) {
-            this.showDocCountError = showDocCountError;
             this.format = formatter;
             this.docCount = docCount;
             this.aggregations = aggregations;
-            this.docCountError = docCountError;
+            this.docCountError = showDocCountError ? docCountError : -1;
         }
 
         /**
          * Read from a stream.
          */
         protected Bucket(StreamInput in, DocValueFormat formatter, boolean showDocCountError) throws IOException {
-            this.showDocCountError = showDocCountError;
             this.format = formatter;
             docCount = in.readVLong();
-            docCountError = -1;
-            if (showDocCountError) {
-                docCountError = in.readLong();
-            }
+            docCountError = showDocCountError ? in.readLong() : -1;
             aggregations = InternalAggregations.readFrom(in);
         }
 
-        @Override
-        public final void writeTo(StreamOutput out) throws IOException {
+        final void writeTo(StreamOutput out, boolean showDocCountError) throws IOException {
             out.writeVLong(getDocCount());
             if (showDocCountError) {
                 out.writeLong(docCountError);
@@ -105,9 +98,6 @@ public void setBucketOrd(long bucketOrd) {
 
         @Override
         public long getDocCountError() {
-            if (showDocCountError == false) {
-                throw new IllegalStateException("show_terms_doc_count_error is false");
-            }
             return docCountError;
         }
 
@@ -121,11 +111,6 @@ protected void updateDocCountError(long docCountErrorDiff) {
             this.docCountError += docCountErrorDiff;
         }
 
-        @Override
-        protected boolean getShowDocCountError() {
-            return showDocCountError;
-        }
-
         @Override
         public InternalAggregations getAggregations() {
             return aggregations;
@@ -155,23 +140,15 @@ public boolean equals(Object obj) {
                 return false;
             }
             Bucket<?> that = (Bucket<?>) obj;
-            if (showDocCountError && docCountError != that.docCountError) {
-                /*
-                 * docCountError doesn't matter if not showing it and
-                 * serialization sets it to -1 no matter what it was
-                 * before.
-                 */
-                return false;
-            }
-            return Objects.equals(docCount, that.docCount)
-                && Objects.equals(showDocCountError, that.showDocCountError)
+            return Objects.equals(docCountError, that.docCountError)
+                && Objects.equals(docCount, that.docCount)
                 && Objects.equals(format, that.format)
                 && Objects.equals(aggregations, that.aggregations);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(getClass(), docCount, format, showDocCountError, showDocCountError ? docCountError : -1, aggregations);
+            return Objects.hash(getClass(), docCount, format, docCountError, aggregations);
         }
     }
 
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/LongTerms.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/LongTerms.java
@@ -178,8 +178,8 @@ public Bucket createBucket(InternalAggregations aggregations, Bucket prototype)
             prototype.term,
             prototype.docCount,
             aggregations,
-            prototype.showDocCountError,
-            prototype.docCountError,
+            showTermDocCountError,
+            prototype.getDocCountError(),
             prototype.format
         );
     }
@@ -260,7 +260,7 @@ public InternalAggregation get() {
 
     @Override
     protected Bucket createBucket(long docCount, InternalAggregations aggs, long docCountError, LongTerms.Bucket prototype) {
-        return new Bucket(prototype.term, docCount, aggs, prototype.showDocCountError, docCountError, format);
+        return new Bucket(prototype.term, docCount, aggs, showTermDocCountError, docCountError, format);
     }
 
     /**
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/StringTerms.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/StringTerms.java
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/UnmappedTerms.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/UnmappedTerms.java
diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketHelpersTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketHelpersTests.java
diff --git a/x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/multiterms/InternalMultiTerms.java b/x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/multiterms/InternalMultiTerms.java
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/InternalCategorizationAggregation.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/InternalCategorizationAggregation.java
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/changepoint/ChangePointBucket.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/changepoint/ChangePointBucket.java

Original file line number	Diff line number	Diff line change
`@@ -220,7 +220,7 @@ private List<B> reducePipelineBuckets(AggregationReduceContext reduceContext, Pi`
`220`	`220`	`return reducedBuckets;`
`221`	`221`	`}`
`222`	`222`
`223`		`- public abstract static class InternalBucket implements Bucket, Writeable {`
	`223`	`+ public abstract static class InternalBucket implements Bucket {`
`224`	`224`
`225`	`225`	`public Object getProperty(String containingAggName, List<String> path) {`
`226`	`226`	`if (path.isEmpty()) {`
`@@ -248,4 +248,8 @@ public Object getProperty(String containingAggName, List<String> path) {`
`248`	`248`	`return aggregation.getProperty(path.subList(1, path.size()));`
`249`	`249`	`}`
`250`	`250`	`}`
	`251`	`+`
	`252`	`+ /** A {@link InternalBucket} that implements the {@link Writeable} interface. Most implementation might want`
	`253`	`+ * to use this one except when specific logic is need to write into the stream. */`
	`254`	`+ public abstract static class InternalBucketWritable extends InternalBucket implements Writeable {}`
`251`	`255`	`}`
Original file line number	Diff line number	Diff line change
`@@ -320,7 +320,9 @@ public int hashCode() {`
`320`	`320`	`return Objects.hash(super.hashCode(), size, buckets, afterKey, Arrays.hashCode(reverseMuls), Arrays.hashCode(missingOrders));`
`321`	`321`	`}`
`322`	`322`
`323`		`- public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucket implements CompositeAggregation.Bucket {`
	`323`	`+ public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucketWritable`
	`324`	`+ implements`
	`325`	`+ CompositeAggregation.Bucket {`
`324`	`326`
`325`	`327`	`private final CompositeKey key;`
`326`	`328`	`private final long docCount;`