Skip to content

Commit

Permalink
Add visitFirstChild helper method
Browse files Browse the repository at this point in the history 
Signed-off-by: currantw <[email protected]>
  • Loading branch information
@currantw
currantw committed Dec 5, 2024
1 parent 3634175 commit 5aca49d
Showing 1 changed file with 28 additions and 17 deletions.
45 changes: 28 additions & 17 deletions core/src/main/java/org/opensearch/sql/analysis/Analyzer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
import org.opensearch.sql.analysis.symbol.Namespace;
import org.opensearch.sql.analysis.symbol.Symbol;
import org.opensearch.sql.ast.AbstractNodeVisitor;
import org.opensearch.sql.ast.Node;
import org.opensearch.sql.ast.expression.Argument;
import org.opensearch.sql.ast.expression.Field;
import org.opensearch.sql.ast.expression.Let;
Expand Down Expand Up @@ -174,7 +175,7 @@ public LogicalPlan visitRelation(Relation node, AnalysisContext context) {

@Override
public LogicalPlan visitRelationSubquery(RelationSubquery node, AnalysisContext context) {
LogicalPlan subquery = analyze(node.getChild().getFirst(), context);
LogicalPlan subquery = visitFirstChild(node, context);
// inherit the parent environment to keep the subquery fields in current environment
TypeEnvironment curEnv = context.peek();

Expand Down Expand Up @@ -227,13 +228,13 @@ public LogicalPlan visitTableFunction(TableFunction node, AnalysisContext contex

@Override
public LogicalPlan visitLimit(Limit node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
return new LogicalLimit(child, node.getLimit(), node.getOffset());
}

@Override
public LogicalPlan visitFilter(Filter node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
Expression condition = expressionAnalyzer.analyze(node.getCondition(), context);

ExpressionReferenceOptimizer optimizer =
Expand Down Expand Up @@ -266,7 +267,7 @@ private void verifySupportsCondition(Expression condition) {
/** Build {@link LogicalRename}. */
@Override
public LogicalPlan visitRename(Rename node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
ImmutableMap.Builder<ReferenceExpression, ReferenceExpression> renameMapBuilder =
new ImmutableMap.Builder<>();
for (Map renameMap : node.getRenameList()) {
Expand All @@ -293,7 +294,7 @@ public LogicalPlan visitRename(Rename node, AnalysisContext context) {
/** Build {@link LogicalAggregation}. */
@Override
public LogicalPlan visitAggregation(Aggregation node, AnalysisContext context) {
final LogicalPlan child = node.getChild().getFirst().accept(this, context);
final LogicalPlan child = visitFirstChild(node, context);
ImmutableList.Builder<NamedAggregator> aggregatorBuilder = new ImmutableList.Builder<>();
for (UnresolvedExpression expr : node.getAggExprList()) {
NamedExpression aggExpr = namedExpressionAnalyzer.analyze(expr, context);
Expand Down Expand Up @@ -331,7 +332,7 @@ public LogicalPlan visitAggregation(Aggregation node, AnalysisContext context) {
/** Build {@link LogicalRareTopN}. */
@Override
public LogicalPlan visitRareTopN(RareTopN node, AnalysisContext context) {
final LogicalPlan child = node.getChild().getFirst().accept(this, context);
final LogicalPlan child = visitFirstChild(node, context);

ImmutableList.Builder<Expression> groupbyBuilder = new ImmutableList.Builder<>();
for (UnresolvedExpression expr : node.getGroupExprList()) {
Expand Down Expand Up @@ -372,7 +373,7 @@ public LogicalPlan visitRareTopN(RareTopN node, AnalysisContext context) {
*/
@Override
public LogicalPlan visitProject(Project node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);

if (node.hasArgument()) {
Argument argument = node.getArgExprList().getFirst();
Expand Down Expand Up @@ -426,7 +427,7 @@ public LogicalPlan visitProject(Project node, AnalysisContext context) {
/** Build {@link LogicalEval}. */
@Override
public LogicalPlan visitEval(Eval node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
ImmutableList.Builder<Pair<ReferenceExpression, Expression>> expressionsBuilder =
new Builder<>();
for (Let let : node.getExpressionList()) {
Expand All @@ -443,7 +444,7 @@ public LogicalPlan visitEval(Eval node, AnalysisContext context) {
/** Build {@link ParseExpression} to context and skip to child nodes. */
@Override
public LogicalPlan visitParse(Parse node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
Expression sourceField = expressionAnalyzer.analyze(node.getSourceField(), context);
ParseMethod parseMethod = node.getParseMethod();
java.util.Map<String, Literal> arguments = node.getArguments();
Expand All @@ -466,7 +467,7 @@ public LogicalPlan visitParse(Parse node, AnalysisContext context) {
/** Build {@link LogicalSort}. */
@Override
public LogicalPlan visitSort(Sort node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
ExpressionReferenceOptimizer optimizer =
new ExpressionReferenceOptimizer(expressionAnalyzer.getRepository(), child);

Expand All @@ -489,7 +490,7 @@ public LogicalPlan visitSort(Sort node, AnalysisContext context) {
/** Build {@link LogicalDedupe}. */
@Override
public LogicalPlan visitDedupe(Dedupe node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
List<Argument> options = node.getOptions();
// Todo, refactor the option.
Integer allowedDuplication = (Integer) options.get(0).getValue().getValue();
Expand All @@ -508,7 +509,7 @@ public LogicalPlan visitDedupe(Dedupe node, AnalysisContext context) {

/** Logical head is identical to {@link LogicalLimit}. */
public LogicalPlan visitHead(Head node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
return new LogicalLimit(child, node.getSize(), node.getFrom());
}

Expand All @@ -528,7 +529,7 @@ public LogicalPlan visitValues(Values node, AnalysisContext context) {
/** Build {@link LogicalMLCommons} for Kmeans command. */
@Override
public LogicalPlan visitKmeans(Kmeans node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
java.util.Map<String, Literal> options = node.getArguments();

TypeEnvironment currentEnv = context.peek();
Expand All @@ -540,7 +541,7 @@ public LogicalPlan visitKmeans(Kmeans node, AnalysisContext context) {
/** Build {@link LogicalAD} for AD command. */
@Override
public LogicalPlan visitAD(AD node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
java.util.Map<String, Literal> options = node.getArguments();

TypeEnvironment currentEnv = context.peek();
Expand All @@ -560,7 +561,7 @@ public LogicalPlan visitAD(AD node, AnalysisContext context) {
/** Build {@link LogicalML} for ml command. */
@Override
public LogicalPlan visitML(ML node, AnalysisContext context) {
LogicalPlan child = node.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(node, context);
TypeEnvironment currentEnv = context.peek();
node.getOutputSchema(currentEnv)
.forEach((key, value) -> currentEnv.define(new Symbol(Namespace.FIELD_NAME, key), value));
Expand All @@ -570,7 +571,7 @@ public LogicalPlan visitML(ML node, AnalysisContext context) {

@Override
public LogicalPlan visitPaginate(Paginate paginate, AnalysisContext context) {
LogicalPlan child = paginate.getChild().getFirst().accept(this, context);
LogicalPlan child = visitFirstChild(paginate, context);
return new LogicalPaginate(paginate.getPageSize(), List.of(child));
}

Expand All @@ -583,7 +584,8 @@ public LogicalPlan visitFetchCursor(FetchCursor cursor, AnalysisContext context)

@Override
public LogicalPlan visitCloseCursor(CloseCursor closeCursor, AnalysisContext context) {
return new LogicalCloseCursor(closeCursor.getChild().getFirst().accept(this, context));
LogicalPlan child = visitFirstChild(closeCursor, context);
return new LogicalCloseCursor(child);
}

/**
Expand All @@ -601,4 +603,13 @@ private SortOption analyzeSortOption(List<Argument> fieldArgs) {
}
return asc ? SortOption.DEFAULT_ASC : SortOption.DEFAULT_DESC;
}

/**
* Visit the first child of the specified node with the given analysis context.
*
* @return the resulting logical plan
*/
private LogicalPlan visitFirstChild(Node node, AnalysisContext context) {
return node.getChild().getFirst().accept(this, context);
}
}

0 comments on commit 5aca49d

Please sign in to comment.