Implement PredictLinear function

Author: pabloem

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/DerivDoubleGroupingAggregatorFunction.java

@@ -6,19 +6,25 @@
  */
 package org.elasticsearch.compute.aggregation;
 
-class SimpleLinearRegressionWithTimeseries {
+public class SimpleLinearRegressionWithTimeseries {
     long count;
     double sumVal;
     long sumTs;
     double sumTsVal;
     long sumTsSq;
+    long maxTs;
+
+    public interface SimpleLinearModelFunction {
+        double predict(SimpleLinearRegressionWithTimeseries model);
+    }
 
     SimpleLinearRegressionWithTimeseries() {
         this.count = 0;
         this.sumVal = 0.0;
         this.sumTs = 0;
         this.sumTsVal = 0.0;
         this.sumTsSq = 0;
+        this.maxTs = Long.MIN_VALUE;
     }
 
     void add(long ts, double val) {
@@ -27,9 +33,16 @@ void add(long ts, double val) {
         sumTs += ts;
         sumTsVal += ts * val;
         sumTsSq += ts * ts;
+        if (ts > maxTs) {
+            maxTs = ts;
+        }
+    }
+
+    public double lastTimestamp() {
+        return maxTs;
     }
 
-    double slope() {
+    public double slope() {
         if (count <= 1) {
             return Double.NaN;
         }
@@ -41,7 +54,7 @@ void add(long ts, double val) {
         return numerator / denominator;
     }
 
-    double intercept() {
+    public double intercept() {
         if (count == 0) {
             return 0.0; // or handle as needed
         }

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/DerivIntGroupingAggregatorFunction.java

@@ -40,15 +40,23 @@ public class Deriv$Type$GroupingAggregatorFunction implements GroupingAggregator
     private final List<Integer> channels;
     private final DriverContext driverContext;
     private ObjectArray<SimpleLinearRegressionWithTimeseries> states;
+    private SimpleLinearRegressionWithTimeseries.SimpleLinearModelFunction fn;
 
-    public Deriv$Type$GroupingAggregatorFunction(List<Integer> channels, DriverContext driverContext) {
+    public Deriv$Type$GroupingAggregatorFunction(List<Integer> channels, DriverContext driverContext, SimpleLinearRegressionWithTimeseries.SimpleLinearModelFunction fn) {
         this.states = driverContext.bigArrays().newObjectArray(256);
         this.channels = channels;
         this.driverContext = driverContext;
+        this.fn = fn;
     }
 
     public static class Supplier implements AggregatorFunctionSupplier {
 
+        SimpleLinearRegressionWithTimeseries.SimpleLinearModelFunction fn;
+
+        public Supplier(SimpleLinearRegressionWithTimeseries.SimpleLinearModelFunction fn) {
+            this.fn = fn;
+        }
+
         @Override
         public List<IntermediateStateDesc> nonGroupingIntermediateStateDesc() {
             throw new UnsupportedOperationException("DerivGroupingAggregatorFunction does not support non-grouping aggregation");
@@ -66,7 +74,7 @@ public class Deriv$Type$GroupingAggregatorFunction implements GroupingAggregator
 
         @Override
         public GroupingAggregatorFunction groupingAggregator(DriverContext driverContext, List<Integer> channels) {
-            return new Deriv$Type$GroupingAggregatorFunction(channels, driverContext);
+            return new Deriv$Type$GroupingAggregatorFunction(channels, driverContext, fn);
         }
 
         @Override
@@ -103,7 +111,7 @@ public class Deriv$Type$GroupingAggregatorFunction implements GroupingAggregator
                         state = new SimpleLinearRegressionWithTimeseries();
                         states.set(groupId, state);
                     }
-                    state.add(ts, (double) vValue);  // TODO - value needs to be converted to double
+                    state.add(ts, (double) vValue);
                 }
             }
 
@@ -283,8 +291,7 @@ public class Deriv$Type$GroupingAggregatorFunction implements GroupingAggregator
                 if (state == null) {
                     resultBuilder.appendNull();
                 } else {
-                    double deriv = state.slope();
-                    resultBuilder.appendDouble(deriv);
+                    resultBuilder.appendDouble(fn.predict(state));
                 }
             }
             blocks[offset] = resultBuilder.build();

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/DerivLongGroupingAggregatorFunction.java

@@ -763,3 +763,57 @@ max_deriv:double | max_rate:double | time_bucket:datetime     | cluster:keyword
 0.0              | 6.980661        | 2024-05-10T00:20:00.000Z | prod
 
 ;
+
+predict_linear_with_long
+required_capability: ts_command_v0
+
+TS k8s
+| STATS predicted_cost = values(round(predict_linear(to_double(network.total_cost), 10), 5)) BY time_bucket = bucket(@timestamp,5minute), cluster
+| KEEP predicted_cost, time_bucket, cluster
+| SORT cluster, time_bucket
+| LIMIT 5
+;
+
+predicted_cost:double                                                  | time_bucket:datetime     | cluster:keyword
+[-6.602588168740215E14, -8.877497236448354E15, 2.2311938303322465E15]  | 2024-05-10T00:00:00.000Z | prod
+[-8.155468708896549E14, -9.51830721515923E14, -1.145502894094415E15]   | 2024-05-10T00:05:00.000Z | prod
+[5.463777793231174E14, -2.0359960651035278E14, -1.8388879144239432E15] | 2024-05-10T00:10:00.000Z | prod
+[-1.3995586522213082E15, 1.0825172194782352E15, 4.0520717676617645E15] | 2024-05-10T00:15:00.000Z | prod
+[-8.47847162491646E14, -9.533695970876474E14, 0.0]                     | 2024-05-10T00:20:00.000Z | prod
+;
+
+predict_linear_with_timedelta
+required_capability: ts_command_v0
+
+TS k8s
+| STATS predicted_cost = values(round(predict_linear(to_double(network.total_cost), 10 minutes), 5)) BY time_bucket = bucket(@timestamp,5minute), cluster
+| KEEP predicted_cost, time_bucket, cluster
+| SORT cluster, time_bucket
+| LIMIT 5
+;
+
+predicted_cost:double                                                  | time_bucket:datetime     | cluster:keyword
+[-6.602588168740214E14, -8.877497236448353E15, 2.2311938303322462E15]  | 2024-05-10T00:00:00.000Z | prod
+[-8.155468708896546E14, -9.51830721515923E14, -1.1455028940944148E15]  | 2024-05-10T00:05:00.000Z | prod
+[5.463777793231174E14, -2.0359960651035275E14, -1.8388879144239428E15] | 2024-05-10T00:10:00.000Z | prod
+[-1.399558652221308E15, 1.082517219478235E15, 4.0520717676617635E15]   | 2024-05-10T00:15:00.000Z | prod
+[-8.478471624916458E14, -9.533695970876471E14, 0.0]                    | 2024-05-10T00:20:00.000Z | prod
+;
+
+predict_linear_with_int_and_timedelta
+required_capability: ts_command_v0
+
+TS k8s
+| STATS predicted_cost_int = values(round(predict_linear(to_long(network.total_bytes_in), 10), 5)) BY time_bucket = bucket(@timestamp,5minute), cluster
+| KEEP predicted_cost_int, time_bucket, cluster
+| SORT cluster, time_bucket
+| LIMIT 5
+;
+
+predicted_cost_int:double                                                | time_bucket:datetime     | cluster:keyword
+[-4.773636702420472E16, -7.885984558091773E17, -2.2361970129287098E17]   | 2024-05-10T00:00:00.000Z | prod
+[-2.8368940048287676E16, -1.9090238696317244E16, -4.5500448741014824E16] | 2024-05-10T00:05:00.000Z | prod
+[-8.229547268699878E16, -4.188543503897301E16, -8.065656914863902E16]    | 2024-05-10T00:10:00.000Z | prod
+[-1.5332165779386784E17, -1.0946919353828525E17, -1.2534486684657331E17] | 2024-05-10T00:15:00.000Z | prod
+[8.510411223069503E17, -8.325264742678522E16, 0.0]                       | 2024-05-10T00:20:00.000Z | prod
+;

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/SimpleLinearRegressionWithTimeseries.java

@@ -1485,6 +1485,7 @@ public enum Cap {
         PERCENTILE_OVER_TIME,
         VARIANCE_STDDEV_OVER_TIME,
         TS_LINREG_DERIVATIVE,
+        TS_LINREG_PREDICT,
         /**
          * INLINE STATS fix incorrect prunning of null filtering
          * https://github.com/elastic/elasticsearch/pull/135011

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/X-DerivGroupingAggregatorFunction.java.st

@@ -44,6 +44,7 @@
 import org.elasticsearch.xpack.esql.expression.function.aggregate.MinOverTime;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Percentile;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.PercentileOverTime;
+import org.elasticsearch.xpack.esql.expression.function.aggregate.PredictLinear;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Present;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.PresentOverTime;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Rate;
@@ -539,6 +540,7 @@ private static FunctionDefinition[][] functions() {
                 defTS(Delta.class, bi(Delta::new), "delta"),
                 defTS(Increase.class, bi(Increase::new), "increase"),
                 def(Deriv.class, uni(Deriv::new), "deriv"),
+                def(PredictLinear.class, bi(PredictLinear::new), "predict_linear"),
                 def(MaxOverTime.class, uni(MaxOverTime::new), "max_over_time"),
                 def(MinOverTime.class, uni(MinOverTime::new), "min_over_time"),
                 def(SumOverTime.class, uni(SumOverTime::new), "sum_over_time"),

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

@@ -30,6 +30,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
             Increase.ENTRY,
             Delta.ENTRY,
             Deriv.ENTRY,
+            PredictLinear.ENTRY,
             Sample.ENTRY,
             SpatialCentroid.ENTRY,
             SpatialExtent.ENTRY,

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.compute.aggregation.DerivDoubleGroupingAggregatorFunction;
 import org.elasticsearch.compute.aggregation.DerivIntGroupingAggregatorFunction;
 import org.elasticsearch.compute.aggregation.DerivLongGroupingAggregatorFunction;
+import org.elasticsearch.compute.aggregation.SimpleLinearRegressionWithTimeseries;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.expression.UnresolvedAttribute;
@@ -99,10 +100,12 @@ public String getWriteableName() {
     @Override
     public AggregatorFunctionSupplier supplier() {
         final DataType type = field().dataType();
+
+        SimpleLinearRegressionWithTimeseries.SimpleLinearModelFunction fn = (SimpleLinearRegressionWithTimeseries model) -> model.slope();
         return switch (type) {
-            case INTEGER -> new DerivIntGroupingAggregatorFunction.Supplier();
-            case LONG -> new DerivLongGroupingAggregatorFunction.Supplier();
-            case DOUBLE -> new DerivDoubleGroupingAggregatorFunction.Supplier();
+            case INTEGER -> new DerivIntGroupingAggregatorFunction.Supplier(fn);
+            case LONG -> new DerivLongGroupingAggregatorFunction.Supplier(fn);
+            case DOUBLE -> new DerivDoubleGroupingAggregatorFunction.Supplier(fn);
             default -> throw new IllegalArgumentException("Unsupported data type for deriv aggregation: " + type);
         };
     }