ES|QL async queries: Partial result on demand

docs/changelog/118122.yaml

@@ -0,0 +1,5 @@
+pr: 118122
+summary: "ES|QL: Partial result on demand for async queries"
+area: ES|QL
+type: enhancement
+issues: []

docs/reference/esql/esql-across-clusters.asciidoc

@@ -231,6 +231,7 @@ Which returns:
     "skipped": 0,
     "partial": 0,
     "failed": 0,
+    "is_partial": false, <7>
     "details": { <3>
       "(local)": { <4>
         "status": "successful",
@@ -275,8 +276,9 @@ Which returns:
 <2> This section of counters shows all possible cluster search states and how many cluster
 searches are currently in that state. The clusters can have one of the following statuses: *running*,
 *successful* (searches on all shards were successful), *skipped* (the search
-failed on a cluster marked with `skip_unavailable`=`true`) or *failed* (the search
-failed on a cluster marked with `skip_unavailable`=`false`).
+failed on a cluster marked with `skip_unavailable`=`true`), *failed* (the search
+failed on a cluster marked with `skip_unavailable`=`false`) or **partial** (the search was
+<<esql-async-query-stop-api, interrupted>> before finishing).
 <3> The `_clusters/details` section shows metadata about the search on each cluster.
 <4> If you included indices from the local cluster you sent the request to in your {ccs},
 it is identified as "(local)".
@@ -285,6 +287,8 @@ which clusters have slower response times than others.
 <6> The shard details for the search on that cluster, including a count of shards that were
 skipped due to the can-match phase results. Shards are skipped when they cannot have any matching data
 and therefore are not included in the full ES|QL query.
+<7> The `is_partial` field is set to `true` if the search was interrupted before finishing using
+<<esql-async-query-stop-api,async stop API>>.
 
 
 The cross-cluster metadata can be used to determine whether any data came back from a cluster.
@@ -327,6 +331,7 @@ Which returns:
     "skipped": 0,
     "partial": 0,
     "failed": 0,
+    "is_partial": false,
     "details": {
       "cluster_one": {
         "status": "successful",

docs/reference/esql/esql-apis.asciidoc

@@ -17,6 +17,7 @@ overview of {esql} and related tutorials, see <<esql>>.
 * <<esql-async-query-api>>
 * <<esql-async-query-get-api>>
 * <<esql-async-query-delete-api>>
+* <<esql-async-query-stop-api>>
 
 
 include::esql-query-api.asciidoc[]
@@ -26,3 +27,5 @@ include::esql-async-query-api.asciidoc[]
 include::esql-async-query-get-api.asciidoc[]
 
 include::esql-async-query-delete-api.asciidoc[]
+
+include::esql-async-query-stop-api.asciidoc[]

docs/reference/esql/esql-async-query-api.asciidoc

@@ -170,3 +170,10 @@ API>> to get the current status and available results for the query.
 (Boolean)
 If `true`, the query request is still executing.
 --
+
+`is_partial`::
++
+--
+(Boolean)
+If `true`, the query has partial results - for example, as a result of using the <<esql-async-query-stop-api, query stop API>>.
+--

docs/reference/esql/esql-async-query-stop-api.asciidoc

@@ -0,0 +1,49 @@
+[[esql-async-query-stop-api]]
+=== {esql} async query stop API
+++++
+<titleabbrev>{esql} async query stop API</titleabbrev>
+++++
+
+.New API reference
+[sidebar]
+--
+For the most up-to-date API details, refer to {api-es}/group/endpoint-esql[ES|QL APIs].
+--
+
+The <<esql,{esql}>> async query stop API is used to manually stop an async query. Once stop command is issued,
+the query stops processing new data and returns the results that have been already processed. Note that due to the pipelined
+nature of {esql} queries, the stop operation is not immediate and may take time to return results.
+
+The results are returned in <<esql-query-api-response-body,the same format>> as for
+<<esql-async-query-get-api,{esql} async query get API>>.
+If the query has been finished by the time the stop command is issued, the results are returned immediately.
+
+If the query processing has not been finished by the time the stop command is issued, the result data will have
+`is_partial` field set to `true`.
+
+[source,console]
+----
+POST /query/async/FkpMRkJGS1gzVDRlM3g4ZzMyRGlLbkEaTXlJZHdNT09TU2VTZVBoNDM3cFZMUToxMDM=/stop
+----
+// TEST[skip: no access to query ID]
+
+[[esql-async-query-stop-api-request]]
+==== {api-request-title}
+
+`POST /_query/async/<query_id>/stop`
+
+[[esql-async-query-stop-api-prereqs]]
+==== {api-prereq-title}
+
+* If the {es} {security-features} are enabled, only the authenticated user that submitted the original query request
+can stop the query.
+
+[[esql-async-query-stop-api-path-params]]
+==== {api-path-parms-title}
+
+`<query_id>`::
+(Required, string)
+Identifier for the query to stop.
++
+A query ID is provided in the <<esql-async-query-api,{esql} async query API>>'s
+response for a query that does not complete in the awaited time.

rest-api-spec/src/main/resources/rest-api-spec/api/esql.async_query_stop.json

@@ -0,0 +1,31 @@
+{
+  "esql.async_query_stop": {
+    "documentation": {
+      "url": "https://www.elastic.co/guide/en/elasticsearch/reference/master/esql-async-query-stop-api.html",
+      "description": "Stops a previously submitted async query request given its ID and collects the results."
+    },
+    "stability": "stable",
+    "visibility": "public",
+    "headers": {
+      "accept": [
+        "application/json"
+      ]
+    },
+    "url": {
+      "paths": [
+        {
+          "path": "/_query/async/{id}/stop",
+          "methods": [
+            "POST"
+          ],
+          "parts": {
+            "id": {
+              "type": "string",
+              "description": "The async query ID"
+            }
+          }
+        }
+      ]
+    }
+  }
+}

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -157,6 +157,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ELASTIC_INFERENCE_SERVICE_UNIFIED_CHAT_COMPLETIONS_INTEGRATION = def(8_822_00_0);
     public static final TransportVersion KQL_QUERY_TECH_PREVIEW = def(8_823_00_0);
     public static final TransportVersion ESQL_PROFILE_ROWS_PROCESSED = def(8_824_00_0);
+    public static final TransportVersion ESQL_RESPONSE_PARTIAL = def(8_825_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/async/AsyncStopRequest.java

@@ -0,0 +1,67 @@
+/*
+ * 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.async;
+
+import org.elasticsearch.action.ActionRequest;
+import org.elasticsearch.action.ActionRequestValidationException;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+import java.util.Objects;
+
+/**
+ * Request for TransportEsqlAsyncStopAction action.
+ */
+public class AsyncStopRequest extends ActionRequest {
+    private final String id;
+
+    /**
+     * Creates a new request
+     *
+     * @param id The id of the search progress request.
+     */
+    public AsyncStopRequest(String id) {
+        this.id = id;
+    }
+
+    public AsyncStopRequest(StreamInput in) throws IOException {
+        super(in);
+        this.id = in.readString();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeString(id);
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        return null;
+    }
+
+    /**
+     * Returns the id of the async search.
+     */
+    public String getId() {
+        return id;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        AsyncStopRequest request = (AsyncStopRequest) o;
+        return Objects.equals(id, request.id);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(id);
+    }
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/esql/EsqlAsyncActionNames.java

@@ -12,4 +12,5 @@
  */
 public class EsqlAsyncActionNames {
     public static final String ESQL_ASYNC_GET_RESULT_ACTION_NAME = "indices:data/read/esql/async/get";
+    public static final String ESQL_ASYNC_STOP_ACTION_NAME = "indices:data/read/esql/async/stop";
 }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeService.java

@@ -76,6 +76,7 @@ public final class ExchangeService extends AbstractLifecycleComponent {
     private final BlockFactory blockFactory;
 
     private final Map<String, ExchangeSinkHandler> sinks = ConcurrentCollections.newConcurrentMap();
+    private final Map<String, ExchangeSourceHandler> exchangeSources = ConcurrentCollections.newConcurrentMap();
 
     public ExchangeService(Settings settings, ThreadPool threadPool, String executorName, BlockFactory blockFactory) {
         this.threadPool = threadPool;
@@ -172,6 +173,32 @@ public static void openExchange(
         );
     }
 
+    /**
+     * Remember the exchange source handler for the given session ID.
+     * This can be used for async/stop requests.
+     */
+    public void addExchangeSourceHandler(String sessionId, ExchangeSourceHandler sourceHandler) {
+        exchangeSources.put(sessionId, sourceHandler);
+    }
+
+    public ExchangeSourceHandler removeExchangeSourceHandler(String sessionId) {
+        return exchangeSources.remove(sessionId);
+    }
+
+    /**
+     * Finishes the session early, i.e., before all sources are finished.
+     * It is called by async/stop API and should be called on the node that coordinates the async request.
+     * It will close all sources and return the results - unlike cancel, this does not discard the results.
+     */
+    public void finishSessionEarly(String sessionId, ActionListener<Void> listener) {
+        ExchangeSourceHandler exchangeSource = removeExchangeSourceHandler(sessionId);
+        if (exchangeSource != null) {
+            exchangeSource.finishEarly(false, listener);
+        } else {
+            listener.onResponse(null);
+        }
+    }
+
     private static class OpenExchangeRequest extends TransportRequest {
         private final String sessionId;
         private final int exchangeBuffer;

x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClustersIT.java

@@ -42,6 +42,7 @@
 import static org.hamcrest.Matchers.any;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.hasItems;
 import static org.hamcrest.Matchers.hasKey;
 
 @ThreadLeakFilters(filters = TestClustersThreadFilter.class)
@@ -260,8 +261,8 @@ public void testUngroupedAggs() throws Exception {
     private void assertClusterDetailsMap(Map<String, Object> result, boolean remoteOnly) {
         @SuppressWarnings("unchecked")
         Map<String, Object> clusters = (Map<String, Object>) result.get("_clusters");
-        assertThat(clusters.size(), equalTo(7));
-        assertThat(clusters.keySet(), equalTo(Set.of("total", "successful", "running", "skipped", "partial", "failed", "details")));
+        assertThat(clusters.size(), greaterThanOrEqualTo(7));
+        assertThat(clusters.keySet(), hasItems("total", "successful", "running", "skipped", "partial", "failed", "details"));
         int expectedNumClusters = remoteOnly ? 1 : 2;
         Set<String> expectedClusterAliases = remoteOnly ? Set.of("remote_cluster") : Set.of("remote_cluster", "(local)");
 
@@ -271,6 +272,10 @@ private void assertClusterDetailsMap(Map<String, Object> result, boolean remoteO
         assertThat(clusters.get("skipped"), equalTo(0));
         assertThat(clusters.get("partial"), equalTo(0));
         assertThat(clusters.get("failed"), equalTo(0));
+        if (clusters.containsKey("is_partial")) {
+            // for some BWC tests, the is_partial key may not be present
+            assertThat(clusters.get("is_partial"), equalTo(false));
+        }
 
         @SuppressWarnings("unchecked")
         Map<String, Object> details = (Map<String, Object>) clusters.get("details");

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractCrossClustersUsageTelemetryIT.java

@@ -127,7 +127,7 @@ protected CCSTelemetrySnapshot getTelemetryFromFailedQuery(String query) throws
         return getTelemetrySnapshot(queryNode);
     }
 
-    private CCSTelemetrySnapshot getTelemetrySnapshot(String nodeName) {
+    protected CCSTelemetrySnapshot getTelemetrySnapshot(String nodeName) {
         var usage = cluster(LOCAL_CLUSTER).getInstance(UsageService.class, nodeName);
         return usage.getEsqlUsageHolder().getCCSTelemetrySnapshot();
     }

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractPauseFieldPlugin.java

@@ -34,12 +34,16 @@ protected void onStartExecute() {}
     // Called when the engine needs to wait for further execution to be allowed.
     protected abstract boolean onWait() throws InterruptedException;
 
+    protected String scriptTypeName() {
+        return "pause";
+    }
+
     @Override
     public ScriptEngine getScriptEngine(Settings settings, Collection<ScriptContext<?>> contexts) {
         return new ScriptEngine() {
             @Override
             public String getType() {
-                return "pause";
+                return scriptTypeName();
             }
 
             @Override