[Rollup] Move toAggCap() methods out of rollup config objects (#32583)

Author: tlrx

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/rollup/action/RollupJobCaps.java

@@ -11,15 +11,26 @@
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.xcontent.ToXContentObject;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder;
+import org.elasticsearch.search.aggregations.bucket.histogram.HistogramAggregationBuilder;
+import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
+import org.elasticsearch.xpack.core.rollup.job.DateHistogramGroupConfig;
+import org.elasticsearch.xpack.core.rollup.job.GroupConfig;
+import org.elasticsearch.xpack.core.rollup.job.HistogramGroupConfig;
+import org.elasticsearch.xpack.core.rollup.job.MetricConfig;
 import org.elasticsearch.xpack.core.rollup.job.RollupJobConfig;
+import org.elasticsearch.xpack.core.rollup.job.TermsGroupConfig;
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.stream.Collectors;
+
+import static java.util.Collections.singletonMap;
 
 /**
  * Represents the Rollup capabilities for a specific job on a single rollup index
@@ -42,52 +53,7 @@ public RollupJobCaps(RollupJobConfig job) {
         jobID = job.getId();
         rollupIndex = job.getRollupIndex();
         indexPattern = job.getIndexPattern();
-        Map<String, Object> dateHistoAggCap = job.getGroupConfig().getDateHistogram().toAggCap();
-        String dateField = job.getGroupConfig().getDateHistogram().getField();
-        RollupFieldCaps fieldCaps = fieldCapLookup.get(dateField);
-        if (fieldCaps == null) {
-            fieldCaps = new RollupFieldCaps();
-        }
-        fieldCaps.addAgg(dateHistoAggCap);
-        fieldCapLookup.put(dateField, fieldCaps);
-
-        if (job.getGroupConfig().getHistogram() != null) {
-            Map<String, Object> histoAggCap = job.getGroupConfig().getHistogram().toAggCap();
-            Arrays.stream(job.getGroupConfig().getHistogram().getFields()).forEach(field -> {
-                RollupFieldCaps caps = fieldCapLookup.get(field);
-                if (caps == null) {
-                    caps = new RollupFieldCaps();
-                }
-                caps.addAgg(histoAggCap);
-                fieldCapLookup.put(field, caps);
-            });
-        }
-
-        if (job.getGroupConfig().getTerms() != null) {
-            Map<String, Object> histoAggCap = job.getGroupConfig().getTerms().toAggCap();
-            Arrays.stream(job.getGroupConfig().getTerms().getFields()).forEach(field -> {
-                RollupFieldCaps caps = fieldCapLookup.get(field);
-                if (caps == null) {
-                    caps = new RollupFieldCaps();
-                }
-                caps.addAgg(histoAggCap);
-                fieldCapLookup.put(field, caps);
-            });
-        }
-
-        if (job.getMetricsConfig().size() > 0) {
-            job.getMetricsConfig().forEach(metricConfig -> {
-                List<Map<String, Object>> metrics = metricConfig.toAggCap();
-                metrics.forEach(m -> {
-                    RollupFieldCaps caps = fieldCapLookup.get(metricConfig.getField());
-                    if (caps == null) {
-                        caps = new RollupFieldCaps();
-                    }
-                    caps.addAgg(m);
-                    fieldCapLookup.put(metricConfig.getField(), caps);
-                });
-            });
-        }
+        fieldCapLookup = createRollupFieldCaps(job);
     }
 
     public RollupJobCaps(StreamInput in) throws IOException {
@@ -149,15 +115,86 @@ public boolean equals(Object other) {
         RollupJobCaps that = (RollupJobCaps) other;
 
         return Objects.equals(this.jobID, that.jobID)
-                && Objects.equals(this.rollupIndex, that.rollupIndex)
-                && Objects.equals(this.fieldCapLookup, that.fieldCapLookup);
+            && Objects.equals(this.rollupIndex, that.rollupIndex)
+            && Objects.equals(this.fieldCapLookup, that.fieldCapLookup);
     }
 
     @Override
     public int hashCode() {
         return Objects.hash(jobID, rollupIndex, fieldCapLookup);
     }
 
+    static Map<String, RollupFieldCaps> createRollupFieldCaps(final RollupJobConfig rollupJobConfig) {
+        final Map<String, RollupFieldCaps> fieldCapLookup = new HashMap<>();
+
+        final GroupConfig groupConfig = rollupJobConfig.getGroupConfig();
+        if (groupConfig != null) {
+            // Create RollupFieldCaps for the date histogram
+            final DateHistogramGroupConfig dateHistogram = groupConfig.getDateHistogram();
+            final Map<String, Object> dateHistogramAggCap = new HashMap<>();
+            dateHistogramAggCap.put("agg", DateHistogramAggregationBuilder.NAME);
+            dateHistogramAggCap.put(DateHistogramGroupConfig.INTERVAL, dateHistogram.getInterval().toString());
+            if (dateHistogram.getDelay() != null) {
+                dateHistogramAggCap.put(DateHistogramGroupConfig.DELAY, dateHistogram.getDelay().toString());
+            }
+            dateHistogramAggCap.put(DateHistogramGroupConfig.TIME_ZONE, dateHistogram.getTimeZone());
+
+            final RollupFieldCaps dateHistogramFieldCaps = new RollupFieldCaps();
+            dateHistogramFieldCaps.addAgg(dateHistogramAggCap);
+            fieldCapLookup.put(dateHistogram.getField(), dateHistogramFieldCaps);
+
+            // Create RollupFieldCaps for the histogram
+            final HistogramGroupConfig histogram = groupConfig.getHistogram();
+            if (histogram != null) {
+                final Map<String, Object> histogramAggCap = new HashMap<>();
+                histogramAggCap.put("agg", HistogramAggregationBuilder.NAME);
+                histogramAggCap.put(HistogramGroupConfig.INTERVAL, histogram.getInterval());
+                for (String field : histogram.getFields()) {
+                    RollupFieldCaps caps = fieldCapLookup.get(field);
+                    if (caps == null) {
+                        caps = new RollupFieldCaps();
+                    }
+                    caps.addAgg(histogramAggCap);
+                    fieldCapLookup.put(field, caps);
+                }
+            }
+
+            // Create RollupFieldCaps for the term
+            final TermsGroupConfig terms = groupConfig.getTerms();
+            if (terms != null) {
+                final Map<String, Object> termsAggCap = singletonMap("agg", TermsAggregationBuilder.NAME);
+                for (String field : terms.getFields()) {
+                    RollupFieldCaps caps = fieldCapLookup.get(field);
+                    if (caps == null) {
+                        caps = new RollupFieldCaps();
+                    }
+                    caps.addAgg(termsAggCap);
+                    fieldCapLookup.put(field, caps);
+                }
+            }
+        }
+
+        // Create RollupFieldCaps for the metrics
+        final List<MetricConfig> metricsConfig = rollupJobConfig.getMetricsConfig();
+            if (metricsConfig.size() > 0) {
+            metricsConfig.forEach(metricConfig -> {
+                final List<Map<String, Object>> metrics = metricConfig.getMetrics().stream()
+                    .map(metric -> singletonMap("agg", (Object) metric))
+                    .collect(Collectors.toList());
+
+                metrics.forEach(m -> {
+                    RollupFieldCaps caps = fieldCapLookup.get(metricConfig.getField());
+                    if (caps == null) {
+                        caps = new RollupFieldCaps();
+                    }
+                    caps.addAgg(m);
+                    fieldCapLookup.put(metricConfig.getField(), caps);
+                });
+            });
+        }
+        return Collections.unmodifiableMap(fieldCapLookup);
+    }
+
     public static class RollupFieldCaps implements Writeable, ToXContentObject {
         private List<Map<String, Object>> aggs = new ArrayList<>();
 

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/rollup/job/DateHistogramGroupConfig.java

@@ -29,7 +29,6 @@
 
 import java.io.IOException;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -55,10 +54,10 @@
 public class DateHistogramGroupConfig implements Writeable, ToXContentObject {
 
     static final String NAME = "date_histogram";
-    private static final String INTERVAL = "interval";
+    public static final String INTERVAL = "interval";
     private static final String FIELD = "field";
     public static final String TIME_ZONE = "time_zone";
-    private static final String DELAY = "delay";
+    public static final String DELAY = "delay";
     private static final String DEFAULT_TIMEZONE = "UTC";
     private static final ConstructingObjectParser<DateHistogramGroupConfig, Void> PARSER;
     static {
@@ -196,21 +195,6 @@ public List<CompositeValuesSourceBuilder<?>> toBuilders() {
         return Collections.singletonList(vsBuilder);
     }
 
-    /**
-     * @return A map representing this config object as a RollupCaps aggregation object
-     */
-    public Map<String, Object> toAggCap() {
-        Map<String, Object> map = new HashMap<>(3);
-        map.put("agg", DateHistogramAggregationBuilder.NAME);
-        map.put(INTERVAL, interval.toString());
-        if (delay != null) {
-            map.put(DELAY, delay.toString());
-        }
-        map.put(TIME_ZONE, timeZone);
-
-        return map;
-    }
-
     public void validateMappings(Map<String, Map<String, FieldCapabilities>> fieldCapsResponse,
                                                              ActionRequestValidationException validationException) {
 

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/rollup/job/HistogramGroupConfig.java

@@ -24,11 +24,9 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Set;
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.common.xcontent.ConstructingObjectParser.constructorArg;
@@ -48,7 +46,7 @@
 public class HistogramGroupConfig implements Writeable, ToXContentObject {
 
     static final String NAME = "histogram";
-    private static final String INTERVAL = "interval";
+    public static final String INTERVAL = "interval";
     private static final String FIELDS = "fields";
     private static final ConstructingObjectParser<HistogramGroupConfig, Void> PARSER;
     static {
@@ -106,20 +104,6 @@ public List<CompositeValuesSourceBuilder<?>> toBuilders() {
         }).collect(Collectors.toList());
     }
 
-    /**
-     * @return A map representing this config object as a RollupCaps aggregation object
-     */
-    public Map<String, Object> toAggCap() {
-        Map<String, Object> map = new HashMap<>(2);
-        map.put("agg", HistogramAggregationBuilder.NAME);
-        map.put(INTERVAL, interval);
-        return map;
-    }
-
-    public Set<String> getAllFields() {
-        return Arrays.stream(fields).collect(Collectors.toSet());
-    }
-
     public void validateMappings(Map<String, Map<String, FieldCapabilities>> fieldCapsResponse,
                                  ActionRequestValidationException validationException) {
 

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/rollup/job/MetricConfig.java

@@ -31,7 +31,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.stream.Collectors;
 
 import static org.elasticsearch.common.xcontent.ConstructingObjectParser.constructorArg;
 
@@ -152,13 +151,6 @@ public List<ValuesSourceAggregationBuilder.LeafOnly> toBuilders() {
         return aggs;
     }
 
-    /**
-     * @return A map representing this config object as a RollupCaps aggregation object
-     */
-    public List<Map<String, Object>> toAggCap() {
-        return metrics.stream().map(metric -> Collections.singletonMap("agg", (Object)metric)).collect(Collectors.toList());
-    }
-
     public void validateMappings(Map<String, Map<String, FieldCapabilities>> fieldCapsResponse,
                                  ActionRequestValidationException validationException) {
 

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/rollup/job/TermsGroupConfig.java

@@ -25,7 +25,6 @@
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -94,15 +93,6 @@ public List<CompositeValuesSourceBuilder<?>> toBuilders() {
         }).collect(Collectors.toList());
     }
 
-    /**
-     * @return A map representing this config object as a RollupCaps aggregation object
-     */
-    public Map<String, Object> toAggCap() {
-        Map<String, Object> map = new HashMap<>(1);
-        map.put("agg", TermsAggregationBuilder.NAME);
-        return map;
-    }
-
     public void validateMappings(Map<String, Map<String, FieldCapabilities>> fieldCapsResponse,
                                  ActionRequestValidationException validationException) {