fix build tests

apache · Jan 23, 2025 · 9f066dd · 9f066dd
1 parent 109b17f
commit 9f066dd
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 14 deletions.
diff --git a/...nsions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/...nsions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
@@ -538,13 +538,15 @@ private MSQTaskReportPayload runInternal(final QueryListener queryListener, fina
 
   private void emitSummaryMetrics(final MSQTaskReportPayload msqTaskReportPayload, final MSQSpec querySpec)
   {
-    final List<Integer> stagesToInclude = new ArrayList<>();
-    for (MSQStagesReport.Stage stage : msqTaskReportPayload.getStages().getStages()) {
-      boolean hasParentStage = stage.getStageDefinition().getInputSpecs().stream()
-          .anyMatch(stageInput -> stageInput instanceof StageInputSpec);
-
-      if (!hasParentStage) {
-        stagesToInclude.add(stage.getStageNumber());
+    final Set<Integer> stagesToInclude = new HashSet<>();
+    final MSQStagesReport stagesReport = msqTaskReportPayload.getStages();
+    if (stagesReport != null) {
+      for (MSQStagesReport.Stage stage : stagesReport.getStages()) {
+        boolean hasParentStage = stage.getStageDefinition().getInputSpecs().stream()
+            .anyMatch(stageInput -> stageInput instanceof StageInputSpec);
+        if (!hasParentStage) {
+          stagesToInclude.add(stage.getStageNumber());
+        }
       }
     }
     long totalProcessedBytes = 0;
@@ -554,7 +556,7 @@ private void emitSummaryMetrics(final MSQTaskReportPayload msqTaskReportPayload,
           .copyMap()
           .entrySet()
           .stream()
-          .filter(entry -> stagesToInclude.contains(entry.getKey())) // Filter by stagesToInclude
+          .filter(entry -> stagesReport == null || stagesToInclude.contains(entry.getKey()))
           .flatMap(counterSnapshotsMap -> counterSnapshotsMap.getValue().values().stream())
           .flatMap(counterSnapshots -> counterSnapshots.getMap().entrySet().stream())
           .filter(entry -> entry.getKey().startsWith("input"))

diff --git a/...ava/org/apache/druid/indexing/common/task/batch/parallel/ParallelIndexSupervisorTask.java b/...ava/org/apache/druid/indexing/common/task/batch/parallel/ParallelIndexSupervisorTask.java
@@ -1252,12 +1252,6 @@ private TaskStatus runSequential(TaskToolbox toolbox) throws Exception
    */
   private TaskReport.ReportMap getTaskCompletionReports(TaskStatus taskStatus)
   {
-    final Number totalProcessedBytes = (Number) indexGenerateRowStats.lhs.get("processedBytes");
-    // Emit the processed bytes metric
-    final ServiceMetricEvent.Builder metricBuilder = new ServiceMetricEvent.Builder();
-    IndexTaskUtils.setTaskDimensions(metricBuilder, this);
-    toolbox.getEmitter().emit(
-        metricBuilder.setMetric("ingest/input/bytes", (Number) totalProcessedBytes));
     return buildIngestionStatsAndContextReport(
         IngestionState.COMPLETED,
         taskStatus.getErrorMsg(),
@@ -1286,11 +1280,29 @@ private void updateAndWriteCompletionReports(TaskStatus status)
 
   private void writeCompletionReports()
   {
+    emitCompletionMetrics();
     if (!isCompactionTask) {
       toolbox.getTaskReportFileWriter().write(getId(), completionReports);
     }
   }
 
+  private void emitCompletionMetrics()
+  {
+    final Map<String, Object> rowStats = getTaskCompletionRowStats();
+    if (rowStats == null) {
+      return;
+    }
+
+    final Number totalProcessedBytes = (Number) rowStats.get("processedBytes");
+    if (totalProcessedBytes == null) {
+      return;
+    }
+
+    final ServiceMetricEvent.Builder metricBuilder = new ServiceMetricEvent.Builder();
+    IndexTaskUtils.setTaskDimensions(metricBuilder, this);
+    toolbox.getEmitter().emit(metricBuilder.setMetric("ingest/input/bytes", totalProcessedBytes));
+  }
+
   private static IndexTuningConfig convertToIndexTuningConfig(ParallelIndexTuningConfig tuningConfig)
   {
     return new IndexTuningConfig(