addressed most comments minus tests

Author: limotova

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/DataType.java

@@ -309,7 +309,7 @@ public enum DataType {
  */
  SEMANTIC_TEXT(builder().esType("semantic_text").unknownSize()),
 
- AGGREGATE_METRIC_DOUBLE(builder().esType("aggregate_metric_double").unknownSize());
+ AGGREGATE_METRIC_DOUBLE(builder().esType("aggregate_metric_double").estimatedSize(Double.BYTES * 3 + Integer.BYTES));
 
  /**
  * Types that are actively being built. These types are not returned

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/AggregateDoubleMetricBlockBuilder.java

@@ -7,21 +7,36 @@
 
 package org.elasticsearch.compute.data;
 
+import org.elasticsearch.common.breaker.CircuitBreakingException;
+import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.mapper.BlockLoader;
 
 public class AggregateDoubleMetricBlockBuilder extends AbstractBlockBuilder implements BlockLoader.AggregateDoubleMetricBuilder {
 
- private final DoubleBlockBuilder minBuilder;
- private final DoubleBlockBuilder maxBuilder;
- private final DoubleBlockBuilder sumBuilder;
- private final IntBlockBuilder countBuilder;
+ private DoubleBlockBuilder minBuilder;
+ private DoubleBlockBuilder maxBuilder;
+ private DoubleBlockBuilder sumBuilder;
+ private IntBlockBuilder countBuilder;
 
  public AggregateDoubleMetricBlockBuilder(int estimatedSize, BlockFactory blockFactory) {
  super(blockFactory);
- minBuilder = new DoubleBlockBuilder(estimatedSize, blockFactory);
- maxBuilder = new DoubleBlockBuilder(estimatedSize, blockFactory);
- sumBuilder = new DoubleBlockBuilder(estimatedSize, blockFactory);
- countBuilder = new IntBlockBuilder(estimatedSize, blockFactory);
+ minBuilder = null;
+ maxBuilder = null;
+ sumBuilder = null;
+ countBuilder = null;
+ try {
+ minBuilder = new DoubleBlockBuilder(estimatedSize, blockFactory);
+ maxBuilder = new DoubleBlockBuilder(estimatedSize, blockFactory);
+ sumBuilder = new DoubleBlockBuilder(estimatedSize, blockFactory);
+ countBuilder = new IntBlockBuilder(estimatedSize, blockFactory);
+ } finally {
+ if (countBuilder == null) {
+ Releasables.closeWhileHandlingException(minBuilder);
+ Releasables.closeWhileHandlingException(maxBuilder);
+ Releasables.closeWhileHandlingException(sumBuilder);
+ Releasables.closeWhileHandlingException(countBuilder);
+ }
+ }
  }
 
  @Override
@@ -45,10 +60,10 @@ protected int elementSize() {
  @Override
  public Block.Builder copyFrom(Block block, int beginInclusive, int endExclusive) {
  CompositeBlock composite = (CompositeBlock) block;
- minBuilder.copyFrom(composite.getBlock(Metric.MIN.ordinal()), beginInclusive, endExclusive);
- maxBuilder.copyFrom(composite.getBlock(Metric.MAX.ordinal()), beginInclusive, endExclusive);
- sumBuilder.copyFrom(composite.getBlock(Metric.SUM.ordinal()), beginInclusive, endExclusive);
- countBuilder.copyFrom(composite.getBlock(Metric.COUNT.ordinal()), beginInclusive, endExclusive);
+ minBuilder.copyFrom(composite.getBlock(Metric.MIN.getIndex()), beginInclusive, endExclusive);
+ maxBuilder.copyFrom(composite.getBlock(Metric.MAX.getIndex()), beginInclusive, endExclusive);
+ sumBuilder.copyFrom(composite.getBlock(Metric.SUM.getIndex()), beginInclusive, endExclusive);
+ countBuilder.copyFrom(composite.getBlock(Metric.COUNT.getIndex()), beginInclusive, endExclusive);
  return this;
  }
 
@@ -64,11 +79,19 @@ public Block.Builder mvOrdering(Block.MvOrdering mvOrdering) {
  @Override
  public Block build() {
  Block[] blocks = new Block[4];
- blocks[Metric.MIN.ordinal()] = minBuilder.build();
- blocks[Metric.MAX.ordinal()] = maxBuilder.build();
- blocks[Metric.SUM.ordinal()] = sumBuilder.build();
- blocks[Metric.COUNT.ordinal()] = countBuilder.build();
- return new CompositeBlock(blocks);
+ try {
+ finish();
+ blocks[Metric.MIN.getIndex()] = minBuilder.build();
+ blocks[Metric.MAX.getIndex()] = maxBuilder.build();
+ blocks[Metric.SUM.getIndex()] = sumBuilder.build();
+ blocks[Metric.COUNT.getIndex()] = countBuilder.build();
+ return new CompositeBlock(blocks);
+ } catch (CircuitBreakingException e) {
+ for (Block block : blocks) {
+ block.close();
+ }
+ throw e;
+ }
  }
 
  @Override
@@ -81,9 +104,19 @@ public BlockLoader.AggregateDoubleMetricBuilder append(double min, double max, d
  }
 
  public enum Metric {
- MIN,
- MAX,
- SUM,
- COUNT;
+ MIN(0),
+ MAX(1),
+ SUM(2),
+ COUNT(3);
+
+ private final int index;
+
+ Metric(int index) {
+ this.index = index;
+ }
+
+ public int getIndex() {
+ return index;
+ }
  }
 }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/CompositeBlock.java

@@ -86,29 +86,21 @@ public int getPositionCount() {
 
  @Override
  public int getTotalValueCount() {
- // TODO: this works for aggregate metric double fields, because the four blocks are guarenteed to have the same length / positions.
- // So check just for one sub block is sufficient.
  return blocks[0].getTotalValueCount();
  }
 
  @Override
  public int getFirstValueIndex(int position) {
- // TODO: this works for aggregate metric double fields, because the four blocks are guarenteed to have the same length / positions.
- // So check just for one sub block is sufficient.
  return blocks[0].getFirstValueIndex(position);
  }
 
  @Override
  public int getValueCount(int position) {
- // TODO: this works for aggregate metric double fields, because the four blocks are guarenteed to have the same length / positions.
- // So check just for one sub block is sufficient.
  return blocks[0].getValueCount(position);
  }
 
  @Override
  public boolean isNull(int position) {
- // TODO: this works for aggregate metric double fields, because the four blocks are guarenteed to have the same length / positions.
- // So check just for one sub block is sufficient.
  return blocks[0].isNull(position);
  }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Count.java

@@ -35,7 +35,6 @@
 import static java.util.Collections.emptyList;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.DEFAULT;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
-import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER;
 
 public class Count extends AggregateFunction implements ToAggregator, SurrogateExpression {
  public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Count", Count::new);
@@ -146,14 +145,7 @@ public Expression surrogate() {
  var s = source();
  var field = field();
  if (field.dataType() == DataType.AGGREGATE_METRIC_DOUBLE) {
- return new Sum(
- s,
- new FromAggregateDoubleMetric(
- source(),
- field,
- new Literal(s, AggregateDoubleMetricBlockBuilder.Metric.COUNT.ordinal(), INTEGER)
- )
- );
+ return new Sum(s, new FromAggregateDoubleMetric(source(), field, AggregateDoubleMetricBlockBuilder.Metric.COUNT));
  }
 
  if (field.foldable()) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Max.java

@@ -153,14 +153,7 @@ public final AggregatorFunctionSupplier supplier(List<Integer> inputChannels) {
  @Override
  public Expression surrogate() {
  if (field().dataType() == DataType.AGGREGATE_METRIC_DOUBLE) {
- return new Max(
- source(),
- new FromAggregateDoubleMetric(
- source(),
- field(),
- new Literal(source(), AggregateDoubleMetricBlockBuilder.Metric.MAX.ordinal(), DataType.INTEGER)
- )
- );
+ return new Max(source(), new FromAggregateDoubleMetric(source(), field(), AggregateDoubleMetricBlockBuilder.Metric.MAX));
  }
  return field().foldable() ? new MvMax(source(), field()) : null;
  }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Min.java

@@ -153,14 +153,7 @@ public final AggregatorFunctionSupplier supplier(List<Integer> inputChannels) {
  @Override
  public Expression surrogate() {
  if (field().dataType() == DataType.AGGREGATE_METRIC_DOUBLE) {
- return new Min(
- source(),
- new FromAggregateDoubleMetric(
- source(),
- field(),
- new Literal(source(), AggregateDoubleMetricBlockBuilder.Metric.MIN.ordinal(), DataType.INTEGER)
- )
- );
+ return new Min(source(), new FromAggregateDoubleMetric(source(), field(), AggregateDoubleMetricBlockBuilder.Metric.MIN));
  }
  return field().foldable() ? new MvMin(source(), field()) : null;
  }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Sum.java

@@ -36,7 +36,6 @@
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
 import static org.elasticsearch.xpack.esql.core.type.DataType.AGGREGATE_METRIC_DOUBLE;
 import static org.elasticsearch.xpack.esql.core.type.DataType.DOUBLE;
-import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER;
 import static org.elasticsearch.xpack.esql.core.type.DataType.LONG;
 import static org.elasticsearch.xpack.esql.core.type.DataType.UNSIGNED_LONG;
 
@@ -139,14 +138,7 @@ public Expression surrogate() {
  var s = source();
  var field = field();
  if (field.dataType() == AGGREGATE_METRIC_DOUBLE) {
- return new Sum(
- s,
- new FromAggregateDoubleMetric(
- source(),
- field,
- new Literal(s, AggregateDoubleMetricBlockBuilder.Metric.SUM.ordinal(), INTEGER)
- )
- );
+ return new Sum(s, new FromAggregateDoubleMetric(source(), field, AggregateDoubleMetricBlockBuilder.Metric.SUM));
  }
 
  // SUM(const) is equivalent to MV_SUM(const)*COUNT(*).

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/FromAggregateDoubleMetric.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
+import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
@@ -61,6 +62,10 @@ public FromAggregateDoubleMetric(
  this.subfieldIndex = subfieldIndex;
  }
 
+ public FromAggregateDoubleMetric(Source source, Expression field, AggregateDoubleMetricBlockBuilder.Metric metric) {
+ this(source, field, new Literal(source, metric.getIndex(), INTEGER));
+ }
+
  private FromAggregateDoubleMetric(StreamInput in) throws IOException {
  this(Source.readFrom((PlanStreamInput) in), in.readNamedWriteable(Expression.class), in.readNamedWriteable(Expression.class));
  }
@@ -84,7 +89,7 @@ public DataType dataType() {
  }
 
  var subfield = ((Number) subfieldIndex.fold(FoldContext.small())).intValue();
- if (subfield == AggregateDoubleMetricBlockBuilder.Metric.COUNT.ordinal()) {
+ if (subfield == AggregateDoubleMetricBlockBuilder.Metric.COUNT.getIndex()) {
  return INTEGER;
  }
  return DOUBLE;
@@ -121,13 +126,10 @@ public EvalOperator.ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvalua
  return new EvalOperator.ExpressionEvaluator() {
  @Override
  public Block eval(Page page) {
- CompositeBlock compositeBlock = (CompositeBlock) eval.eval(page);
- try {
+ try (CompositeBlock compositeBlock = (CompositeBlock) eval.eval(page)) {
  Block block = compositeBlock.getBlock(((Number) subfieldIndex.fold(FoldContext.small())).intValue());
  block.incRef();
  return block;
- } finally {
- compositeBlock.close();
  }
  }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AggregateMapper.java

@@ -213,7 +213,7 @@ private static Stream<AggDef> groupingAndNonGrouping(Tuple<Class<?>, Tuple<Strin
  // rate doesn't support non-grouping aggregations
  return Stream.of(new AggDef(tuple.v1(), tuple.v2().v1(), tuple.v2().v2(), true));
  } else if (tuple.v2().v1().equals("AggregatedMetricDouble")) {
- // not supporting grouping aggregations yet
+ // TODO: support grouping aggregations for aggregate metric double
  return Stream.of(new AggDef(tuple.v1(), tuple.v2().v1(), tuple.v2().v2(), false));
  } else {
  return Stream.of(

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/FromAggregateDoubleMetricTests.java

@@ -0,0 +1,4 @@
+package org.elasticsearch.xpack.esql.expression.function.scalar.convert;
+
+public class FromAggregateDoubleMetricTests {
+}