POC - Semantic search CCS support with ccs_minimize_roundtrips=false

server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ResolvedIndices;
 import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.client.internal.RemoteClusterClient;
 import org.elasticsearch.cluster.metadata.DataStream;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.routing.allocation.DataTier;
@@ -36,11 +37,13 @@
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.search.builder.PointInTimeBuilder;
 import org.elasticsearch.transport.RemoteClusterAware;
+import org.elasticsearch.transport.RemoteClusterService;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -52,6 +55,8 @@
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
+import static org.elasticsearch.common.util.concurrent.EsExecutors.DIRECT_EXECUTOR_SERVICE;
+
 /**
  * Context object used to rewrite {@link QueryBuilder} instances into simplified version.
  */
@@ -72,6 +77,7 @@ public class QueryRewriteContext {
     protected final Client client;
     protected final LongSupplier nowInMillis;
     private final List<BiConsumer<Client, ActionListener<?>>> asyncActions = new ArrayList<>();
+    private final Map<String, List<BiConsumer<RemoteClusterClient, ActionListener<?>>>> remoteAsyncActions = new HashMap<>();
     protected boolean allowUnmappedFields;
     protected boolean mapUnmappedFieldAsString;
     protected Predicate<String> allowedFields;
@@ -346,22 +352,35 @@ public void registerAsyncAction(BiConsumer<Client, ActionListener<?>> asyncActio
         asyncActions.add(asyncAction);
     }
 
+    public void registerRemoteAsyncAction(String clusterAlias, BiConsumer<RemoteClusterClient, ActionListener<?>> asyncAction) {
+        List<BiConsumer<RemoteClusterClient, ActionListener<?>>> asyncActions = remoteAsyncActions.computeIfAbsent(
+            clusterAlias,
+            k -> new ArrayList<>()
+        );
+        asyncActions.add(asyncAction);
+    }
+
     /**
      * Returns <code>true</code> if there are any registered async actions.
      */
     public boolean hasAsyncActions() {
-        return asyncActions.isEmpty() == false;
+        return asyncActions.isEmpty() == false || remoteAsyncActions.isEmpty() == false;
     }
 
     /**
      * Executes all registered async actions and notifies the listener once it's done. The value that is passed to the listener is always
      * <code>null</code>. The list of registered actions is cleared once this method returns.
      */
     public void executeAsyncActions(ActionListener<Void> listener) {
-        if (asyncActions.isEmpty()) {
+        if (asyncActions.isEmpty() && remoteAsyncActions.isEmpty()) {
             listener.onResponse(null);
         } else {
-            CountDown countDown = new CountDown(asyncActions.size());
+            int actionCount = asyncActions.size();
+            for (var remoteAsyncActionList : remoteAsyncActions.values()) {
+                actionCount += remoteAsyncActionList.size();
+            }
+
+            CountDown countDown = new CountDown(actionCount);
             ActionListener<?> internalListener = new ActionListener<>() {
                 @Override
                 public void onResponse(Object o) {
@@ -377,12 +396,28 @@ public void onFailure(Exception e) {
                     }
                 }
             };
+
             // make a copy to prevent concurrent modification exception
             List<BiConsumer<Client, ActionListener<?>>> biConsumers = new ArrayList<>(asyncActions);
             asyncActions.clear();
             for (BiConsumer<Client, ActionListener<?>> action : biConsumers) {
                 action.accept(client, internalListener);
             }
+
+            for (var entry : remoteAsyncActions.entrySet()) {
+                String clusterAlias = entry.getKey();
+                List<BiConsumer<RemoteClusterClient, ActionListener<?>>> remoteBiConsumers = entry.getValue();
+
+                RemoteClusterClient remoteClient = client.getRemoteClusterClient(
+                    clusterAlias,
+                    DIRECT_EXECUTOR_SERVICE,
+                    RemoteClusterService.DisconnectedStrategy.RECONNECT_UNLESS_SKIP_UNAVAILABLE
+                );
+                for (BiConsumer<RemoteClusterClient, ActionListener<?>> action : remoteBiConsumers) {
+                    action.accept(remoteClient, internalListener);
+                }
+            }
+            remoteAsyncActions.clear();
         }
     }
 

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/GetInferenceFieldsAction.java

@@ -0,0 +1,162 @@
+/*
+ * 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.core.inference.action;
+
+import org.elasticsearch.action.ActionRequest;
+import org.elasticsearch.action.ActionRequestValidationException;
+import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.action.ActionType;
+import org.elasticsearch.action.RemoteClusterActionType;
+import org.elasticsearch.cluster.metadata.InferenceFieldMetadata;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.inference.InferenceResults;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+public class GetInferenceFieldsAction extends ActionType<GetInferenceFieldsAction.Response> {
+    public static final GetInferenceFieldsAction INSTANCE = new GetInferenceFieldsAction();
+    public static final RemoteClusterActionType<Response> REMOTE_TYPE = new RemoteClusterActionType<>(INSTANCE.name(), Response::new);
+
+    public static final String NAME = "cluster:monitor/xpack/inference_fields/get";
+
+    public GetInferenceFieldsAction() {
+        super(NAME);
+    }
+
+    public static class Request extends ActionRequest {
+        private final List<String> indices;
+        private final List<String> fields;
+        private final boolean resolveWildcards;
+        private final boolean useDefaultFields;
+        private final String query;
+
+        public Request(
+            List<String> indices,
+            List<String> fields,
+            boolean resolveWildcards,
+            boolean useDefaultFields,
+            @Nullable String query
+        ) {
+            this.indices = indices;
+            this.fields = fields;
+            this.resolveWildcards = resolveWildcards;
+            this.useDefaultFields = useDefaultFields;
+            this.query = query;
+        }
+
+        public Request(StreamInput in) throws IOException {
+            super(in);
+            this.indices = in.readCollectionAsList(StreamInput::readString);
+            this.fields = in.readCollectionAsList(StreamInput::readString);
+            this.resolveWildcards = in.readBoolean();
+            this.useDefaultFields = in.readBoolean();
+            this.query = in.readOptionalString();
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            super.writeTo(out);
+            out.writeStringCollection(indices);
+            out.writeStringCollection(fields);
+            out.writeBoolean(resolveWildcards);
+            out.writeBoolean(useDefaultFields);
+            out.writeOptionalString(query);
+        }
+
+        @Override
+        public ActionRequestValidationException validate() {
+            return null;
+        }
+
+        public List<String> getIndices() {
+            return Collections.unmodifiableList(indices);
+        }
+
+        public List<String> getFields() {
+            return Collections.unmodifiableList(fields);
+        }
+
+        public boolean resolveWildcards() {
+            return resolveWildcards;
+        }
+
+        public boolean useDefaultFields() {
+            return useDefaultFields;
+        }
+
+        public String getQuery() {
+            return query;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            Request request = (Request) o;
+            return Objects.equals(indices, request.indices)
+                && Objects.equals(fields, request.fields)
+                && resolveWildcards == request.resolveWildcards
+                && useDefaultFields == request.useDefaultFields
+                && Objects.equals(query, request.query);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(indices, fields, resolveWildcards, useDefaultFields, query);
+        }
+    }
+
+    public static class Response extends ActionResponse {
+        private final Map<String, List<InferenceFieldMetadata>> inferenceFieldsMap;
+        private final Map<String, InferenceResults> inferenceResultsMap;
+
+        public Response(Map<String, List<InferenceFieldMetadata>> inferenceFieldsMap, Map<String, InferenceResults> inferenceResultsMap) {
+            this.inferenceFieldsMap = inferenceFieldsMap;
+            this.inferenceResultsMap = inferenceResultsMap;
+        }
+
+        public Response(StreamInput in) throws IOException {
+            this.inferenceFieldsMap = in.readImmutableMap(i -> i.readCollectionAsImmutableList(InferenceFieldMetadata::new));
+            this.inferenceResultsMap = in.readImmutableMap(i -> i.readNamedWriteable(InferenceResults.class));
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            out.writeMap(inferenceFieldsMap, StreamOutput::writeCollection);
+            out.writeMap(inferenceResultsMap, StreamOutput::writeNamedWriteable);
+        }
+
+        public Map<String, List<InferenceFieldMetadata>> getInferenceFieldsMap() {
+            return Collections.unmodifiableMap(this.inferenceFieldsMap);
+        }
+
+        public Map<String, InferenceResults> getInferenceResultsMap() {
+            return Collections.unmodifiableMap(this.inferenceResultsMap);
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            Response response = (Response) o;
+            return Objects.equals(inferenceFieldsMap, response.inferenceFieldsMap)
+                && Objects.equals(inferenceResultsMap, response.inferenceResultsMap);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(inferenceFieldsMap, inferenceResultsMap);
+        }
+    }
+}

x-pack/plugin/esql/build.gradle

@@ -71,6 +71,9 @@ dependencies {
   testImplementation('org.webjars.npm:fontsource__roboto-mono:4.5.7')
 
   internalClusterTestImplementation project(":modules:mapper-extras")
+  internalClusterTestImplementation project(xpackModule('inference'))
+  internalClusterTestImplementation testArtifact(project(xpackModule('inference')))
+  internalClusterTestImplementation testArtifact(project(xpackModule('inference')), 'internalClusterTest')
 }
 
 tasks.named("dependencyLicenses").configure {