Feature: PPL command - WMA Trendline

core/src/main/java/org/opensearch/sql/ast/tree/Trendline.java

@@ -66,6 +66,7 @@ public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
   }
 
   public enum TrendlineType {
-    SMA
+    SMA,
+    WMA
   }
 }

core/src/main/java/org/opensearch/sql/planner/physical/TrendlineOperator.java

@@ -11,6 +11,7 @@
 import com.google.common.collect.ImmutableMap.Builder;
 import java.time.Instant;
 import java.time.LocalTime;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -21,6 +22,7 @@
 import lombok.ToString;
 import org.apache.commons.lang3.tuple.Pair;
 import org.opensearch.sql.ast.tree.Trendline;
+import org.opensearch.sql.data.model.ExprDoubleValue;
 import org.opensearch.sql.data.model.ExprTupleValue;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
@@ -106,9 +108,11 @@ private Map<String, ExprValue> consumeInputTuple(ExprValue inputValue) {
 
   private static TrendlineAccumulator createAccumulator(
       Pair<Trendline.TrendlineComputation, ExprCoreType> computation) {
-    // Add a switch statement based on computation type to choose the accumulator when more
-    // types of computations are supported.
-    return new SimpleMovingAverageAccumulator(computation.getKey(), computation.getValue());
+    return switch (computation.getKey().getComputationType()) {
+      case SMA -> new SimpleMovingAverageAccumulator(computation.getKey(), computation.getValue());
+      case WMA -> new WeightedMovingAverageAccumulator(
+          computation.getKey(), computation.getValue());
+    };
   }
 
   /** Maintains stateful information for calculating the trendline. */
@@ -187,6 +191,103 @@ public ExprValue calculate() {
     }
   }
 
+  private static class WeightedMovingAverageAccumulator implements TrendlineAccumulator {
+    private final LiteralExpression dataPointsNeeded;
+    private final ArrayList<ExprValue> receivedValues;
+    private final WmaTrendlineEvaluator evaluator;
+
+    public WeightedMovingAverageAccumulator(
+        Trendline.TrendlineComputation computation, ExprCoreType type) {
+      this.dataPointsNeeded = DSL.literal(computation.getNumberOfDataPoints().doubleValue());
+      this.receivedValues = new ArrayList<>(computation.getNumberOfDataPoints());
+      this.evaluator = getEvaluator(type);
+    }
+
+    static WmaTrendlineEvaluator getEvaluator(ExprCoreType type) {
+      return switch (type) {
+        case DOUBLE -> NumericWmaEvaluator.INSTANCE;
+        case DATE, TIMESTAMP -> TimeStampWmaEvaluator.INSTANCE;
+        case TIME -> TimeWmaEvaluator.INSTANCE;
+        default -> throw new IllegalArgumentException(
+            String.format("Invalid type %s used for weighted moving average.", type.typeName()));
+      };
+    }
+
+    @Override
+    public void accumulate(ExprValue value) {
+      receivedValues.add(value);
+      if (receivedValues.size() > dataPointsNeeded.valueOf().integerValue()) {
+        receivedValues.removeFirst();
+      }
+    }
+
+    @Override
+    public ExprValue calculate() {
+      if (receivedValues.size() < dataPointsNeeded.valueOf().integerValue()) {
+        return null;
+      } else if (dataPointsNeeded.valueOf().integerValue() == 1) {
+        return receivedValues.getFirst();
+      }
+      return evaluator.evaluate(receivedValues);
+    }
+
+    private static class NumericWmaEvaluator implements WmaTrendlineEvaluator {
+
+      private static final NumericWmaEvaluator INSTANCE = new NumericWmaEvaluator();
+
+      @Override
+      public ExprValue evaluate(ArrayList<ExprValue> receivedValues) {
+        double sum = 0D;
+        int totalWeight = (receivedValues.size() * (receivedValues.size() + 1)) / 2;
+        for (int i = 0; i < receivedValues.size(); i++) {
+          sum += receivedValues.get(i).doubleValue() * ((i + 1D) / totalWeight);
+        }
+        return new ExprDoubleValue(sum);
+      }
+    }
+
+    private static class TimeStampWmaEvaluator implements WmaTrendlineEvaluator {
+
+      private static final TimeStampWmaEvaluator INSTANCE = new TimeStampWmaEvaluator();
+
+      @Override
+      public ExprValue evaluate(ArrayList<ExprValue> receivedValues) {
+        long sum = 0L;
+        int totalWeight = (receivedValues.size() * (receivedValues.size() + 1)) / 2;
+        for (int i = 0; i < receivedValues.size(); i++) {
+          sum +=
+              (long)
+                  (receivedValues.get(i).timestampValue().toEpochMilli()
+                      * ((i + 1D) / totalWeight));
+        }
+
+        return ExprValueUtils.timestampValue(Instant.ofEpochMilli((sum)));
+      }
+    }
+
+    private static class TimeWmaEvaluator implements WmaTrendlineEvaluator {
+
+      private static final TimeWmaEvaluator INSTANCE = new TimeWmaEvaluator();
+
+      @Override
+      public ExprValue evaluate(ArrayList<ExprValue> receivedValues) {
+        long sum = 0L;
+        int totalWeight = (receivedValues.size() * (receivedValues.size() + 1)) / 2;
+        for (int i = 0; i < receivedValues.size(); i++) {
+          sum +=
+              (long)
+                  (MILLIS.between(LocalTime.MIN, receivedValues.get(i).timeValue())
+                      * ((i + 1D) / totalWeight));
+        }
+        return ExprValueUtils.timeValue(LocalTime.MIN.plus(sum, MILLIS));
+      }
+    }
+
+    private interface WmaTrendlineEvaluator {
+      ExprValue evaluate(ArrayList<ExprValue> receivedValues);
+    }
+  }
+
   private interface ArithmeticEvaluator {
     Expression calculateFirstTotal(List<ExprValue> dataPoints);