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
   }
 }

docs/user/ppl/cmd/trendline.rst

@@ -15,25 +15,24 @@ Description
 
 Syntax
 ============
-`TRENDLINE [sort <[+|-] sort-field>] SMA(number-of-datapoints, field) [AS alias] [SMA(number-of-datapoints, field) [AS alias]]...`
+`TRENDLINE [sort <[+|-] sort-field>] <trendline-type>(number-of-datapoints, field) [AS alias] [<trendline-type>(number-of-datapoints, field) [AS alias]]...`
 
 * [+|-]: optional. The plus [+] stands for ascending order and NULL/MISSING first and a minus [-] stands for descending order and NULL/MISSING last. **Default:** ascending order and NULL/MISSING first.
 * sort-field: mandatory when sorting is used. The field used to sort.
+* trendline-type: mandatory. The type of algorithm being used for the calculation, only SMA (simple moving average) or WMA (weighted moving average) are supported at the moment.
 * number-of-datapoints: mandatory. The number of datapoints to calculate the moving average (must be greater than zero).
 * field: mandatory. The name of the field the moving average should be calculated for.
-* alias: optional. The name of the resulting column containing the moving average (defaults to the field name with "_trendline").
+* alias: optional. The name of the resulting column containing the moving average (defaults to the field name with "_<trendline-type>_trendline").
 
-At the moment only the Simple Moving Average (SMA) type is supported.
-
-It is calculated like
+In the case of Simple Moving Average - SMA, result will be calculated as per the below formula.
 
     f[i]: The value of field 'f' in the i-th data-point
     n: The number of data-points in the moving window (period)
     t: The current time index
 
     SMA(t) = (1/n) * Σ(f[i]), where i = t-n+1 to t
 
-Example 1: Calculate the moving average on one field.
+Example 1: Calculate the simple moving average on one field.
 =====================================================
 
 The example shows how to calculate the moving average on one field.
@@ -52,7 +51,7 @@ PPL query::
     +------+
 
 
-Example 2: Calculate the moving average on multiple fields.
+Example 2: Calculate the simple moving average on multiple fields.
 ===========================================================
 
 The example shows how to calculate the moving average on multiple fields.
@@ -70,10 +69,10 @@ PPL query::
     | 15.5 | 30.5      |
     +------+-----------+
 
-Example 4: Calculate the moving average on one field without specifying an alias.
+Example 3: Calculate the simple moving average on one field without specifying an alias.
 =================================================================================
 
-The example shows how to calculate the moving average on one field.
+The example shows how to calculate the moving average on one field without specifying an alias.
 
 PPL query::
 
@@ -88,3 +87,69 @@ PPL query::
     | 15.5                     |
     +--------------------------+
 
+
+
+In the case of Weighted Moving Average - WMA, result will be calculated as per the below formula.
+
+    f[i]: The value of field 'f' in the i-th data point
+    n: The number of data points in the moving window (period)
+    t: The current time index
+    w[i]: The weight of the i-th data point, increasing by one per step to prioritize recent points.
+
+    WMA(t) = ( Σ from i=t−n+1 to t of (w[i] * f[i]) ) / ( Σ from i=t−n+1 to t of w[i] )
+
+Example 1: Calculate the weighted moving average on one field.
+=====================================================
+
+The example shows how to calculate the weighted moving average on one field.
+
+PPL query::
+
+    os> source=accounts | trendline wma(2, account_number) as an | fields an;
+    fetched rows / total rows = 4/4
+    +--------------------+
+    | an                 |
+    |--------------------|
+    | null               |
+    | 4.333333333333333  |
+    | 10.666666666666666 |
+    | 16.333333333333332 |
+    +--------------------+
+
+Example 2: Calculate the weighted moving average on multiple fields.
+===========================================================
+
+The example shows how to calculate the weighted moving average on multiple fields.
+
+PPL query::
+
+    os> source=accounts | trendline wma(2, account_number) as an sma(2, age) as age_trend | fields an, age_trend ;
+    fetched rows / total rows = 4/4
+    +--------------------+-----------+
+    | an                 | age_trend |
+    |--------------------+-----------|
+    | null               | null      |
+    | 4.333333333333333  | 34.0      |
+    | 10.666666666666666 | 32.0      |
+    | 16.333333333333332 | 30.5      |
+    +--------------------+-----------+
+
+
+Example 3: Calculate the weighted moving average on one field without specifying an alias.
+=================================================================================
+
+The example shows how to calculate the weighted moving average on one field without specifying an alias.
+
+PPL query::
+
+    os> source=accounts | trendline wma(2, account_number) | fields account_number_trendline;
+    fetched rows / total rows = 4/4
+    +--------------------------+
+    | account_number_trendline |
+    |--------------------------|
+    | null                     |
+    | 4.333333333333333        |
+    | 10.666666666666666       |
+    | 16.333333333333332       |
+    +--------------------------+
+

integ-test/src/test/java/org/opensearch/sql/ppl/TrendlineCommandIT.java

@@ -75,4 +75,68 @@ public void testTrendlineWithSort() throws IOException {
                 TEST_INDEX_BANK));
     verifyDataRows(result, rows(new Object[] {null}), rows(44313.0), rows(39882.5));
   }
+
+  @Test
+  public void testTrendlineWma() throws IOException {
+    final JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | where balance > 39000 | sort balance | trendline wma(2, balance) as"
+                    + " balance_trend | fields balance_trend",
+                TEST_INDEX_BANK));
+    verifyDataRows(
+        result, rows(new Object[] {null}), rows(45570.666666666664), rows(40101.666666666664));
+  }
+
+  @Test
+  public void testTrendlineMultipleFieldsWma() throws IOException {
+    final JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | where balance > 39000 | sort balance | trendline wma(2, balance) as"
+                    + " balance_trend wma(2, account_number) as account_number_trend | fields"
+                    + " balance_trend, account_number_trend",
+                TEST_INDEX_BANK));
+    verifyDataRows(
+        result,
+        rows(null, null),
+        rows(40101.666666666664, 16.999999999999996),
+        rows(45570.666666666664, 29.666666666666664));
+  }
+
+  @Test
+  public void testTrendlineOverwritesExistingFieldWma() throws IOException {
+    final JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | where balance > 39000 | sort balance | trendline wma(2, balance) as"
+                    + " age | fields age",
+                TEST_INDEX_BANK));
+    verifyDataRows(
+        result, rows(new Object[] {null}), rows(40101.666666666664), rows(45570.666666666664));
+  }
+
+  @Test
+  public void testTrendlineNoAliasWma() throws IOException {
+    final JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | where balance > 39000 | sort balance | trendline wma(2, balance) |"
+                    + " fields balance_trendline",
+                TEST_INDEX_BANK));
+    verifyDataRows(
+        result, rows(new Object[] {null}), rows(40101.666666666664), rows(45570.666666666664));
+  }
+
+  @Test
+  public void testTrendlineWithSortWma() throws IOException {
+    final JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | where balance > 39000 | trendline sort balance wma(2, balance) |"
+                    + " fields balance_trendline",
+                TEST_INDEX_BANK));
+    verifyDataRows(
+        result, rows(new Object[] {null}), rows(40101.666666666664), rows(45570.666666666664));
+  }
 }

ppl/src/main/antlr/OpenSearchPPLLexer.g4

@@ -60,6 +60,7 @@ NUM:                                'NUM';
 
 // TRENDLINE KEYWORDS
 SMA:                                'SMA';
+WMA:                                'WMA';
 
 // ARGUMENT KEYWORDS
 KEEPEMPTY:                          'KEEPEMPTY';

ppl/src/main/antlr/OpenSearchPPLParser.g4

@@ -156,6 +156,7 @@ trendlineClause
 
 trendlineType
    : SMA
+   | WMA
    ;
 
 kmeansCommand

ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java

@@ -89,13 +89,12 @@ public Trendline.TrendlineComputation visitTrendlineClause(
     }
 
     final Field dataField = (Field) this.visitFieldExpression(ctx.field);
+    final Trendline.TrendlineType computationType =
+        Trendline.TrendlineType.valueOf(ctx.trendlineType().getText().toUpperCase(Locale.ROOT));
     final String alias =
         ctx.alias != null
             ? ctx.alias.getText()
-            : dataField.getChild().get(0).toString() + "_trendline";
-
-    final Trendline.TrendlineType computationType =
-        Trendline.TrendlineType.valueOf(ctx.trendlineType().getText().toUpperCase(Locale.ROOT));
+            : dataField.getChild().get(0).toString() + "_" + computationType.name() + "_trendline";
     return new Trendline.TrendlineComputation(
         numberOfDataPoints, dataField, alias, computationType);
   }