Remove the need to have toQuery() on ShardContext

Author: carlosdelest

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneQueryExpressionEvaluator.java

@@ -29,11 +29,9 @@
 import org.elasticsearch.compute.operator.EvalOperator;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
-import org.elasticsearch.index.query.QueryBuilder;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
-import java.util.List;
 
 /**
  * {@link EvalOperator.ExpressionEvaluator} to run a Lucene {@link Query} during
@@ -50,23 +48,6 @@ public record ShardConfig(Query query, IndexSearcher searcher) {}
 
     private ShardState[] perShardState = EMPTY_SHARD_STATES;
 
-    public LuceneQueryExpressionEvaluator(
-        BlockFactory blockFactory,
-        List<? extends ShardContext> shardContexts,
-        QueryBuilder queryBuilder
-    ) {
-        assert shardContexts.isEmpty() == false;
-        assert queryBuilder != null;
-
-        this.blockFactory = blockFactory;
-        this.shards = new ShardConfig[shardContexts.size()];
-
-        int i = 0;
-        for (ShardContext shardContext : shardContexts) {
-            this.shards[i++] = new ShardConfig(shardContext.toQuery(queryBuilder), shardContext.searcher());
-        }
-    }
-
     public LuceneQueryExpressionEvaluator(BlockFactory blockFactory, ShardConfig[] shards) {
         this.blockFactory = blockFactory;
         this.shards = shards;
@@ -362,17 +343,15 @@ public void close() {
     }
 
     public static class Factory implements EvalOperator.ExpressionEvaluator.Factory {
-        private final List<? extends ShardContext> shardContexts;
-        private final QueryBuilder queryBuilder;
+        private final ShardConfig[] shardConfigs;
 
-        public Factory(List<? extends ShardContext> shardContexts, QueryBuilder queryBuilder) {
-            this.shardContexts = shardContexts;
-            this.queryBuilder = queryBuilder;
+        public Factory(ShardConfig[] shardConfigs) {
+            this.shardConfigs = shardConfigs;
         }
 
         @Override
         public EvalOperator.ExpressionEvaluator get(DriverContext context) {
-            return new LuceneQueryExpressionEvaluator(context.blockFactory(), shardContexts, queryBuilder);
+            return new LuceneQueryExpressionEvaluator(context.blockFactory(), shardConfigs);
         }
     }
 }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/ShardContext.java

@@ -8,8 +8,6 @@
 package org.elasticsearch.compute.lucene;
 
 import org.apache.lucene.search.IndexSearcher;
-import org.apache.lucene.search.Query;
-import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.sort.SortAndFormats;
 import org.elasticsearch.search.sort.SortBuilder;
 
@@ -41,9 +39,4 @@ public interface ShardContext {
      * {@code _cat/shards}.
      */
     String shardIdentifier();
-
-    /**
-     * Converts a {@link QueryBuilder} to a {@link Query} in the shard
-     */
-    Query toQuery(QueryBuilder queryBuilder);
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/EvalMapper.java

@@ -16,7 +16,6 @@
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.data.Vector;
-import org.elasticsearch.compute.lucene.ShardContext;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.EvalOperator;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
@@ -34,6 +33,7 @@
 import org.elasticsearch.xpack.esql.evaluator.mapper.ExpressionMapper;
 import org.elasticsearch.xpack.esql.expression.function.fulltext.Match;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.InsensitiveEqualsMapper;
+import org.elasticsearch.xpack.esql.planner.EsPhysicalOperationProviders.ShardContext;
 import org.elasticsearch.xpack.esql.planner.Layout;
 
 import java.util.List;
@@ -61,10 +61,10 @@ public static ExpressionEvaluator.Factory toEvaluator(
         FoldContext foldCtx,
         Expression exp,
         Layout layout,
-        List<? extends ShardContext> shardContexts
+        List<ShardContext> shardContexts
     ) {
         if (exp instanceof Match m) {
-            return m.toEvaluator(shardContexts, m.queryBuilder());
+            return m.toEvaluator(shardContexts);
         } else if (exp instanceof EvaluatorMapper m) {
             return m.toEvaluator(new EvaluatorMapper.ToEvaluator() {
                 @Override
@@ -88,12 +88,7 @@ public FoldContext foldCtx() {
 
     static class BooleanLogic extends ExpressionMapper<BinaryLogic> {
         @Override
-        public ExpressionEvaluator.Factory map(
-            FoldContext foldCtx,
-            BinaryLogic bc,
-            Layout layout,
-            List<? extends ShardContext> shardContexts
-        ) {
+        public ExpressionEvaluator.Factory map(FoldContext foldCtx, BinaryLogic bc, Layout layout, List<ShardContext> shardContexts) {
             var leftEval = toEvaluator(foldCtx, bc.left(), layout, shardContexts);
             var rightEval = toEvaluator(foldCtx, bc.right(), layout, shardContexts);
             /**
@@ -171,7 +166,7 @@ public void close() {
 
     static class Nots extends ExpressionMapper<Not> {
         @Override
-        public ExpressionEvaluator.Factory map(FoldContext foldCtx, Not not, Layout layout, List<? extends ShardContext> shardContexts) {
+        public ExpressionEvaluator.Factory map(FoldContext foldCtx, Not not, Layout layout, List<ShardContext> shardContexts) {
             var expEval = toEvaluator(foldCtx, not.field(), layout);
             return dvrCtx -> new org.elasticsearch.xpack.esql.evaluator.predicate.operator.logical.NotEvaluator(
                 not.source(),
@@ -183,12 +178,7 @@ public ExpressionEvaluator.Factory map(FoldContext foldCtx, Not not, Layout layo
 
     static class Attributes extends ExpressionMapper<Attribute> {
         @Override
-        public ExpressionEvaluator.Factory map(
-            FoldContext foldCtx,
-            Attribute attr,
-            Layout layout,
-            List<? extends ShardContext> shardContexts
-        ) {
+        public ExpressionEvaluator.Factory map(FoldContext foldCtx, Attribute attr, Layout layout, List<ShardContext> shardContexts) {
             record Attribute(int channel) implements ExpressionEvaluator {
                 @Override
                 public Block eval(Page page) {
@@ -223,12 +213,7 @@ public boolean eagerEvalSafeInLazy() {
     static class Literals extends ExpressionMapper<Literal> {
 
         @Override
-        public ExpressionEvaluator.Factory map(
-            FoldContext foldCtx,
-            Literal lit,
-            Layout layout,
-            List<? extends ShardContext> shardContexts
-        ) {
+        public ExpressionEvaluator.Factory map(FoldContext foldCtx, Literal lit, Layout layout, List<ShardContext> shardContexts) {
             record LiteralsEvaluator(DriverContext context, Literal lit) implements ExpressionEvaluator {
                 @Override
                 public Block eval(Page page) {
@@ -285,12 +270,7 @@ private static Block block(Literal lit, BlockFactory blockFactory, int positions
     static class IsNulls extends ExpressionMapper<IsNull> {
 
         @Override
-        public ExpressionEvaluator.Factory map(
-            FoldContext foldCtx,
-            IsNull isNull,
-            Layout layout,
-            List<? extends ShardContext> shardContexts
-        ) {
+        public ExpressionEvaluator.Factory map(FoldContext foldCtx, IsNull isNull, Layout layout, List<ShardContext> shardContexts) {
             var field = toEvaluator(foldCtx, isNull.field(), layout);
             return new IsNullEvaluatorFactory(field);
         }
@@ -338,12 +318,7 @@ public String toString() {
     static class IsNotNulls extends ExpressionMapper<IsNotNull> {
 
         @Override
-        public ExpressionEvaluator.Factory map(
-            FoldContext foldCtx,
-            IsNotNull isNotNull,
-            Layout layout,
-            List<? extends ShardContext> shardContexts
-        ) {
+        public ExpressionEvaluator.Factory map(FoldContext foldCtx, IsNotNull isNotNull, Layout layout, List<ShardContext> shardContexts) {
             return new IsNotNullEvaluatorFactory(toEvaluator(foldCtx, isNotNull.field(), layout));
         }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/mapper/ExpressionMapper.java

@@ -7,11 +7,11 @@
 
 package org.elasticsearch.xpack.esql.evaluator.mapper;
 
-import org.elasticsearch.compute.lucene.ShardContext;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.util.ReflectionUtils;
+import org.elasticsearch.xpack.esql.planner.EsPhysicalOperationProviders.ShardContext;
 import org.elasticsearch.xpack.esql.planner.Layout;
 
 import java.util.List;
@@ -23,10 +23,5 @@ public ExpressionMapper() {
         typeToken = ReflectionUtils.detectSuperTypeForRuleLike(getClass());
     }
 
-    public abstract ExpressionEvaluator.Factory map(
-        FoldContext foldCtx,
-        E expression,
-        Layout layout,
-        List<? extends ShardContext> shardContexts
-    );
+    public abstract ExpressionEvaluator.Factory map(FoldContext foldCtx, E expression, Layout layout, List<ShardContext> shardContexts);
 }

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

@@ -13,7 +13,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.compute.lucene.LuceneQueryExpressionEvaluator;
-import org.elasticsearch.compute.lucene.ShardContext;
+import org.elasticsearch.compute.lucene.LuceneQueryExpressionEvaluator.ShardConfig;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.xpack.esql.capabilities.PostOptimizationVerificationAware;
 import org.elasticsearch.xpack.esql.common.Failure;
@@ -33,8 +33,8 @@
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.AbstractConvertFunction;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.planner.EsPhysicalOperationProviders.ShardContext;
 import org.elasticsearch.xpack.esql.planner.EsqlExpressionTranslators;
-import org.elasticsearch.xpack.esql.planner.Layout;
 import org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter;
 
 import java.io.IOException;
@@ -296,10 +296,12 @@ private boolean isOperator() {
         return isOperator;
     }
 
-    public LuceneQueryExpressionEvaluator.Factory toEvaluator(
-        List<? extends ShardContext> shardContexts,
-        QueryBuilder queryBuilder
-    ) {
-        return new LuceneQueryExpressionEvaluator.Factory(shardContexts, queryBuilder);
+    public LuceneQueryExpressionEvaluator.Factory toEvaluator(List<ShardContext> shardContexts) {
+        ShardConfig[] shardConfigs = new ShardConfig[shardContexts.size()];
+        int i = 0;
+        for (ShardContext shardContext : shardContexts) {
+            shardConfigs[i++] = new ShardConfig(shardContext.toQuery(queryBuilder()), shardContext.searcher());
+        }
+        return new LuceneQueryExpressionEvaluator.Factory(shardConfigs);
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/operator/comparison/InsensitiveEqualsMapper.java

@@ -12,13 +12,13 @@
 import org.apache.lucene.util.automaton.ByteRunAutomaton;
 import org.elasticsearch.common.TriFunction;
 import org.elasticsearch.common.lucene.BytesRefs;
-import org.elasticsearch.compute.lucene.ShardContext;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.evaluator.mapper.ExpressionMapper;
+import org.elasticsearch.xpack.esql.planner.EsPhysicalOperationProviders.ShardContext;
 import org.elasticsearch.xpack.esql.planner.Layout;
 
 import java.util.List;
@@ -35,7 +35,7 @@ public final ExpressionEvaluator.Factory map(
         FoldContext foldCtx,
         InsensitiveEquals bc,
         Layout layout,
-        List<? extends ShardContext> shardContexts
+        List<ShardContext> shardContexts
     ) {
         DataType leftType = bc.left().dataType();
         DataType rightType = bc.right().dataType();

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java

@@ -17,7 +17,6 @@
 import org.elasticsearch.compute.data.LocalCircuitBreaker;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.lucene.LuceneOperator;
-import org.elasticsearch.compute.lucene.ShardContext;
 import org.elasticsearch.compute.operator.ColumnExtractOperator;
 import org.elasticsearch.compute.operator.ColumnLoadOperator;
 import org.elasticsearch.compute.operator.Driver;
@@ -91,6 +90,7 @@
 import org.elasticsearch.xpack.esql.plan.physical.ProjectExec;
 import org.elasticsearch.xpack.esql.plan.physical.ShowExec;
 import org.elasticsearch.xpack.esql.plan.physical.TopNExec;
+import org.elasticsearch.xpack.esql.planner.EsPhysicalOperationProviders.ShardContext;
 import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
 import org.elasticsearch.xpack.esql.session.Configuration;
 
@@ -130,7 +130,7 @@ public class LocalExecutionPlanner {
     private final EnrichLookupService enrichLookupService;
     private final LookupFromIndexService lookupFromIndexService;
     private final PhysicalOperationProviders physicalOperationProviders;
-    private final List<? extends ShardContext> shardContexts;
+    private final List<ShardContext> shardContexts;
 
     public LocalExecutionPlanner(
         String sessionId,
@@ -145,7 +145,7 @@ public LocalExecutionPlanner(
         EnrichLookupService enrichLookupService,
         LookupFromIndexService lookupFromIndexService,
         PhysicalOperationProviders physicalOperationProviders,
-        List<? extends ShardContext> shardContexts
+        List<ShardContext> shardContexts
     ) {
 
         this.sessionId = sessionId;