Allow partial results in ES|QL

Author: dnhatn

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

@@ -189,6 +189,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_PROFILE_ASYNC_NANOS = def(9_007_00_0);
     public static final TransportVersion ESQL_LOOKUP_JOIN_SOURCE_TEXT = def(9_008_0_00);
     public static final TransportVersion REMOVE_ALL_APPLICABLE_SELECTOR = def(9_009_0_00);
+    public static final TransportVersion ESQL_SUPPORT_PARTIAL_RESULTS = def(9_010_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/ConfigurationTestUtils.java

@@ -71,7 +71,8 @@ public static Configuration randomConfiguration(String query, Map<String, Map<St
             query,
             profile,
             tables,
-            System.nanoTime()
+            System.nanoTime(),
+            false
         );
     }
 

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -388,7 +388,8 @@ public static Configuration configuration(QueryPragmas pragmas, String query) {
             query,
             false,
             TABLES,
-            System.nanoTime()
+            System.nanoTime(),
+            false
         );
     }
 

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

@@ -96,7 +96,16 @@ private EsqlQueryResponse runWithBreaking(EsqlQueryRequest request) throws Circu
     @Override
     protected EsqlQueryResponse run(EsqlQueryRequest request) {
         try {
-            return runWithBreaking(request);
+            if (randomBoolean()) {
+                request.allowPartialResults(randomBoolean());
+            }
+            var resp = runWithBreaking(request);
+            if (resp.isPartial() == false) {
+                return resp;
+            }
+            try (resp) {
+                assertTrue(request.allowPartialResults());
+            }
         } catch (Exception e) {
             try (EsqlQueryResponse resp = super.run(request)) {
                 assertThat(e, instanceOf(CircuitBreakingException.class));
@@ -105,6 +114,7 @@ protected EsqlQueryResponse run(EsqlQueryRequest request) {
                 return resp;
             }
         }
+        return super.run(request);
     }
 
     /**

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

@@ -83,16 +83,31 @@ private EsqlQueryResponse runQueryWithDisruption(EsqlQueryRequest request) {
         logger.info("--> start disruption scheme [{}]", disruptionScheme);
         disruptionScheme.startDisrupting();
         logger.info("--> executing esql query with disruption {} ", request.query());
+        if (randomBoolean()) {
+            request.allowPartialResults(randomBoolean());
+        }
         ActionFuture<EsqlQueryResponse> future = client().execute(EsqlQueryAction.INSTANCE, request);
         try {
-            return future.actionGet(2, TimeUnit.MINUTES);
+            var resp = future.actionGet(2, TimeUnit.MINUTES);
+            if (resp.isPartial() == false) {
+                return resp;
+            }
+            try (resp) {
+                assertTrue(request.allowPartialResults());
+            }
         } catch (Exception ignored) {
 
         } finally {
             clearDisruption();
         }
         try {
-            return future.actionGet(2, TimeUnit.MINUTES);
+            var resp = future.actionGet(2, TimeUnit.MINUTES);
+            if (resp.isPartial() == false) {
+                return resp;
+            }
+            try (resp) {
+                assertTrue(request.allowPartialResults());
+            }
         } catch (Exception e) {
             logger.info(
                 "running tasks: {}",
@@ -113,8 +128,8 @@ private EsqlQueryResponse runQueryWithDisruption(EsqlQueryRequest request) {
             ensureBlocksReleased();
             logger.info("--> failed to execute esql query with disruption; retrying...", e);
             EsqlTestUtils.assertEsqlFailure(e);
-            return client().execute(EsqlQueryAction.INSTANCE, request).actionGet(2, TimeUnit.MINUTES);
         }
+        return super.run(request);
     }
 
     private ServiceDisruptionScheme addRandomDisruptionScheme() {

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

@@ -16,13 +16,17 @@
 import org.elasticsearch.test.FailingFieldPlugin;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.json.JsonXContent;
+import org.elasticsearch.xpack.esql.EsqlTestUtils;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.in;
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 /**
  * Make sure the failures on the data node come back as failures over the wire.
@@ -48,10 +52,7 @@ protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
         return settings;
     }
 
-    /**
-     * Use a runtime field that fails when loading field values to fail the entire query.
-     */
-    public void testFailureLoadingFields() throws IOException {
+    public Set<String> populateIndices() throws Exception {
         XContentBuilder mapping = JsonXContent.contentBuilder().startObject();
         mapping.startObject("runtime");
         {
@@ -63,17 +64,62 @@ public void testFailureLoadingFields() throws IOException {
             mapping.endObject();
         }
         mapping.endObject();
-        client().admin().indices().prepareCreate("fail").setSettings(indexSettings(1, 0)).setMapping(mapping.endObject()).get();
-
-        int docCount = 50;
-        List<IndexRequestBuilder> docs = new ArrayList<>(docCount);
-        for (int d = 0; d < docCount; d++) {
-            docs.add(client().prepareIndex("ok").setSource("foo", d));
+        client().admin().indices().prepareCreate("fail").setMapping(mapping.endObject()).get();
+        int okCount = between(1, 50);
+        Set<String> okIds = new HashSet<>();
+        List<IndexRequestBuilder> docs = new ArrayList<>(okCount);
+        for (int d = 0; d < okCount; d++) {
+            String id = "ok-" + d;
+            okIds.add(id);
+            docs.add(client().prepareIndex("ok").setId(id).setSource("foo", d));
+        }
+        int failCount = between(1, 50);
+        for (int d = 0; d < failCount; d++) {
+            docs.add(client().prepareIndex("fail").setId("fail-" + d).setSource("foo", d));
         }
-        docs.add(client().prepareIndex("fail").setSource("foo", 0));
         indexRandom(true, docs);
+        return okIds;
+    }
 
+    /**
+     * Use a runtime field that fails when loading field values to fail the entire query.
+     */
+    public void testFailureLoadingFields() throws Exception {
+        populateIndices();
         IllegalStateException e = expectThrows(IllegalStateException.class, () -> run("FROM fail,ok | LIMIT 100").close());
         assertThat(e.getMessage(), equalTo("Accessing failing field"));
     }
+
+    public void testPartialResults() throws Exception {
+        Set<String> okIds = populateIndices();
+        {
+            EsqlQueryRequest request = new EsqlQueryRequest();
+            request.query("FROM fail,ok | LIMIT 100");
+            request.allowPartialResults(true);
+            request.pragmas(randomPragmas());
+            try (EsqlQueryResponse resp = run(request)) {
+                assertTrue(resp.isPartial());
+                List<List<Object>> rows = EsqlTestUtils.getValuesList(resp);
+                assertThat(rows.size(), lessThanOrEqualTo(okIds.size()));
+            }
+        }
+        {
+            EsqlQueryRequest request = new EsqlQueryRequest();
+            request.query("FROM fail,ok METADATA _id | KEEP _id, fail_me | LIMIT 100");
+            request.allowPartialResults(true);
+            request.pragmas(randomPragmas());
+            try (EsqlQueryResponse resp = run(request)) {
+                assertTrue(resp.isPartial());
+                List<List<Object>> rows = EsqlTestUtils.getValuesList(resp);
+                assertThat(rows.size(), lessThanOrEqualTo(okIds.size()));
+                Set<String> actualIds = new HashSet<>();
+                for (List<Object> row : rows) {
+                    assertThat(row.size(), equalTo(2));
+                    String id = (String) row.getFirst();
+                    assertThat(id, in(okIds));
+                    assertTrue(actualIds.add(id));
+                }
+            }
+        }
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -809,7 +809,12 @@ public enum Cap {
          * Fixes a series of issues with inlinestats which had an incomplete implementation after lookup and inlinestats
          * were refactored.
          */
-        INLINESTATS_V3(EsqlPlugin.INLINESTATS_FEATURE_FLAG);
+        INLINESTATS_V3(EsqlPlugin.INLINESTATS_FEATURE_FLAG),
+
+        /**
+         * Support partial_results
+         */
+        SUPPORT_PARTIAL_RESULTS;
 
         private final boolean enabled;
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlExecutionInfo.java

@@ -174,6 +174,9 @@ public TimeValue planningTookTime() {
     public void markEndQuery() {
         assert relativeStartNanos != null : "Relative start time must be set when markEndQuery is called";
         overallTook = new TimeValue(System.nanoTime() - relativeStartNanos, TimeUnit.NANOSECONDS);
+        if (isPartial == false) {
+            isPartial = clusterInfo.values().stream().anyMatch(c -> c.failedShards != null && c.failedShards > 0);
+        }
     }
 
     // for testing only - use markEndQuery in production code

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryRequest.java

@@ -52,6 +52,7 @@ public class EsqlQueryRequest extends org.elasticsearch.xpack.core.esql.action.E
     private boolean keepOnCompletion;
     private boolean onSnapshotBuild = Build.current().isSnapshot();
     private boolean acceptedPragmaRisks = false;
+    private boolean allowPartialResults = false;
 
     /**
      * "Tables" provided in the request for use with things like {@code LOOKUP}.
@@ -231,6 +232,14 @@ public Map<String, Map<String, Column>> tables() {
         return tables;
     }
 
+    public boolean allowPartialResults() {
+        return allowPartialResults;
+    }
+
+    public void allowPartialResults(boolean allowPartialResults) {
+        this.allowPartialResults = allowPartialResults;
+    }
+
     @Override
     public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
         // Pass the query as the description

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponse.java

@@ -181,6 +181,10 @@ public boolean isAsync() {
         return isRunning;
     }
 
+    public boolean isPartial() {
+        return executionInfo != null && executionInfo.isPartial();
+    }
+
     public EsqlExecutionInfo getExecutionInfo() {
         return executionInfo;
     }