SQL: Streamline declaration of LeafAggs (#55380)

Avoid repetition of the aggregation builder setup

Relates #55241

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/AvgAgg.java

@@ -5,18 +5,19 @@
  */
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
-import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
-import static org.elasticsearch.search.aggregations.AggregationBuilders.avg;
+import java.util.function.Function;
 
-public class AvgAgg extends LeafAgg {
+public class AvgAgg extends DefaultAggSourceLeafAgg {
 
     public AvgAgg(String id, AggSource source) {
         super(id, source);
     }
 
     @Override
-    AggregationBuilder toBuilder() {
-        return addAggSource(avg(id()));
+    Function<String, ValuesSourceAggregationBuilder<?>> builder() {
+        return AggregationBuilders::avg;
     }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/CardinalityAgg.java

@@ -5,18 +5,19 @@
  */
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
-import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
-import static org.elasticsearch.search.aggregations.AggregationBuilders.cardinality;
+import java.util.function.Function;
 
-public class CardinalityAgg extends LeafAgg {
+public class CardinalityAgg extends DefaultAggSourceLeafAgg {
 
     public CardinalityAgg(String id, AggSource source) {
         super(id, source);
     }
 
     @Override
-    AggregationBuilder toBuilder() {
-        return addAggSource(cardinality(id()));
+    Function<String, ValuesSourceAggregationBuilder<?>> builder() {
+        return AggregationBuilders::cardinality;
     }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/DefaultAggSourceLeafAgg.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.sql.querydsl.agg;
+
+import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
+
+import java.util.function.Function;
+
+abstract class DefaultAggSourceLeafAgg extends LeafAgg {
+
+    DefaultAggSourceLeafAgg(String id, AggSource source) {
+        super(id, source);
+    }
+
+    @Override
+    AggregationBuilder toBuilder() {
+        return source().with(builder().apply(id()));
+    }
+
+    abstract Function<String, ValuesSourceAggregationBuilder<?>> builder();
+}

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/ExtendedStatsAgg.java

@@ -5,18 +5,19 @@
  */
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
-import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
-import static org.elasticsearch.search.aggregations.AggregationBuilders.extendedStats;
+import java.util.function.Function;
 
-public class ExtendedStatsAgg extends LeafAgg {
+public class ExtendedStatsAgg extends DefaultAggSourceLeafAgg {
 
     public ExtendedStatsAgg(String id, AggSource source) {
         super(id, source);
     }
 
     @Override
-    AggregationBuilder toBuilder() {
-        return addAggSource(extendedStats(id()));
+    Function<String, ValuesSourceAggregationBuilder<?>> builder() {
+        return AggregationBuilders::extendedStats;
     }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/FilterExistsAgg.java

@@ -5,6 +5,7 @@
  */
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.search.aggregations.AggregationBuilder;
 import org.elasticsearch.xpack.ql.expression.gen.script.ScriptTemplate;
@@ -27,11 +28,15 @@ public class FilterExistsAgg extends LeafAgg {
 
     @Override
     AggregationBuilder toBuilder() {
+        QueryBuilder qb;
+        
         if (source().fieldName() != null) {
-            return filter(id(), QueryBuilders.existsQuery(source().fieldName()));
+            qb = QueryBuilders.existsQuery(source().fieldName());
         } else {
-            return filter(id(), QueryBuilders.scriptQuery(wrapWithIsNotNull(source().script()).toPainless()));
+            qb = QueryBuilders.scriptQuery(wrapWithIsNotNull(source().script()).toPainless());
         }
+
+        return filter(id(), qb);
     }
 
     private static ScriptTemplate wrapWithIsNotNull(ScriptTemplate script) {

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/LeafAgg.java

@@ -6,7 +6,6 @@
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
 import org.elasticsearch.search.aggregations.AggregationBuilder;
-import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
 public abstract class LeafAgg extends Agg {
 
@@ -15,9 +14,4 @@ public abstract class LeafAgg extends Agg {
     }
 
     abstract AggregationBuilder toBuilder();
-
-    @SuppressWarnings("rawtypes")
-    protected ValuesSourceAggregationBuilder addAggSource(ValuesSourceAggregationBuilder builder) {
-        return source().with(builder);
-    }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/MaxAgg.java

@@ -5,18 +5,19 @@
  */
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
-import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
-import static org.elasticsearch.search.aggregations.AggregationBuilders.max;
+import java.util.function.Function;
 
-public class MaxAgg extends LeafAgg {
+public class MaxAgg extends DefaultAggSourceLeafAgg {
 
     public MaxAgg(String id, AggSource source) {
         super(id, source);
     }
 
     @Override
-    AggregationBuilder toBuilder() {
-        return addAggSource(max(id()));
+    Function<String, ValuesSourceAggregationBuilder<?>> builder() {
+        return AggregationBuilders::max;
     }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/MedianAbsoluteDeviationAgg.java

@@ -6,18 +6,19 @@
 
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
-import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
-import static org.elasticsearch.search.aggregations.AggregationBuilders.medianAbsoluteDeviation;
+import java.util.function.Function;
 
-public class MedianAbsoluteDeviationAgg extends LeafAgg {
+public class MedianAbsoluteDeviationAgg extends DefaultAggSourceLeafAgg {
 
     public MedianAbsoluteDeviationAgg(String id, AggSource source) {
         super(id, source);
     }
 
     @Override
-    AggregationBuilder toBuilder() {
-        return addAggSource(medianAbsoluteDeviation(id()));
+    Function<String, ValuesSourceAggregationBuilder<?>> builder() {
+        return AggregationBuilders::medianAbsoluteDeviation;
     }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/MinAgg.java

@@ -5,18 +5,19 @@
  */
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
-import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
-import static org.elasticsearch.search.aggregations.AggregationBuilders.min;
+import java.util.function.Function;
 
-public class MinAgg extends LeafAgg {
+public class MinAgg extends DefaultAggSourceLeafAgg {
 
     public MinAgg(String id, AggSource source) {
         super(id, source);
     }
 
     @Override
-    AggregationBuilder toBuilder() {
-        return addAggSource(min(id()));
+    Function<String, ValuesSourceAggregationBuilder<?>> builder() {
+        return AggregationBuilders::min;
     }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/PercentileRanksAgg.java

@@ -5,13 +5,14 @@
  */
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
-import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
 import java.util.List;
+import java.util.function.Function;
 
 import static org.elasticsearch.search.aggregations.AggregationBuilders.percentileRanks;
 
-public class PercentileRanksAgg extends LeafAgg {
+public class PercentileRanksAgg extends DefaultAggSourceLeafAgg {
 
     private final List<Double> values;
 
@@ -21,7 +22,7 @@ public class PercentileRanksAgg extends LeafAgg {
     }
 
     @Override
-    AggregationBuilder toBuilder() {
-        return addAggSource(percentileRanks(id(), values.stream().mapToDouble(Double::doubleValue).toArray()));
+    Function<String, ValuesSourceAggregationBuilder<?>> builder() {
+        return s -> percentileRanks(s, values.stream().mapToDouble(Double::doubleValue).toArray());
     }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/PercentilesAgg.java

@@ -6,13 +6,14 @@
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
 import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
 import org.elasticsearch.search.aggregations.metrics.PercentilesAggregationBuilder;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
 import java.util.List;
+import java.util.function.Function;
 
-import static org.elasticsearch.search.aggregations.AggregationBuilders.percentiles;
-
-public class PercentilesAgg extends LeafAgg {
+public class PercentilesAgg extends DefaultAggSourceLeafAgg {
 
     private final List<Double> percents;
 
@@ -20,11 +21,16 @@ public class PercentilesAgg extends LeafAgg {
         super(id, source);
         this.percents = percents;
     }
-
+    
     @Override
     AggregationBuilder toBuilder() {
         // TODO: look at keyed
-        PercentilesAggregationBuilder builder = (PercentilesAggregationBuilder) addAggSource(percentiles(id()));
+        PercentilesAggregationBuilder builder = (PercentilesAggregationBuilder) super.toBuilder();
         return builder.percentiles(percents.stream().mapToDouble(Double::doubleValue).toArray());
     }
-}
+
+    @Override
+    Function<String, ValuesSourceAggregationBuilder<?>> builder() {
+        return AggregationBuilders::percentiles;
+    }
+}

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/StatsAgg.java

@@ -5,18 +5,19 @@
  */
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
-import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
-import static org.elasticsearch.search.aggregations.AggregationBuilders.stats;
+import java.util.function.Function;
 
-public class StatsAgg extends LeafAgg {
+public class StatsAgg extends DefaultAggSourceLeafAgg {
 
     public StatsAgg(String id, AggSource source) {
         super(id, source);
     }
 
     @Override
-    AggregationBuilder toBuilder() {
-        return addAggSource(stats(id()));
+    Function<String, ValuesSourceAggregationBuilder<?>> builder() {
+        return AggregationBuilders::stats;
     }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/SumAgg.java

@@ -5,18 +5,19 @@
  */
 package org.elasticsearch.xpack.sql.querydsl.agg;
 
-import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 
-import static org.elasticsearch.search.aggregations.AggregationBuilders.sum;
+import java.util.function.Function;
 
-public class SumAgg extends LeafAgg {
+public class SumAgg extends DefaultAggSourceLeafAgg {
 
     public SumAgg(String id, AggSource source) {
         super(id, source);
     }
 
     @Override
-    AggregationBuilder toBuilder() {
-        return addAggSource(sum(id()));
+    Function<String, ValuesSourceAggregationBuilder<?>> builder() {
+        return AggregationBuilders::sum;
     }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/agg/TopHitsAgg.java

@@ -84,10 +84,12 @@ public class TopHitsAgg extends LeafAgg {
 
         TopHitsAggregationBuilder builder = topHits(id());
         if (source().fieldName() != null) {
-            return builder.docValueField(source().fieldName(), SqlDataTypes.format(fieldDataType)).sorts(sortBuilderList).size(1);
+            builder.docValueField(source().fieldName(), SqlDataTypes.format(fieldDataType));
         } else {
-            return builder.scriptField(id(), source().script().toPainless()).sorts(sortBuilderList).size(1);
+            builder.scriptField(id(), source().script().toPainless());
         }
+
+        return builder.sorts(sortBuilderList).size(1);
     }
 
     @Override