Add Trendline parser test

Signed-off-by: James Duong <[email protected]>
Bit-Quill · Oct 22, 2024 · ef30957 · ef30957
1 parent d006675
commit ef30957
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 2 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java b/core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java
@@ -59,6 +59,7 @@
 import org.opensearch.sql.ast.tree.Sort;
 import org.opensearch.sql.ast.tree.Sort.SortOption;
 import org.opensearch.sql.ast.tree.TableFunction;
+import org.opensearch.sql.ast.tree.Trendline;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.ast.tree.Values;
 
@@ -463,6 +464,18 @@ public static Limit limit(UnresolvedPlan input, Integer limit, Integer offset) {
     return new Limit(limit, offset).attach(input);
   }
 
+  public static Trendline trendline(UnresolvedPlan input, Trendline.TrendlineComputation... computations) {
+    return new Trendline(Arrays.asList(computations)).attach(input);
+  }
+
+  public static Trendline.TrendlineComputation computation(
+      Integer numDataPoints,
+      UnresolvedExpression dataField,
+      String alias,
+      String type) {
+    return new Trendline.TrendlineComputation(numDataPoints, dataField, alias, type);
+  }
+
   public static Parse parse(
       UnresolvedPlan input,
       ParseMethod parseMethod,

diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Trendline.java b/core/src/main/java/org/opensearch/sql/ast/tree/Trendline.java
@@ -7,6 +7,8 @@
 
 import com.google.common.collect.ImmutableList;
 import java.util.List;
+import java.util.Locale;
+
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
@@ -25,7 +27,7 @@ public class Trendline extends UnresolvedPlan {
   private final List<UnresolvedExpression> computations;
 
   @Override
-  public UnresolvedPlan attach(UnresolvedPlan child) {
+  public Trendline attach(UnresolvedPlan child) {
     this.child = child;
     return this;
   }
@@ -56,7 +58,7 @@ public TrendlineComputation(
       this.numberOfDataPoints = numberOfDataPoints;
       this.dataField = dataField;
       this.alias = alias;
-      this.computationType = Trendline.TrendlineType.valueOf(computationType.toUpperCase());
+      this.computationType = Trendline.TrendlineType.valueOf(computationType.toUpperCase(Locale.ROOT));
     }
 
     @Override

diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstBuilderTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstBuilderTest.java
@@ -13,6 +13,7 @@
 import static org.opensearch.sql.ast.dsl.AstDSL.argument;
 import static org.opensearch.sql.ast.dsl.AstDSL.booleanLiteral;
 import static org.opensearch.sql.ast.dsl.AstDSL.compare;
+import static org.opensearch.sql.ast.dsl.AstDSL.computation;
 import static org.opensearch.sql.ast.dsl.AstDSL.dedupe;
 import static org.opensearch.sql.ast.dsl.AstDSL.defaultDedupArgs;
 import static org.opensearch.sql.ast.dsl.AstDSL.defaultFieldsArgs;
@@ -38,6 +39,7 @@
 import static org.opensearch.sql.ast.dsl.AstDSL.span;
 import static org.opensearch.sql.ast.dsl.AstDSL.stringLiteral;
 import static org.opensearch.sql.ast.dsl.AstDSL.tableFunction;
+import static org.opensearch.sql.ast.dsl.AstDSL.trendline;
 import static org.opensearch.sql.ast.dsl.AstDSL.unresolvedArg;
 import static org.opensearch.sql.utils.SystemIndexUtils.DATASOURCES_TABLE_NAME;
 import static org.opensearch.sql.utils.SystemIndexUtils.mappingTable;
@@ -660,6 +662,17 @@ public void testMLCommand() {
                 .build()));
   }
 
+  @Test
+  public void testTrendline() {
+    assertEqual(
+        "source=t | trendline sma(5, test_field) as test_field_alias sma(1, test_field_2) as test_field_alias_2",
+        trendline(
+            relation("t"),
+            computation(5, field("test_field"), "test_field_alias", "sma"),
+            computation(1, field("test_field)2"), "test_field_alias_2", "sma")
+    ));
+  }
+
   @Test
   public void testDescribeCommand() {
     assertEqual("describe t", relation(mappingTable("t")));