diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/capabilities/RewriteableAware.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/capabilities/RewriteableAware.java new file mode 100644 index 0000000000000..41886573fdbb2 --- /dev/null +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/capabilities/RewriteableAware.java @@ -0,0 +1,31 @@ +/* + * 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.capabilities; + +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.xpack.esql.core.expression.Expression; + +/** + * Defines objects that need to go through the rewrite phase. + */ +public interface RewriteableAware extends TranslationAware { + + /** + * @return The current active query builder. + */ + QueryBuilder queryBuilder(); + + /** + * Replaces the current query builder with a rewritten iteration. This happens multiple times through the rewrite phase until + * the final iteration of the query builder is stored. + * @param queryBuilder QueryBuilder + * @return Expression defining the active QueryBuilder + */ + Expression replaceQueryBuilder(QueryBuilder queryBuilder); + +} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java index 452b18e0d7185..c273da317dec2 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java @@ -17,6 +17,7 @@ import org.elasticsearch.xpack.esql.action.EsqlCapabilities; import org.elasticsearch.xpack.esql.capabilities.PostAnalysisPlanVerificationAware; import org.elasticsearch.xpack.esql.capabilities.PostOptimizationVerificationAware; +import org.elasticsearch.xpack.esql.capabilities.RewriteableAware; import org.elasticsearch.xpack.esql.capabilities.TranslationAware; import org.elasticsearch.xpack.esql.common.Failures; import org.elasticsearch.xpack.esql.core.expression.Expression; @@ -71,7 +72,8 @@ public abstract class FullTextFunction extends Function PostAnalysisPlanVerificationAware, EvaluatorMapper, ExpressionScoreMapper, - PostOptimizationVerificationAware { + PostOptimizationVerificationAware, + RewriteableAware { private final Expression query; private final QueryBuilder queryBuilder; @@ -164,14 +166,13 @@ public Query asQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHand return queryBuilder != null ? new TranslationAwareExpressionQuery(source(), queryBuilder) : translate(pushdownPredicates, handler); } + @Override public QueryBuilder queryBuilder() { return queryBuilder; } protected abstract Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler); - public abstract Expression replaceQueryBuilder(QueryBuilder queryBuilder); - @Override public BiConsumer postAnalysisPlanVerification() { return FullTextFunction::checkFullTextQueryFunctions; diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryBuilderResolver.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryBuilderResolver.java index ef3828a3f2fbb..4f924f4ad0419 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryBuilderResolver.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryBuilderResolver.java @@ -12,6 +12,8 @@ import org.elasticsearch.index.query.QueryBuilder; import org.elasticsearch.index.query.QueryRewriteContext; import org.elasticsearch.index.query.Rewriteable; +import org.elasticsearch.xpack.esql.capabilities.RewriteableAware; +import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.util.Holder; import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates; import org.elasticsearch.xpack.esql.plan.logical.EsRelation; @@ -25,24 +27,28 @@ import java.util.Set; /** - * Some {@link FullTextFunction} implementations such as {@link org.elasticsearch.xpack.esql.expression.function.fulltext.Match} + * Some {@link RewriteableAware} 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 QueryBuilderResolver#resolveQueryBuilders(LogicalPlan, TransportActionServices, ActionListener)} will rewrite the plan by - * replacing {@link FullTextFunction} expression with new ones that hold rewritten {@link QueryBuilder}s. + * replacing {@link RewriteableAware} expression with new ones that hold rewritten {@link QueryBuilder}s. */ public final class QueryBuilderResolver { private QueryBuilderResolver() {} public static void resolveQueryBuilders(LogicalPlan plan, TransportActionServices services, ActionListener listener) { - var hasFullTextFunctions = plan.anyMatch(p -> { - Holder hasFullTextFunction = new Holder<>(false); - p.forEachExpression(FullTextFunction.class, unused -> hasFullTextFunction.set(true)); - return hasFullTextFunction.get(); + var hasRewriteableAwareFunctions = plan.anyMatch(p -> { + Holder hasRewriteable = new Holder<>(false); + p.forEachExpression(expr -> { + if (expr instanceof RewriteableAware) { + hasRewriteable.set(true); + } + }); + return hasRewriteable.get(); }); - if (hasFullTextFunctions) { + if (hasRewriteableAwareFunctions) { Rewriteable.rewriteAndFetch( - new FullTextFunctionsRewritable(plan), + new FunctionsRewriteable(plan), queryRewriteContext(services, indexNames(plan)), listener.delegateFailureAndWrap((l, r) -> l.onResponse(r.plan)) ); @@ -70,29 +76,33 @@ private static Set indexNames(LogicalPlan plan) { return indexNames; } - private record FullTextFunctionsRewritable(LogicalPlan plan) implements Rewriteable { + private record FunctionsRewriteable(LogicalPlan plan) implements Rewriteable { @Override - public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOException { + public FunctionsRewriteable rewrite(QueryRewriteContext ctx) throws IOException { Holder exceptionHolder = new Holder<>(); Holder updated = new Holder<>(false); - LogicalPlan newPlan = plan.transformExpressionsDown(FullTextFunction.class, f -> { - QueryBuilder builder = f.queryBuilder(), initial = builder; - builder = builder == null - ? f.asQuery(LucenePushdownPredicates.DEFAULT, TranslatorHandler.TRANSLATOR_HANDLER).toQueryBuilder() - : builder; - try { - builder = builder.rewrite(ctx); - } catch (IOException e) { - exceptionHolder.setIfAbsent(e); + LogicalPlan newPlan = plan.transformExpressionsDown(Expression.class, expr -> { + Expression finalExpression = expr; + if (expr instanceof RewriteableAware rewriteableAware) { + QueryBuilder builder = rewriteableAware.queryBuilder(), initial = builder; + builder = builder == null + ? rewriteableAware.asQuery(LucenePushdownPredicates.DEFAULT, TranslatorHandler.TRANSLATOR_HANDLER).toQueryBuilder() + : builder; + try { + builder = builder.rewrite(ctx); + } catch (IOException e) { + exceptionHolder.setIfAbsent(e); + } + var rewritten = builder != initial; + updated.set(updated.get() || rewritten); + finalExpression = rewritten ? rewriteableAware.replaceQueryBuilder(builder) : finalExpression; } - var rewritten = builder != initial; - updated.set(updated.get() || rewritten); - return rewritten ? f.replaceQueryBuilder(builder) : f; + return finalExpression; }); if (exceptionHolder.get() != null) { throw exceptionHolder.get(); } - return updated.get() ? new FullTextFunctionsRewritable(newPlan) : this; + return updated.get() ? new FunctionsRewriteable(newPlan) : this; } } }