return partial_results

dnhatn · dnhatn · commit 2e40166aca3d · 2025-02-10T14:40:46.000-08:00
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java
@@ -389,7 +389,7 @@ public static Configuration configuration(QueryPragmas pragmas, String query) {
             false,
             TABLES,
             System.nanoTime(),
-          false
+            false
         );
     }
 
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractEsqlIntegTestCase.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractEsqlIntegTestCase.java
@@ -150,9 +150,6 @@ protected final EsqlQueryResponse run(String esqlCommands, QueryPragmas pragmas)
     protected EsqlQueryResponse run(String esqlCommands, QueryPragmas pragmas, QueryBuilder filter) {
         EsqlQueryRequest request = EsqlQueryRequest.syncEsqlQueryRequest();
         request.query(esqlCommands);
-        if (randomBoolean()) {
-            request.allowPartialResults(randomBoolean());
-        }
         if (pragmas != null) {
             request.pragmas(pragmas);
         }
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionBreakerIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionBreakerIT.java
@@ -96,7 +96,17 @@ private EsqlQueryResponse runWithBreaking(EsqlQueryRequest request) throws Circu
     @Override
     protected EsqlQueryResponse run(EsqlQueryRequest request) {
         try {
-            return runWithBreaking(request);
+            if (randomBoolean()) {
+                request.allowPartialResults(randomBoolean());
+            }
+            try (EsqlQueryResponse resp = runWithBreaking(request)) {
+                if (resp.isPartial()) {
+                    assertTrue(request.allowPartialResults());
+                } else {
+                    resp.incRef();
+                    return resp;
+                }
+            }
         } catch (Exception e) {
             try (EsqlQueryResponse resp = super.run(request)) {
                 assertThat(e, instanceOf(CircuitBreakingException.class));
@@ -105,6 +115,7 @@ protected EsqlQueryResponse run(EsqlQueryRequest request) {
                 return resp;
             }
         }
+        return super.run(request);
     }
 
     /**
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlDisruptionIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlDisruptionIT.java
@@ -84,16 +84,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);
+            try (var resp = future.actionGet(2, TimeUnit.MINUTES)) {
+                if (resp.isPartial()) {
+                    assertTrue(request.allowPartialResults());
+                } else {
+                    resp.incRef();
+                    return resp;
+                }
+            }
         } catch (Exception ignored) {
 
         } finally {
             clearDisruption();
         }
-        try {
-            return future.actionGet(2, TimeUnit.MINUTES);
+        try (EsqlQueryResponse resp = future.actionGet(2, TimeUnit.MINUTES)) {
+            if (resp.isPartial()) {
+                assertTrue(request.allowPartialResults());
+            } else {
+                resp.incRef();
+                return resp;
+            }
         } catch (Exception e) {
             logger.info(
                 "running tasks: {}",
@@ -114,9 +129,9 @@ private EsqlQueryResponse runQueryWithDisruption(EsqlQueryRequest request) {
             ensureBlocksReleased();
             logger.info("--> failed to execute esql query with disruption; retrying...", e);
             EsqlTestUtils.assertEsqlFailure(e);
-            request.allowPartialResults(false);
             return client().execute(EsqlQueryAction.INSTANCE, request).actionGet(2, TimeUnit.MINUTES);
         }
+        return super.run(request);
     }
 
     private ServiceDisruptionScheme addRandomDisruptionScheme() {
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlNodeFailureIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlNodeFailureIT.java
@@ -16,13 +16,15 @@
 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.List;
 
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 /**
  * Make sure the failures on the data node come back as failures over the wire.
@@ -73,7 +75,23 @@ public void testFailureLoadingFields() throws IOException {
         docs.add(client().prepareIndex("fail").setSource("foo", 0));
         indexRandom(true, docs);
 
-        IllegalStateException e = expectThrows(IllegalStateException.class, () -> run("FROM fail,ok | LIMIT 100").close());
-        assertThat(e.getMessage(), equalTo("Accessing failing field"));
+        // without partial results
+        String query = "FROM fail,ok | LIMIT 100";
+        {
+            IllegalStateException e = expectThrows(IllegalStateException.class, () -> run(query).close());
+            assertThat(e.getMessage(), equalTo("Accessing failing field"));
+        }
+        // allow partial results
+        {
+            EsqlQueryRequest request = new EsqlQueryRequest();
+            request.query(query);
+            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(50));
+            }
+        }
     }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlExecutionInfo.java b/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
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponse.java b/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;
     }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java
@@ -209,11 +209,11 @@ public void execute(
                         transportService.getThreadPool(),
                         cancelQueryOnFailure,
                         computeListener.acquireCompute().delegateFailure((l, profiles) -> {
-                            if (execInfo.isCrossClusterSearch() && execInfo.clusterAliases().contains(LOCAL_CLUSTER)) {
-                                var tookTime = TimeValue.timeValueNanos(System.nanoTime() - execInfo.getRelativeStartNanos());
-                                var status = localClusterWasInterrupted.get()
-                                    ? EsqlExecutionInfo.Cluster.Status.PARTIAL
-                                    : EsqlExecutionInfo.Cluster.Status.SUCCESSFUL;
+                            var tookTime = TimeValue.timeValueNanos(System.nanoTime() - execInfo.getRelativeStartNanos());
+                            var status = localClusterWasInterrupted.get()
+                                ? EsqlExecutionInfo.Cluster.Status.PARTIAL
+                                : EsqlExecutionInfo.Cluster.Status.SUCCESSFUL;
+                            if (execInfo.clusterInfo.containsKey(LOCAL_CLUSTER)) {
                                 execInfo.swapCluster(
                                     LOCAL_CLUSTER,
                                     (k, v) -> new EsqlExecutionInfo.Cluster.Builder(v).setStatus(status).setTook(tookTime).build()
@@ -252,16 +252,14 @@ public void execute(
                             cancelQueryOnFailure,
                             localListener.acquireCompute().map(r -> {
                                 localClusterWasInterrupted.set(execInfo.isPartial());
-                                if (execInfo.isCrossClusterSearch() && execInfo.clusterAliases().contains(LOCAL_CLUSTER)) {
-                                    execInfo.swapCluster(
-                                        LOCAL_CLUSTER,
-                                        (k, v) -> new EsqlExecutionInfo.Cluster.Builder(v).setTotalShards(r.getTotalShards())
-                                            .setSuccessfulShards(r.getSuccessfulShards())
-                                            .setSkippedShards(r.getSkippedShards())
-                                            .setFailedShards(r.getFailedShards())
-                                            .build()
-                                    );
-                                }
+                                execInfo.swapCluster(
+                                    LOCAL_CLUSTER,
+                                    (k, v) -> new EsqlExecutionInfo.Cluster.Builder(v).setTotalShards(r.getTotalShards())
+                                        .setSuccessfulShards(r.getSuccessfulShards())
+                                        .setSkippedShards(r.getSkippedShards())
+                                        .setFailedShards(r.getFailedShards())
+                                        .build()
+                                );
                                 return r.getProfiles();
                             })
                         );
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java
@@ -127,13 +127,32 @@ protected void sendRequest(
                     esqlExecutor,
                     listener.delegateFailureAndWrap((l, unused) -> {
                         final AtomicReference<DataNodeComputeResponse> nodeResponseRef = new AtomicReference<>();
+                        final Runnable onGroupFailure;
+                        if (allowPartialResults) {
+                            CancellableTask groupTask = RemoteListenerGroup.createGroupTask(
+                                transportService,
+                                parentTask,
+                                () -> "compute group, data-node [" + node.getName() + "], " + shardIds + " [" + shardIds + "]"
+                            );
+                            onGroupFailure = computeService.cancelQueryOnFailure(groupTask);
+                            l = ActionListener.runAfter(l, () -> transportService.getTaskManager().unregister(groupTask));
+                        } else {
+                            onGroupFailure = () -> {
+                                if (nodeResponseRef.get() != null) {
+                                    exchangeService.finishSinkHandler(
+                                        childSessionId,
+                                        new TaskCancelledException(parentTask.getReasonCancelled())
+                                    );
+                                }
+                            };
+                        }
                         try (
-                            var computeListener = new ComputeListener(threadPool, runOnTaskFailure, l.map(ignored -> nodeResponseRef.get()))
+                            var computeListener = new ComputeListener(threadPool, onGroupFailure, l.map(ignored -> nodeResponseRef.get()))
                         ) {
                             final var remoteSink = exchangeService.newRemoteSink(parentTask, childSessionId, transportService, connection);
                             exchangeSource.addRemoteSink(
                                 remoteSink,
-                                true,
+                                allowPartialResults == false,
                                 pagesFetched::incrementAndGet,
                                 queryPragmas.concurrentExchangeClients(),
                                 computeListener.acquireAvoid()
@@ -239,6 +258,7 @@ public void onFailure(Exception e) {
                         }
                         onResponse(List.of());
                     } else {
+                        // TODO: fatal vs non-fatal
                         try {
                             exchangeService.finishSinkHandler(request.sessionId(), e);
                         } finally {
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestSender.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestSender.java
@@ -82,13 +82,14 @@ final void startComputeOnDataNodes(
         searchShards(rootTask, clusterAlias, requestFilter, concreteIndices, originalIndices, ActionListener.wrap(targetShards -> {
             try (var computeListener = new ComputeListener(transportService.getThreadPool(), runOnTaskFailure, listener.map(profiles -> {
                 TimeValue took = TimeValue.timeValueNanos(System.nanoTime() - startTimeInNanos);
+                final int failedShards = shardFailures.size();
                 return new ComputeResponse(
                     profiles,
                     took,
                     targetShards.totalShards(),
-                    targetShards.totalShards(),
+                    targetShards.totalShards() - failedShards,
                     targetShards.skippedShards(),
-                    0
+                    failedShards
                 );
             }))) {
                 for (TargetShard shard : targetShards.shards.values()) {
@@ -111,21 +112,26 @@ private void trySendingRequestsForPendingShards(TargetShards targetShards, Compu
                     if (changed.compareAndSet(true, false) == false) {
                         break;
                     }
+                    for (ShardId shardId : pendingShardIds) {
+                        if (targetShards.getShard(shardId).remainingNodes.isEmpty()) shardFailures.compute(
+                            shardId,
+                            (k, v) -> new ShardFailure(
+                                true,
+                                v == null ? new NoShardAvailableActionException(shardId, "no shard copies found") : v.failure
+                            )
+                        );
+                    }
+                    // remove shards with fatal failure from pending shards
                     final Iterator<ShardId> shardIts = pendingShardIds.iterator();
                     while (shardIts.hasNext()) {
                         final ShardId shardId = shardIts.next();
-                        if (targetShards.getShard(shardId).remainingNodes.isEmpty()) {
-                            shardFailures.compute(
-                                shardId,
-                                (k, v) -> new ShardFailure(
-                                    true,
-                                    v == null ? new NoShardAvailableActionException(shardId, "no shard copies found") : v.failure
-                                )
-                            );
+                        final ShardFailure failure = shardFailures.get(shardId);
+                        if (failure != null && failure.fatal) {
                             shardIts.remove();
                         }
                     }
-                    if (reportedFailure || shardFailures.values().stream().anyMatch(shardFailure -> shardFailure.fatal)) {
+                    if (reportedFailure
+                        || (allowPartialResults == false && shardFailures.values().stream().anyMatch(shardFailure -> shardFailure.fatal))) {
                         reportedFailure = true;
                         reportFailures(computeListener);
                     } else {
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/RemoteListenerGroup.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/RemoteListenerGroup.java
@@ -50,7 +50,7 @@ class RemoteListenerGroup {
         this.taskManager = transportService.getTaskManager();
         this.clusterAlias = clusterAlias;
         this.executionInfo = executionInfo;
-        groupTask = createGroupTask(rootTask, () -> rootTask.getDescription() + "[" + clusterAlias + "]");
+        groupTask = createGroupTask(transportService, rootTask, () -> rootTask.getDescription() + "[" + clusterAlias + "]");
         CountDown countDown = new CountDown(2);
         // The group is done when both the sink and the cluster request are done
         Runnable finishGroup = () -> {
@@ -92,7 +92,8 @@ public ActionListener<List<DriverProfile>> getClusterRequestListener() {
         return clusterRequestListener;
     }
 
-    private CancellableTask createGroupTask(Task parentTask, Supplier<String> description) {
+    public static CancellableTask createGroupTask(TransportService transportService, Task parentTask, Supplier<String> description) {
+        final TaskManager taskManager = transportService.getTaskManager();
         return (CancellableTask) taskManager.register(
             "transport",
             "esql_compute_group",

Original file line number	Diff line number	Diff line change
`@@ -389,7 +389,7 @@ public static Configuration configuration(QueryPragmas pragmas, String query) {`
`389`	`389`	`false,`
`390`	`390`	`TABLES,`
`391`	`391`	`System.nanoTime(),`
`392`		`- false`
	`392`	`+ false`
`393`	`393`	`);`
`394`	`394`	`}`
`395`	`395`
Original file line number	Diff line number	Diff line change
`@@ -174,6 +174,9 @@ public TimeValue planningTookTime() {`
`174`	`174`	`public void markEndQuery() {`
`175`	`175`	`assert relativeStartNanos != null : "Relative start time must be set when markEndQuery is called";`
`176`	`176`	`overallTook = new TimeValue(System.nanoTime() - relativeStartNanos, TimeUnit.NANOSECONDS);`
	`177`	`+ if (isPartial == false) {`
	`178`	`+ isPartial = clusterInfo.values().stream().anyMatch(c -> c.failedShards != null && c.failedShards > 0);`
	`179`	`+ }`
`177`	`180`	`}`
`178`	`181`
`179`	`182`	`// for testing only - use markEndQuery in production code`
Original file line number	Diff line number	Diff line change
`@@ -181,6 +181,10 @@ public boolean isAsync() {`
`181`	`181`	`return isRunning;`
`182`	`182`	`}`
`183`	`183`
	`184`	`+ public boolean isPartial() {`
	`185`	`+ return executionInfo != null && executionInfo.isPartial();`
	`186`	`+ }`
	`187`	`+`
`184`	`188`	`public EsqlExecutionInfo getExecutionInfo() {`
`185`	`189`	`return executionInfo;`
`186`	`190`	`}`