Make POS count failed task info and perform dedup on them

Author: tanjialiang

presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java

@@ -679,7 +679,7 @@ else if (preparedQuery.isExplainTypeValidate()) {
                 planAndMore = queryPlanner.createQueryPlan(session, preparedQuery, warningCollector, variableAllocator, planNodeIdAllocator, sparkContext, sql);
                 JavaSparkContext javaSparkContext = new JavaSparkContext(sparkContext);
                 CollectionAccumulator<SerializedTaskInfo> taskInfoCollector = new CollectionAccumulator<>();
-                taskInfoCollector.register(sparkContext, Option.empty(), false);
+                taskInfoCollector.register(sparkContext, Option.empty(), true);
                 CollectionAccumulator<PrestoSparkShuffleStats> shuffleStatsCollector = new CollectionAccumulator<>();
                 shuffleStatsCollector.register(sparkContext, Option.empty(), false);
                 TempStorage tempStorage = tempStorageManager.getTempStorage(storageBasedBroadcastJoinStorage);

presto-spark-base/src/main/java/com/facebook/presto/spark/execution/AbstractPrestoSparkQueryExecution.java

@@ -29,7 +29,9 @@
 import com.facebook.presto.execution.QueryState;
 import com.facebook.presto.execution.QueryStateTimer;
 import com.facebook.presto.execution.StageInfo;
+import com.facebook.presto.execution.TaskId;
 import com.facebook.presto.execution.TaskInfo;
+import com.facebook.presto.execution.TaskState;
 import com.facebook.presto.execution.scheduler.ExecutionWriterTarget;
 import com.facebook.presto.execution.scheduler.StreamingPlanSection;
 import com.facebook.presto.execution.scheduler.StreamingSubPlan;
@@ -116,6 +118,7 @@
 import java.util.TreeMap;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 import static com.facebook.airlift.units.DataSize.Unit.BYTE;
@@ -572,23 +575,78 @@ protected void validateStorageCapabilities(TempStorage tempStorage)
         }
     }
 
+    /**
+     * Updates the taskInfoMap to ensure it stores the most relevant {@link TaskInfo} for each
+     * logical task, identified by task ID (excluding attempt number).
+     * <p>
+     * This method ensures that, for each logical task, the map retains the latest successful
+     * attempt if available, or otherwise the most recent attempt based on attempt number. Warnings
+     * are logged in cases of unexpected duplicate or multiple successful attempts.
+     *
+     * @param taskInfoMap the map from logical task ID (taskId excluding attempt number) to
+     * {@link TaskInfo}
+     * @param taskInfo the {@link TaskInfo} to consider for updating the map
+     */
+    private void updateTaskInfoMap(HashMap<String, TaskInfo> taskInfoMap, TaskInfo taskInfo)
+    {
+        TaskId newTaskId = taskInfo.getTaskId();
+        String taskIdWithoutAttemptId = new StringBuilder()
+                .append(newTaskId.getStageExecutionId().toString())
+                .append(".")
+                .append(newTaskId.getId())
+                .toString();
+        if (!taskInfoMap.containsKey(taskIdWithoutAttemptId)) {
+            taskInfoMap.put(taskIdWithoutAttemptId, taskInfo);
+            return;
+        }
+
+        TaskInfo storedTaskInfo = taskInfoMap.get(taskIdWithoutAttemptId);
+        TaskId storedTaskId = storedTaskInfo.getTaskId();
+        TaskState storedTaskState = storedTaskInfo.getTaskStatus().getState();
+        TaskState newTaskState = taskInfo.getTaskStatus().getState();
+        if (storedTaskState == TaskState.FINISHED) {
+            if (newTaskState == TaskState.FINISHED) {
+                log.warn("Multiple attempts of the same task have succeeded %s vs %s",
+                        storedTaskId.toString(), newTaskId.toString());
+            }
+            // Successful one has been stored. Nothing needs to be done.
+            return;
+        }
+
+        int storedAttemptNumber = storedTaskId.getAttemptNumber();
+        int newAttemptNumber = newTaskId.getAttemptNumber();
+        if (newTaskState == TaskState.FINISHED || storedAttemptNumber < newAttemptNumber) {
+            taskInfoMap.put(taskIdWithoutAttemptId, taskInfo);
+        }
+        if (storedAttemptNumber == newAttemptNumber) {
+            log.warn("Received multiple identical TaskId %s vs %s",
+                    storedTaskId.toString(), newTaskId.toString());
+        }
+    }
+
     protected void queryCompletedEvent(Optional<ExecutionFailureInfo> failureInfo, OptionalLong updateCount)
     {
         List<SerializedTaskInfo> serializedTaskInfos = taskInfoCollector.value();
-        ImmutableList.Builder<TaskInfo> taskInfos = ImmutableList.builder();
+        HashMap<String, TaskInfo> taskInfoMap = new HashMap<>();
         long totalSerializedTaskInfoSizeInBytes = 0;
         for (SerializedTaskInfo serializedTaskInfo : serializedTaskInfos) {
             byte[] bytes = serializedTaskInfo.getBytesAndClear();
             totalSerializedTaskInfoSizeInBytes += bytes.length;
             TaskInfo taskInfo = deserializeZstdCompressed(taskInfoCodec, bytes);
-            taskInfos.add(taskInfo);
+            updateTaskInfoMap(taskInfoMap, taskInfo);
         }
         taskInfoCollector.reset();
 
-        log.info("Total serialized task info size: %s", DataSize.succinctBytes(totalSerializedTaskInfoSizeInBytes));
+        log.info("Total serialized task info count %s size: %s. Total deduped task info count %s",
+                serializedTaskInfos.size(),
+                DataSize.succinctBytes(totalSerializedTaskInfoSizeInBytes),
+                taskInfoMap.size());
 
         Optional<StageInfo> stageInfoOptional = getFinalFragmentedPlan().map(finalFragmentedPlan ->
-                PrestoSparkQueryExecutionFactory.createStageInfo(session.getQueryId(), finalFragmentedPlan, taskInfos.build()));
+                PrestoSparkQueryExecutionFactory.createStageInfo(
+                        session.getQueryId(),
+                        finalFragmentedPlan,
+                        taskInfoMap.values().stream().collect(Collectors.toList())));
         QueryState queryState = failureInfo.isPresent() ? FAILED : FINISHED;
 
         QueryInfo queryInfo = PrestoSparkQueryExecutionFactory.createQueryInfo(