SED-4848 Make time-series aggregations configurable on both axes

step-framework-timeseries/src/main/java/step/core/timeseries/aggregation/TimeSeriesAggregationPipeline.java

@@ -9,7 +9,12 @@
 import step.core.timeseries.bucket.BucketAttributes;
 import step.core.timeseries.bucket.BucketBuilder;
 
-import java.util.*;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeMap;
 import java.util.concurrent.atomic.LongAdder;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -82,35 +87,15 @@ public TimeSeriesAggregationResponse collect(TimeSeriesAggregationQuery query) {
         long sourceResolution = idealAvailableCollection.getResolutionMs();
         TimeSeriesProcessedParams finalParams = processQueryParams(query, sourceResolution);
 
-        Map<BucketAttributes, Map<Long, BucketBuilder>> seriesBuilder = new HashMap<>();
-
-        LongAdder bucketCount = new LongAdder();
-        long t1 = System.currentTimeMillis();
-        try (Stream<Bucket> stream = idealAvailableCollection.queryTimeSeries(finalParams)) {
-            stream.forEach(bucket -> {
-                bucketCount.increment();
-                BucketAttributes bucketAttributes = bucket.getAttributes() != null ? bucket.getAttributes() : new BucketAttributes();
-
-                BucketAttributes seriesKey;
-                if (CollectionUtils.isNotEmpty(finalParams.getGroupDimensions())) {
-                    seriesKey = bucketAttributes.subset(finalParams.getGroupDimensions());
-                } else {
-                    seriesKey = new BucketAttributes();
-                }
-                Map<Long, BucketBuilder> series = seriesBuilder.computeIfAbsent(seriesKey, a -> new TreeMap<>());
-
-                long index = calculateBucketBeginAnchor(bucket.getBegin(), finalParams);
-                series.computeIfAbsent(index, i -> new BucketBuilder(i, i + getBucketSize(finalParams.getFrom(), finalParams.getTo(), finalParams.isShrink(), finalParams.getResolution()))
-                    .withAccumulateAttributes(query.getCollectAttributeKeys(), query.getCollectAttributesValuesLimit())).accumulate(bucket);
-            });
-        }
-        long t2 = System.currentTimeMillis();
-        if (logger.isDebugEnabled()) {
-            logger.debug("Performed query in " + (t2 - t1) + "ms. Number of buckets processed: " + bucketCount.longValue());
+        Map<BucketAttributes, Map<Long, BucketBuilder>> resultBuilder;
+        if (query.getTimeAggregation().isMerge() && query.getGroupAggregation().isMerge()) {
+            resultBuilder = collectByMerging(query, finalParams, idealAvailableCollection);
+        } else {
+            resultBuilder = collectByAggregating(query, finalParams, idealAvailableCollection);
         }
 
-        Map<BucketAttributes, Map<Long, Bucket>> result = seriesBuilder.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e ->
-            e.getValue().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, i -> i.getValue().build()))));
+        Map<BucketAttributes, Map<Long, Bucket>> result = resultBuilder.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e ->
+            e.getValue().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, i -> i.getValue().buildAggregate()))));
 
         return new TimeSeriesAggregationResponseBuilder()
             .setSeries(result)
@@ -124,6 +109,112 @@ public TimeSeriesAggregationResponse collect(TimeSeriesAggregationQuery query) {
             .build();
     }
 
+    /**
+     * Collects the source buckets when both axes merge, i.e. when the whole aggregation amounts to one single merge
+     * of all the source buckets of a group falling into the same time bucket. A merge being associative and
+     * commutative, the source buckets can be merged directly into their resulting bucket, without materializing the
+     * series they belong to. The memory footprint is therefore driven by the number of groups and time buckets of the
+     * response, and not by the cardinality of the attributes of the source buckets.
+     */
+    private Map<BucketAttributes, Map<Long, BucketBuilder>> collectByMerging(TimeSeriesAggregationQuery query, TimeSeriesProcessedParams finalParams, TimeSeriesCollection collection) {
+        Map<BucketAttributes, Map<Long, BucketBuilder>> resultBuilder = new HashMap<>();
+        LongAdder bucketCount = new LongAdder();
+        long t1 = System.currentTimeMillis();
+        try (Stream<Bucket> stream = collection.queryTimeSeries(finalParams)) {
+            stream.forEach(bucket -> {
+                bucketCount.increment();
+                BucketAttributes groupAttributes = getGroupAttributes(bucket, finalParams.getGroupDimensions());
+                long timeSliceIndex = calculateBucketBeginAnchor(bucket.getBegin(), finalParams);
+
+                Map<Long, BucketBuilder> resultSeriesBuilder = resultBuilder.computeIfAbsent(groupAttributes, a -> new TreeMap<>());
+                resultSeriesBuilder.computeIfAbsent(timeSliceIndex, i -> newGroupBucketBuilder(query, finalParams, groupAttributes, i))
+                    .merge(bucket);
+            });
+        }
+        logAggregationDuration("merge aggregation", t1, bucketCount);
+        return resultBuilder;
+    }
+
+    /**
+     * Collects the source buckets when at least one of the axes reduces its inputs to a scalar. The time-window
+     * aggregation has then to be applied per series, before the group-by aggregation, which requires the series of
+     * each group to be materialized.
+     */
+    private Map<BucketAttributes, Map<Long, BucketBuilder>> collectByAggregating(TimeSeriesAggregationQuery query, TimeSeriesProcessedParams finalParams, TimeSeriesCollection collection) {
+        // Perform time-window aggregation and partition the time series:
+        // Aggregate each source series into the requested aligned time buckets and assign the resulting series to their respective groups (defined by the group dimensions).
+        // Do not perform any cross-series aggregation at this stage.
+        Map<Long, Map<BucketAttributes, Map<BucketAttributes, BucketBuilder>>> timeBucketedSeriesByGroup = new HashMap<>();
+        LongAdder bucketCount = new LongAdder();
+        long t1 = System.currentTimeMillis();
+        try (Stream<Bucket> stream = collection.queryTimeSeries(finalParams)) {
+            stream.forEach(bucket -> {
+                bucketCount.increment();
+                // The attributes of the source series
+                BucketAttributes bucketAttributes = bucket.getAttributes() != null ? bucket.getAttributes() : new BucketAttributes();
+                // The subset of attributes corresponding to the requested group dimensions (group by)
+                BucketAttributes groupAttributes = getGroupAttributes(bucket, finalParams.getGroupDimensions());
+                // The time slice index on the result series
+                long timeSliceIndex = calculateBucketBeginAnchor(bucket.getBegin(), finalParams);
+                // Get or create the time slice corresponding to the time index of the current bucket on the result series (aligned)
+                Map<BucketAttributes, Map<BucketAttributes, BucketBuilder>> timeSlice = timeBucketedSeriesByGroup.computeIfAbsent(timeSliceIndex, a -> new HashMap<>());
+                // Get or create the group of builders corresponding to the current group (defined by the group dimensions)
+                Map<BucketAttributes, BucketBuilder> indexSeriesBuckets = timeSlice.computeIfAbsent(groupAttributes, a -> new HashMap<>());
+                // Get the builder for the attributes of the current bucket. The full attributes of the series are kept
+                // at this stage, so that the attribute collection can be performed on them during the group-by aggregation
+                BucketBuilder bucketBuilder = indexSeriesBuckets.computeIfAbsent(bucketAttributes, a -> new BucketBuilder(query.getTimeAggregation(), timeSliceIndex, getBucketEnd(timeSliceIndex, finalParams)).withAttributes(bucketAttributes));
+                // Merge the current source bucket into the builder. The configured time-window aggregation is
+                // applied when the builder is reduced, at the group-by stage
+                bucketBuilder.merge(bucket);
+            });
+        }
+        logAggregationDuration("time-window aggregation", t1, bucketCount);
+
+        // Aggregate the grouped series:
+        // For each time bucket, apply the configured group-by aggregation across the aligned series in each group.
+        Map<BucketAttributes, Map<Long, BucketBuilder>> resultBuilder = new HashMap<>();
+        // For each time slice
+        timeBucketedSeriesByGroup.forEach((timeSliceIndex, timeSlice) -> {
+            // For each group
+            timeSlice.forEach((groupAttributes, group) -> {
+                // For each series of the group
+                Map<Long, BucketBuilder> resultSeriesBuilder = resultBuilder.computeIfAbsent(groupAttributes, a -> new TreeMap<>());
+                group.forEach((seriesAttributes, series) -> {
+                    BucketBuilder bucketBuilder = resultSeriesBuilder.computeIfAbsent(timeSliceIndex, i -> newGroupBucketBuilder(query, finalParams, groupAttributes, i));
+                    // Aggregate the series into the group. How the series contributes is defined by the
+                    // time-window aggregation it was built with
+                    bucketBuilder.aggregate(series);
+                });
+            });
+        });
+        return resultBuilder;
+    }
+
+    private BucketAttributes getGroupAttributes(Bucket bucket, Set<String> groupDimensions) {
+        BucketAttributes bucketAttributes = bucket.getAttributes();
+        if (bucketAttributes == null || CollectionUtils.isEmpty(groupDimensions)) {
+            return new BucketAttributes();
+        }
+        return bucketAttributes.subset(groupDimensions);
+    }
+
+    private BucketBuilder newGroupBucketBuilder(TimeSeriesAggregationQuery query, TimeSeriesProcessedParams finalParams, BucketAttributes groupAttributes, long timeSliceIndex) {
+        return new BucketBuilder(query.getGroupAggregation(), timeSliceIndex, getBucketEnd(timeSliceIndex, finalParams))
+            // The group attributes are copied, so that collecting the attributes doesn't mutate the key of the response
+            .withAttributes(new BucketAttributes(groupAttributes))
+            .withAccumulateAttributes(query.getCollectAttributeKeys(), query.getCollectAttributesValuesLimit());
+    }
+
+    private void logAggregationDuration(String aggregationName, long startTime, LongAdder bucketCount) {
+        if (logger.isDebugEnabled()) {
+            logger.debug("Performed " + aggregationName + " in " + (System.currentTimeMillis() - startTime) + "ms. Number of buckets processed: " + bucketCount.longValue());
+        }
+    }
+
+    private long getBucketEnd(long i, TimeSeriesProcessedParams finalParams) {
+        return i + getBucketSize(finalParams.getFrom(), finalParams.getTo(), finalParams.isShrink(), finalParams.getResolution());
+    }
+
     private long roundDownToAvailableResolution(long targetResolution) {
         List<Long> availableResolutions = getAvailableResolutions();
         for (int i = 1; i < availableResolutions.size(); i++) {

step-framework-timeseries/src/main/java/step/core/timeseries/aggregation/TimeSeriesAggregationQuery.java

@@ -1,6 +1,7 @@
 package step.core.timeseries.aggregation;
 
 import step.core.collections.Filter;
+import step.core.timeseries.bucket.Aggregation;
 import step.core.timeseries.query.TimeSeriesQuery;
 
 import javax.annotation.Nullable;
@@ -13,6 +14,12 @@ public class TimeSeriesAggregationQuery extends TimeSeriesQuery {
     // Group
     private Set<String> groupDimensions;
 
+    // Aggregation applied across the series of one group
+    private Aggregation groupAggregation;
+
+    // Aggregation applied across the successive buckets of one series falling into the same time window
+    private Aggregation timeAggregation;
+
     // Ideal number of buckets the interval will be split into
     @Nullable
     private Integer bucketsCount;
@@ -31,6 +38,8 @@ protected TimeSeriesAggregationQuery(
         Filter filter,
         TimeSeriesOptimizationType optimizationType,
         Set<String> groupDimensions,
+        Aggregation groupAggregation,
+        Aggregation timeAggregation,
         Long bucketsResolution,
         Long from,
         Long to,
@@ -42,6 +51,8 @@ protected TimeSeriesAggregationQuery(
         this.optimizationType = optimizationType;
         this.bucketsCount = bucketsCount;
         this.groupDimensions = groupDimensions;
+        this.groupAggregation = groupAggregation;
+        this.timeAggregation = timeAggregation;
         this.bucketsResolution = bucketsResolution;
         this.collectAttributeKeys = collectAttributeKeys;
         this.collectAttributesValuesLimit = collectAttributesValuesLimit;
@@ -51,6 +62,14 @@ public Set<String> getGroupDimensions() {
         return groupDimensions;
     }
 
+    public Aggregation getGroupAggregation() {
+        return groupAggregation;
+    }
+
+    public Aggregation getTimeAggregation() {
+        return timeAggregation;
+    }
+
     public Filter getFilter() {
         return filter;
     }

step-framework-timeseries/src/main/java/step/core/timeseries/aggregation/TimeSeriesAggregationQueryBuilder.java

@@ -3,8 +3,8 @@
 import step.core.collections.Filter;
 import step.core.collections.Filters;
 import step.core.timeseries.TimeSeriesFilterBuilder;
+import step.core.timeseries.bucket.Aggregation;
 
-import java.util.Date;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
@@ -13,6 +13,8 @@ public class TimeSeriesAggregationQueryBuilder {
 
     private TimeSeriesOptimizationType optimizationType = TimeSeriesOptimizationType.MOST_EFFICIENT;
     private Set<String> groupDimensions = new HashSet<>();
+    private Aggregation groupAggregation = Aggregation.MERGE;
+    private Aggregation timeAggregation = Aggregation.MERGE;
     private Filter filter = Filters.empty();
     private Long from;
     private Long to;
@@ -32,6 +34,31 @@ public TimeSeriesAggregationQueryBuilder withGroupDimensions(Set<String> groupDi
         return this;
     }
 
+    /**
+     * Defines the group-by (series) aggregation, i.e. how the aligned buckets of the series belonging to the same
+     * group are reduced into one single bucket. Defaults to {@link Aggregation#AVG}.
+     *
+     * @param groupDimensions    the attribute keys defining the groups
+     * @param groupAggregation the aggregation applied across the series of each group
+     */
+    public TimeSeriesAggregationQueryBuilder groupBy(Set<String> groupDimensions, Aggregation groupAggregation) {
+        this.groupDimensions = groupDimensions;
+        this.groupAggregation = groupAggregation;
+        return this;
+    }
+
+    /**
+     * Defines the time-window aggregation, i.e. how the successive buckets of one single series falling into the same
+     * time window are reduced into one single bucket. Independent of the group-by aggregation defined by
+     * {@link #groupBy(Set, Aggregation)}. Defaults to {@link Aggregation#AVG}.
+     *
+     * @param timeAggregation the aggregation applied across the buckets of one series within a time window
+     */
+    public TimeSeriesAggregationQueryBuilder withTimeAggregation(Aggregation timeAggregation) {
+        this.timeAggregation = timeAggregation;
+        return this;
+    }
+
     public TimeSeriesAggregationQueryBuilder withFilter(Filter filter) {
         this.filter = filter;
         return this;
@@ -88,7 +115,7 @@ public TimeSeriesAggregationQueryBuilder split(int targetNumberOfBuckets) {
      */
     public TimeSeriesAggregationQuery build() {
         this.from = this.from != null ? this.from : 0;
-        return new TimeSeriesAggregationQuery(filter, optimizationType, groupDimensions, this.proposedResolution, this.from, this.to, shrink, this.bucketsCount, collectAttributeKeys, collectAttributesValuesLimit);
+        return new TimeSeriesAggregationQuery(filter, optimizationType, groupDimensions, groupAggregation, timeAggregation, this.proposedResolution, this.from, this.to, shrink, this.bucketsCount, collectAttributeKeys, collectAttributesValuesLimit);
     }