Add count() support for PPL (#894)

opensearch-project · Dec 4, 2020 · 41d48dc · 41d48dc
1 parent 7dc8699
commit 41d48dc
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 15 deletions.
diff --git a/docs/experiment/ppl/cmd/stats.rst b/docs/experiment/ppl/cmd/stats.rst
@@ -38,8 +38,23 @@ stats <aggregation>... [by-clause]...
 * aggregation: mandatory. A statistical aggregation function. The argument of aggregation must be field.
 * by-clause: optional. The one or more fields to group the results by. **Default**: If no <by-clause> is specified, the stats command returns only one row, which is the aggregation over the entire result set.
 
+Example 1: Calculate the count of events
+========================================
 
-Example 1: Calculate the average of a field
+The example show calculate the count of events in the accounts.
+
+PPL query::
+
+    od> source=accounts | stats count();
+    fetched rows / total rows = 1/1
+    +-----------+
+    | count()   |
+    |-----------|
+    | 4         |
+    +-----------+
+
+
+Example 2: Calculate the average of a field
 ===========================================
 
 The example show calculate the average age of all the accounts.
@@ -55,7 +70,7 @@ PPL query::
     +------------+
 
 
-Example 2: Calculate the average of a field by group
+Example 3: Calculate the average of a field by group
 ====================================================
 
 The example show calculate the average age of all the accounts group by gender.
@@ -72,23 +87,23 @@ PPL query::
     +----------+--------------------+
 
 
-Example 3: Calculate the average and sum of a field by group
-============================================================
+Example 4: Calculate the average, sum and count of a field by group
+===================================================================
 
-The example show calculate the average age and sum age of all the accounts group by gender.
+The example show calculate the average age, sum age and count of events of all the accounts group by gender.
 
 PPL query::
 
-    od> source=accounts | stats avg(age), sum(age) by gender;
+    od> source=accounts | stats avg(age), sum(age), count() by gender;
     fetched rows / total rows = 2/2
-    +----------+--------------------+------------+
-    | gender   | avg(age)           | sum(age)   |
-    |----------+--------------------+------------|
-    | F        | 28.0               | 28         |
-    | M        | 33.666666666666664 | 101        |
-    +----------+--------------------+------------+
-
-Example 4: Calculate the maximum of a field
+    +----------+--------------------+------------+-----------+
+    | gender   | avg(age)           | sum(age)   | count()   |
+    |----------+--------------------+------------+-----------|
+    | F        | 28.0               | 28         | 1         |
+    | M        | 33.666666666666664 | 101        | 3         |
+    +----------+--------------------+------------+-----------+
+
+Example 5: Calculate the maximum of a field
 ===========================================
 
 The example calculates the max age of all the accounts.
@@ -103,7 +118,7 @@ PPL query::
     | 36         |
     +------------+
 
-Example 5: Calculate the maximum and minimum of a field by group
+Example 6: Calculate the maximum and minimum of a field by group
 ================================================================
 
 The example calculates the max and min age values of all the accounts group by gender.

diff --git a/integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/ppl/StatsCommandIT.java b/integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/ppl/StatsCommandIT.java
@@ -58,6 +58,14 @@ public void testStatsCount() throws IOException {
     verifyDataRows(response, rows(1000));
   }
 
+  @Test
+  public void testStatsCountAll() throws IOException {
+    JSONObject response =
+        executeQuery(String.format("source=%s | stats count()", TEST_INDEX_ACCOUNT));
+    verifySchema(response, schema("count()", null, "integer"));
+    verifyDataRows(response, rows(1000));
+  }
+
   @Test
   public void testStatsMin() throws IOException {
     JSONObject response = executeQuery(String.format(

diff --git a/ppl/src/main/antlr/OpenDistroPPLParser.g4 b/ppl/src/main/antlr/OpenDistroPPLParser.g4
@@ -125,6 +125,7 @@ statsAggTerm
 /** aggregation functions */
 statsFunction
     : statsFunctionName LT_PRTHS valueExpression RT_PRTHS           #statsFunctionCall
+    | COUNT LT_PRTHS RT_PRTHS                                       #countAllFunctionCall
     | percentileAggFunction                                         #percentileAggFunctionCall
     ;
 

diff --git a/.../main/java/com/amazon/opendistroforelasticsearch/sql/ppl/parser/AstExpressionBuilder.java b/.../main/java/com/amazon/opendistroforelasticsearch/sql/ppl/parser/AstExpressionBuilder.java
@@ -41,6 +41,7 @@
 import static com.amazon.opendistroforelasticsearch.sql.ppl.antlr.parser.OpenDistroPPLParser.WcFieldExpressionContext;
 
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.AggregateFunction;
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.AllFields;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.And;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Argument;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Compare;
@@ -169,6 +170,12 @@ public UnresolvedExpression visitStatsFunctionCall(StatsFunctionCallContext ctx)
     return new AggregateFunction(ctx.statsFunctionName().getText(), visit(ctx.valueExpression()));
   }
 
+  @Override
+  public UnresolvedExpression visitCountAllFunctionCall(
+      OpenDistroPPLParser.CountAllFunctionCallContext ctx) {
+    return new AggregateFunction("count", AllFields.of());
+  }
+
   @Override
   public UnresolvedExpression visitPercentileAggFunction(PercentileAggFunctionContext ctx) {
     return new AggregateFunction(ctx.PERCENTILE().getText(), visit(ctx.aggField),

diff --git a/...t/java/com/amazon/opendistroforelasticsearch/sql/ppl/parser/AstExpressionBuilderTest.java b/...t/java/com/amazon/opendistroforelasticsearch/sql/ppl/parser/AstExpressionBuilderTest.java
@@ -47,6 +47,7 @@
 import static com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL.xor;
 import static java.util.Collections.emptyList;
 
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.AllFields;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.DataType;
 import org.junit.Ignore;
 import org.junit.Test;
@@ -324,6 +325,27 @@ public void testPercentileAggFuncExpr() {
         ));
   }
 
+  @Test
+  public void testCountFuncCallExpr() {
+    assertEqual("source=t | stats count() by b",
+        agg(
+            relation("t"),
+            exprList(
+                alias(
+                    "count()",
+                    aggregate("count", AllFields.of())
+                )
+            ),
+            emptyList(),
+            exprList(
+                alias(
+                    "b",
+                    field("b")
+                )),
+            defaultStatsArgs()
+        ));
+  }
+
   @Test
   public void testEvalFuncCallExpr() {
     assertEqual("source=t | eval f=abs(a)",