Add shortcut for count() (opensearch-project#4142)

Author: noCharger

docs/user/ppl/cmd/stats.rst

@@ -80,7 +80,7 @@ COUNT
 Description
 >>>>>>>>>>>
 
-Usage: Returns a count of the number of expr in the rows retrieved by a SELECT statement.
+Usage: Returns a count of the number of expr in the rows retrieved by a SELECT statement. The ``C()`` function can be used as an abbreviation for ``COUNT()``.
 
 Example::
 
@@ -92,6 +92,14 @@ Example::
     | 4       |
     +---------+
 
+    os> source=accounts | stats c();
+    fetched rows / total rows = 1/1
+    +-----+
+    | c() |
+    |-----|
+    | 4   |
+    +-----+
+
 SUM
 ---
 

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLAggregationIT.java

@@ -60,6 +60,10 @@ public void testSimpleCount() throws IOException {
         executeQuery(String.format("source=%s | stats count() as c", TEST_INDEX_BANK));
     verifySchema(actual, schema("c", "bigint"));
     verifyDataRows(actual, rows(7));
+
+    actual = executeQuery(String.format("source=%s | stats c() as count_emp", TEST_INDEX_BANK));
+    verifySchema(actual, schema("count_emp", "bigint"));
+    verifyDataRows(actual, rows(7));
   }
 
   @Test

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

@@ -88,6 +88,14 @@ public void testStatsCountAll() throws IOException {
       verifySchema(response, schema("count()", null, "int"));
     }
     verifyDataRows(response, rows(1000));
+
+    response = executeQuery(String.format("source=%s | stats c()", TEST_INDEX_ACCOUNT));
+    if (isCalciteEnabled()) {
+      verifySchema(response, schema("c()", null, "bigint"));
+    } else {
+      verifySchema(response, schema("c()", null, "int"));
+    }
+    verifyDataRows(response, rows(1000));
   }
 
   @Test

ppl/src/main/antlr/OpenSearchPPLParser.g4

@@ -459,7 +459,7 @@ statsAggTerm
 // aggregation functions
 statsFunction
    : statsFunctionName LT_PRTHS valueExpression RT_PRTHS        # statsFunctionCall
-   | COUNT LT_PRTHS RT_PRTHS                                    # countAllFunctionCall
+   | (COUNT | C) LT_PRTHS RT_PRTHS                              # countAllFunctionCall
    | PERCENTILE_SHORTCUT LT_PRTHS valueExpression RT_PRTHS      # percentileShortcutFunctionCall
    | (DISTINCT_COUNT | DC | DISTINCT_COUNT_APPROX) LT_PRTHS valueExpression RT_PRTHS    # distinctCountFunctionCall
    | takeAggFunction                                            # takeAggFunctionCall

ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLAggregationTest.java

@@ -30,6 +30,19 @@ public void testSimpleCount() {
 
     String expectedSparkSql = "" + "SELECT COUNT(*) `c`\n" + "FROM `scott`.`EMP`";
     verifyPPLToSparkSQL(root, expectedSparkSql);
+
+    ppl = "source=EMP | stats c() as count_emp";
+    root = getRelNode(ppl);
+    expectedLogical =
+        ""
+            + "LogicalAggregate(group=[{}], count_emp=[COUNT()])\n"
+            + "  LogicalTableScan(table=[[scott, EMP]])\n";
+    verifyLogical(root, expectedLogical);
+    expectedResult = "count_emp=14\n";
+    verifyResult(root, expectedResult);
+
+    expectedSparkSql = "" + "SELECT COUNT(*) `count_emp`\n" + "FROM `scott`.`EMP`";
+    verifyPPLToSparkSQL(root, expectedSparkSql);
   }
 
   @Test
@@ -86,6 +99,27 @@ public void testMultipleAggregatesWithAliases() {
     verifyPPLToSparkSQL(root, expectedSparkSql);
   }
 
+  @Test
+  public void testMultipleAggregatesWithCountAbbreviation() {
+    String ppl =
+        "source=EMP | stats avg(SAL) as avg_sal, max(SAL) as max_sal, min(SAL) as min_sal, c()"
+            + " as cnt";
+    RelNode root = getRelNode(ppl);
+    String expectedLogical =
+        "LogicalAggregate(group=[{}], avg_sal=[AVG($0)], max_sal=[MAX($0)], min_sal=[MIN($0)],"
+            + " cnt=[COUNT()])\n"
+            + "  LogicalProject(SAL=[$5])\n"
+            + "    LogicalTableScan(table=[[scott, EMP]])\n";
+    verifyLogical(root, expectedLogical);
+    String expectedResult = "avg_sal=2073.214285; max_sal=5000.00; min_sal=800.00; cnt=14\n";
+    verifyResult(root, expectedResult);
+
+    String expectedSparkSql =
+        "SELECT AVG(`SAL`) `avg_sal`, MAX(`SAL`) `max_sal`, MIN(`SAL`) `min_sal`, COUNT(*) `cnt`\n"
+            + "FROM `scott`.`EMP`";
+    verifyPPLToSparkSQL(root, expectedSparkSql);
+  }
+
   @Test
   public void testMultipleAggregatesWithAliasesByClause() {
     String ppl =

ppl/src/test/java/org/opensearch/sql/ppl/parser/AstBuilderTest.java

@@ -240,6 +240,18 @@ public void testStatsCommand() {
             defaultStatsArgs()));
   }
 
+  @Test
+  public void testStatsCommandWithCountAbbreviation() {
+    assertEqual(
+        "source=t | stats c()",
+        agg(
+            relation("t"),
+            exprList(alias("c()", aggregate("count", AstDSL.allFields()))),
+            emptyList(),
+            emptyList(),
+            defaultStatsArgs()));
+  }
+
   @Test
   public void testStatsCommandWithByClause() {
     assertEqual(