Adapt the plugin to Elastic 8

Thanks to @nicenatalia

Co-authored-by: Natalia Escalera <[email protected]>

Author: garaud

src/main/java/org/opendatasoft/elasticsearch/search/aggregations/bucket/DateHierarchyAggregationBuilder.java

@@ -1,6 +1,7 @@
 package org.opendatasoft.elasticsearch.search.aggregations.bucket;
 
-import org.elasticsearch.Version;
+import org.elasticsearch.TransportVersion;
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.Rounding;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -40,6 +41,11 @@
  * The builder of the aggregatorFactory. Also implements the parsing of the request.
  */
 public class DateHierarchyAggregationBuilder extends ValuesSourceAggregationBuilder<DateHierarchyAggregationBuilder> {
+    @Override
+    public TransportVersion getMinimalSupportedVersion() {
+        return TransportVersions.V_8_0_0;
+    }
+
     public static final String NAME = "date_hierarchy";
     public static final ValuesSourceRegistry.RegistryKey<DateHierarchyAggregationSupplier> REGISTRY_KEY =
         new ValuesSourceRegistry.RegistryKey<>(NAME, DateHierarchyAggregationSupplier.class);
@@ -140,8 +146,8 @@ public void writeTo(StreamOutput out) throws IOException {
         DateHierarchyAggregationBuilder::new
     );
     static {
-
-        ValuesSourceAggregationBuilder.declareFields(PARSER, true, true, true);
+        // ES 8.x introduces field validation. Setting timezoneAware to false to avoid duplication of the timezone field
+        ValuesSourceAggregationBuilder.declareFields(PARSER, true, true, false);
 
         PARSER.declareString(DateHierarchyAggregationBuilder::interval, INTERVAL_FIELD);
 
@@ -180,7 +186,7 @@ private DateHierarchyAggregationBuilder(String name) {
     }
 
     @Override
-    protected boolean serializeTargetValueType(Version version) {
+    protected boolean serializeTargetValueType(TransportVersion version) {
         return true;
     }
 
@@ -412,9 +418,4 @@ public boolean equals(Object obj) {
     public String getType() {
         return NAME;
     }
-
-    @Override
-    protected ValuesSourceRegistry.RegistryKey<?> getRegistryKey() {
-        return REGISTRY_KEY;
-    }
 }

src/main/java/org/opendatasoft/elasticsearch/search/aggregations/bucket/DateHierarchyAggregator.java

@@ -1,15 +1,16 @@
 package org.opendatasoft.elasticsearch.search.aggregations.bucket;
 
-import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.ElasticsearchException;
-import org.elasticsearch.common.Rounding;
 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.BytesRefHash;
+import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.core.Releasables;
+import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.BucketOrder;
@@ -25,7 +26,6 @@
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -54,7 +54,6 @@ public DateHierarchyAggregator(
         this.bucketCountThresholds = bucketCountThresholds;
         order.validate(this);
         this.order = order;
-        this.partiallyBuiltBucketComparator = order == null ? null : order.partiallyBuiltBucketComparator(b -> b.bucketOrd, this);
     }
 
     public static class BucketCountThresholds implements Writeable, ToXContentFragment {
@@ -144,7 +143,6 @@ public boolean equals(Object obj) {
     private final long minDocCount;
     private final BucketCountThresholds bucketCountThresholds;
     private final List<DateHierarchyAggregationBuilder.PreparedRounding> preparedRoundings;
-    protected final Comparator<InternalPathHierarchy.InternalBucket> partiallyBuiltBucketComparator;
 
     /**
      * The collector collects the docs, including or not some score (depending of the including of a Scorer) in the
@@ -153,11 +151,11 @@ public boolean equals(Object obj) {
      * The LeafBucketCollector is a "Per-leaf bucket collector". It collects docs for the account of buckets.
      */
     @Override
-    public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
+    public LeafBucketCollector getLeafCollector(AggregationExecutionContext ctx, LeafBucketCollector sub) throws IOException {
         if (valuesSource == null) {
             return LeafBucketCollector.NO_OP_COLLECTOR;
         }
-        final SortedNumericDocValues values = valuesSource.longValues(ctx);
+        final SortedNumericDocValues values = valuesSource.longValues(ctx.getLeafReaderContext());
 
         return new LeafBucketCollectorBase(sub, values) {
 
@@ -189,63 +187,76 @@ public void collect(int doc, long bucket) throws IOException {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrdinals) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrdinals) throws IOException {
 
-        InternalDateHierarchy.InternalBucket[][] topBucketsPerOrd = new InternalDateHierarchy.InternalBucket[owningBucketOrdinals.length][];
-        InternalDateHierarchy[] results = new InternalDateHierarchy[owningBucketOrdinals.length];
+        // InternalDateHierarchy.InternalBucket[][] topBucketsPerOrd = new
+        // InternalDateHierarchy.InternalBucket[owningBucketOrdinals.length][];
+        // InternalDateHierarchy[] results = new InternalDateHierarchy[owningBucketOrdinals.length];
 
-        for (int ordIdx = 0; ordIdx < owningBucketOrdinals.length; ordIdx++) {
-            assert owningBucketOrdinals[ordIdx] == 0;
+        try (
+            ObjectArray<InternalDateHierarchy.InternalBucket[]> topBucketsPerOrd = bigArrays().newObjectArray(owningBucketOrdinals.size())
+        ) {
 
-            // build buckets and store them sorted
-            final int size = (int) Math.min(bucketOrds.size(), bucketCountThresholds.getShardSize());
+            InternalDateHierarchy[] results = new InternalDateHierarchy[Math.toIntExact(owningBucketOrdinals.size())];
 
-            PathSortedTree<String, InternalDateHierarchy.InternalBucket> pathSortedTree = new PathSortedTree<>(order.comparator(), size);
+            for (int ordIdx = 0; ordIdx < owningBucketOrdinals.size(); ordIdx++) {
+                // assert owningBucketOrdinals[ordIdx] == 0;
+                assert owningBucketOrdinals.get(ordIdx) == 0;
 
-            InternalDateHierarchy.InternalBucket spare;
-            for (int i = 0; i < bucketOrds.size(); i++) {
-                spare = new InternalDateHierarchy.InternalBucket(0, null, null, null, 0, null);
+                // build buckets and store them sorted
+                final int size = (int) Math.min(bucketOrds.size(), bucketCountThresholds.getShardSize());
 
-                BytesRef term = new BytesRef();
-                bucketOrds.get(i, term);
-                String[] paths = term.utf8ToString().split("/", -1);
+                PathSortedTree<String, InternalDateHierarchy.InternalBucket> pathSortedTree = new PathSortedTree<>(
+                    order.comparator(),
+                    size
+                );
 
-                spare.paths = paths;
-                spare.key = term;
-                spare.level = paths.length - 1;
-                spare.name = paths[spare.level];
-                spare.docCount = bucketDocCount(i);
-                spare.bucketOrd = i;
+                InternalDateHierarchy.InternalBucket spare;
+                for (int i = 0; i < bucketOrds.size(); i++) {
+                    spare = new InternalDateHierarchy.InternalBucket(0, null, null, null, 0, null);
 
-                pathSortedTree.add(spare.paths, spare);
-            }
+                    BytesRef term = new BytesRef();
+                    bucketOrds.get(i, term);
+                    String[] paths = term.utf8ToString().split("/", -1);
 
-            // Get the top buckets
-            topBucketsPerOrd[ordIdx] = new InternalDateHierarchy.InternalBucket[size];
-            long otherHierarchyNodes = pathSortedTree.getFullSize();
-            Iterator<InternalDateHierarchy.InternalBucket> iterator = pathSortedTree.consumer();
-            for (int i = 0; i < size; i++) {
-                final InternalDateHierarchy.InternalBucket bucket = iterator.next();
-                topBucketsPerOrd[ordIdx][i] = bucket;
-                otherHierarchyNodes -= 1;
-            }
+                    spare.paths = paths;
+                    spare.key = term;
+                    spare.level = paths.length - 1;
+                    spare.name = paths[spare.level];
+                    spare.docCount = bucketDocCount(i);
+                    spare.bucketOrd = i;
 
-            results[ordIdx] = new InternalDateHierarchy(
-                name,
-                Arrays.asList(topBucketsPerOrd[ordIdx]),
-                order,
-                minDocCount,
-                bucketCountThresholds.getRequiredSize(),
-                bucketCountThresholds.getShardSize(),
-                otherHierarchyNodes,
-                metadata()
-            );
-        }
+                    pathSortedTree.add(spare.paths, spare);
+                }
 
-        // Build sub-aggregations for pruned buckets
-        buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, aggregations) -> b.aggregations = aggregations);
+                // Get the top buckets
+                topBucketsPerOrd.set(ordIdx, new InternalDateHierarchy.InternalBucket[size]);
+                // topBucketsPerOrd[ordIdx] = new InternalDateHierarchy.InternalBucket[size];
+                long otherHierarchyNodes = pathSortedTree.getFullSize();
+                Iterator<InternalDateHierarchy.InternalBucket> iterator = pathSortedTree.consumer();
+                for (int i = 0; i < size; i++) {
+                    final InternalDateHierarchy.InternalBucket bucket = iterator.next();
+                    topBucketsPerOrd.get(ordIdx)[i] = bucket;
+                    otherHierarchyNodes -= 1;
+                }
 
-        return results;
+                results[ordIdx] = new InternalDateHierarchy(
+                    name,
+                    Arrays.asList(topBucketsPerOrd.get(ordIdx)),
+                    order,
+                    minDocCount,
+                    bucketCountThresholds.getRequiredSize(),
+                    bucketCountThresholds.getShardSize(),
+                    otherHierarchyNodes,
+                    metadata()
+                );
+            }
+
+            // Build sub-aggregations for pruned buckets
+            buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, aggregations) -> b.aggregations = aggregations);
+
+            return results;
+        }
     }
 
     @Override

src/main/java/org/opendatasoft/elasticsearch/search/aggregations/bucket/InternalDateHierarchy.java

@@ -3,7 +3,8 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.search.aggregations.Aggregations;
+import org.elasticsearch.search.aggregations.AggregationReduceContext;
+import org.elasticsearch.search.aggregations.AggregatorReducer;
 import org.elasticsearch.search.aggregations.BucketOrder;
 import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.InternalAggregations;
@@ -15,11 +16,13 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.TreeMap;
 
 /**
  * An internal implementation of {@link InternalMultiBucketAggregation}
@@ -28,11 +31,67 @@
  */
 public class InternalDateHierarchy extends InternalMultiBucketAggregation<InternalDateHierarchy, InternalDateHierarchy.InternalBucket> {
 
+    @Override
+    protected AggregatorReducer getLeaderReducer(AggregationReduceContext reduceContext, int size) {
+        Map<BytesRef, List<InternalBucket>> buckets = new LinkedHashMap<>();
+
+        return new AggregatorReducer() {
+            private long otherHierarchyNodes = 0;
+
+            @Override
+            public void accept(InternalAggregation aggregation) {
+                InternalDateHierarchy dateHierarchy = (InternalDateHierarchy) aggregation;
+
+                otherHierarchyNodes += dateHierarchy.getSumOtherHierarchyNodes();
+
+                for (InternalBucket bucket : dateHierarchy.buckets) {
+                    List<InternalBucket> existingBuckets = buckets.get(bucket.key);
+                    if (existingBuckets == null) {
+                        existingBuckets = new ArrayList<>(size);
+                        buckets.put(bucket.key, existingBuckets);
+                    }
+                    existingBuckets.add(bucket);
+                }
+            }
+
+            @Override
+            public InternalAggregation get() {
+                final int size = !reduceContext.isFinalReduce() ? buckets.size() : Math.min(requiredSize, buckets.size());
+                PathSortedTree<String, InternalBucket> ordered = new PathSortedTree<>(order.comparator(), size);
+
+                for (List<InternalBucket> sameTermBuckets : buckets.values()) {
+                    final InternalBucket b = reduceBucket(sameTermBuckets, reduceContext);
+                    if (b.getDocCount() >= minDocCount || !reduceContext.isFinalReduce()) {
+                        reduceContext.consumeBucketsAndMaybeBreak(1);
+                        ordered.add(b.paths, b);
+                    } else {
+                        reduceContext.consumeBucketsAndMaybeBreak(-countInnerBucket(b));
+                    }
+                }
+
+                long sum_other_hierarchy_nodes = ordered.getFullSize() - size + otherHierarchyNodes;
+
+                return new InternalDateHierarchy(
+                    getName(),
+                    ordered.getAsList(),
+                    order,
+                    minDocCount,
+                    requiredSize,
+                    shardSize,
+                    sum_other_hierarchy_nodes,
+                    getMetadata()
+                );
+            }
+        };
+    }
+
     /**
      * The bucket class of InternalDateHierarchy.
      * @see MultiBucketsAggregation.Bucket
      */
-    public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucket implements KeyComparable<InternalBucket> {
+    public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucketWritable
+        implements
+            KeyComparable<InternalBucket> {
 
         BytesRef key;
         String name;
@@ -104,11 +163,10 @@ public long getDocCount() {
         }
 
         @Override
-        public Aggregations getAggregations() {
+        public InternalAggregations getAggregations() {
             return aggregations;
         }
 
-        @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
             builder.startObject();
             builder.field(CommonFields.DOC_COUNT.getPreferredName(), docCount);
@@ -214,62 +272,7 @@ public List<InternalBucket> getBuckets() {
         return buckets;
     }
 
-    /**
-     * Reduces the given aggregations to a single one and returns it.
-     */
-    @Override
-    public InternalDateHierarchy reduce(List<InternalAggregation> aggregations, ReduceContext reduceContext) {
-        Map<BytesRef, List<InternalBucket>> buckets = null;
-        long otherHierarchyNodes = 0;
-
-        // extract buckets from aggregations
-        for (InternalAggregation aggregation : aggregations) {
-            InternalDateHierarchy dateHierarchy = (InternalDateHierarchy) aggregation;
-            if (buckets == null) {
-                buckets = new LinkedHashMap<>();
-            }
-
-            otherHierarchyNodes += dateHierarchy.getSumOtherHierarchyNodes();
-
-            for (InternalBucket bucket : dateHierarchy.buckets) {
-                List<InternalBucket> existingBuckets = buckets.get(bucket.key);
-                if (existingBuckets == null) {
-                    existingBuckets = new ArrayList<>(aggregations.size());
-                    buckets.put(bucket.key, existingBuckets);
-                }
-                existingBuckets.add(bucket);
-            }
-        }
-
-        // reduce and sort buckets depending of ordering rules
-        final int size = !reduceContext.isFinalReduce() ? buckets.size() : Math.min(requiredSize, buckets.size());
-        PathSortedTree<String, InternalBucket> ordered = new PathSortedTree<>(order.comparator(), size);
-        for (List<InternalBucket> sameTermBuckets : buckets.values()) {
-
-            final InternalBucket b = reduceBucket(sameTermBuckets, reduceContext);
-            if (b.getDocCount() >= minDocCount || !reduceContext.isFinalReduce()) {
-                reduceContext.consumeBucketsAndMaybeBreak(1);
-                ordered.add(b.paths, b);
-            } else {
-                reduceContext.consumeBucketsAndMaybeBreak(-countInnerBucket(b));
-            }
-        }
-
-        long sum_other_hierarchy_nodes = ordered.getFullSize() - size + otherHierarchyNodes;
-        return new InternalDateHierarchy(
-            getName(),
-            ordered.getAsList(),
-            order,
-            minDocCount,
-            requiredSize,
-            shardSize,
-            sum_other_hierarchy_nodes,
-            getMetadata()
-        );
-    }
-
-    @Override
-    protected InternalBucket reduceBucket(List<InternalBucket> buckets, ReduceContext context) {
+    protected InternalBucket reduceBucket(List<InternalBucket> buckets, AggregationReduceContext context) {
         List<InternalAggregations> aggregationsList = new ArrayList<>(buckets.size());
         InternalBucket reduced = null;
         for (InternalBucket bucket : buckets) {