`Sqrt` function for ESQL (#98449)

* Sqrt function for ESQL

Introduces a unary scalar function for square root, which is a thin
wrapper over the Java.Math implementation.

* Fix area for ESQL integration changelog.

* Restore changelog.

* Restore area in changelog.

docs/reference/esql/esql-functions.asciidoc

@@ -102,6 +102,7 @@ include::functions/round.asciidoc[]
 include::functions/sin.asciidoc[]
 include::functions/sinh.asciidoc[]
 include::functions/split.asciidoc[]
+include::functions/sqrt.asciidoc[]
 include::functions/starts_with.asciidoc[]
 include::functions/substring.asciidoc[]
 include::functions/tan.asciidoc[]

docs/reference/esql/functions/sqrt.asciidoc

@@ -0,0 +1,15 @@
+[[esql-sqrt]]
+=== `SQRT`
+Returns the square root of a number.  The input can be any numeric value, the return value
+is always a double.
+
+Square roots of negative numbers are NaN. Square roots of infinites are infinite.
+
+[source.merge.styled,esql]
+----
+include::{esql-specs}/math.csv-spec[tag=sqrt]
+----
+[%header.monospaced.styled,format=dsv,separator=|]
+|===
+include::{esql-specs}/math.csv-spec[tag=sqrt-result]
+|===

x-pack/plugin/esql/qa/testFixtures/src/main/resources/math.csv-spec

@@ -900,4 +900,56 @@ a:double
 // end::floor-result[]
 ;
 
+sqrt
+// tag::sqrt[]
+ROW d = 100.0
+| EVAL s = SQRT(d);
+// end::sqrt[]
 
+// tag::sqrt-result[]
+d: double | s:double
+100.0    | 10.0
+// end::sqrt-result[]
+;
+
+sqrtOfInteger
+row i = 81 | eval s = sqrt(i);
+
+i:integer | s:double
+81        | 9
+;
+
+sqrtOfNegative
+row d = -1.0 | eval s = is_nan(sqrt(d));
+
+d:double | s:boolean
+-1.0     | true
+;
+
+sqrtOfNan
+row d = 0.0/0.0 | eval s = is_nan(sqrt(d));
+
+d:double | s:boolean
+NaN      | true
+;
+
+sqrtOfZero
+row d = 0.0 |eval s = sqrt(d);
+
+d:double | s:double
+0.0     | 0.0
+;
+
+sqrtOfNegativeZero
+row d = -0.0 |eval s = sqrt(d);
+
+d:double | s:double
+-0.0     | -0.0
+;
+
+sqrtOfInfinite
+row d = 1/0.0 | eval s = is_infinite(sqrt(d));
+
+d:double | s:boolean
+Infinity | true
+;

x-pack/plugin/esql/qa/testFixtures/src/main/resources/show.csv-spec

@@ -55,6 +55,7 @@ round                    |round(arg1, arg2)
 sin                      |sin(arg1)
 sinh                     |sinh(arg1)
 split                    |split(arg1, arg2)
+sqrt                     |sqrt(arg1)
 starts_with              |starts_with(arg1, arg2)
 substring                |substring(arg1, arg2, arg3)
 sum                      |sum(arg1)

x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/math/SqrtDoubleEvaluator.java

@@ -0,0 +1,64 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License
+// 2.0; you may not use this file except in compliance with the Elastic License
+// 2.0.
+package org.elasticsearch.xpack.esql.expression.function.scalar.math;
+
+import java.lang.Override;
+import java.lang.String;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.DoubleVector;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.EvalOperator;
+
+/**
+ * {@link EvalOperator.ExpressionEvaluator} implementation for {@link Sqrt}.
+ * This class is generated. Do not edit it.
+ */
+public final class SqrtDoubleEvaluator implements EvalOperator.ExpressionEvaluator {
+  private final EvalOperator.ExpressionEvaluator val;
+
+  public SqrtDoubleEvaluator(EvalOperator.ExpressionEvaluator val) {
+    this.val = val;
+  }
+
+  @Override
+  public Block eval(Page page) {
+    Block valUncastBlock = val.eval(page);
+    if (valUncastBlock.areAllValuesNull()) {
+      return Block.constantNullBlock(page.getPositionCount());
+    }
+    DoubleBlock valBlock = (DoubleBlock) valUncastBlock;
+    DoubleVector valVector = valBlock.asVector();
+    if (valVector == null) {
+      return eval(page.getPositionCount(), valBlock);
+    }
+    return eval(page.getPositionCount(), valVector).asBlock();
+  }
+
+  public DoubleBlock eval(int positionCount, DoubleBlock valBlock) {
+    DoubleBlock.Builder result = DoubleBlock.newBlockBuilder(positionCount);
+    position: for (int p = 0; p < positionCount; p++) {
+      if (valBlock.isNull(p) || valBlock.getValueCount(p) != 1) {
+        result.appendNull();
+        continue position;
+      }
+      result.appendDouble(Sqrt.process(valBlock.getDouble(valBlock.getFirstValueIndex(p))));
+    }
+    return result.build();
+  }
+
+  public DoubleVector eval(int positionCount, DoubleVector valVector) {
+    DoubleVector.Builder result = DoubleVector.newVectorBuilder(positionCount);
+    position: for (int p = 0; p < positionCount; p++) {
+      result.appendDouble(Sqrt.process(valVector.getDouble(p)));
+    }
+    return result.build();
+  }
+
+  @Override
+  public String toString() {
+    return "SqrtDoubleEvaluator[" + "val=" + val + "]";
+  }
+}

x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/math/SqrtIntEvaluator.java

@@ -0,0 +1,66 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License
+// 2.0; you may not use this file except in compliance with the Elastic License
+// 2.0.
+package org.elasticsearch.xpack.esql.expression.function.scalar.math;
+
+import java.lang.Override;
+import java.lang.String;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.DoubleVector;
+import org.elasticsearch.compute.data.IntBlock;
+import org.elasticsearch.compute.data.IntVector;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.EvalOperator;
+
+/**
+ * {@link EvalOperator.ExpressionEvaluator} implementation for {@link Sqrt}.
+ * This class is generated. Do not edit it.
+ */
+public final class SqrtIntEvaluator implements EvalOperator.ExpressionEvaluator {
+  private final EvalOperator.ExpressionEvaluator val;
+
+  public SqrtIntEvaluator(EvalOperator.ExpressionEvaluator val) {
+    this.val = val;
+  }
+
+  @Override
+  public Block eval(Page page) {
+    Block valUncastBlock = val.eval(page);
+    if (valUncastBlock.areAllValuesNull()) {
+      return Block.constantNullBlock(page.getPositionCount());
+    }
+    IntBlock valBlock = (IntBlock) valUncastBlock;
+    IntVector valVector = valBlock.asVector();
+    if (valVector == null) {
+      return eval(page.getPositionCount(), valBlock);
+    }
+    return eval(page.getPositionCount(), valVector).asBlock();
+  }
+
+  public DoubleBlock eval(int positionCount, IntBlock valBlock) {
+    DoubleBlock.Builder result = DoubleBlock.newBlockBuilder(positionCount);
+    position: for (int p = 0; p < positionCount; p++) {
+      if (valBlock.isNull(p) || valBlock.getValueCount(p) != 1) {
+        result.appendNull();
+        continue position;
+      }
+      result.appendDouble(Sqrt.process(valBlock.getInt(valBlock.getFirstValueIndex(p))));
+    }
+    return result.build();
+  }
+
+  public DoubleVector eval(int positionCount, IntVector valVector) {
+    DoubleVector.Builder result = DoubleVector.newVectorBuilder(positionCount);
+    position: for (int p = 0; p < positionCount; p++) {
+      result.appendDouble(Sqrt.process(valVector.getInt(p)));
+    }
+    return result.build();
+  }
+
+  @Override
+  public String toString() {
+    return "SqrtIntEvaluator[" + "val=" + val + "]";
+  }
+}

x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/math/SqrtLongEvaluator.java

@@ -0,0 +1,66 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License
+// 2.0; you may not use this file except in compliance with the Elastic License
+// 2.0.
+package org.elasticsearch.xpack.esql.expression.function.scalar.math;
+
+import java.lang.Override;
+import java.lang.String;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.DoubleVector;
+import org.elasticsearch.compute.data.LongBlock;
+import org.elasticsearch.compute.data.LongVector;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.EvalOperator;
+
+/**
+ * {@link EvalOperator.ExpressionEvaluator} implementation for {@link Sqrt}.
+ * This class is generated. Do not edit it.
+ */
+public final class SqrtLongEvaluator implements EvalOperator.ExpressionEvaluator {
+  private final EvalOperator.ExpressionEvaluator val;
+
+  public SqrtLongEvaluator(EvalOperator.ExpressionEvaluator val) {
+    this.val = val;
+  }
+
+  @Override
+  public Block eval(Page page) {
+    Block valUncastBlock = val.eval(page);
+    if (valUncastBlock.areAllValuesNull()) {
+      return Block.constantNullBlock(page.getPositionCount());
+    }
+    LongBlock valBlock = (LongBlock) valUncastBlock;
+    LongVector valVector = valBlock.asVector();
+    if (valVector == null) {
+      return eval(page.getPositionCount(), valBlock);
+    }
+    return eval(page.getPositionCount(), valVector).asBlock();
+  }
+
+  public DoubleBlock eval(int positionCount, LongBlock valBlock) {
+    DoubleBlock.Builder result = DoubleBlock.newBlockBuilder(positionCount);
+    position: for (int p = 0; p < positionCount; p++) {
+      if (valBlock.isNull(p) || valBlock.getValueCount(p) != 1) {
+        result.appendNull();
+        continue position;
+      }
+      result.appendDouble(Sqrt.process(valBlock.getLong(valBlock.getFirstValueIndex(p))));
+    }
+    return result.build();
+  }
+
+  public DoubleVector eval(int positionCount, LongVector valVector) {
+    DoubleVector.Builder result = DoubleVector.newVectorBuilder(positionCount);
+    position: for (int p = 0; p < positionCount; p++) {
+      result.appendDouble(Sqrt.process(valVector.getLong(p)));
+    }
+    return result.build();
+  }
+
+  @Override
+  public String toString() {
+    return "SqrtLongEvaluator[" + "val=" + val + "]";
+  }
+}

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

@@ -53,6 +53,7 @@ import org.elasticsearch.xpack.esql.expression.function.scalar.math.Pow;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Round;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Sin;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Sinh;
+import org.elasticsearch.xpack.esql.expression.function.scalar.math.Sqrt;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Tan;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Tanh;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Tau;
@@ -120,6 +121,7 @@ public class EsqlFunctionRegistry extends FunctionRegistry {
                 def(Round.class, Round::new, "round"),
                 def(Sin.class, Sin::new, "sin"),
                 def(Sinh.class, Sinh::new, "sinh"),
+                def(Sqrt.class, Sqrt::new, "sqrt"),
                 def(Tan.class, Tan::new, "tan"),
                 def(Tanh.class, Tanh::new, "tanh"),
                 def(Tau.class, Tau::new, "tau") },

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

@@ -0,0 +1,96 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.expression.function.scalar.math;
+
+import org.elasticsearch.compute.ann.Evaluator;
+import org.elasticsearch.compute.operator.EvalOperator;
+import org.elasticsearch.xpack.esql.expression.function.scalar.UnaryScalarFunction;
+import org.elasticsearch.xpack.esql.planner.Mappable;
+import org.elasticsearch.xpack.ql.expression.Expression;
+import org.elasticsearch.xpack.ql.tree.NodeInfo;
+import org.elasticsearch.xpack.ql.tree.Source;
+import org.elasticsearch.xpack.ql.type.DataType;
+import org.elasticsearch.xpack.ql.type.DataTypes;
+
+import java.util.List;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+import static org.elasticsearch.xpack.ql.expression.TypeResolutions.ParamOrdinal.DEFAULT;
+import static org.elasticsearch.xpack.ql.expression.TypeResolutions.isNumeric;
+
+public class Sqrt extends UnaryScalarFunction implements Mappable {
+    public Sqrt(Source source, Expression field) {
+        super(source, field);
+    }
+
+    @Override
+    public Supplier<EvalOperator.ExpressionEvaluator> toEvaluator(
+        Function<Expression, Supplier<EvalOperator.ExpressionEvaluator>> toEvaluator
+    ) {
+        Supplier<EvalOperator.ExpressionEvaluator> field = toEvaluator.apply(field());
+        var fieldType = field().dataType();
+        var eval = field.get();
+
+        if (fieldType == DataTypes.DOUBLE) {
+            return () -> new SqrtDoubleEvaluator(eval);
+        }
+        if (fieldType == DataTypes.INTEGER) {
+            return () -> new SqrtIntEvaluator(eval);
+        }
+        if (fieldType == DataTypes.LONG) {
+            return () -> new SqrtLongEvaluator(eval);
+        }
+
+        throw new UnsupportedOperationException("Unsupported type " + fieldType);
+    }
+
+    @Evaluator(extraName = "Double")
+    static double process(double val) {
+        return Math.sqrt(val);
+    }
+
+    @Evaluator(extraName = "Long")
+    static double process(long val) {
+        return Math.sqrt(val);
+    }
+
+    @Evaluator(extraName = "Int")
+    static double process(int val) {
+        return Math.sqrt(val);
+    }
+
+    @Override
+    public final Expression replaceChildren(List<Expression> newChildren) {
+        return new Sqrt(source(), newChildren.get(0));
+    }
+
+    @Override
+    protected NodeInfo<? extends Expression> info() {
+        return NodeInfo.create(this, Sqrt::new, field());
+    }
+
+    @Override
+    public DataType dataType() {
+        return DataTypes.DOUBLE;
+    }
+
+    @Override
+    public Object fold() {
+        return Mappable.super.fold();
+    }
+
+    @Override
+    protected TypeResolution resolveType() {
+        if (childrenResolved() == false) {
+            return new TypeResolution("Unresolved children");
+        }
+
+        return isNumeric(field, sourceText(), DEFAULT);
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java

@@ -63,6 +63,7 @@ import org.elasticsearch.xpack.esql.expression.function.scalar.math.Pow;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Round;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Sin;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Sinh;
+import org.elasticsearch.xpack.esql.expression.function.scalar.math.Sqrt;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Tan;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Tanh;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Tau;
@@ -298,6 +299,7 @@ public final class PlanNamedTypes {
             of(ESQL_UNARY_SCLR_CLS, Log10.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar),
             of(ESQL_UNARY_SCLR_CLS, Sin.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar),
             of(ESQL_UNARY_SCLR_CLS, Sinh.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar),
+            of(ESQL_UNARY_SCLR_CLS, Sqrt.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar),
             of(ESQL_UNARY_SCLR_CLS, Tan.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar),
             of(ESQL_UNARY_SCLR_CLS, Tanh.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar),
             of(ESQL_UNARY_SCLR_CLS, ToBoolean.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar),
@@ -1023,6 +1025,7 @@ public final class PlanNamedTypes {
         entry(name(Log10.class), Log10::new),
         entry(name(Sin.class), Sin::new),
         entry(name(Sinh.class), Sinh::new),
+        entry(name(Sqrt.class), Sqrt::new),
         entry(name(Tan.class), Tan::new),
         entry(name(Tanh.class), Tanh::new),
         entry(name(ToBoolean.class), ToBoolean::new),

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

@@ -0,0 +1,62 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.expression.function.scalar.math;
+
+import com.carrotsearch.randomizedtesting.annotations.Name;
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
+import org.elasticsearch.xpack.esql.expression.function.scalar.AbstractScalarFunctionTestCase;
+import org.elasticsearch.xpack.ql.expression.Expression;
+import org.elasticsearch.xpack.ql.tree.Source;
+import org.elasticsearch.xpack.ql.type.DataType;
+import org.elasticsearch.xpack.ql.type.DataTypes;
+import org.hamcrest.Matcher;
+
+import java.util.List;
+import java.util.function.Supplier;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class SqrtTests extends AbstractScalarFunctionTestCase {
+    public SqrtTests(@Name("TestCase") Supplier<TestCase> testCaseSupplier) {
+        this.testCase = testCaseSupplier.get();
+    }
+
+    @ParametersFactory
+    public static Iterable<Object[]> parameters() {
+        return parameterSuppliersFromTypedData(List.of(new TestCaseSupplier("Sqrt of Double", () -> {
+            // TODO: include larger values here
+            double arg = randomDouble();
+            return new TestCase(
+                List.of(new TypedData(arg, DataTypes.DOUBLE, "arg")),
+                "SqrtDoubleEvaluator[val=Attribute[channel=0]]",
+                DataTypes.DOUBLE,
+                equalTo(Math.sqrt(arg))
+            );
+        })));
+    }
+
+    private Matcher<Object> resultsMatcher(List<TypedData> typedData) {
+        return equalTo(Math.sqrt((Double) typedData.get(0).data()));
+    }
+
+    @Override
+    protected Expression build(Source source, List<Expression> args) {
+        return new Sqrt(source, args.get(0));
+    }
+
+    @Override
+    protected List<ArgumentSpec> argSpec() {
+        return List.of(required(numerics()));
+    }
+
+    @Override
+    protected DataType expectedType(List<DataType> argTypes) {
+        return DataTypes.DOUBLE;
+    }
+}