ESQL: Add MV_PSERIES_WEIGHTED_SUM for score calculations used by security solution

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestsDataLoader.java

@@ -54,6 +54,7 @@ public class CsvTestsDataLoader {
     private static final TestsDataset HOSTS = new TestsDataset("hosts", "mapping-hosts.json", "hosts.csv");
     private static final TestsDataset APPS = new TestsDataset("apps", "mapping-apps.json", "apps.csv");
     private static final TestsDataset LANGUAGES = new TestsDataset("languages", "mapping-languages.json", "languages.csv");
+    private static final TestsDataset ALERTS = new TestsDataset("alerts", "mapping-alerts.json", "alerts.csv");
     private static final TestsDataset UL_LOGS = new TestsDataset("ul_logs", "mapping-ul_logs.json", "ul_logs.csv");
     private static final TestsDataset SAMPLE_DATA = new TestsDataset("sample_data", "mapping-sample_data.json", "sample_data.csv");
     private static final TestsDataset CLIENT_IPS = new TestsDataset("clientips", "mapping-clientips.json", "clientips.csv");
@@ -94,6 +95,7 @@ public class CsvTestsDataLoader {
         Map.entry(LANGUAGES.indexName, LANGUAGES),
         Map.entry(UL_LOGS.indexName, UL_LOGS),
         Map.entry(SAMPLE_DATA.indexName, SAMPLE_DATA),
+        Map.entry(ALERTS.indexName, ALERTS),
         Map.entry(CLIENT_IPS.indexName, CLIENT_IPS),
         Map.entry(CLIENT_CIDR.indexName, CLIENT_CIDR),
         Map.entry(AGES.indexName, AGES),

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

@@ -0,0 +1,11 @@
+host.name:keyword,kibana.alert.risk_score:double
+test-host-1,21.0
+test-host-2,17.0
+test-host-2,23.0
+test-host-1,45.0
+test-host-2,12.0
+test-host-2,16.0
+test-host-1,21.0
+test-host-1,70.0
+test-host-1,21.0
+test-host-2,5.0

x-pack/plugin/esql/qa/testFixtures/src/main/resources/mapping-alerts.json

@@ -0,0 +1,10 @@
+{
+  "properties": {
+    "host.name": {
+      "type": "keyword"
+    },
+    "kibana.alert.risk_score": {
+      "type": "double"
+    }
+  }
+}

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

@@ -0,0 +1,24 @@
+multivalue
+ROW data = [21.0, 45.0, 21.0, 70.0, 21.0]
+| EVAL sorted = MV_SORT(data, "desc")
+| EVAL score = MV_RIEMANN_ZETA(sorted, 1.5)
+| EVAL normalized_score = ROUND(100 * score / 261.2, 2)
+| KEEP normalized_score, score;
+
+normalized_score:double|score:double
+36.16                  |94.45465156212452
+;
+
+
+aggregation
+from alerts
+| where host.name is not null
+| SORT host.name, kibana.alert.risk_score
+| stats score = RIEMANN_ZETA(kibana.alert.risk_score, 1.5) by host.name
+| EVAL normalized_score = ROUND(100 * score / 261.2, 2)
+| keep host.name, normalized_score, score;
+
+host.name:keyword|normalized_score:double|score:double
+test-host-1      |36.16                  |94.45465156212452
+test-host-2      |36.16                  |94.45465156212452
+;

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

@@ -21,6 +21,7 @@
 import org.elasticsearch.xpack.esql.expression.function.aggregate.MedianAbsoluteDeviation;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Min;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Percentile;
+import org.elasticsearch.xpack.esql.expression.function.aggregate.RiemannZeta;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialCentroid;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Sum;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Values;
@@ -87,6 +88,7 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvSlice;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvSort;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvSum;
+import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvRiemannZeta;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvZip;
 import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
 import org.elasticsearch.xpack.esql.expression.function.scalar.spatial.SpatialContains;
@@ -189,6 +191,7 @@ private FunctionDefinition[][] functions() {
                 def(MedianAbsoluteDeviation.class, MedianAbsoluteDeviation::new, "median_absolute_deviation"),
                 def(Min.class, Min::new, "min"),
                 def(Percentile.class, Percentile::new, "percentile"),
+                def(RiemannZeta.class, RiemannZeta::new, "riemann_zeta"),
                 def(Sum.class, Sum::new, "sum"),
                 def(Values.class, Values::new, "values") },
             // math
@@ -291,7 +294,8 @@ private FunctionDefinition[][] functions() {
                 def(MvSlice.class, MvSlice::new, "mv_slice"),
                 def(MvZip.class, MvZip::new, "mv_zip"),
                 def(MvSum.class, MvSum::new, "mv_sum"),
-                def(Split.class, Split::new, "split") } };
+                def(Split.class, Split::new, "split"),
+                def(MvRiemannZeta.class, MvRiemannZeta::new, "mv_riemann_zeta") } };
     }
 
     @Override

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

@@ -0,0 +1,86 @@
+/*
+ * 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.aggregate;
+
+import org.elasticsearch.compute.aggregation.AggregatorFunctionSupplier;
+import org.elasticsearch.compute.aggregation.SumDoubleAggregatorFunctionSupplier;
+import org.elasticsearch.compute.aggregation.SumIntAggregatorFunctionSupplier;
+import org.elasticsearch.compute.aggregation.SumLongAggregatorFunctionSupplier;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+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;
+import org.elasticsearch.xpack.esql.core.type.DataTypes;
+import org.elasticsearch.xpack.esql.core.util.StringUtils;
+import org.elasticsearch.xpack.esql.expression.SurrogateExpression;
+import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
+import org.elasticsearch.xpack.esql.expression.function.Param;
+import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvSum;
+import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Mul;
+
+import java.util.List;
+
+import static org.elasticsearch.xpack.esql.core.type.DataTypes.*;
+
+/**
+ * Sum all values of a field in matching documents.
+ */
+public class RiemannZeta extends NumericAggregate implements SurrogateExpression {
+    private final Expression p;
+
+    @FunctionInfo(returnType = "double", description = "todo", isAggregation = true)
+    public RiemannZeta(
+        Source source,
+        @Param(name = "number", type = { "double" }) Expression field,
+        @Param(name = "p", type = { "double" }) Expression p
+    ) {
+        super(source, field, List.of(p));
+        this.p = p;
+    }
+
+    @Override
+    protected NodeInfo<RiemannZeta> info() {
+        return NodeInfo.create(this, RiemannZeta::new, field(), p);
+    }
+
+    @Override
+    public RiemannZeta replaceChildren(List<Expression> newChildren) {
+        return new RiemannZeta(source(), newChildren.get(0), newChildren.get(1));
+    }
+
+    @Override
+    public DataType dataType() {
+        return field().dataType();
+    }
+
+    @Override
+    protected AggregatorFunctionSupplier longSupplier(List<Integer> inputChannels) {
+        return new SumLongAggregatorFunctionSupplier(inputChannels);
+    }
+
+    @Override
+    protected AggregatorFunctionSupplier intSupplier(List<Integer> inputChannels) {
+        return new SumIntAggregatorFunctionSupplier(inputChannels);
+    }
+
+    @Override
+    protected AggregatorFunctionSupplier doubleSupplier(List<Integer> inputChannels) {
+        return new SumDoubleAggregatorFunctionSupplier(inputChannels);
+    }
+
+    @Override
+    public Expression surrogate() {
+        var s = source();
+        var field = field();
+
+        // SUM(const) is equivalent to MV_SUM(const)*COUNT(*).
+        return field.foldable()
+            ? new Mul(s, new MvSum(s, field), new Count(s, new Literal(s, StringUtils.WILDCARD, DataTypes.KEYWORD)))
+            : null;
+    }
+}