Skip to content

Commit 7220ed0

Browse files
yuancuyuancu
committed
WIP: Implementing expand command
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 112d5ce commit 7220ed0

11 files changed

Lines changed: 127 additions & 115 deletions

File tree

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,10 @@ public T visitRelationSubquery(RelationSubquery node, C context) {
8585
return visitChildren(node, context);
8686
}
8787

88+
public T visitExpand(Expand expand, C context) {
89+
return visitChildren(expand, context);
90+
}
91+
8892
public T visitTableFunction(TableFunction node, C context) {
8993
return visitChildren(node, context);
9094
}
@@ -317,10 +321,6 @@ public T visitFillNull(FillNull fillNull, C context) {
317321
return visitChildren(fillNull, context);
318322
}
319323

320-
public T visitExpand(Expand expand, C context) {
321-
return visitChildren(expand, context);
322-
}
323-
324324
public T visitPatterns(Patterns patterns, C context) {
325325
return visitChildren(patterns, context);
326326
}

‎core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java‎

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -48,29 +48,9 @@
4848
import org.opensearch.sql.ast.expression.When;
4949
import org.opensearch.sql.ast.expression.WindowFunction;
5050
import org.opensearch.sql.ast.expression.Xor;
51-
import org.opensearch.sql.ast.tree.Aggregation;
52-
import org.opensearch.sql.ast.tree.Dedupe;
53-
import org.opensearch.sql.ast.tree.DescribeRelation;
54-
import org.opensearch.sql.ast.tree.Eval;
55-
import org.opensearch.sql.ast.tree.FillNull;
56-
import org.opensearch.sql.ast.tree.Filter;
57-
import org.opensearch.sql.ast.tree.Head;
58-
import org.opensearch.sql.ast.tree.Limit;
59-
import org.opensearch.sql.ast.tree.Parse;
60-
import org.opensearch.sql.ast.tree.Patterns;
61-
import org.opensearch.sql.ast.tree.Project;
62-
import org.opensearch.sql.ast.tree.RareTopN;
51+
import org.opensearch.sql.ast.tree.*;
6352
import org.opensearch.sql.ast.tree.RareTopN.CommandType;
64-
import org.opensearch.sql.ast.tree.Relation;
65-
import org.opensearch.sql.ast.tree.RelationSubquery;
66-
import org.opensearch.sql.ast.tree.Rename;
67-
import org.opensearch.sql.ast.tree.Sort;
6853
import org.opensearch.sql.ast.tree.Sort.SortOption;
69-
import org.opensearch.sql.ast.tree.SubqueryAlias;
70-
import org.opensearch.sql.ast.tree.TableFunction;
71-
import org.opensearch.sql.ast.tree.Trendline;
72-
import org.opensearch.sql.ast.tree.UnresolvedPlan;
73-
import org.opensearch.sql.ast.tree.Values;
7454

7555
/** Class of static methods to create specific node instances. */
7656
@UtilityClass
@@ -117,6 +97,10 @@ public static Eval eval(UnresolvedPlan input, Let... projectList) {
11797
return new Eval(Arrays.asList(projectList)).attach(input);
11898
}
11999

100+
public Expand expand(UnresolvedPlan input, Field field) {
101+
return new Expand(field).attach(input);
102+
}
103+
120104
public static UnresolvedPlan projectWithArg(
121105
UnresolvedPlan input, List<Argument> argList, UnresolvedExpression... projectList) {
122106
return new Project(Arrays.asList(projectList), argList).attach(input);

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

Lines changed: 24 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,48 +1,40 @@
11
/*
2-
*
3-
* * Copyright OpenSearch Contributors
4-
* * SPDX-License-Identifier: Apache-2.0
5-
*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
64
*/
75

86
package org.opensearch.sql.ast.tree;
97

10-
import lombok.EqualsAndHashCode;
8+
import com.google.common.collect.ImmutableList;
9+
import java.util.List;
1110
import lombok.Getter;
1211
import lombok.RequiredArgsConstructor;
1312
import lombok.ToString;
1413
import org.opensearch.sql.ast.AbstractNodeVisitor;
15-
import org.opensearch.sql.ast.Node;
1614
import org.opensearch.sql.ast.expression.Field;
17-
import org.opensearch.sql.ast.expression.UnresolvedExpression;
18-
19-
import java.util.List;
20-
import java.util.Optional;
2115

22-
/** AST node represent Expand operation. */
23-
@RequiredArgsConstructor
24-
@EqualsAndHashCode(callSuper = false)
16+
/** AST node representing an {@code expand <field>} operation. */
17+
@Getter
2518
@ToString
26-
public class Expand extends UnresolvedPlan{
27-
private UnresolvedPlan child;
28-
@Getter
29-
private final Field field;
30-
@Getter
31-
private final Optional<UnresolvedExpression> alias;
19+
@RequiredArgsConstructor
20+
public class Expand extends UnresolvedPlan {
21+
22+
private UnresolvedPlan child;
23+
@Getter private final Field field;
3224

33-
@Override
34-
public Expand attach(UnresolvedPlan child) {
35-
this.child = child;
36-
return this;
37-
}
25+
@Override
26+
public Expand attach(UnresolvedPlan child) {
27+
this.child = child;
28+
return this;
29+
}
3830

39-
@Override
40-
public List<? extends Node> getChild() {
41-
return child == null ? List.of() : List.of(child);
42-
}
31+
@Override
32+
public List<UnresolvedPlan> getChild() {
33+
return ImmutableList.of(child);
34+
}
4335

44-
@Override
45-
public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
46-
return nodeVisitor.visitExpand(this, context);
47-
}
36+
@Override
37+
public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
38+
return nodeVisitor.visitExpand(this, context);
39+
}
4840
}

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

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.apache.calcite.plan.ViewExpanders;
2929
import org.apache.calcite.rel.RelNode;
3030
import org.apache.calcite.rel.core.Aggregate;
31+
import org.apache.calcite.rel.core.JoinRelType;
3132
import org.apache.calcite.rel.type.RelDataTypeField;
3233
import org.apache.calcite.rex.RexCall;
3334
import org.apache.calcite.rex.RexCorrelVariable;
@@ -164,7 +165,7 @@ public RelNode visitProject(Project node, CalcitePlanContext context) {
164165
}
165166

166167
/** See logic in {@link org.opensearch.sql.analysis.symbol.SymbolTable#lookupAllFields} */
167-
private void tryToRemoveNestedFields(CalcitePlanContext context) {
168+
private static void tryToRemoveNestedFields(CalcitePlanContext context) {
168169
Set<String> allFields = new HashSet<>(context.relBuilder.peek().getRowType().getFieldNames());
169170
List<RexNode> duplicatedNestedFields =
170171
allFields.stream()
@@ -824,8 +825,52 @@ public RelNode visitTrendline(Trendline node, CalcitePlanContext context) {
824825
throw new CalciteUnsupportedException("Trendline command is unsupported in Calcite");
825826
}
826827

828+
/**
829+
* Expand command visitor to handle array field expansion. 1. Unnest 2. Join with the original
830+
* table to get all fields
831+
*
832+
* <p>S = π_{field, other_fields}(R ⨝ UNNEST_field(R))
833+
*
834+
* @param expand Expand command to be visited
835+
* @param context CalcitePlanContext containing the RelBuilder and other context
836+
* @return RelNode representing records with the expanded array field
837+
*/
827838
@Override
828839
public RelNode visitExpand(Expand expand, CalcitePlanContext context) {
829-
throw new CalciteUnsupportedException("Expand command is unsupported in Calcite");
840+
// 1. Visit Children
841+
visitChildren(expand, context);
842+
843+
var relBuilder = context.relBuilder;
844+
845+
// 3. Get the field to expand
846+
Field arrayField = expand.getField();
847+
848+
// 5. Unnest the array field
849+
// Analyze the array field to get its RexNode
850+
RexNode arrayFieldRex = rexVisitor.analyze(arrayField, context);
851+
852+
// Push the original table to the RelBuilder stack
853+
RelNode originalTable = relBuilder.peek();
854+
// No alias is provided in the expand command, so we remove the original array field,
855+
// then replace it with the unnest result.
856+
relBuilder.projectExcept(arrayFieldRex);
857+
relBuilder.push(originalTable);
858+
859+
// Join on ROW_NUMBER_COLUMN_NAME
860+
Holder<RexCorrelVariable> correlVariable = Holder.empty();
861+
relBuilder.variable(correlVariable::set);
862+
863+
relBuilder.project(List.of(arrayFieldRex), List.of(), false, List.of(correlVariable.get().id));
864+
// Alias is not supported in expand yet, we pass in an empty list
865+
relBuilder.uncollect(List.of(), false);
866+
867+
List<RexNode> allFields =
868+
relBuilder.peek().getRowType().getFieldList().stream()
869+
.map(f -> (RexNode) relBuilder.field(f.getName()))
870+
.toList();
871+
872+
relBuilder.correlate(JoinRelType.INNER, correlVariable.get().id, relBuilder.fields());
873+
874+
return relBuilder.peek();
830875
}
831876
}

‎core/src/main/java/org/opensearch/sql/planner/logical/LogicalExpand.java‎

Lines changed: 11 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,30 +7,25 @@
77

88
package org.opensearch.sql.planner.logical;
99

10+
import java.util.Collections;
1011
import lombok.EqualsAndHashCode;
1112
import lombok.Getter;
1213
import lombok.ToString;
1314
import org.opensearch.sql.expression.Expression;
1415

15-
import java.util.Collections;
16-
1716
@ToString
1817
@EqualsAndHashCode(callSuper = true)
19-
public class LogicalExpand extends LogicalPlan{
18+
public class LogicalExpand extends LogicalPlan {
2019

21-
@Getter
22-
private final Expression field;
23-
@Getter
24-
private final Expression alias;
20+
@Getter private final Expression field;
2521

26-
public LogicalExpand(LogicalPlan child, Expression field, Expression alias) {
27-
super(Collections.singletonList(child));
28-
this.field = field;
29-
this.alias = alias;
30-
}
22+
public LogicalExpand(LogicalPlan child, Expression field) {
23+
super(Collections.singletonList(child));
24+
this.field = field;
25+
}
3126

32-
@Override
33-
public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
34-
return visitor.visitExpand(this, context);
35-
}
27+
@Override
28+
public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
29+
return visitor.visitExpand(this, context);
30+
}
3631
}

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,10 @@
1010
import org.opensearch.sql.ppl.ExpandCommandIT;
1111

1212
public class CalciteExpandIT extends ExpandCommandIT {
13-
@Override
14-
public void init() throws Exception {
15-
super.init();
16-
enableCalcite();
17-
disallowCalciteFallback();
18-
}
13+
@Override
14+
public void init() throws Exception {
15+
super.init();
16+
enableCalcite();
17+
disallowCalciteFallback();
18+
}
1919
}

‎integ-test/src/test/java/org/opensearch/sql/ppl/ExpandCommandIT.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import static org.opensearch.sql.util.MatcherUtils.verifySchema;
1414

1515
import org.json.JSONObject;
16+
import org.junit.Ignore;
1617
import org.junit.jupiter.api.Test;
1718

1819
public class ExpandCommandIT extends PPLIntegTestCase {
@@ -27,10 +28,16 @@ public void testExpand() throws Exception {
2728
JSONObject response =
2829
executeQuery(String.format("source=%s | expand address", TEST_INDEX_NESTED_SIMPLE));
2930
verifySchema(
30-
response, schema("name", "string"), schema("age", "integer"), schema("id", "integer"), schema("address", "object"));
31+
response,
32+
schema("name", "string"),
33+
schema("age", "integer"),
34+
schema("id", "integer"),
35+
schema("address", "object"));
3136
verifyNumOfRows(response, 11);
3237
}
3338

39+
// TODO: double check if expand with alias is supported
40+
@Ignore
3441
@Test
3542
public void testExpandWithAlias() throws Exception {
3643
JSONObject response =

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,8 @@ KMEANS: 'KMEANS';
3737
AD: 'AD';
3838
ML: 'ML';
3939
FILLNULL: 'FILLNULL';
40-
EXPAND: 'EXPAND';
4140
TRENDLINE: 'TRENDLINE';
41+
EXPAND: 'EXPAND';
4242
SIMPLE_PATTERN: 'SIMPLE_PATTERN';
4343
BRAIN: 'BRAIN';
4444
VARIABLE_COUNT_THRESHOLD: 'VARIABLE_COUNT_THRESHOLD';

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,8 @@ commands
6969
| adCommand
7070
| mlCommand
7171
| fillnullCommand
72-
| expandCommand
7372
| trendlineCommand
73+
| expandCommand
7474
;
7575

7676
commandName
@@ -223,10 +223,6 @@ fillNullUsing
223223
: USING replacementPair (COMMA replacementPair)*
224224
;
225225

226-
expandCommand
227-
: EXPAND fieldExpression (AS alias = qualifiedName)?
228-
;
229-
230226
replacementPair
231227
: fieldExpression EQUAL replacement = valueExpression
232228
;
@@ -243,6 +239,10 @@ trendlineType
243239
: SMA
244240
;
245241

242+
expandCommand
243+
: EXPAND fieldExpression
244+
;
245+
246246
kmeansCommand
247247
: KMEANS (kmeansParameter)*
248248
;

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -415,6 +415,13 @@ public UnresolvedPlan visitTopCommand(OpenSearchPPLParser.TopCommandContext ctx)
415415
groupList);
416416
}
417417

418+
/** expand command. */
419+
@Override
420+
public UnresolvedPlan visitExpandCommand(OpenSearchPPLParser.ExpandCommandContext ctx) {
421+
Field fieldExpression = (Field) internalVisitExpression(ctx.fieldExpression());
422+
return new Expand(fieldExpression);
423+
}
424+
418425
@Override
419426
public UnresolvedPlan visitGrokCommand(OpenSearchPPLParser.GrokCommandContext ctx) {
420427
UnresolvedExpression sourceField = internalVisitExpression(ctx.source_field);
@@ -628,13 +635,6 @@ public UnresolvedPlan visitFillNullUsing(OpenSearchPPLParser.FillNullUsingContex
628635
return FillNull.ofVariousValue(replacementsBuilder.build());
629636
}
630637

631-
/** expand command. */
632-
@Override
633-
public UnresolvedPlan visitExpandCommand(OpenSearchPPLParser.ExpandCommandContext ctx) {
634-
return new Expand((Field) internalVisitExpression(ctx.fieldExpression()),
635-
ctx.alias!=null ? Optional.of(internalVisitExpression(ctx.alias)) : Optional.empty());
636-
}
637-
638638
/** trendline command. */
639639
@Override
640640
public UnresolvedPlan visitTrendlineCommand(OpenSearchPPLParser.TrendlineCommandContext ctx) {

0 commit comments

Comments
 (0)