Skip to content

Commit 170d83c

Browse files
smalyshevsmalyshev
committed
refactor
1 parent 2c051b5 commit 170d83c

File tree

5 files changed

+39
-42
lines changed

5 files changed

+39
-42
lines changed

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Enrich.java‎

Lines changed: 24 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import org.elasticsearch.index.IndexMode;
1818
import org.elasticsearch.transport.RemoteClusterAware;
1919
import org.elasticsearch.xpack.core.enrich.EnrichPolicy;
20-
import org.elasticsearch.xpack.esql.capabilities.PostAnalysisPlanVerificationAware;
20+
import org.elasticsearch.xpack.esql.capabilities.PostAnalysisVerificationAware;
2121
import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
2222
import org.elasticsearch.xpack.esql.common.Failures;
2323
import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
@@ -35,7 +35,6 @@
3535
import org.elasticsearch.xpack.esql.index.EsIndex;
3636
import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
3737
import org.elasticsearch.xpack.esql.plan.GeneratingPlan;
38-
import org.elasticsearch.xpack.esql.plan.logical.join.LookupJoin;
3938

4039
import java.io.IOException;
4140
import java.util.ArrayList;
@@ -45,13 +44,18 @@
4544
import java.util.Map;
4645
import java.util.Objects;
4746
import java.util.Set;
48-
import java.util.function.BiConsumer;
4947

5048
import static org.elasticsearch.xpack.esql.common.Failure.fail;
5149
import static org.elasticsearch.xpack.esql.core.expression.Expressions.asAttributes;
5250
import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
5351

54-
public class Enrich extends UnaryPlan implements GeneratingPlan<Enrich>, PostAnalysisPlanVerificationAware, TelemetryAware, SortAgnostic {
52+
public class Enrich extends UnaryPlan
53+
implements
54+
GeneratingPlan<Enrich>,
55+
PostAnalysisVerificationAware,
56+
TelemetryAware,
57+
SortAgnostic,
58+
ExecutesOn {
5559
public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
5660
LogicalPlan.class,
5761
"Enrich",
@@ -68,6 +72,16 @@ public class Enrich extends UnaryPlan implements GeneratingPlan<Enrich>, PostAna
6872

6973
private final Mode mode;
7074

75+
@Override
76+
public ExecuteLocation executesOn() {
77+
if (mode == Mode.REMOTE) {
78+
return ExecuteLocation.REMOTE;
79+
} else if (mode == Mode.COORDINATOR) {
80+
return ExecuteLocation.COORDINATOR;
81+
}
82+
return ExecuteLocation.ANY;
83+
}
84+
7185
public enum Mode {
7286
ANY,
7387
COORDINATOR,
@@ -279,11 +293,6 @@ public int hashCode() {
279293
return Objects.hash(super.hashCode(), mode, policyName, matchField, policy, concreteIndices, enrichFields);
280294
}
281295

282-
@Override
283-
public BiConsumer<LogicalPlan, Failures> postAnalysisPlanVerification() {
284-
return Enrich::checkRemoteEnrich;
285-
}
286-
287296
/**
288297
* Ensure that no remote enrich is allowed after a reduction or an enrich with coordinator mode.
289298
* <p>
@@ -296,15 +305,6 @@ public BiConsumer<LogicalPlan, Failures> postAnalysisPlanVerification() {
296305
* We might consider implementing the actual remote enrich on the coordinating cluster, however, this requires
297306
* retaining the originating cluster and restructing pages for routing, which might be complicated.
298307
*/
299-
private static void checkRemoteEnrich(LogicalPlan plan, Failures failures) {
300-
// First look for remote ENRICH, and then look at its children. Going over the whole plan once is trickier as remote ENRICHs can be
301-
// in separate FORK branches which are valid by themselves.
302-
plan.forEachUp(Enrich.class, enrich -> checkForPlansForbiddenBeforeRemoteEnrich(enrich, failures));
303-
}
304-
305-
/**
306-
* For a given remote {@link Enrich}, check if there are any forbidden plans upstream.
307-
*/
308308
private static void checkForPlansForbiddenBeforeRemoteEnrich(Enrich enrich, Failures failures) {
309309
if (enrich.mode != Mode.REMOTE) {
310310
return;
@@ -313,25 +313,18 @@ private static void checkForPlansForbiddenBeforeRemoteEnrich(Enrich enrich, Fail
313313
Set<Source> badCommands = new HashSet<>();
314314

315315
enrich.forEachUp(LogicalPlan.class, u -> {
316-
if (u instanceof ExecutesOn.Coordinator
317-
|| u instanceof Enrich upstreamEnrich && upstreamEnrich.mode() == Enrich.Mode.COORDINATOR
318-
|| u instanceof LookupJoin) {
316+
if (u instanceof ExecutesOn ex && ex.executesOn() == ExecuteLocation.COORDINATOR) {
319317
badCommands.add(u.source());
320318
}
321-
322-
// if (u instanceof Aggregate) {
323-
// badCommands.add("STATS");
324-
// } else if (u instanceof Enrich upstreamEnrich && upstreamEnrich.mode() == Enrich.Mode.COORDINATOR) {
325-
// badCommands.add("another ENRICH with coordinator policy");
326-
// } else if (u instanceof LookupJoin) {
327-
// badCommands.add("LOOKUP JOIN");
328-
// } else if (u instanceof Fork) {
329-
// badCommands.add("FORK");
330-
// }
331319
});
332320

333321
badCommands.forEach(
334322
f -> failures.add(fail(enrich, "ENRICH with remote policy can't be executed after [" + f.text() + "]" + f.source()))
335323
);
336324
}
325+
326+
@Override
327+
public void postAnalysisVerification(Failures failures) {
328+
checkForPlansForbiddenBeforeRemoteEnrich(this, failures);
329+
}
337330
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/ExecutesOn.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@
1313
public interface ExecutesOn {
1414
enum ExecuteLocation {
1515
COORDINATOR,
16-
REMOTE
16+
REMOTE,
17+
ANY; // Can be executed on either coordinator or remote nodes
1718
}
1819

1920
ExecuteLocation executesOn();

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/Join.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.elasticsearch.xpack.esql.core.type.DataType;
2222
import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
2323
import org.elasticsearch.xpack.esql.plan.logical.BinaryPlan;
24+
import org.elasticsearch.xpack.esql.plan.logical.ExecutesOn;
2425
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
2526
import org.elasticsearch.xpack.esql.plan.logical.SortAgnostic;
2627

@@ -56,7 +57,7 @@
5657
import static org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes.LEFT;
5758
import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.commonType;
5859

59-
public class Join extends BinaryPlan implements PostAnalysisVerificationAware, SortAgnostic {
60+
public class Join extends BinaryPlan implements PostAnalysisVerificationAware, SortAgnostic, ExecutesOn {
6061
public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Join", Join::new);
6162
public static final DataType[] UNSUPPORTED_TYPES = {
6263
TEXT,
@@ -309,4 +310,9 @@ private static boolean comparableTypes(Attribute left, Attribute right) {
309310
public boolean isRemote() {
310311
return isRemote;
311312
}
313+
314+
@Override
315+
public ExecuteLocation executesOn() {
316+
return isRemote ? ExecuteLocation.REMOTE : ExecuteLocation.COORDINATOR;
317+
}
312318
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/LookupJoin.java‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,16 @@
1313
import org.elasticsearch.xpack.esql.core.expression.Attribute;
1414
import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
1515
import org.elasticsearch.xpack.esql.core.tree.Source;
16-
import org.elasticsearch.xpack.esql.plan.logical.Enrich;
1716
import org.elasticsearch.xpack.esql.plan.logical.ExecutesOn;
1817
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
1918
import org.elasticsearch.xpack.esql.plan.logical.PipelineBreaker;
2019
import org.elasticsearch.xpack.esql.plan.logical.SurrogateLogicalPlan;
2120
import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
2221
import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes.UsingJoinType;
2322

24-
import java.util.LinkedList;
23+
import java.util.HashSet;
2524
import java.util.List;
25+
import java.util.Set;
2626

2727
import static java.util.Collections.emptyList;
2828
import static org.elasticsearch.xpack.esql.common.Failure.fail;
@@ -105,13 +105,10 @@ public void postAnalysisVerification(Failures failures) {
105105
}
106106

107107
private void checkRemoteJoin(Failures failures) {
108-
List<Source> fails = new LinkedList<>();
108+
Set<Source> fails = new HashSet<>();
109109

110110
this.forEachUp(UnaryPlan.class, u -> {
111-
if (u instanceof PipelineBreaker || u instanceof ExecutesOn.Coordinator) {
112-
fails.add(u.source());
113-
}
114-
if (u instanceof Enrich enrich && enrich.mode() == Enrich.Mode.COORDINATOR) {
111+
if (u instanceof PipelineBreaker || (u instanceof ExecutesOn ex && ex.executesOn() == ExecuteLocation.COORDINATOR)) {
115112
fails.add(u.source());
116113
}
117114
});

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7004,7 +7004,7 @@ public void testAggThenEnrichRemote() {
70047004
| eval employee_id = to_str(emp_no)
70057005
| ENRICH _remote:departments
70067006
"""));
7007-
assertThat(error.getMessage(), containsString("line 4:3: ENRICH with remote policy can't be executed after STATS"));
7007+
assertThat(error.getMessage(), containsString("line 4:3: ENRICH with remote policy can't be executed after [STATS size=count(*) BY emp_no]@2:3"));
70087008
}
70097009

70107010
public void testEnrichBeforeLimit() {
@@ -7354,7 +7354,7 @@ public void testRejectRemoteEnrichAfterCoordinatorEnrich() {
73547354
"""));
73557355
assertThat(
73567356
error.getMessage(),
7357-
containsString("ENRICH with remote policy can't be executed after another ENRICH with coordinator policy")
7357+
containsString("ENRICH with remote policy can't be executed after [ENRICH _coordinator:departments]@3:3")
73587358
);
73597359
}
73607360

0 commit comments

Comments
 (0)