Implement logging and explain for trendline

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

Hide file tree

Showing 4 changed files with 46 additions and 2 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/executor/Explain.java b/core/src/main/java/org/opensearch/sql/executor/Explain.java
@@ -8,12 +8,14 @@
 import com.google.common.collect.ImmutableMap;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.tuple.Pair;
 import org.opensearch.sql.ast.tree.Sort;
+import org.opensearch.sql.ast.tree.Trendline;
 import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse;
 import org.opensearch.sql.executor.ExecutionEngine.ExplainResponseNode;
 import org.opensearch.sql.expression.Expression;
@@ -31,6 +33,7 @@
 import org.opensearch.sql.planner.physical.RenameOperator;
 import org.opensearch.sql.planner.physical.SortOperator;
 import org.opensearch.sql.planner.physical.TakeOrderedOperator;
+import org.opensearch.sql.planner.physical.TrendlineOperator;
 import org.opensearch.sql.planner.physical.ValuesOperator;
 import org.opensearch.sql.planner.physical.WindowOperator;
 import org.opensearch.sql.storage.TableScanOperator;
@@ -211,6 +214,15 @@ public ExplainResponseNode visitNested(NestedOperator node, Object context) {
         explanNode -> explanNode.setDescription(ImmutableMap.of("nested", node.getFields())));
   }
 
+  @Override
+  public ExplainResponseNode visitTrendline(TrendlineOperator node, Object context) {
+    return explain(
+        node,
+        context,
+        explainNode -> explainNode.setDescription(
+            ImmutableMap.of("computations", describeTrendlineComputations(node.getComputations()))));
+  }
+
   protected ExplainResponseNode explain(
       PhysicalPlan node, Object context, Consumer<ExplainResponseNode> doExplain) {
     ExplainResponseNode explainNode = new ExplainResponseNode(getOperatorName(node));
@@ -245,4 +257,17 @@ private Map<String, Map<String, String>> describeSortList(
                         "sortOrder", p.getLeft().getSortOrder().toString(),
                         "nullOrder", p.getLeft().getNullOrder().toString())));
   }
+
+  private List<Map<String, String>> describeTrendlineComputations(
+      List<Trendline.TrendlineComputation> computations) {
+    return computations.stream()
+        .map(computation ->
+                ImmutableMap.of(
+                    "computationType", computation.getComputationType().name().toLowerCase(Locale.ROOT),
+                    "numberOfDataPoints", computation.getNumberOfDataPoints().toString(),
+                    "dataField", computation.getDataField().getChild().getFirst().toString(),
+                    "alias", computation.getAlias() != null ? computation.getAlias() : ""))
+        .collect(Collectors.toList());
+  }
+
 }
diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java b/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java
@@ -5,6 +5,7 @@
 
 package org.opensearch.sql.planner.physical;
 
+import org.opensearch.sql.ast.tree.Trendline;
 import org.opensearch.sql.storage.TableScanOperator;
 import org.opensearch.sql.storage.write.TableWriteOperator;
 
@@ -96,7 +97,7 @@ public R visitML(PhysicalPlan node, C context) {
     return visitNode(node, context);
   }
 
-  public R visitTrendline(PhysicalPlan node, C context) {
+  public R visitTrendline(TrendlineOperator node, C context) {
     return visitNode(node, context);
   }
 

diff --git a/...n/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java b/...n/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtector.java
@@ -189,7 +189,7 @@ public PhysicalPlan visitML(PhysicalPlan node, Object context) {
   }
 
   @Override
-  public PhysicalPlan visitTrendline(PhysicalPlan node, Object context) {
+  public PhysicalPlan visitTrendline(TrendlineOperator node, Object context) {
     TrendlineOperator trendlineOperator = (TrendlineOperator) node;
     return doProtect(
         new TrendlineOperator(

diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizer.java b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizer.java
@@ -9,6 +9,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
@@ -42,6 +43,7 @@
 import org.opensearch.sql.ast.tree.Rename;
 import org.opensearch.sql.ast.tree.Sort;
 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.common.utils.StringUtils;
 import org.opensearch.sql.planner.logical.LogicalAggregation;
@@ -220,6 +222,13 @@ public String visitHead(Head node, String context) {
     return StringUtils.format("%s | head %d", child, size);
   }
 
+  @Override
+  public String visitTrendline(Trendline node, String context) {
+    String child = node.getChild().getFirst().accept(this, context);
+    String computations = visitExpressionList(node.getComputations());
+    return StringUtils.format("%s | trendline %s", child, computations);
+  }
+
   private String visitFieldList(List<Field> fieldList) {
     return fieldList.stream().map(this::visitExpression).collect(Collectors.joining(","));
   }
@@ -316,5 +325,14 @@ public String visitAlias(Alias node, String context) {
       String expr = node.getDelegated().accept(this, context);
       return StringUtils.format("%s", expr);
     }
+
+    @Override
+    public String visitTrendlineComputation(Trendline.TrendlineComputation node, String context) {
+      final String dataField = node.getDataField().accept(this, context);
+      final String aliasOrEmpty = node.getAlias() != null ? " as " + node.getAlias() : "";
+      final String computationType = node.getComputationType().name().toLowerCase(Locale.ROOT);
+      return StringUtils.format("%s(%d, %s)%s",
+          computationType, node.getNumberOfDataPoints(), dataField, aliasOrEmpty);
+    }
   }
 }