From 67250cf438497ac9422a51a848355fa31754a151 Mon Sep 17 00:00:00 2001 From: Kathleen DeRusso Date: Mon, 18 Aug 2025 10:32:25 -0400 Subject: [PATCH 1/3] Refactor: Add RewriteableAware interface for functions that require rewriting --- .../esql/capabilities/RewriteableAware.java | 31 ++++++++++ .../function/fulltext/FullTextFunction.java | 7 ++- .../fulltext/QueryBuilderResolver.java | 56 +++++++++++-------- 3 files changed, 68 insertions(+), 26 deletions(-) create mode 100644 x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/capabilities/RewriteableAware.java 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..b8e1de0b8bfad 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 FunctionsRewritable(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 FunctionsRewritable(LogicalPlan plan) implements Rewriteable { @Override - public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOException { + public FunctionsRewritable 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 = Rewriteable.rewrite(builder, 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 FunctionsRewritable(newPlan) : this; } } } From e28a10a2e6b11d0c5596ff2e5e5ba9ddf223aadc Mon Sep 17 00:00:00 2001 From: Kathleen DeRusso Date: Tue, 19 Aug 2025 10:10:45 -0400 Subject: [PATCH 2/3] Refactor/rename - fix typo --- .../function/fulltext/QueryBuilderResolver.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 b8e1de0b8bfad..3060caf1ff9fd 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 @@ -48,7 +48,7 @@ public static void resolveQueryBuilders(LogicalPlan plan, TransportActionService }); if (hasRewriteableAwareFunctions) { Rewriteable.rewriteAndFetch( - new FunctionsRewritable(plan), + new FunctionsRewriteable(plan), queryRewriteContext(services, indexNames(plan)), listener.delegateFailureAndWrap((l, r) -> l.onResponse(r.plan)) ); @@ -76,9 +76,9 @@ private static Set indexNames(LogicalPlan plan) { return indexNames; } - private record FunctionsRewritable(LogicalPlan plan) implements Rewriteable { + private record FunctionsRewriteable(LogicalPlan plan) implements Rewriteable { @Override - public FunctionsRewritable 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(Expression.class, expr -> { @@ -102,7 +102,7 @@ public FunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOException { if (exceptionHolder.get() != null) { throw exceptionHolder.get(); } - return updated.get() ? new FunctionsRewritable(newPlan) : this; + return updated.get() ? new FunctionsRewriteable(newPlan) : this; } } } From ff31d3b047a43beea52e03b2a6e99f11d4bce931 Mon Sep 17 00:00:00 2001 From: Kathleen DeRusso Date: Tue, 19 Aug 2025 10:12:09 -0400 Subject: [PATCH 3/3] Revert rewrite change --- .../esql/expression/function/fulltext/QueryBuilderResolver.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 3060caf1ff9fd..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 @@ -89,7 +89,7 @@ public FunctionsRewriteable rewrite(QueryRewriteContext ctx) throws IOException ? rewriteableAware.asQuery(LucenePushdownPredicates.DEFAULT, TranslatorHandler.TRANSLATOR_HANDLER).toQueryBuilder() : builder; try { - builder = Rewriteable.rewrite(builder, ctx); + builder = builder.rewrite(ctx); } catch (IOException e) { exceptionHolder.setIfAbsent(e); }