ESQL: Add a pre-mapping logical plan processing step #120368
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| pr: 120368 | ||
| summary: Add a pre-mapping logical plan processing step | ||
| area: ES|QL | ||
| type: enhancement | ||
| issues: [] |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,129 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.esql.planner.mapper.preprocessor; | ||
|
|
||
| import org.elasticsearch.action.ActionListener; | ||
| import org.elasticsearch.action.ResolvedIndices; | ||
| import org.elasticsearch.index.query.QueryBuilder; | ||
| import org.elasticsearch.index.query.QueryRewriteContext; | ||
| import org.elasticsearch.index.query.Rewriteable; | ||
| import org.elasticsearch.xpack.esql.expression.function.fulltext.FullTextFunction; | ||
| import org.elasticsearch.xpack.esql.plan.logical.EsRelation; | ||
| import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; | ||
| import org.elasticsearch.xpack.esql.planner.TranslatorHandler; | ||
| import org.elasticsearch.xpack.esql.plugin.TransportActionServices; | ||
| import org.elasticsearch.xpack.esql.session.IndexResolver; | ||
|
|
||
| import java.io.IOException; | ||
| import java.util.HashMap; | ||
| import java.util.HashSet; | ||
| import java.util.Map; | ||
| import java.util.Set; | ||
|
|
||
| /** | ||
| * Some {@link FullTextFunction} implementations such as {@link org.elasticsearch.xpack.esql.expression.function.fulltext.Match} | ||
| * will be translated to a {@link QueryBuilder} that require a rewrite phase on the coordinator. | ||
| * {@link FullTextFunctionMapperPreprocessor#preprocess(LogicalPlan, TransportActionServices, ActionListener)} will rewrite the plan by | ||
| * replacing {@link FullTextFunction} expression with new ones that hold rewritten {@link QueryBuilder}s. | ||
| */ | ||
| public class FullTextFunctionMapperPreprocessor implements MappingPreProcessor { | ||
|
There was a problem hiding this comment. This uses a pull approach where a built-in pre-processor searches for FullTextFunctions before rewriting the query. There was a problem hiding this comment. I've reworked this part, to have the |
||
|
|
||
| @Override | ||
| public void preprocess(LogicalPlan plan, TransportActionServices services, ActionListener<LogicalPlan> listener) { | ||
| Set<FullTextFunction> unresolved = fullTextFunctions(plan); | ||
| Set<String> indexNames = indexNames(plan); | ||
|
|
||
| if (indexNames == null || indexNames.isEmpty() || unresolved.isEmpty()) { | ||
| listener.onResponse(plan); | ||
| return; | ||
| } | ||
| QueryRewriteContext ctx = queryRewriteContext(services, indexNames); | ||
| FullTextFunctionsRewritable rewritable = new FullTextFunctionsRewritable(unresolved); | ||
| Rewriteable.rewriteAndFetch( | ||
| rewritable, | ||
| ctx, | ||
| listener.delegateFailureAndWrap((l, r) -> l.onResponse(updateQueryBuilders(plan, r.results()))) | ||
| ); | ||
| } | ||
|
|
||
| public LogicalPlan updateQueryBuilders(LogicalPlan plan, Map<FullTextFunction, QueryBuilder> newQueryBuilders) { | ||
|
There was a problem hiding this comment. Why isn't the rewrite applied on the spot and instead an associated map is used? There was a problem hiding this comment. I've reworked this bit. But my guess is that the intention was to avoid multiple tree passes. |
||
| LogicalPlan newPlan = plan.transformExpressionsDown(FullTextFunction.class, m -> { | ||
| if (newQueryBuilders.containsKey(m)) { | ||
| return m.replaceQueryBuilder(newQueryBuilders.get(m)); | ||
| } | ||
| return m; | ||
| }); | ||
| // The given plan was already analyzed and optimized, so we set the resulted plan to optimized as well. | ||
| newPlan.setOptimized(); | ||
| return newPlan; | ||
| } | ||
|
|
||
| private static QueryRewriteContext queryRewriteContext(TransportActionServices services, Set<String> indexNames) { | ||
| ResolvedIndices resolvedIndices = ResolvedIndices.resolveWithIndexNamesAndOptions( | ||
| indexNames.toArray(String[]::new), | ||
| IndexResolver.FIELD_CAPS_INDICES_OPTIONS, | ||
| services.clusterService().state(), | ||
| services.indexNameExpressionResolver(), | ||
| services.transportService().getRemoteClusterService(), | ||
| System.currentTimeMillis() | ||
| ); | ||
|
There was a problem hiding this comment. The indices have already been resolved through field-caps; double check if this information can be reused instead of going back to string format. There was a problem hiding this comment.
|
||
|
|
||
| return services.searchService().getRewriteContext(System::currentTimeMillis, resolvedIndices, null); | ||
| } | ||
|
|
||
| private static Set<FullTextFunction> fullTextFunctions(LogicalPlan plan) { | ||
| Set<FullTextFunction> functions = new HashSet<>(); | ||
| plan.forEachExpressionDown(FullTextFunction.class, functions::add); | ||
| return functions; | ||
| } | ||
|
|
||
| public Set<String> indexNames(LogicalPlan plan) { | ||
| Set<String> indexNames = new HashSet<>(); | ||
| plan.forEachDown(EsRelation.class, esRelation -> indexNames.addAll(esRelation.index().concreteIndices())); | ||
| return indexNames; | ||
| } | ||
|
|
||
| private static class FullTextFunctionsRewritable | ||
| implements | ||
| Rewriteable<FullTextFunctionMapperPreprocessor.FullTextFunctionsRewritable> { | ||
|
|
||
| private final Map<FullTextFunction, QueryBuilder> queryBuilderMap; | ||
|
|
||
| FullTextFunctionsRewritable(Map<FullTextFunction, QueryBuilder> queryBuilderMap) { | ||
| this.queryBuilderMap = queryBuilderMap; | ||
| } | ||
|
|
||
| FullTextFunctionsRewritable(Set<FullTextFunction> functions) { | ||
| this.queryBuilderMap = new HashMap<>(); | ||
|
|
||
| for (FullTextFunction func : functions) { | ||
| queryBuilderMap.put(func, func.asQuery(TranslatorHandler.TRANSLATOR_HANDLER).asBuilder()); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public FullTextFunctionMapperPreprocessor.FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOException { | ||
| Map<FullTextFunction, QueryBuilder> results = new HashMap<>(); | ||
|
|
||
| boolean hasChanged = false; | ||
| for (var entry : queryBuilderMap.entrySet()) { | ||
| var initial = entry.getValue(); | ||
| var rewritten = initial.rewrite(ctx); | ||
| hasChanged |= rewritten != initial; | ||
|
|
||
| results.put(entry.getKey(), rewritten); | ||
| } | ||
|
|
||
| return hasChanged ? new FullTextFunctionMapperPreprocessor.FullTextFunctionsRewritable(results) : this; | ||
| } | ||
|
|
||
| public Map<FullTextFunction, QueryBuilder> results() { | ||
| return queryBuilderMap; | ||
| } | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.esql.planner.mapper.preprocessor; | ||
|
|
||
| import org.elasticsearch.action.ActionListener; | ||
| import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; | ||
| import org.elasticsearch.xpack.esql.plugin.TransportActionServices; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.Collection; | ||
| import java.util.List; | ||
|
|
||
| public class MapperPreprocessorExecutor { | ||
|
|
||
| private final TransportActionServices services; | ||
| private final List<MappingPreProcessor> proprocessors = new ArrayList<>(); | ||
|
|
||
| public MapperPreprocessorExecutor(TransportActionServices services) { | ||
| this.services = services; | ||
| } | ||
|
|
||
| public MapperPreprocessorExecutor addPreprocessor(MappingPreProcessor preProcessor) { | ||
| proprocessors.add(preProcessor); | ||
| return this; | ||
| } | ||
|
|
||
| public MapperPreprocessorExecutor addPreprocessors(Collection<MappingPreProcessor> preProcessors) { | ||
| proprocessors.addAll(preProcessors); | ||
| return this; | ||
| } | ||
|
|
||
| public void execute(LogicalPlan plan, ActionListener<LogicalPlan> listener) { | ||
| execute(plan, 0, listener); | ||
| } | ||
|
|
||
| private void execute(LogicalPlan plan, int index, ActionListener<LogicalPlan> listener) { | ||
| if (index == proprocessors.size()) { | ||
| listener.onResponse(plan); | ||
| } else { | ||
| proprocessors.get(index).preprocess(plan, services, listener.delegateFailureAndWrap((l, p) -> execute(p, index + 1, l))); | ||
| } | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,30 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.esql.planner.mapper.preprocessor; | ||
|
|
||
| import org.elasticsearch.action.ActionListener; | ||
| import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; | ||
| import org.elasticsearch.xpack.esql.plugin.TransportActionServices; | ||
|
|
||
| /** | ||
| * Interface for a LogicalPlan processing rule occurring after the optimization, but before mapping to a physical plan. | ||
| * This step occurs on the coordinator. The rule may use services provided to the transport action and thus can resolve indices, rewrite | ||
| * queries, perform substitutions, etc. | ||
| * Note that the LogicalPlan following the rules' changes will not undergo another logical optimization round. The changes these rules | ||
| * should apply are only those that require access to services that need to be performed asynchronously. | ||
| */ | ||
| public interface MappingPreProcessor { | ||
|
There was a problem hiding this comment. Will there be multiple preprocessors? There was a problem hiding this comment.
That was my assumption, yes.
There are some particularities to how queries are repeatedly rewritten (and stopping that is done by identity checking of a return object) -- these should remain particular to But the execution of the preprocessor is done at a dedicated stage. Not sure if the current format meets your expectation, but happy to shape it further if it can be done better. |
||
|
|
||
| /** | ||
| * Process a logical plan making use of the available services and provide the updated plan to the provided listener. | ||
| * @param plan the logical plan to process | ||
| * @param services the services available from the transport action | ||
| * @param listener the listener to notify when processing is complete | ||
| */ | ||
| void preprocess(LogicalPlan plan, TransportActionServices services, ActionListener<LogicalPlan> listener); | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,24 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.esql.plugin; | ||
|
|
||
| import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver; | ||
| import org.elasticsearch.cluster.service.ClusterService; | ||
| import org.elasticsearch.compute.operator.exchange.ExchangeService; | ||
| import org.elasticsearch.search.SearchService; | ||
| import org.elasticsearch.transport.TransportService; | ||
| import org.elasticsearch.usage.UsageService; | ||
|
|
||
| public record TransportActionServices( | ||
| TransportService transportService, | ||
| SearchService searchService, | ||
| ExchangeService exchangeService, | ||
| ClusterService clusterService, | ||
| IndexNameExpressionResolver indexNameExpressionResolver, | ||
| UsageService usageService | ||
| ) {} |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Good catch