Skip to content

Commit 4e13896

Browse files
ahkcsahkcs
committed
Add UT and optimize code
Signed-off-by: Kai Huang <ahkcs@amazon.com>
1 parent e6ca895 commit 4e13896

6 files changed

Lines changed: 59 additions & 34 deletions

File tree

‎core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -366,6 +366,8 @@ public enum BuiltinFunctionName {
366366
// .put("earliest", BuiltinFunctionName.EARLIEST)
367367
// .put("latest", BuiltinFunctionName.LATEST)
368368
.put("distinct_count_approx", BuiltinFunctionName.DISTINCT_COUNT_APPROX)
369+
.put("dc", BuiltinFunctionName.DISTINCT_COUNT_APPROX)
370+
.put("distinct_count", BuiltinFunctionName.DISTINCT_COUNT_APPROX)
369371
.put("pattern", BuiltinFunctionName.INTERNAL_PATTERN)
370372
.build();
371373

‎core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
import static org.opensearch.sql.expression.function.BuiltinFunctionName.DAY_OF_WEEK;
6262
import static org.opensearch.sql.expression.function.BuiltinFunctionName.DAY_OF_YEAR;
6363
import static org.opensearch.sql.expression.function.BuiltinFunctionName.DEGREES;
64+
import static org.opensearch.sql.expression.function.BuiltinFunctionName.DISTINCT_COUNT_APPROX;
6465
import static org.opensearch.sql.expression.function.BuiltinFunctionName.DIVIDE;
6566
import static org.opensearch.sql.expression.function.BuiltinFunctionName.DIVIDEFUNCTION;
6667
import static org.opensearch.sql.expression.function.BuiltinFunctionName.E;
@@ -1122,6 +1123,15 @@ void populate() {
11221123
wrapSqlOperandTypeChecker(
11231124
SqlStdOperatorTable.COUNT.getOperandTypeChecker(), COUNT.name(), false));
11241125

1126+
register(
1127+
DISTINCT_COUNT_APPROX,
1128+
(distinct, field, argList, ctx) ->
1129+
ctx.relBuilder.aggregateCall(SqlStdOperatorTable.APPROX_COUNT_DISTINCT, field),
1130+
wrapSqlOperandTypeChecker(
1131+
SqlStdOperatorTable.APPROX_COUNT_DISTINCT.getOperandTypeChecker(),
1132+
DISTINCT_COUNT_APPROX.name(),
1133+
false));
1134+
11251135
register(
11261136
VARSAMP,
11271137
(distinct, field, argList, ctx) -> ctx.relBuilder.aggregateCall(VAR_SAMP_NULLABLE, field),

‎docs/user/ppl/cmd/eventstats.rst‎

Lines changed: 9 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -287,29 +287,16 @@ Usage: DISTINCT_COUNT(expr), DC(expr). Returns the approximate number of distinc
287287

288288
Example::
289289

290-
PPL> source=accounts | eventstats dc(state) as distinct_states;
290+
PPL> source=accounts | eventstats dc(state) as distinct_states, distinct_count(state) as dc_states_alt by gender;
291291
fetched rows / total rows = 4/4
292-
+----------------+-----------+----------------------+---------+--------+--------+----------+-------+-----+-----------------------+----------+-----------------+
293-
| account_number | firstname | address | balance | gender | city | employer | state | age | email | lastname | distinct_states |
294-
|----------------+-----------+----------------------+---------+--------+--------+----------+-------+-----+-----------------------+----------+-----------------|
295-
| 1 | Amber | 880 Holmes Lane | 39225 | M | Brogan | Pyrami | IL | 32 | amberduke@pyrami.com | Duke | 4 |
296-
| 6 | Hattie | 671 Bristol Street | 5686 | M | Dante | Netagy | TN | 36 | hattiebond@netagy.com | Bond | 4 |
297-
| 13 | Nanette | 789 Madison Street | 32838 | F | Nogal | Quility | VA | 28 | null | Bates | 4 |
298-
| 18 | Dale | 467 Hutchinson Court | 4180 | M | Orick | null | MD | 33 | daleadams@boink.com | Adams | 4 |
299-
+----------------+-----------+----------------------+---------+--------+--------+----------+-------+-----+-----------------------+----------+-----------------+
300-
301-
Example with BY clause::
302-
303-
PPL> source=accounts | eventstats distinct_count(state) as distinct_states by gender;
304-
fetched rows / total rows = 4/4
305-
+----------------+-----------+----------------------+---------+--------+--------+----------+-------+-----+-----------------------+----------+-----------------+
306-
| account_number | firstname | address | balance | gender | city | employer | state | age | email | lastname | distinct_states |
307-
|----------------+-----------+----------------------+---------+--------+--------+----------+-------+-----+-----------------------+----------+-----------------|
308-
| 13 | Nanette | 789 Madison Street | 32838 | F | Nogal | Quility | VA | 28 | null | Bates | 1 |
309-
| 1 | Amber | 880 Holmes Lane | 39225 | M | Brogan | Pyrami | IL | 32 | amberduke@pyrami.com | Duke | 3 |
310-
| 6 | Hattie | 671 Bristol Street | 5686 | M | Dante | Netagy | TN | 36 | hattiebond@netagy.com | Bond | 3 |
311-
| 18 | Dale | 467 Hutchinson Court | 4180 | M | Orick | null | MD | 33 | daleadams@boink.com | Adams | 3 |
312-
+----------------+-----------+----------------------+---------+--------+--------+----------+-------+-----+-----------------------+----------+-----------------+
292+
+----------------+-----------+----------------------+---------+--------+--------+----------+-------+-----+-----------------------+----------+-----------------+-----------------+
293+
| account_number | firstname | address | balance | gender | city | employer | state | age | email | lastname | distinct_states | dc_states_alt |
294+
|----------------+-----------+----------------------+---------+--------+--------+----------+-------+-----+-----------------------+----------+-----------------|-----------------|
295+
| 13 | Nanette | 789 Madison Street | 32838 | F | Nogal | Quility | VA | 28 | null | Bates | 1 | 1 |
296+
| 1 | Amber | 880 Holmes Lane | 39225 | M | Brogan | Pyrami | IL | 32 | amberduke@pyrami.com | Duke | 3 | 3 |
297+
| 6 | Hattie | 671 Bristol Street | 5686 | M | Dante | Netagy | TN | 36 | hattiebond@netagy.com | Bond | 3 | 3 |
298+
| 18 | Dale | 467 Hutchinson Court | 4180 | M | Orick | null | MD | 33 | daleadams@boink.com | Adams | 3 | 3 |
299+
+----------------+-----------+----------------------+---------+--------+--------+----------+-------+-----+-----------------------+----------+-----------------+-----------------+
313300

314301

315302
Configuration

‎ppl/src/main/antlr/OpenSearchPPLParser.g4‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -387,7 +387,6 @@ windowFunction
387387
windowFunctionName
388388
: statsFunctionName
389389
| scalarWindowFunctionName
390-
| distinctCountWindowFunction
391390
;
392391

393392
scalarWindowFunctionName
@@ -400,10 +399,7 @@ scalarWindowFunctionName
400399
| LAST
401400
| NTH
402401
| NTILE
403-
;
404-
405-
distinctCountWindowFunction
406-
: DISTINCT_COUNT
402+
| DISTINCT_COUNT
407403
| DC
408404
;
409405

‎ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -558,13 +558,8 @@ public UnresolvedExpression visitBetween(OpenSearchPPLParser.BetweenContext ctx)
558558

559559
@Override
560560
public UnresolvedExpression visitWindowFunction(OpenSearchPPLParser.WindowFunctionContext ctx) {
561-
Function f;
562-
if (ctx.windowFunctionName().distinctCountWindowFunction() != null) {
563-
// All distinct count variants (dc, distinct_count, distinct_count_approx) use the UDAF
564-
f = buildFunction("distinct_count_approx", ctx.functionArgs().functionArg());
565-
} else {
566-
f = buildFunction(ctx.windowFunctionName().getText(), ctx.functionArgs().functionArg());
567-
}
561+
Function f =
562+
buildFunction(ctx.windowFunctionName().getText(), ctx.functionArgs().functionArg());
568563
// In PPL eventstats command, all window functions have the same partition and order spec.
569564
return new WindowFunction(f);
570565
}

‎ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLEventstatsTest.java‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,4 +70,39 @@ public void testEventstatsAvg() {
7070
+ "FROM `scott`.`EMP`";
7171
verifyPPLToSparkSQL(root, expectedSparkSql);
7272
}
73+
74+
@Test
75+
public void testEventstatsDistinctCount() {
76+
String ppl = "source=EMP | eventstats dc(DEPTNO)";
77+
RelNode root = getRelNode(ppl);
78+
String expectedLogical =
79+
"LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5],"
80+
+ " COMM=[$6], DEPTNO=[$7], dc(DEPTNO)=[APPROX_COUNT_DISTINCT($7) OVER ()])\n"
81+
+ " LogicalTableScan(table=[[scott, EMP]])\n";
82+
verifyLogical(root, expectedLogical);
83+
}
84+
85+
@Test
86+
public void testEventstatsDistinctCountFunction() {
87+
String ppl = "source=EMP | eventstats distinct_count(DEPTNO)";
88+
RelNode root = getRelNode(ppl);
89+
String expectedLogical =
90+
"LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5],"
91+
+ " COMM=[$6], DEPTNO=[$7], distinct_count(DEPTNO)=[APPROX_COUNT_DISTINCT($7) OVER"
92+
+ " ()])\n"
93+
+ " LogicalTableScan(table=[[scott, EMP]])\n";
94+
verifyLogical(root, expectedLogical);
95+
}
96+
97+
@Test
98+
public void testEventstatsDistinctCountWithPartition() {
99+
String ppl = "source=EMP | eventstats dc(JOB) by DEPTNO";
100+
RelNode root = getRelNode(ppl);
101+
String expectedLogical =
102+
"LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5],"
103+
+ " COMM=[$6], DEPTNO=[$7], dc(JOB)=[APPROX_COUNT_DISTINCT($2) OVER (PARTITION BY"
104+
+ " $7)])\n"
105+
+ " LogicalTableScan(table=[[scott, EMP]])\n";
106+
verifyLogical(root, expectedLogical);
107+
}
73108
}

0 commit comments

Comments
 (0)