Skip to content

Commit b692c60

Browse files
ioanatiaioanatia
committed
Make Fork n-ary
1 parent fc933d4 commit b692c60

File tree

13 files changed

+91
-265
lines changed

13 files changed

+91
-265
lines changed

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java‎

Lines changed: 2 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -154,15 +154,7 @@ public class Analyzer extends ParameterizedRuleExecutor<LogicalPlan, AnalyzerCon
154154
);
155155

156156
private static final List<Batch<LogicalPlan>> RULES = List.of(
157-
new Batch<>(
158-
"Initialize",
159-
Limiter.ONCE,
160-
new ResolveTable(),
161-
new ResolveEnrich(),
162-
new ResolveLookupTables(),
163-
new ResolveFunctions(),
164-
new ResolveForkFunctions()
165-
),
157+
new Batch<>("Initialize", Limiter.ONCE, new ResolveTable(), new ResolveEnrich(), new ResolveLookupTables(), new ResolveFunctions()),
166158
new Batch<>(
167159
"Resolution",
168160
/*
@@ -171,7 +163,6 @@ public class Analyzer extends ParameterizedRuleExecutor<LogicalPlan, AnalyzerCon
171163
* attempts to resolve this reference.
172164
*/
173165
new ImplicitCasting(),
174-
new ImplicitForkCasting(),
175166
new ResolveRefs(),
176167
new ResolveUnionTypes() // Must be after ResolveRefs, so union types can be found
177168
),
@@ -492,10 +483,6 @@ protected LogicalPlan rule(LogicalPlan plan, AnalyzerContext context) {
492483
return resolveInsist(i, childrenOutput, context.indexResolution());
493484
}
494485

495-
if (plan instanceof Fork f) {
496-
return resolveFork(f, context);
497-
}
498-
499486
if (plan instanceof Dedup dedup) {
500487
return resolveDedup(dedup, childrenOutput);
501488
}
@@ -683,16 +670,6 @@ private Join resolveLookupJoin(LookupJoin join) {
683670
return join;
684671
}
685672

686-
private LogicalPlan resolveFork(Fork fork, AnalyzerContext context) {
687-
List<LogicalPlan> subPlans = fork.subPlans();
688-
689-
List<LogicalPlan> newSubPlans = new ArrayList<>();
690-
for (var logicalPlan : subPlans) {
691-
newSubPlans.add(logicalPlan.transformUp(LogicalPlan.class, p -> p.childrenResolved() == false ? p : rule(p, context)));
692-
}
693-
return new Fork(fork.source(), fork.child(), newSubPlans);
694-
}
695-
696673
private List<Attribute> resolveUsingColumns(List<Attribute> cols, List<Attribute> output, String side) {
697674
List<Attribute> resolved = new ArrayList<>(cols.size());
698675
for (Attribute col : cols) {
@@ -1154,23 +1131,6 @@ public static org.elasticsearch.xpack.esql.core.expression.function.Function res
11541131
}
11551132
}
11561133

1157-
private static class ResolveForkFunctions extends ParameterizedAnalyzerRule<LogicalPlan, AnalyzerContext> {
1158-
private final ResolveFunctions resolveFunctions = new ResolveFunctions();
1159-
1160-
@Override
1161-
protected LogicalPlan rule(LogicalPlan plan, AnalyzerContext context) {
1162-
return plan.transformUp(Fork.class, fork -> resolveFunctionsInForkSubQueries(fork, context));
1163-
}
1164-
1165-
private LogicalPlan resolveFunctionsInForkSubQueries(Fork fork, AnalyzerContext ctx) {
1166-
List<LogicalPlan> newSubPlans = new ArrayList<>();
1167-
for (var subPlan : fork.subPlans()) {
1168-
newSubPlans.add(resolveFunctions.apply(subPlan, ctx));
1169-
}
1170-
return fork.replaceSubPlans(newSubPlans);
1171-
}
1172-
}
1173-
11741134
private static class AddImplicitLimit extends ParameterizedRule<LogicalPlan, LogicalPlan, AnalyzerContext> {
11751135
@Override
11761136
public LogicalPlan apply(LogicalPlan logicalPlan, AnalyzerContext context) {
@@ -1200,7 +1160,7 @@ public LogicalPlan apply(LogicalPlan logicalPlan, AnalyzerContext context) {
12001160

12011161
private LogicalPlan addImplicitLimitToForkSubQueries(Fork fork, AnalyzerContext ctx) {
12021162
List<LogicalPlan> newSubPlans = new ArrayList<>();
1203-
for (var subPlan : fork.subPlans()) {
1163+
for (var subPlan : fork.children()) {
12041164
newSubPlans.add(addImplicitLimit.apply(subPlan, ctx));
12051165
}
12061166
return fork.replaceSubPlans(newSubPlans);
@@ -1458,23 +1418,6 @@ private static Expression castStringLiteral(Expression from, DataType target) {
14581418
}
14591419
}
14601420

1461-
private static class ImplicitForkCasting extends ParameterizedRule<LogicalPlan, LogicalPlan, AnalyzerContext> {
1462-
private final ImplicitCasting implicitCasting = new ImplicitCasting();
1463-
1464-
@Override
1465-
public LogicalPlan apply(LogicalPlan logicalPlan, AnalyzerContext context) {
1466-
return logicalPlan.transformUp(Fork.class, fork -> implicitCastForkSubQueries(fork, context));
1467-
}
1468-
1469-
private LogicalPlan implicitCastForkSubQueries(Fork fork, AnalyzerContext ctx) {
1470-
List<LogicalPlan> newSubPlans = new ArrayList<>();
1471-
for (var subPlan : fork.subPlans()) {
1472-
newSubPlans.add(implicitCasting.apply(subPlan, ctx));
1473-
}
1474-
return fork.replaceSubPlans(newSubPlans);
1475-
}
1476-
}
1477-
14781421
/**
14791422
* The EsqlIndexResolver will create InvalidMappedField instances for fields that are ambiguous (i.e. have multiple mappings).
14801423
* During {@link ResolveRefs} we do not convert these to UnresolvedAttribute instances, as we want to first determine if they can

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/PreAnalyzer.java‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,10 @@
1313
import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
1414

1515
import java.util.ArrayList;
16+
import java.util.Collection;
17+
import java.util.HashSet;
1618
import java.util.List;
19+
import java.util.Set;
1720

1821
import static java.util.Collections.emptyList;
1922

@@ -45,19 +48,20 @@ public PreAnalysis preAnalyze(LogicalPlan plan) {
4548
}
4649

4750
protected PreAnalysis doPreAnalyze(LogicalPlan plan) {
48-
List<TableInfo> indices = new ArrayList<>();
51+
Set<TableInfo> indices = new HashSet<>();
4952
List<Enrich> unresolvedEnriches = new ArrayList<>();
5053
List<TableInfo> lookupIndices = new ArrayList<>();
5154

5255
plan.forEachUp(UnresolvedRelation.class, p -> {
53-
List<TableInfo> list = p.indexMode() == IndexMode.LOOKUP ? lookupIndices : indices;
54-
list.add(new TableInfo(p.indexPattern()));
56+
Collection<TableInfo> collection = p.indexMode() == IndexMode.LOOKUP ? lookupIndices : indices;
57+
collection.add(new TableInfo(p.indexPattern()));
5558
});
59+
5660
plan.forEachUp(Enrich.class, unresolvedEnriches::add);
5761

5862
// mark plan as preAnalyzed (if it were marked, there would be no analysis)
5963
plan.forEachUp(LogicalPlan::setPreAnalyzed);
6064

61-
return new PreAnalysis(indices, unresolvedEnriches, lookupIndices);
65+
return new PreAnalysis(indices.stream().toList(), unresolvedEnriches, lookupIndices);
6266
}
6367
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/TableInfo.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99

1010
import org.elasticsearch.xpack.esql.plan.IndexPattern;
1111

12+
import java.util.Objects;
13+
1214
public class TableInfo {
1315

1416
private final IndexPattern id;
@@ -20,4 +22,23 @@ public TableInfo(IndexPattern id) {
2022
public IndexPattern id() {
2123
return id;
2224
}
25+
26+
@Override
27+
public int hashCode() {
28+
return Objects.hash(id);
29+
}
30+
31+
@Override
32+
public boolean equals(Object obj) {
33+
if (this == obj) {
34+
return true;
35+
}
36+
37+
if (obj == null || getClass() != obj.getClass()) {
38+
return false;
39+
}
40+
41+
TableInfo other = (TableInfo) obj;
42+
return Objects.equals(id, other.id);
43+
}
2344
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@
3030
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.NotEquals;
3131
import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
3232
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
33-
import org.elasticsearch.xpack.esql.plan.logical.Fork;
3433
import org.elasticsearch.xpack.esql.plan.logical.Insist;
3534
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
3635
import org.elasticsearch.xpack.esql.plan.logical.Lookup;
@@ -180,11 +179,6 @@ else if (p instanceof Lookup lookup) {
180179
else {
181180
lookup.matchFields().forEach(unresolvedExpressions);
182181
}
183-
} else if (p instanceof Fork fork) {
184-
var subPlans = fork.subPlans();
185-
for (var subPlan : subPlans) {
186-
checkUnresolvedAttributes(subPlan, failures);
187-
}
188182
}
189183

190184
else {

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryBuilderResolver.java‎

Lines changed: 6 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,8 @@
1212
import org.elasticsearch.index.query.QueryBuilder;
1313
import org.elasticsearch.index.query.QueryRewriteContext;
1414
import org.elasticsearch.index.query.Rewriteable;
15-
import org.elasticsearch.xpack.esql.core.expression.Expression;
1615
import org.elasticsearch.xpack.esql.core.util.Holder;
1716
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
18-
import org.elasticsearch.xpack.esql.plan.logical.Fork;
1917
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
2018
import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
2119
import org.elasticsearch.xpack.esql.plugin.TransportActionServices;
@@ -24,8 +22,6 @@
2422
import java.io.IOException;
2523
import java.util.HashSet;
2624
import java.util.Set;
27-
import java.util.function.Function;
28-
import java.util.stream.Collectors;
2925

3026
/**
3127
* Some {@link FullTextFunction} implementations such as {@link org.elasticsearch.xpack.esql.expression.function.fulltext.Match}
@@ -38,7 +34,11 @@ public final class QueryBuilderResolver {
3834
private QueryBuilderResolver() {}
3935

4036
public static void resolveQueryBuilders(LogicalPlan plan, TransportActionServices services, ActionListener<LogicalPlan> listener) {
41-
var hasFullTextFunctions = hasFullTextFunctions(plan);
37+
var hasFullTextFunctions = plan.anyMatch(p -> {
38+
Holder<Boolean> hasFullTextFunction = new Holder<>(false);
39+
p.forEachExpression(FullTextFunction.class, unused -> hasFullTextFunction.set(true));
40+
return hasFullTextFunction.get();
41+
});
4242
if (hasFullTextFunctions) {
4343
Rewriteable.rewriteAndFetch(
4444
new FullTextFunctionsRewritable(plan),
@@ -69,29 +69,12 @@ private static Set<String> indexNames(LogicalPlan plan) {
6969
return indexNames;
7070
}
7171

72-
private static boolean hasFullTextFunctions(LogicalPlan plan) {
73-
return plan.anyMatch(p -> {
74-
Holder<Boolean> hasFullTextFunction = new Holder<>(false);
75-
p.forEachExpression(FullTextFunction.class, unused -> hasFullTextFunction.set(true));
76-
77-
if (p instanceof Fork fork) {
78-
fork.subPlans().forEach(subPlan -> {
79-
if (hasFullTextFunctions(subPlan)) {
80-
hasFullTextFunction.set(true);
81-
}
82-
});
83-
}
84-
85-
return hasFullTextFunction.get();
86-
});
87-
}
88-
8972
private record FullTextFunctionsRewritable(LogicalPlan plan) implements Rewriteable<QueryBuilderResolver.FullTextFunctionsRewritable> {
9073
@Override
9174
public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOException {
9275
Holder<IOException> exceptionHolder = new Holder<>();
9376
Holder<Boolean> updated = new Holder<>(false);
94-
LogicalPlan newPlan = transformPlan(plan, f -> {
77+
LogicalPlan newPlan = plan.transformExpressionsDown(FullTextFunction.class, f -> {
9578
QueryBuilder builder = f.queryBuilder(), initial = builder;
9679
builder = builder == null ? f.asQuery(TranslatorHandler.TRANSLATOR_HANDLER).toQueryBuilder() : builder;
9780
try {
@@ -108,15 +91,5 @@ public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOExc
10891
}
10992
return updated.get() ? new FullTextFunctionsRewritable(newPlan) : this;
11093
}
111-
112-
private LogicalPlan transformPlan(LogicalPlan plan, Function<FullTextFunction, ? extends Expression> rule) {
113-
return plan.transformExpressionsDown(FullTextFunction.class, rule).transformDown(Fork.class, fork -> {
114-
var subPlans = fork.subPlans()
115-
.stream()
116-
.map(subPlan -> subPlan.transformExpressionsDown(FullTextFunction.class, rule))
117-
.collect(Collectors.toList());
118-
return new Fork(fork.source(), fork.child(), subPlans);
119-
});
120-
}
12194
}
12295
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,6 @@
6868
import org.elasticsearch.xpack.esql.plan.logical.RrfScoreEval;
6969
import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
7070
import org.elasticsearch.xpack.esql.plan.logical.join.LookupJoin;
71-
import org.elasticsearch.xpack.esql.plan.logical.join.StubRelation;
7271
import org.elasticsearch.xpack.esql.plan.logical.show.ShowInfo;
7372
import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
7473
import org.joni.exception.SyntaxException;
@@ -630,9 +629,8 @@ public PlanFactory visitForkCommand(EsqlBaseParser.ForkCommandContext ctx) {
630629
throw new ParsingException(source(ctx), "Fork requires at least two branches");
631630
}
632631
return input -> {
633-
var stub = StubRelation.EMPTY;
634-
List<LogicalPlan> subPlans = subQueries.stream().map(planFactory -> planFactory.apply(stub)).toList();
635-
return new Fork(source(ctx), input, subPlans);
632+
List<LogicalPlan> subPlans = subQueries.stream().map(planFactory -> planFactory.apply(input)).toList();
633+
return new Fork(source(ctx), subPlans);
636634
};
637635
}
638636

0 commit comments

Comments
 (0)