add support for agg sample context

Signed-off-by: YANGDB <[email protected]>

ppl-spark-integration/src/main/java/org/opensearch/sql/ast/tree/Aggregation.java

@@ -6,6 +6,7 @@
 package org.opensearch.sql.ast.tree;
 
 import com.google.common.collect.ImmutableList;
+import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.Setter;
@@ -16,6 +17,7 @@
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 
 /** Logical plan node of Aggregation, the interface for building aggregation actions in queries. */
 @Getter
@@ -29,7 +31,8 @@ public class Aggregation extends UnresolvedPlan {
   private UnresolvedExpression span;
   private List<Argument> argExprList;
   private UnresolvedPlan child;
-
+  private Optional<TablesampleContext> sample;
+
   /** Aggregation Constructor without span and argument. */
   public Aggregation(
       List<UnresolvedExpression> aggExprList,
@@ -71,4 +74,14 @@ public List<UnresolvedPlan> getChild() {
   public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
     return nodeVisitor.visitAggregation(this, context);
   }
+
+  @Getter
+  @Setter
+  @ToString
+  @EqualsAndHashCode(callSuper = false)
+  @AllArgsConstructor
+  public static class TablesampleContext {
+    public int percentage;
+  }
+
 }

ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java

@@ -21,6 +21,7 @@
 import org.opensearch.sql.ast.expression.DataType;
 import org.opensearch.sql.ast.expression.EqualTo;
 import org.opensearch.sql.ast.expression.Field;
+import org.opensearch.sql.ast.tree.Aggregation.TablesampleContext;
 import org.opensearch.sql.ast.tree.FieldSummary;
 import org.opensearch.sql.ast.expression.FieldsMapping;
 import org.opensearch.sql.ast.expression.Let;
@@ -448,6 +449,10 @@ public UnresolvedPlan visitTopCommand(OpenSearchPPLParser.TopCommandContext ctx)
                     aggListBuilder.build(),
                     aggListBuilder.build(),
                     groupListBuilder.build());
+    if(ctx.sampleClause()!=null) {
+      int percentage = Integer.parseInt(ctx.sampleClause().percentage.getText());
+      aggregation.setSample(Optional.of(new TablesampleContext(percentage)));
+    }
     return aggregation;
   }
 
@@ -493,6 +498,10 @@ public UnresolvedPlan visitRareCommand(OpenSearchPPLParser.RareCommandContext ct
                     aggListBuilder.build(),
                     aggListBuilder.build(),
                     groupListBuilder.build());
+    if(ctx.sampleClause()!=null) {
+      int percentage = Integer.parseInt(ctx.sampleClause().percentage.getText());
+      aggregation.setSample(Optional.of(new TablesampleContext(percentage)));
+    }
     return aggregation;
   }