Skip to content

Commit 68ec263

Browse files
yuancuyuancu
committed
WIP: Make poc implementation for chart command
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent fdb09e8 commit 68ec263

8 files changed

Lines changed: 285 additions & 38 deletions

File tree

‎core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
import org.opensearch.sql.ast.tree.Append;
5050
import org.opensearch.sql.ast.tree.AppendCol;
5151
import org.opensearch.sql.ast.tree.Bin;
52+
import org.opensearch.sql.ast.tree.Chart;
5253
import org.opensearch.sql.ast.tree.CloseCursor;
5354
import org.opensearch.sql.ast.tree.Dedupe;
5455
import org.opensearch.sql.ast.tree.Eval;
@@ -269,6 +270,10 @@ public T visitReverse(Reverse node, C context) {
269270
return visitChildren(node, context);
270271
}
271272

273+
public T visitChart(Chart node, C context) {
274+
return visitChildren(node, context);
275+
}
276+
272277
public T visitTimechart(Timechart node, C context) {
273278
return visitChildren(node, context);
274279
}

‎core/src/main/java/org/opensearch/sql/ast/tree/Chart.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.ast.tree;
7+
8+
import com.google.common.collect.ImmutableList;
9+
import java.util.List;
10+
import lombok.AllArgsConstructor;
11+
import lombok.EqualsAndHashCode;
12+
import lombok.Getter;
13+
import lombok.ToString;
14+
import org.opensearch.sql.ast.AbstractNodeVisitor;
15+
import org.opensearch.sql.ast.expression.Argument;
16+
import org.opensearch.sql.ast.expression.UnresolvedExpression;
17+
18+
/** AST node represent chart command. */
19+
@Getter
20+
@ToString
21+
@EqualsAndHashCode(callSuper = false)
22+
@AllArgsConstructor
23+
@lombok.Builder(toBuilder = true)
24+
public class Chart extends UnresolvedPlan {
25+
private UnresolvedPlan child;
26+
private UnresolvedExpression rowSplit;
27+
private UnresolvedExpression columnSplit;
28+
private List<UnresolvedExpression> aggregationFunctions;
29+
private List<Argument> arguments;
30+
31+
@Override
32+
public UnresolvedPlan attach(UnresolvedPlan child) {
33+
this.child = child;
34+
return this;
35+
}
36+
37+
@Override
38+
public List<UnresolvedPlan> getChild() {
39+
return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child);
40+
}
41+
42+
@Override
43+
public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
44+
return nodeVisitor.visitChart(this, context);
45+
}
46+
}

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java‎

Lines changed: 92 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@
9898
import org.opensearch.sql.ast.tree.Append;
9999
import org.opensearch.sql.ast.tree.AppendCol;
100100
import org.opensearch.sql.ast.tree.Bin;
101+
import org.opensearch.sql.ast.tree.Chart;
101102
import org.opensearch.sql.ast.tree.CloseCursor;
102103
import org.opensearch.sql.ast.tree.Dedupe;
103104
import org.opensearch.sql.ast.tree.Eval;
@@ -1023,6 +1024,11 @@ private Pair<List<RexNode>, List<AggCall>> resolveAttributesForAggregation(
10231024

10241025
@Override
10251026
public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
1027+
return visitAggregationAndReturnProjection(node, context).getLeft();
1028+
}
1029+
1030+
private Pair<RelNode, List<RexNode>> visitAggregationAndReturnProjection(
1031+
Aggregation node, CalcitePlanContext context) {
10261032
visitChildren(node, context);
10271033

10281034
List<UnresolvedExpression> aggExprList = node.getAggExprList();
@@ -1100,14 +1106,14 @@ public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
11001106
aggregationAttributes.getLeft().stream()
11011107
.map(this::extractAliasLiteral)
11021108
.flatMap(Optional::stream)
1103-
.map(ref -> ((RexLiteral) ref).getValueAs(String.class))
1109+
.map(ref -> ref.getValueAs(String.class))
11041110
.map(context.relBuilder::field)
11051111
.map(f -> (RexNode) f)
11061112
.toList();
11071113
reordered.addAll(aliasedGroupByList);
11081114
context.relBuilder.project(reordered);
11091115

1110-
return context.relBuilder.peek();
1116+
return Pair.of(context.relBuilder.peek(), reordered);
11111117
}
11121118

11131119
private Optional<UnresolvedExpression> getTimeSpanField(UnresolvedExpression expr) {
@@ -1947,6 +1953,90 @@ private String getValueFunctionName(UnresolvedExpression aggregateFunction) {
19471953
return sb.toString();
19481954
}
19491955

1956+
@Override
1957+
public RelNode visitChart(Chart node, CalcitePlanContext context) {
1958+
visitChildren(node, context);
1959+
ArgumentMap argMap = ArgumentMap.of(node.getArguments());
1960+
List<UnresolvedExpression> groupExprList = new ArrayList<>();
1961+
UnresolvedExpression span;
1962+
if (node.getColumnSplit() instanceof Span && node.getRowSplit() instanceof Span) {
1963+
throw new UnsupportedOperationException("It is not supported to have two span splits");
1964+
} else if (node.getRowSplit() instanceof Span) {
1965+
if (node.getColumnSplit() != null) {
1966+
groupExprList.add(node.getColumnSplit());
1967+
}
1968+
span = node.getRowSplit();
1969+
} else if (node.getColumnSplit() instanceof Span) {
1970+
if (node.getRowSplit() != null) {
1971+
groupExprList.add(node.getRowSplit());
1972+
}
1973+
span = node.getColumnSplit();
1974+
} else {
1975+
groupExprList.addAll(
1976+
Stream.of(node.getRowSplit(), node.getColumnSplit()).filter(Objects::nonNull).toList());
1977+
span = null;
1978+
}
1979+
Aggregation aggregation =
1980+
new Aggregation(node.getAggregationFunctions(), List.of(), groupExprList, span, List.of());
1981+
Pair<RelNode, List<RexNode>> aggregated =
1982+
visitAggregationAndReturnProjection(aggregation, context);
1983+
// If row or column split does not present or limit equals 0, this is the same as `stats agg
1984+
// [group by col]`
1985+
1986+
Integer limit =
1987+
Optional.ofNullable(argMap.get("limit")).map(l -> (Integer) l.getValue()).orElse(10);
1988+
Boolean top =
1989+
Optional.ofNullable(argMap.get("top")).map(t -> (Boolean) t.getValue()).orElse(true);
1990+
if (node.getRowSplit() == null || node.getColumnSplit() == null || Objects.equals(limit, 0)) {
1991+
return aggregated.getLeft();
1992+
}
1993+
List<RexNode> projected = aggregated.getRight();
1994+
String columSplitName = aggregated.getLeft().getRowType().getFieldNames().getLast();
1995+
RelBuilder relBuilder = context.relBuilder;
1996+
// 0: agg; 2: column-split
1997+
relBuilder.project(relBuilder.field(0), relBuilder.field(2));
1998+
relBuilder.filter(relBuilder.isNotNull(relBuilder.field(1)));
1999+
// 1: column split; 0: agg
2000+
relBuilder.aggregate(
2001+
relBuilder.groupKey(relBuilder.field(1)),
2002+
relBuilder.sum(relBuilder.field(0)).as("__grand_total__")); // results: group key, agg calls
2003+
RexNode grandTotal = relBuilder.field("__grand_total__");
2004+
if (top) {
2005+
grandTotal = relBuilder.desc(grandTotal);
2006+
}
2007+
RexNode rowNum =
2008+
PlanUtils.makeOver(
2009+
context,
2010+
BuiltinFunctionName.ROW_NUMBER,
2011+
relBuilder.literal(1),
2012+
List.of(),
2013+
List.of(),
2014+
List.of(grandTotal),
2015+
WindowFrame.toCurrentRow());
2016+
relBuilder.projectPlus(relBuilder.alias(rowNum, "__row_number__"));
2017+
RelNode ranked = relBuilder.build();
2018+
2019+
relBuilder.push(aggregated.getLeft());
2020+
relBuilder.push(ranked);
2021+
2022+
// on column-split = group key
2023+
relBuilder.join(
2024+
JoinRelType.INNER, relBuilder.equals(relBuilder.field(2, 0, 2), relBuilder.field(2, 1, 0)));
2025+
RexNode caseExpr =
2026+
relBuilder.alias(
2027+
relBuilder.call(
2028+
SqlStdOperatorTable.CASE,
2029+
relBuilder.call(
2030+
SqlStdOperatorTable.LESS_THAN_OR_EQUAL,
2031+
relBuilder.field("__row_number__"),
2032+
relBuilder.literal(limit)),
2033+
relBuilder.field(2),
2034+
relBuilder.literal("OTHER")),
2035+
columSplitName);
2036+
relBuilder.project(relBuilder.field(0), relBuilder.field(1), caseExpr);
2037+
return relBuilder.peek();
2038+
}
2039+
19502040
/** Transforms timechart command into SQL-based operations. */
19512041
@Override
19522042
public RelNode visitTimechart(
@@ -2064,7 +2154,6 @@ private RelNode buildTopCategoriesQuery(
20642154
if (limit > 0) {
20652155
context.relBuilder.limit(0, limit);
20662156
}
2067-
20682157
return context.relBuilder.build();
20692158
}
20702159

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ ML: 'ML';
4545
FILLNULL: 'FILLNULL';
4646
FLATTEN: 'FLATTEN';
4747
TRENDLINE: 'TRENDLINE';
48+
CHART: 'CHART';
4849
TIMECHART: 'TIMECHART';
4950
APPENDCOL: 'APPENDCOL';
5051
EXPAND: 'EXPAND';
@@ -76,6 +77,7 @@ RIGHT_HINT: 'HINT.RIGHT';
7677
// COMMAND ASSIST KEYWORDS
7778
AS: 'AS';
7879
BY: 'BY';
80+
OVER: 'OVER';
7981
SOURCE: 'SOURCE';
8082
INDEX: 'INDEX';
8183
A: 'A';
@@ -92,6 +94,7 @@ COST: 'COST';
9294
EXTENDED: 'EXTENDED';
9395
OVERRIDE: 'OVERRIDE';
9496
OVERWRITE: 'OVERWRITE';
97+
BOTTOM: 'BOTTOM';
9598

9699
// SORT FIELD KEYWORDS
97100
// TODO #3180: Fix broken sort functionality

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ commands
7777
| flattenCommand
7878
| reverseCommand
7979
| regexCommand
80+
| chartCommand
8081
| timechartCommand
8182
| rexCommand
8283
;
@@ -248,6 +249,24 @@ reverseCommand
248249
: REVERSE
249250
;
250251

252+
chartCommand
253+
: CHART chartOptions* statsAggTerm (COMMA statsAggTerm)* (OVER rowSplit)? (BY columnSplit)?
254+
| CHART chartOptions* statsAggTerm (COMMA statsAggTerm)* BY rowSplit (COMMA)? columnSplit
255+
;
256+
257+
chartOptions
258+
: LIMIT EQUAL (TOP | BOTTOM)? integerLiteral
259+
| USEOTHER EQUAL booleanLiteral
260+
;
261+
262+
rowSplit
263+
: fieldExpression binOption*
264+
;
265+
266+
columnSplit
267+
: fieldExpression binOption*
268+
;
269+
251270
timechartCommand
252271
: TIMECHART timechartParameter* statsFunction (BY fieldExpression)?
253272
;

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

Lines changed: 50 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
import org.antlr.v4.runtime.ParserRuleContext;
4444
import org.antlr.v4.runtime.Token;
4545
import org.antlr.v4.runtime.tree.ParseTree;
46+
import org.antlr.v4.runtime.tree.TerminalNode;
4647
import org.apache.commons.lang3.tuple.Pair;
4748
import org.opensearch.sql.ast.EmptySourcePropagateVisitor;
4849
import org.opensearch.sql.ast.dsl.AstDSL;
@@ -51,6 +52,7 @@
5152
import org.opensearch.sql.ast.tree.Aggregation;
5253
import org.opensearch.sql.ast.tree.Append;
5354
import org.opensearch.sql.ast.tree.AppendCol;
55+
import org.opensearch.sql.ast.tree.Chart;
5456
import org.opensearch.sql.ast.tree.CountBin;
5557
import org.opensearch.sql.ast.tree.Dedupe;
5658
import org.opensearch.sql.ast.tree.DefaultBin;
@@ -479,60 +481,39 @@ public UnresolvedPlan visitBinCommand(BinCommandContext ctx) {
479481
UnresolvedExpression aligntime = null;
480482
UnresolvedExpression start = null;
481483
UnresolvedExpression end = null;
482-
484+
String errorFormat = "Duplicate %s parameter in bin command";
483485
// Process each bin option: detect duplicates and assign values in one shot
484486
for (OpenSearchPPLParser.BinOptionContext option : ctx.binOption()) {
487+
UnresolvedExpression resolvedOption = internalVisitExpression(option);
485488
// SPAN parameter
486489
if (option.span != null) {
487-
if (!seenParams.add("SPAN")) {
488-
throw new IllegalArgumentException("Duplicate SPAN parameter in bin command");
489-
}
490-
span = internalVisitExpression(option.span);
490+
checkParamDuplication(seenParams, option.SPAN(), errorFormat);
491+
span = resolvedOption;
491492
}
492-
493493
// BINS parameter
494494
if (option.bins != null) {
495-
if (!seenParams.add("BINS")) {
496-
throw new IllegalArgumentException("Duplicate BINS parameter in bin command");
497-
}
498-
bins = Integer.parseInt(option.bins.getText());
495+
checkParamDuplication(seenParams, option.SPAN(), errorFormat);
496+
bins = (Integer) ((Literal) resolvedOption).getValue();
499497
}
500-
501498
// MINSPAN parameter
502499
if (option.minspan != null) {
503-
if (!seenParams.add("MINSPAN")) {
504-
throw new IllegalArgumentException("Duplicate MINSPAN parameter in bin command");
505-
}
506-
minspan = internalVisitExpression(option.minspan);
500+
checkParamDuplication(seenParams, option.MINSPAN(), errorFormat);
501+
minspan = resolvedOption;
507502
}
508-
509503
// ALIGNTIME parameter
510504
if (option.aligntime != null) {
511-
if (!seenParams.add("ALIGNTIME")) {
512-
throw new IllegalArgumentException("Duplicate ALIGNTIME parameter in bin command");
513-
}
514-
aligntime =
515-
option.aligntime.EARLIEST() != null
516-
? org.opensearch.sql.ast.dsl.AstDSL.stringLiteral("earliest")
517-
: option.aligntime.LATEST() != null
518-
? org.opensearch.sql.ast.dsl.AstDSL.stringLiteral("latest")
519-
: internalVisitExpression(option.aligntime.literalValue());
505+
checkParamDuplication(seenParams, option.ALIGNTIME(), errorFormat);
506+
aligntime = resolvedOption;
520507
}
521-
522508
// START parameter
523509
if (option.start != null) {
524-
if (!seenParams.add("START")) {
525-
throw new IllegalArgumentException("Duplicate START parameter in bin command");
526-
}
527-
start = internalVisitExpression(option.start);
510+
checkParamDuplication(seenParams, option.START(), errorFormat);
511+
start = resolvedOption;
528512
}
529-
530513
// END parameter
531514
if (option.end != null) {
532-
if (!seenParams.add("END")) {
533-
throw new IllegalArgumentException("Duplicate END parameter in bin command");
534-
}
535-
end = internalVisitExpression(option.end);
515+
checkParamDuplication(seenParams, option.END(), errorFormat);
516+
end = resolvedOption;
536517
}
537518
}
538519

@@ -561,6 +542,14 @@ public UnresolvedPlan visitBinCommand(BinCommandContext ctx) {
561542
}
562543
}
563544

545+
private void checkParamDuplication(
546+
Set<String> seenParams, TerminalNode terminalNode, String errorFormat) {
547+
String paramName = terminalNode.getText();
548+
if (!seenParams.add(paramName)) {
549+
throw new IllegalArgumentException(StringUtils.format(errorFormat, paramName));
550+
}
551+
}
552+
564553
/** Sort command. */
565554
@Override
566555
public UnresolvedPlan visitSortCommand(SortCommandContext ctx) {
@@ -596,6 +585,32 @@ public UnresolvedPlan visitReverseCommand(OpenSearchPPLParser.ReverseCommandCont
596585
return new Reverse();
597586
}
598587

588+
/** Chart command. */
589+
@Override
590+
public UnresolvedPlan visitChartCommand(OpenSearchPPLParser.ChartCommandContext ctx) {
591+
UnresolvedExpression rowSplit =
592+
ctx.rowSplit() == null ? null : internalVisitExpression(ctx.rowSplit());
593+
UnresolvedExpression columnSplit =
594+
ctx.columnSplit() == null ? null : internalVisitExpression(ctx.columnSplit());
595+
List<Argument> arguments = ArgumentFactory.getArgumentList(ctx);
596+
ImmutableList.Builder<UnresolvedExpression> aggListBuilder = new ImmutableList.Builder<>();
597+
for (OpenSearchPPLParser.StatsAggTermContext aggCtx : ctx.statsAggTerm()) {
598+
UnresolvedExpression aggExpression = internalVisitExpression(aggCtx.statsFunction());
599+
String name =
600+
aggCtx.alias == null
601+
? getTextInQuery(aggCtx)
602+
: StringUtils.unquoteIdentifier(aggCtx.alias.getText());
603+
Alias alias = new Alias(name, aggExpression);
604+
aggListBuilder.add(alias);
605+
}
606+
return Chart.builder()
607+
.rowSplit(rowSplit)
608+
.columnSplit(columnSplit)
609+
.aggregationFunctions(aggListBuilder.build())
610+
.arguments(arguments)
611+
.build();
612+
}
613+
599614
/** Timechart command. */
600615
@Override
601616
public UnresolvedPlan visitTimechartCommand(OpenSearchPPLParser.TimechartCommandContext ctx) {

0 commit comments

Comments
 (0)