Skip to content

Commit b06667e

Browse files
yuancuyuancu
committed
Support usenull and nullstr (when both row split and col split present)
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent ce237b9 commit b06667e

6 files changed

Lines changed: 53 additions & 23 deletions

File tree

‎core/src/main/java/org/opensearch/sql/analysis/Analyzer.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
import org.opensearch.sql.ast.tree.Append;
6262
import org.opensearch.sql.ast.tree.AppendCol;
6363
import org.opensearch.sql.ast.tree.Bin;
64+
import org.opensearch.sql.ast.tree.Chart;
6465
import org.opensearch.sql.ast.tree.CloseCursor;
6566
import org.opensearch.sql.ast.tree.Dedupe;
6667
import org.opensearch.sql.ast.tree.Eval;
@@ -763,6 +764,11 @@ public LogicalPlan visitSpath(SPath node, AnalysisContext context) {
763764
throw getOnlyForCalciteException("Spath");
764765
}
765766

767+
@Override
768+
public LogicalPlan visitChart(Chart node, AnalysisContext context) {
769+
throw getOnlyForCalciteException("Chart");
770+
}
771+
766772
@Override
767773
public LogicalPlan visitTimechart(Timechart node, AnalysisContext context) {
768774
throw getOnlyForCalciteException("Timechart");

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@
1212
import lombok.Getter;
1313
import lombok.ToString;
1414
import org.opensearch.sql.ast.AbstractNodeVisitor;
15+
import org.opensearch.sql.ast.dsl.AstDSL;
1516
import org.opensearch.sql.ast.expression.Argument;
17+
import org.opensearch.sql.ast.expression.Literal;
1618
import org.opensearch.sql.ast.expression.UnresolvedExpression;
1719

1820
/** AST node represent chart command. */
@@ -22,6 +24,13 @@
2224
@AllArgsConstructor
2325
@lombok.Builder(toBuilder = true)
2426
public class Chart extends UnresolvedPlan {
27+
public static final Literal DEFAULT_USE_OTHER = Literal.TRUE;
28+
public static final Literal DEFAULT_OTHER_STR = AstDSL.stringLiteral("OTHER");
29+
public static final Literal DEFAULT_LIMIT = AstDSL.intLiteral(10);
30+
public static final Literal DEFAULT_USE_NULL = Literal.TRUE;
31+
public static final Literal DEFAULT_NULL_STR = AstDSL.stringLiteral("NULL");
32+
public static final Literal DEFAULT_TOP = Literal.TRUE;
33+
2534
private UnresolvedPlan child;
2635
private UnresolvedExpression rowSplit;
2736
private UnresolvedExpression columnSplit;

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

Lines changed: 30 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1983,23 +1983,25 @@ public RelNode visitChart(Chart node, CalcitePlanContext context) {
19831983
// If row or column split does not present or limit equals 0, this is the same as `stats agg
19841984
// [group by col]`
19851985

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-
Boolean useOther =
1991-
Optional.ofNullable(argMap.get("useother")).map(u -> (Boolean) u.getValue()).orElse(true);
1992-
String otherStr =
1993-
Optional.ofNullable(argMap.get("otherstr")).map(o -> (String) o.getValue()).orElse("OTHER");
1986+
Integer limit = (Integer) argMap.getOrDefault("limit", Chart.DEFAULT_LIMIT).getValue();
19941987
if (node.getRowSplit() == null || node.getColumnSplit() == null || Objects.equals(limit, 0)) {
19951988
return aggregated.getLeft();
19961989
}
1997-
List<RexNode> projected = aggregated.getRight();
1990+
1991+
Boolean top = (Boolean) argMap.getOrDefault("top", Chart.DEFAULT_TOP).getValue();
1992+
Boolean useOther =
1993+
(Boolean) argMap.getOrDefault("useother", Chart.DEFAULT_USE_OTHER).getValue();
1994+
Boolean useNull = (Boolean) argMap.getOrDefault("usenull", Chart.DEFAULT_USE_NULL).getValue();
1995+
String otherStr = (String) argMap.getOrDefault("otherstr", Chart.DEFAULT_OTHER_STR).getValue();
1996+
String nullStr = (String) argMap.getOrDefault("nullstr", Chart.DEFAULT_NULL_STR).getValue();
1997+
19981998
String columSplitName = aggregated.getLeft().getRowType().getFieldNames().getLast();
19991999
RelBuilder relBuilder = context.relBuilder;
20002000
// 0: agg; 2: column-split
20012001
relBuilder.project(relBuilder.field(0), relBuilder.field(2));
2002-
relBuilder.filter(relBuilder.isNotNull(relBuilder.field(1)));
2002+
if (!useNull) {
2003+
relBuilder.filter(relBuilder.isNotNull(relBuilder.field(1)));
2004+
}
20032005
// 1: column split; 0: agg
20042006
relBuilder.aggregate(
20052007
relBuilder.groupKey(relBuilder.field(1)),
@@ -2008,11 +2010,13 @@ public RelNode visitChart(Chart node, CalcitePlanContext context) {
20082010
if (top) {
20092011
grandTotal = relBuilder.desc(grandTotal);
20102012
}
2013+
// Always set it to null last so that it does not interfere with top / bottom calculation
2014+
grandTotal = relBuilder.nullsLast(grandTotal);
20112015
RexNode rowNum =
20122016
PlanUtils.makeOver(
20132017
context,
20142018
BuiltinFunctionName.ROW_NUMBER,
2015-
relBuilder.literal(1),
2019+
relBuilder.literal(1), // dummy expression for row number calculation
20162020
List.of(),
20172021
List.of(),
20182022
List.of(grandTotal),
@@ -2025,26 +2029,29 @@ public RelNode visitChart(Chart node, CalcitePlanContext context) {
20252029

20262030
// on column-split = group key
20272031
relBuilder.join(
2028-
JoinRelType.INNER, relBuilder.equals(relBuilder.field(2, 0, 2), relBuilder.field(2, 1, 0)));
2032+
JoinRelType.LEFT, relBuilder.equals(relBuilder.field(2, 0, 2), relBuilder.field(2, 1, 0)));
20292033

2030-
RexNode condition =
2034+
RexNode colSplitPostJoin = relBuilder.field(2);
2035+
RexNode lteCondition =
20312036
relBuilder.call(
20322037
SqlStdOperatorTable.LESS_THAN_OR_EQUAL,
20332038
relBuilder.field("__row_number__"),
20342039
relBuilder.literal(limit));
2040+
RexNode nullCondition = relBuilder.isNull(colSplitPostJoin);
20352041
RexNode columnSplitExpr;
2036-
if (useOther) {
2037-
columnSplitExpr =
2038-
relBuilder.call(
2039-
SqlStdOperatorTable.CASE,
2040-
condition,
2041-
relBuilder.field(2),
2042-
relBuilder.literal(otherStr));
2043-
} else {
2044-
relBuilder.filter(condition);
2045-
columnSplitExpr = relBuilder.field(2);
2042+
if (!useOther) {
2043+
relBuilder.filter(lteCondition);
20462044
}
20472045

2046+
columnSplitExpr =
2047+
relBuilder.call(
2048+
SqlStdOperatorTable.CASE,
2049+
nullCondition,
2050+
relBuilder.literal(nullStr),
2051+
lteCondition,
2052+
relBuilder.field(2),
2053+
relBuilder.literal(otherStr));
2054+
20482055
relBuilder.project(
20492056
relBuilder.field(0),
20502057
relBuilder.field(1),

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,8 @@ SHOWCOUNT: 'SHOWCOUNT';
135135
LIMIT: 'LIMIT';
136136
USEOTHER: 'USEOTHER';
137137
OTHERSTR: 'OTHERSTR';
138+
USENULL: 'USENULL';
139+
NULLSTR: 'NULLSTR';
138140
INPUT: 'INPUT';
139141
OUTPUT: 'OUTPUT';
140142
PATH: 'PATH';

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,8 @@ chartOptions
258258
: LIMIT EQUAL (TOP | BOTTOM)? integerLiteral
259259
| USEOTHER EQUAL booleanLiteral
260260
| OTHERSTR EQUAL stringLiteral
261+
| USENULL EQUAL booleanLiteral
262+
| NULLSTR EQUAL stringLiteral
261263
;
262264

263265
rowSplit

‎ppl/src/main/java/org/opensearch/sql/ppl/utils/ArgumentFactory.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,10 @@ public static List<Argument> getArgumentList(ChartCommandContext ctx) {
159159
arguments.add(new Argument("useother", getArgumentValue(optionCtx.booleanLiteral())));
160160
} else if (optionCtx.OTHERSTR() != null) {
161161
arguments.add(new Argument("otherstr", getArgumentValue(optionCtx.stringLiteral())));
162+
} else if (optionCtx.USENULL() != null) {
163+
arguments.add(new Argument("usenull", getArgumentValue(optionCtx.booleanLiteral())));
164+
} else if (optionCtx.NULLSTR() != null) {
165+
arguments.add(new Argument("nullstr", getArgumentValue(optionCtx.stringLiteral())));
162166
}
163167
}
164168
return arguments;

0 commit comments

Comments
 (0)