Refactor: Add RewriteableAware interface for functions that require r… (#133075)

* Refactor: Add RewriteableAware interface for functions that require rewriting

* Refactor/rename - fix typo

* Revert rewrite change

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);
+
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java

@@ -17,6 +17,7 @@ import org.elasticsearch.index.query.QueryBuilder;
 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 abstract class FullTextFunction extends Function
         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<LogicalPlan, Failures> postAnalysisPlanVerification() {
         return FullTextFunction::checkFullTextQueryFunctions;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryBuilderResolver.java

@@ -12,6 +12,8 @@ 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.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.HashSet;
 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<LogicalPlan> listener) {
-        var hasFullTextFunctions = plan.anyMatch(p -> {
-            Holder<Boolean> hasFullTextFunction = new Holder<>(false);
-            p.forEachExpression(FullTextFunction.class, unused -> hasFullTextFunction.set(true));
-            return hasFullTextFunction.get();
+        var hasRewriteableAwareFunctions = plan.anyMatch(p -> {
+            Holder<Boolean> 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 @@ public final class QueryBuilderResolver {
         return indexNames;
     }
 
-    private record FullTextFunctionsRewritable(LogicalPlan plan) implements Rewriteable<QueryBuilderResolver.FullTextFunctionsRewritable> {
+    private record FunctionsRewriteable(LogicalPlan plan) implements Rewriteable<FunctionsRewriteable> {
         @Override
-        public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOException {
+        public FunctionsRewriteable rewrite(QueryRewriteContext ctx) throws IOException {
             Holder<IOException> exceptionHolder = new Holder<>();
             Holder<Boolean> 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;
         }
     }
 }