diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/lookup-join.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/lookup-join.csv-spec index 618149f2c3dde..39638899cf6b6 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/lookup-join.csv-spec +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/lookup-join.csv-spec @@ -198,6 +198,51 @@ language_code:integer | language_name:keyword 4 | German ; +dropAllLookedUpFieldsOnTheDataNode-Ignore +// Depends on +// https://github.com/elastic/elasticsearch/issues/118778 +// https://github.com/elastic/elasticsearch/issues/118781 +required_capability: join_lookup_v8 + +FROM employees +| EVAL language_code = emp_no % 10 +| LOOKUP JOIN languages_lookup_non_unique_key ON language_code +| WHERE emp_no == 10001 +| SORT emp_no +| DROP language* +; + +emp_no:integer +10001 +10001 +10001 +10001 +; + +dropAllLookedUpFieldsOnTheCoordinator-Ignore +// Depends on +// https://github.com/elastic/elasticsearch/issues/118778 +// https://github.com/elastic/elasticsearch/issues/118781 +required_capability: join_lookup_v8 + +FROM employees +| SORT emp_no +| LIMIT 2 +| EVAL language_code = emp_no % 10 +| LOOKUP JOIN languages_lookup_non_unique_key ON language_code +| DROP language* +; + +emp_no:integer +10001 +10001 +10001 +10001 +10002 +10002 +10002 +; + ############################################### # Filtering tests with languages_lookup index ############################################### @@ -860,6 +905,26 @@ ignoreOrder:true 2023-10-23T12:15:03.360Z | 172.21.2.162 | 3450233 | Connected to 10.1.0.3 | Success | 172.21.2.162 | null ; +lookupMessageFromIndexTwiceFullyShadowing +required_capability: join_lookup_v8 + +FROM sample_data +| LOOKUP JOIN message_types_lookup ON message +| LOOKUP JOIN message_types_lookup ON message +| KEEP @timestamp, client_ip, event_duration, message, type +; +ignoreOrder:true + +@timestamp:date | client_ip:ip | event_duration:long | message:keyword | type:keyword +2023-10-23T13:55:01.543Z | 172.21.3.15 | 1756467 | Connected to 10.1.0.1 | Success +2023-10-23T13:53:55.832Z | 172.21.3.15 | 5033755 | Connection error | Error +2023-10-23T13:52:55.015Z | 172.21.3.15 | 8268153 | Connection error | Error +2023-10-23T13:51:54.732Z | 172.21.3.15 | 725448 | Connection error | Error +2023-10-23T13:33:34.937Z | 172.21.0.5 | 1232382 | Disconnected | Disconnected +2023-10-23T12:27:28.948Z | 172.21.2.113 | 2764889 | Connected to 10.1.0.2 | Success +2023-10-23T12:15:03.360Z | 172.21.2.162 | 3450233 | Connected to 10.1.0.3 | Success +; + ############################################### # Tests with clientips_lookup and message_types_lookup indexes ############################################### diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java index e01608c546090..1c7e765c6bd59 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java @@ -9,16 +9,16 @@ import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BlockUtils; +import org.elasticsearch.index.IndexMode; import org.elasticsearch.xpack.esql.core.expression.AttributeSet; import org.elasticsearch.xpack.esql.core.expression.EmptyAttribute; import org.elasticsearch.xpack.esql.core.expression.Expressions; import org.elasticsearch.xpack.esql.core.expression.NamedExpression; -import org.elasticsearch.xpack.esql.core.util.Holder; import org.elasticsearch.xpack.esql.plan.logical.Aggregate; +import org.elasticsearch.xpack.esql.plan.logical.EsRelation; import org.elasticsearch.xpack.esql.plan.logical.Eval; import org.elasticsearch.xpack.esql.plan.logical.Limit; import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; -import org.elasticsearch.xpack.esql.plan.logical.Project; import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation; import org.elasticsearch.xpack.esql.plan.logical.local.LocalSupplier; import org.elasticsearch.xpack.esql.planner.PlannerUtils; @@ -34,13 +34,17 @@ public final class PruneColumns extends Rule { @Override public LogicalPlan apply(LogicalPlan plan) { - var used = new AttributeSet(); - // don't remove Evals without any Project/Aggregate (which might not occur as the last node in the plan) - var seenProjection = new Holder<>(Boolean.FALSE); - - // start top-to-bottom - // and track used references + // track used references + var used = plan.outputSet(); + // while going top-to-bottom (upstream) var pl = plan.transformDown(p -> { + // Note: It is NOT required to do anything special for binary plans like JOINs. It is perfectly fine that transformDown descends + // first into the left side, adding all kinds of attributes to the `used` set, and then descends into the right side - even + // though the `used` set will contain stuff only used in the left hand side. That's because any attribute that is used in the + // left hand side must have been created in the left side as well. Even field attributes belonging to the same index fields will + // have different name ids in the left and right hand sides - as in the extreme example + // `FROM lookup_idx | LOOKUP JOIN lookup_idx ON key_field`. + // skip nodes that simply pass the input through if (p instanceof Limit) { return p; @@ -53,7 +57,7 @@ public LogicalPlan apply(LogicalPlan plan) { do { recheck = false; if (p instanceof Aggregate aggregate) { - var remaining = seenProjection.get() ? removeUnused(aggregate.aggregates(), used) : null; + var remaining = removeUnused(aggregate.aggregates(), used); if (remaining != null) { if (remaining.isEmpty()) { @@ -87,10 +91,8 @@ public LogicalPlan apply(LogicalPlan plan) { ); } } - - seenProjection.set(Boolean.TRUE); } else if (p instanceof Eval eval) { - var remaining = seenProjection.get() ? removeUnused(eval.fields(), used) : null; + var remaining = removeUnused(eval.fields(), used); // no fields, no eval if (remaining != null) { if (remaining.isEmpty()) { @@ -100,8 +102,16 @@ public LogicalPlan apply(LogicalPlan plan) { p = new Eval(eval.source(), eval.child(), remaining); } } - } else if (p instanceof Project) { - seenProjection.set(Boolean.TRUE); + } else if (p instanceof EsRelation esRelation && esRelation.indexMode() == IndexMode.LOOKUP) { + // Normally, pruning EsRelation has no effect because InsertFieldExtraction only extracts the required fields, anyway. + // The field extraction for LOOKUP JOIN works differently, however - we extract all fields (other than the join key) + // that the EsRelation has. + var remaining = removeUnused(esRelation.output(), used); + // TODO: LookupFromIndexOperator cannot handle 0 lookup fields, yet. That means 1 field in total (key field + lookup). + // https://github.com/elastic/elasticsearch/issues/118778 + if (remaining != null && remaining.size() > 1) { + p = new EsRelation(esRelation.source(), esRelation.index(), remaining, esRelation.indexMode(), esRelation.frozen()); + } } } while (recheck); diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java index f9d86ecf0f61a..e41e500aad110 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java @@ -104,6 +104,7 @@ else if (plan instanceof Project project) { FieldAttribute.class, // Do not use the attribute name, this can deviate from the field name for union types. // Also skip fields from lookup indices because we do not have stats for these. + // TODO: We do have stats for lookup indices in case they are being used in the FROM clause; this can be refined. f -> stats.exists(f.fieldName()) || lookupFields.contains(f) ? f : Literal.of(f, null) ); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/LookupJoin.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/LookupJoin.java index 57c8cb00baa32..4e009156072df 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/LookupJoin.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/LookupJoin.java @@ -11,7 +11,6 @@ import org.elasticsearch.xpack.esql.core.tree.NodeInfo; import org.elasticsearch.xpack.esql.core.tree.Source; import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; -import org.elasticsearch.xpack.esql.plan.logical.Project; import org.elasticsearch.xpack.esql.plan.logical.SurrogateLogicalPlan; import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes.UsingJoinType; @@ -50,9 +49,8 @@ public LookupJoin(Source source, LogicalPlan left, LogicalPlan right, JoinConfig */ @Override public LogicalPlan surrogate() { - Join normalized = new Join(source(), left(), right(), config()); // TODO: decide whether to introduce USING or just basic ON semantics - keep the ordering out for now - return new Project(source(), normalized, output()); + return new Join(source(), left(), right(), config()); } @Override diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java index 72c771367d6ae..444fb51576da9 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java @@ -749,7 +749,7 @@ public String describe() { return Stream.concat( Stream.concat(Stream.of(sourceOperatorFactory), intermediateOperatorFactories.stream()), Stream.of(sinkOperatorFactory) - ).map(Describable::describe).collect(joining("\n\\_", "\\_", "")); + ).map(describable -> describable == null ? "null" : describable.describe()).collect(joining("\n\\_", "\\_", "")); } @Override diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java index ba0b67c842daf..fdf6fffdc6f03 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java @@ -3011,6 +3011,24 @@ public void testPruneChainedEval() { var source = as(limit.child(), EsRelation.class); } + public void testPruneChainedEvalNoProjection() { + var plan = plan(""" + from test + | eval garbage = salary + 3 + | eval garbage = emp_no / garbage, garbage = garbage + | eval garbage = 1 + """); + var eval = as(plan, Eval.class); + var limit = as(eval.child(), Limit.class); + var source = as(limit.child(), EsRelation.class); + + assertEquals(1, eval.fields().size()); + var alias = as(eval.fields().get(0), Alias.class); + assertEquals(alias.name(), "garbage"); + var literal = as(alias.child(), Literal.class); + assertEquals(1, literal.value()); + } + /** * Expects * Limit[1000[INTEGER]] @@ -4914,8 +4932,7 @@ public void testPlanSanityCheckWithBinaryPlans() throws Exception { | LOOKUP JOIN languages_lookup ON language_code """); - var project = as(plan, Project.class); - var join = as(project.child(), Join.class); + var join = as(plan, Join.class); var joinWithInvalidLeftPlan = join.replaceChildren(join.right(), join.right()); IllegalStateException e = expectThrows(IllegalStateException.class, () -> logicalOptimizer.optimize(joinWithInvalidLeftPlan)); @@ -5921,10 +5938,9 @@ public void testLookupJoinPushDownFilterOnJoinKeyWithRename() { """; var plan = optimizedPlan(query); - var project = as(plan, Project.class); - var join = as(project.child(), Join.class); + var join = as(plan, Join.class); assertThat(join.config().type(), equalTo(JoinTypes.LEFT)); - project = as(join.left(), Project.class); + var project = as(join.left(), Project.class); var limit = as(project.child(), Limit.class); assertThat(limit.limit().fold(), equalTo(1000)); var filter = as(limit.child(), Filter.class); @@ -5965,10 +5981,9 @@ public void testLookupJoinPushDownFilterOnLeftSideField() { var plan = optimizedPlan(query); - var project = as(plan, Project.class); - var join = as(project.child(), Join.class); + var join = as(plan, Join.class); assertThat(join.config().type(), equalTo(JoinTypes.LEFT)); - project = as(join.left(), Project.class); + var project = as(join.left(), Project.class); var limit = as(project.child(), Limit.class); assertThat(limit.limit().fold(), equalTo(1000)); @@ -6009,8 +6024,7 @@ public void testLookupJoinPushDownDisabledForLookupField() { var plan = optimizedPlan(query); - var project = as(plan, Project.class); - var limit = as(project.child(), Limit.class); + var limit = as(plan, Limit.class); assertThat(limit.limit().fold(), equalTo(1000)); var filter = as(limit.child(), Filter.class); @@ -6022,7 +6036,7 @@ public void testLookupJoinPushDownDisabledForLookupField() { var join = as(filter.child(), Join.class); assertThat(join.config().type(), equalTo(JoinTypes.LEFT)); - project = as(join.left(), Project.class); + var project = as(join.left(), Project.class); var leftRel = as(project.child(), EsRelation.class); var rightRel = as(join.right(), EsRelation.class); @@ -6054,8 +6068,7 @@ public void testLookupJoinPushDownSeparatedForConjunctionBetweenLeftAndRightFiel var plan = optimizedPlan(query); - var project = as(plan, Project.class); - var limit = as(project.child(), Limit.class); + var limit = as(plan, Limit.class); assertThat(limit.limit().fold(), equalTo(1000)); // filter kept in place, working on the right side var filter = as(limit.child(), Filter.class); @@ -6067,7 +6080,7 @@ public void testLookupJoinPushDownSeparatedForConjunctionBetweenLeftAndRightFiel var join = as(filter.child(), Join.class); assertThat(join.config().type(), equalTo(JoinTypes.LEFT)); - project = as(join.left(), Project.class); + var project = as(join.left(), Project.class); // filter pushed down filter = as(project.child(), Filter.class); op = as(filter.condition(), GreaterThan.class); @@ -6079,7 +6092,6 @@ public void testLookupJoinPushDownSeparatedForConjunctionBetweenLeftAndRightFiel var leftRel = as(filter.child(), EsRelation.class); var rightRel = as(join.right(), EsRelation.class); - } /** @@ -6107,8 +6119,7 @@ public void testLookupJoinPushDownDisabledForDisjunctionBetweenLeftAndRightField var plan = optimizedPlan(query); - var project = as(plan, Project.class); - var limit = as(project.child(), Limit.class); + var limit = as(plan, Limit.class); assertThat(limit.limit().fold(), equalTo(1000)); var filter = as(limit.child(), Filter.class); @@ -6128,7 +6139,7 @@ public void testLookupJoinPushDownDisabledForDisjunctionBetweenLeftAndRightField var join = as(filter.child(), Join.class); assertThat(join.config().type(), equalTo(JoinTypes.LEFT)); - project = as(join.left(), Project.class); + var project = as(join.left(), Project.class); var leftRel = as(project.child(), EsRelation.class); var rightRel = as(join.right(), EsRelation.class); diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java index 48fc319e6643e..3518beda7dc48 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java @@ -14,8 +14,11 @@ import org.elasticsearch.common.geo.ShapeRelation; import org.elasticsearch.common.lucene.BytesRefs; import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.common.util.set.Sets; import org.elasticsearch.compute.aggregation.AggregatorMode; +import org.elasticsearch.compute.operator.exchange.ExchangeSinkHandler; +import org.elasticsearch.compute.operator.exchange.ExchangeSourceHandler; import org.elasticsearch.core.Tuple; import org.elasticsearch.geometry.Circle; import org.elasticsearch.geometry.Polygon; @@ -37,6 +40,7 @@ import org.elasticsearch.xpack.esql.EsqlTestUtils; import org.elasticsearch.xpack.esql.EsqlTestUtils.TestConfigurableSearchStats; import org.elasticsearch.xpack.esql.EsqlTestUtils.TestConfigurableSearchStats.Config; +import org.elasticsearch.xpack.esql.TestBlockFactory; import org.elasticsearch.xpack.esql.VerificationException; import org.elasticsearch.xpack.esql.action.EsqlCapabilities; import org.elasticsearch.xpack.esql.analysis.Analyzer; @@ -122,6 +126,8 @@ import org.elasticsearch.xpack.esql.plan.physical.ProjectExec; import org.elasticsearch.xpack.esql.plan.physical.TopNExec; import org.elasticsearch.xpack.esql.plan.physical.UnaryExec; +import org.elasticsearch.xpack.esql.planner.EsPhysicalOperationProviders; +import org.elasticsearch.xpack.esql.planner.LocalExecutionPlanner; import org.elasticsearch.xpack.esql.planner.PlannerUtils; import org.elasticsearch.xpack.esql.planner.mapper.Mapper; import org.elasticsearch.xpack.esql.plugin.QueryPragmas; @@ -133,12 +139,14 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Collection; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; +import java.util.regex.Pattern; import java.util.stream.Collectors; import static java.util.Arrays.asList; @@ -6736,6 +6744,299 @@ public void testLookupThenTopN() { ); } + public void testLookupJoinFieldLoading() throws Exception { + assumeTrue("Requires LOOKUP JOIN", EsqlCapabilities.Cap.JOIN_LOOKUP_V8.isEnabled()); + + TestDataSource data = dataSetWithLookupIndices(Map.of("lookup_index", List.of("first_name", "foo", "bar", "baz"))); + + String query = """ + FROM test + | LOOKUP JOIN lookup_index ON first_name + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo", "bar", "baz"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index ON first_name + | KEEP b* + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("bar", "baz"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index ON first_name + | DROP b* + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index ON first_name + | EVAL bar = 10 + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo", "baz"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index ON first_name + | RENAME bar AS foobar + | KEEP f* + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo", "bar"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index ON first_name + | STATS count_distinct(foo) BY bar + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo", "bar")), true); + + query = """ + FROM test + | LOOKUP JOIN lookup_index ON first_name + | MV_EXPAND foo + | KEEP foo + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index ON first_name + | MV_EXPAND foo + | DROP foo + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo", "bar", "baz"))); + + query = """ + FROM lookup_index + | LOOKUP JOIN lookup_index ON first_name + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo", "bar", "baz"))); + + query = """ + FROM lookup_index + | LOOKUP JOIN lookup_index ON first_name + | KEEP foo + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo"))); + } + + public void testLookupJoinFieldLoadingTwoLookups() throws Exception { + assumeTrue("Requires LOOKUP JOIN", EsqlCapabilities.Cap.JOIN_LOOKUP_V8.isEnabled()); + + TestDataSource data = dataSetWithLookupIndices( + Map.of( + "lookup_index1", + List.of("first_name", "foo", "bar", "baz"), + "lookup_index2", + List.of("first_name", "foo", "bar2", "baz2") + ) + ); + + String query = """ + FROM test + | LOOKUP JOIN lookup_index1 ON first_name + | LOOKUP JOIN lookup_index2 ON first_name + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("bar", "baz"), Set.of("foo", "bar2", "baz2"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index1 ON first_name + | LOOKUP JOIN lookup_index2 ON first_name + | DROP foo + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("bar", "baz"), Set.of("bar2", "baz2"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index1 ON first_name + | LOOKUP JOIN lookup_index2 ON first_name + | KEEP b* + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("bar", "baz"), Set.of("bar2", "baz2"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index1 ON first_name + | LOOKUP JOIN lookup_index2 ON first_name + | DROP baz* + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("bar"), Set.of("foo", "bar2"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index1 ON first_name + | EVAL foo = to_upper(foo) + | LOOKUP JOIN lookup_index2 ON first_name + | EVAL foo = to_lower(foo) + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("bar", "baz"), Set.of("foo", "bar2", "baz2"))); + } + + @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/119082") + public void testLookupJoinFieldLoadingTwoLookupsProjectInBetween() throws Exception { + assumeTrue("Requires LOOKUP JOIN", EsqlCapabilities.Cap.JOIN_LOOKUP_V8.isEnabled()); + + TestDataSource data = dataSetWithLookupIndices( + Map.of( + "lookup_index1", + List.of("first_name", "foo", "bar", "baz"), + "lookup_index2", + List.of("first_name", "foo", "bar2", "baz2") + ) + ); + + String query = """ + FROM test + | LOOKUP JOIN lookup_index1 ON first_name + | RENAME foo AS foo1 + | LOOKUP JOIN lookup_index2 ON first_name + | DROP b* + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo"), Set.of("foo"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index1 ON first_name + | DROP bar + | LOOKUP JOIN lookup_index2 ON first_name + | DROP b* + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("foo"), Set.of("foo"))); + + query = """ + FROM test + | LOOKUP JOIN lookup_index1 ON first_name + | KEEP first_name, b* + | LOOKUP JOIN lookup_index2 ON first_name + | DROP bar* + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of("baz"), Set.of("foo", "baz2"))); + } + + @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/118778") + public void testLookupJoinFieldLoadingDropAllFields() throws Exception { + assumeTrue("Requires LOOKUP JOIN", EsqlCapabilities.Cap.JOIN_LOOKUP_V8.isEnabled()); + + TestDataSource data = dataSetWithLookupIndices(Map.of("lookup_index", List.of("first_name", "foo", "bar", "baz"))); + + String query = """ + FROM test + | LOOKUP JOIN lookup_index ON first_name + | DROP foo, b* + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of())); + + query = """ + FROM test + | LOOKUP JOIN lookup_index ON first_name + | LOOKUP JOIN lookup_index ON first_name + """; + assertLookupJoinFieldNames(query, data, List.of(Set.of(), Set.of("foo", "bar", "baz"))); + } + + private void assertLookupJoinFieldNames(String query, TestDataSource data, List> expectedFieldNames) { + assertLookupJoinFieldNames(query, data, expectedFieldNames, false); + } + + private void assertLookupJoinFieldNames( + String query, + TestDataSource data, + List> expectedFieldNames, + boolean useDataNodePlan + ) { + // Do not assert serialization: + // This will have a LookupJoinExec, which is not serializable because it doesn't leave the coordinator. + var plan = physicalPlan(query, data, false); + + var physicalOperations = physicalOperationsFromPhysicalPlan(plan, useDataNodePlan); + + List> fields = findFieldNamesInLookupJoinDescription(physicalOperations); + + assertEquals(expectedFieldNames.size(), fields.size()); + for (int i = 0; i < expectedFieldNames.size(); i++) { + assertEquals(expectedFieldNames.get(i), fields.get(i)); + } + } + + private TestDataSource dataSetWithLookupIndices(Map> indexNameToFieldNames) { + Map lookupIndices = new HashMap<>(); + + for (Map.Entry> entry : indexNameToFieldNames.entrySet()) { + String lookupIndexName = entry.getKey(); + Map lookup_fields = fields(entry.getValue()); + + EsIndex lookupIndex = new EsIndex(lookupIndexName, lookup_fields, Map.of(lookupIndexName, IndexMode.LOOKUP)); + lookupIndices.put(lookupIndexName, IndexResolution.valid(lookupIndex)); + } + + return makeTestDataSource( + "test", + "mapping-basic.json", + new EsqlFunctionRegistry(), + lookupIndices, + setupEnrichResolution(), + TEST_SEARCH_STATS + ); + } + + private Map fields(Collection fieldNames) { + Map fields = new HashMap<>(); + + for (String fieldName : fieldNames) { + fields.put(fieldName, new EsField(fieldName, DataType.KEYWORD, Map.of(), false)); + } + + return fields; + } + + private LocalExecutionPlanner.LocalExecutionPlan physicalOperationsFromPhysicalPlan(PhysicalPlan plan, boolean useDataNodePlan) { + // The TopN needs an estimated row size for the planner to work + var plans = PlannerUtils.breakPlanBetweenCoordinatorAndDataNode(EstimatesRowSize.estimateRowSize(0, plan), config); + plan = useDataNodePlan ? plans.v2() : plans.v1(); + plan = PlannerUtils.localPlan(List.of(), config, plan); + LocalExecutionPlanner planner = new LocalExecutionPlanner( + "test", + "", + null, + BigArrays.NON_RECYCLING_INSTANCE, + TestBlockFactory.getNonBreakingInstance(), + Settings.EMPTY, + config, + new ExchangeSourceHandler(10, null, null), + new ExchangeSinkHandler(null, 10, () -> 10), + null, + null, + new EsPhysicalOperationProviders(List.of(), null) + ); + + return planner.plan(plan); + } + + private List> findFieldNamesInLookupJoinDescription(LocalExecutionPlanner.LocalExecutionPlan physicalOperations) { + + String[] descriptionLines = physicalOperations.describe().split("\\r?\\n|\\r"); + + // Capture the inside of "...load_fields=[field{f}#19, other_field{f}#20]". + String insidePattern = "[^\\]]*"; + Pattern expected = Pattern.compile("\\\\_LookupOperator.*load_fields=\\[(" + insidePattern + ")].*"); + + List> results = new ArrayList<>(); + for (String line : descriptionLines) { + var matcher = expected.matcher(line); + if (matcher.find()) { + String allFields = matcher.group(1); + Set loadedFields = Arrays.stream(allFields.split(",")) + .map(name -> name.trim().split("\\{f}#")[0]) + .collect(Collectors.toSet()); + results.add(loadedFields); + } + } + + return results; + } + public void testScore() { assumeTrue("'METADATA _score' is disabled", EsqlCapabilities.Cap.METADATA_SCORE.isEnabled()); var plan = physicalPlan("""