Skip to content

ESQL: Allow using the same index in FROM and LOOKUP #118768

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Dec 18, 2024
Merged

ESQL: Allow using the same index in FROM and LOOKUP #118768

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
24f81b5
Allow using the same index in FROM and LOOKUP
ivancea Dec 16, 2024
14683f2
Added capability and test
ivancea Dec 16, 2024
c0ff359
Merge branch 'main' into lookup-join-repeated-index
ivancea Dec 16, 2024
b41ed69
Merge branch 'main' into lookup-join-repeated-index
ivancea Dec 17, 2024
721f702
Moved Node method to PlannerUtils
ivancea Dec 17, 2024
1a680c0
Merge branch 'main' into lookup-join-repeated-index
ivancea Dec 17, 2024
d330a3f
Post merge capability fix
ivancea Dec 17, 2024
3c10075
Merge branch 'main' into lookup-join-repeated-index
ivancea Dec 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion x-pack/plugin/esql/qa/testFixtures/src/main/resources/lookup-join.csv-spec
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,22 @@ language_code:integer | language_name:keyword | country:keyword
2 | [German, German, German] | [Austria, Germany, Switzerland]
;

repeatedIndexOnFrom
required_capability: join_lookup_v7
required_capability: join_lookup_repeated_index_from

FROM languages_lookup
| LOOKUP JOIN languages_lookup ON language_code
| SORT language_code
;

language_code:integer | language_name:keyword
1 | English
2 | French
3 | Spanish
4 | German
;

###############################################
# Filtering tests with languages_lookup index
###############################################
Expand Down Expand Up @@ -1061,4 +1077,3 @@ ignoreOrder:true
2023-10-23T12:27:28.948Z | 172.21.2.113 | 2764889 | QA | null
2023-10-23T12:15:03.360Z | 172.21.2.162 | 3450233 | QA | null
;

5 changes: 5 additions & 0 deletions x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -549,6 +549,11 @@ public enum Cap {
*/
JOIN_LOOKUP_V7(Build.current().isSnapshot()),

/**
* LOOKUP JOIN with the same index as the FROM
*/
JOIN_LOOKUP_REPEATED_INDEX_FROM(JOIN_LOOKUP_V7.isEnabled()),

/**
* Fix for https://github.com/elastic/elasticsearch/issues/117054
*/
Expand Down
24 changes: 23 additions & 1 deletion x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
import org.elasticsearch.xpack.esql.core.expression.Expression;
import org.elasticsearch.xpack.esql.core.expression.predicate.Predicates;
import org.elasticsearch.xpack.esql.core.tree.Node;
import org.elasticsearch.xpack.esql.core.tree.Source;
import org.elasticsearch.xpack.esql.core.type.DataType;
import org.elasticsearch.xpack.esql.core.util.Holder;
Expand All @@ -40,6 +41,7 @@
import org.elasticsearch.xpack.esql.plan.physical.ExchangeSinkExec;
import org.elasticsearch.xpack.esql.plan.physical.ExchangeSourceExec;
import org.elasticsearch.xpack.esql.plan.physical.FragmentExec;
import org.elasticsearch.xpack.esql.plan.physical.LookupJoinExec;
import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
import org.elasticsearch.xpack.esql.planner.mapper.LocalMapper;
import org.elasticsearch.xpack.esql.planner.mapper.Mapper;
Expand All @@ -48,9 +50,12 @@
import org.elasticsearch.xpack.esql.stats.SearchStats;

import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;
import java.util.function.Consumer;
import java.util.function.Function;

import static java.util.Arrays.asList;
import static org.elasticsearch.index.mapper.MappedFieldType.FieldExtractPreference.DOC_VALUES;
Expand Down Expand Up @@ -105,10 +110,27 @@ public static Set<String> planConcreteIndices(PhysicalPlan plan) {
return Set.of();
}
var indices = new LinkedHashSet<String>();
plan.forEachUp(FragmentExec.class, f -> f.fragment().forEachUp(EsRelation.class, r -> indices.addAll(r.index().concreteIndices())));
// TODO: This only works for LEFT join, we still need to support RIGHT join
forEachUpWithChildren(plan, node -> {
if (node instanceof FragmentExec f) {
f.fragment().forEachUp(EsRelation.class, r -> indices.addAll(r.index().concreteIndices()));
}
}, node -> node instanceof LookupJoinExec join ? List.of(join.left()) : node.children());
return indices;
}

/**
* Similar to {@link Node#forEachUp(Consumer)}, but with a custom callback to get the node children.
*/
private static <T extends Node<T>> void forEachUpWithChildren(
T node,
Consumer<? super T> action,
Function<? super T, Collection<T>> childrenGetter
) {
childrenGetter.apply(node).forEach(c -> forEachUpWithChildren(c, action, childrenGetter));
action.accept(node);
}

/**
* Returns the original indices specified in the FROM command of the query. We need the original query to resolve alias filters.
*/
Expand Down
49 changes: 3 additions & 46 deletions x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java
View file
Open in desktop
Copy link
Contributor Author

@ivancea ivancea Dec 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Most of the changes here are just reverting to the original logic

Original file line number Diff line number Diff line change
Expand Up @@ -63,12 +63,8 @@
import org.elasticsearch.xpack.esql.core.expression.Attribute;
import org.elasticsearch.xpack.esql.enrich.EnrichLookupService;
import org.elasticsearch.xpack.esql.enrich.LookupFromIndexService;
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
import org.elasticsearch.xpack.esql.plan.logical.join.Join;
import org.elasticsearch.xpack.esql.plan.physical.ExchangeSinkExec;
import org.elasticsearch.xpack.esql.plan.physical.ExchangeSourceExec;
import org.elasticsearch.xpack.esql.plan.physical.FragmentExec;
import org.elasticsearch.xpack.esql.plan.physical.LookupJoinExec;
import org.elasticsearch.xpack.esql.plan.physical.OutputExec;
import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
import org.elasticsearch.xpack.esql.planner.EsPhysicalOperationProviders;
Expand All @@ -81,7 +77,6 @@
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
Expand Down Expand Up @@ -167,11 +162,9 @@ public void execute(
Map<String, OriginalIndices> clusterToConcreteIndices = transportService.getRemoteClusterService()
.groupIndices(SearchRequest.DEFAULT_INDICES_OPTIONS, PlannerUtils.planConcreteIndices(physicalPlan).toArray(String[]::new));
QueryPragmas queryPragmas = configuration.pragmas();
Set<String> lookupIndexNames = findLookupIndexNames(physicalPlan);
Set<String> concreteIndexNames = selectConcreteIndices(clusterToConcreteIndices, lookupIndexNames);
if (dataNodePlan == null) {
if (concreteIndexNames.isEmpty() == false) {
String error = "expected no concrete indices without data node plan; got " + concreteIndexNames;
if (clusterToConcreteIndices.values().stream().allMatch(v -> v.indices().length == 0) == false) {
String error = "expected no concrete indices without data node plan; got " + clusterToConcreteIndices;
assert false : error;
listener.onFailure(new IllegalStateException(error));
return;
Expand All @@ -194,7 +187,7 @@ public void execute(
return;
}
} else {
if (concreteIndexNames.isEmpty()) {
if (clusterToConcreteIndices.values().stream().allMatch(v -> v.indices().length == 0)) {
var error = "expected concrete indices with data node plan but got empty; data node plan " + dataNodePlan;
assert false : error;
listener.onFailure(new IllegalStateException(error));
Expand Down Expand Up @@ -268,42 +261,6 @@ public void execute(
}
}

private Set<String> selectConcreteIndices(Map<String, OriginalIndices> clusterToConcreteIndices, Set<String> indexesToIgnore) {
Set<String> concreteIndexNames = new HashSet<>();
clusterToConcreteIndices.forEach((clusterAlias, concreteIndices) -> {
for (String index : concreteIndices.indices()) {
if (indexesToIgnore.contains(index) == false) {
concreteIndexNames.add(index);
}
}
});
return concreteIndexNames;
}

private Set<String> findLookupIndexNames(PhysicalPlan physicalPlan) {
Set<String> lookupIndexNames = new HashSet<>();
// When planning JOIN on the coordinator node: "LookupJoinExec.lookup()->FragmentExec.fragment()->EsRelation.index()"
physicalPlan.forEachDown(
LookupJoinExec.class,
lookupJoinExec -> lookupJoinExec.lookup()
.forEachDown(
FragmentExec.class,
frag -> frag.fragment().forEachDown(EsRelation.class, esRelation -> lookupIndexNames.add(esRelation.index().name()))
)
);
// When planning JOIN on the data node: "FragmentExec.fragment()->Join.right()->EsRelation.index()"
// TODO this only works for LEFT join, so we still need to support RIGHT join
physicalPlan.forEachDown(
FragmentExec.class,
fragmentExec -> fragmentExec.fragment()
.forEachDown(
Join.class,
join -> join.right().forEachDown(EsRelation.class, esRelation -> lookupIndexNames.add(esRelation.index().name()))
)
);
return lookupIndexNames;
}

// For queries like: FROM logs* | LIMIT 0 (including cross-cluster LIMIT 0 queries)
private static void updateShardCountForCoordinatorOnlyQuery(EsqlExecutionInfo execInfo) {
if (execInfo.isCrossClusterSearch()) {
Expand Down