Serialize Configuration to get the source in the LOOKUP request, + TransportVersion

ivancea · ivancea · commit e627316175ff · 2025-02-04T15:41:16.000+01:00
diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java
@@ -173,6 +173,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INFERENCE_REQUEST_ADAPTIVE_RATE_LIMITING = def(8_839_00_0);
     public static final TransportVersion ML_INFERENCE_IBM_WATSONX_RERANK_ADDED = def(8_840_00_0);
     public static final TransportVersion ELASTICSEARCH_9_0 = def(9_000_00_0);
+    public static final TransportVersion ESQL_LOOKUP_JOIN_CONFIGURATION_IN_REQUEST = def(9_001_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/LookupFromIndexIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/LookupFromIndexIT.java
@@ -46,6 +46,7 @@
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.async.AsyncExecutionId;
+import org.elasticsearch.xpack.esql.ConfigurationTestUtils;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -214,6 +215,7 @@ private void runLookup(DataType keyType, PopulateIndices populateIndices) throws
             final String finalNodeWithShard = nodeWithShard;
             LookupFromIndexOperator.Factory lookup = new LookupFromIndexOperator.Factory(
                 "test",
+                ConfigurationTestUtils.randomConfiguration(),
                 parentTask,
                 QueryPragmas.ENRICH_MAX_WORKERS.get(Settings.EMPTY),
                 1,
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperator.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperator.java
@@ -26,6 +26,7 @@
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.session.Configuration;
 
 import java.io.IOException;
 import java.util.Iterator;
@@ -38,6 +39,7 @@
 public final class LookupFromIndexOperator extends AsyncOperator<LookupFromIndexOperator.OngoingJoin> {
     public record Factory(
         String sessionId,
+        Configuration configuration,
         CancellableTask parentTask,
         int maxOutstandingRequests,
         int inputChannel,
@@ -67,6 +69,7 @@ public String describe() {
         public Operator get(DriverContext driverContext) {
             return new LookupFromIndexOperator(
                 sessionId,
+                configuration,
                 driverContext,
                 parentTask,
                 maxOutstandingRequests,
@@ -83,6 +86,7 @@ public Operator get(DriverContext driverContext) {
 
     private final LookupFromIndexService lookupService;
     private final String sessionId;
+    private final Configuration configuration;
     private final CancellableTask parentTask;
     private final int inputChannel;
     private final DataType inputDataType;
@@ -102,6 +106,7 @@ public Operator get(DriverContext driverContext) {
 
     public LookupFromIndexOperator(
         String sessionId,
+        Configuration configuration,
         DriverContext driverContext,
         CancellableTask parentTask,
         int maxOutstandingRequests,
@@ -115,6 +120,7 @@ public LookupFromIndexOperator(
     ) {
         super(driverContext, maxOutstandingRequests);
         this.sessionId = sessionId;
+        this.configuration = configuration;
         this.parentTask = parentTask;
         this.inputChannel = inputChannel;
         this.lookupService = lookupService;
@@ -131,6 +137,7 @@ protected void performAsync(Page inputPage, ActionListener<OngoingJoin> listener
         totalTerms += inputBlock.getTotalValueCount();
         LookupFromIndexService.Request request = new LookupFromIndexService.Request(
             sessionId,
+            configuration,
             lookupIndex,
             inputDataType,
             matchField,
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexService.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexService.java
@@ -9,6 +9,8 @@
 
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -30,11 +32,14 @@
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamOutput;
+import org.elasticsearch.xpack.esql.session.Configuration;
 
 import java.io.IOException;
 import java.util.List;
 import java.util.Objects;
 
+import static java.lang.System.in;
+
 /**
  * {@link LookupFromIndexService} performs lookup against a Lookup index for
  * a given input page. See {@link AbstractLookupService} for how it works
@@ -66,6 +71,7 @@ public LookupFromIndexService(
     protected TransportRequest transportRequest(LookupFromIndexService.Request request, ShardId shardId) {
         return new TransportRequest(
             request.sessionId,
+            request.configuration,
             shardId,
             request.inputDataType,
             request.inputPage,
@@ -101,10 +107,12 @@ protected AbstractLookupService.LookupResponse readLookupResponse(StreamInput in
     }
 
     public static class Request extends AbstractLookupService.Request {
+        private final Configuration configuration;
         private final String matchField;
 
         Request(
             String sessionId,
+            Configuration configuration,
             String index,
             DataType inputDataType,
             String matchField,
@@ -113,15 +121,18 @@ public static class Request extends AbstractLookupService.Request {
             Source source
         ) {
             super(sessionId, index, inputDataType, inputPage, extractFields, source);
+            this.configuration = configuration;
             this.matchField = matchField;
         }
     }
 
     protected static class TransportRequest extends AbstractLookupService.TransportRequest {
+        private final Configuration configuration;
         private final String matchField;
 
         TransportRequest(
             String sessionId,
+            Configuration configuration,
             ShardId shardId,
             DataType inputDataType,
             Page inputPage,
@@ -131,19 +142,30 @@ protected static class TransportRequest extends AbstractLookupService.TransportR
             Source source
         ) {
             super(sessionId, shardId, inputDataType, inputPage, toRelease, extractFields, source);
+            this.configuration = configuration;
             this.matchField = matchField;
         }
 
         static TransportRequest readFrom(StreamInput in, BlockFactory blockFactory) throws IOException {
             TaskId parentTaskId = TaskId.readFromStream(in);
             String sessionId = in.readString();
+            Configuration configuration = null;
+            if (in.getTransportVersion().onOrAfter(TransportVersions.ESQL_LOOKUP_JOIN_CONFIGURATION_IN_REQUEST)) {
+                configuration = new Configuration(
+                    // TODO make EsqlConfiguration Releasable
+                    new BlockStreamInput(
+                        in,
+                        new BlockFactory(new NoopCircuitBreaker(CircuitBreaker.REQUEST), BigArrays.NON_RECYCLING_INSTANCE)
+                    )
+                );
+            }
             ShardId shardId = new ShardId(in);
             DataType inputDataType = DataType.fromTypeName(in.readString());
             Page inputPage;
             try (BlockStreamInput bsi = new BlockStreamInput(in, blockFactory)) {
                 inputPage = new Page(bsi);
             }
-            PlanStreamInput planIn = new PlanStreamInput(in, in.namedWriteableRegistry(), null);
+            PlanStreamInput planIn = new PlanStreamInput(in, in.namedWriteableRegistry(), configuration);
             List<NamedExpression> extractFields = planIn.readNamedWriteableCollectionAsList(NamedExpression.class);
             String matchField = in.readString();
             var source = Source.EMPTY;
@@ -152,6 +174,7 @@ static TransportRequest readFrom(StreamInput in, BlockFactory blockFactory) thro
             }
             TransportRequest result = new TransportRequest(
                 sessionId,
+                configuration,
                 shardId,
                 inputDataType,
                 inputPage,
@@ -168,6 +191,9 @@ static TransportRequest readFrom(StreamInput in, BlockFactory blockFactory) thro
         public void writeTo(StreamOutput out) throws IOException {
             super.writeTo(out);
             out.writeString(sessionId);
+            if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_LOOKUP_JOIN_CONFIGURATION_IN_REQUEST)) {
+                configuration.writeTo(out);
+            }
             out.writeWriteable(shardId);
             out.writeString(inputDataType.typeName());
             out.writeWriteable(inputPage);
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java
@@ -601,6 +601,7 @@ private PhysicalOperation planLookupJoin(LookupJoinExec join, LocalExecutionPlan
         return source.with(
             new LookupFromIndexOperator.Factory(
                 sessionId,
+                configuration,
                 parentTask,
                 context.queryPragmas().enrichMaxWorkers(),
                 matchConfig.channel(),
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperatorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperatorTests.java
@@ -53,6 +53,7 @@
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
+import org.elasticsearch.xpack.esql.ConfigurationTestUtils;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -137,6 +138,7 @@ protected Operator.OperatorFactory simple() {
         );
         return new LookupFromIndexOperator.Factory(
             sessionId,
+            ConfigurationTestUtils.randomConfiguration(),
             parentTask,
             maxOutstandingRequests,
             inputChannel,
