Skip to content

Adding stage level and task level metrics. #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
prudhvimaharishi wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Adding stage level and task level metrics. #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions src/main/java/com/google/cloud/gcs/CustomMetricListener.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,13 @@
import org.apache.spark.scheduler.SparkListener;
import org.apache.spark.scheduler.SparkListenerJobStart;
import org.apache.spark.scheduler.SparkListenerStageCompleted;
import org.apache.spark.scheduler.SparkListenerTaskEnd;
import org.apache.spark.scheduler.StageInfo;

public class CustomMetricListener extends SparkListener {

private static final Deque<StageInfo> stageInfoDeque = new LinkedBlockingDeque<>();
private static final Deque<SparkListenerTaskEnd> taskEndDeque = new LinkedBlockingDeque<>();
private static final ConcurrentHashMap<Integer, Long> stageToExecutionId =
new ConcurrentHashMap<>();
private CountDownLatch latch = new CountDownLatch(1);
Expand All @@ -24,6 +26,10 @@ public static Deque<StageInfo> getStageInfoDeque() {
return stageInfoDeque;
}

public static Deque<SparkListenerTaskEnd> getTaskEndDeque() {
return taskEndDeque;
}

public static ConcurrentHashMap<Integer, Long> getStageToExecutionId() {
return stageToExecutionId;
}
Expand Down Expand Up @@ -65,6 +71,14 @@ public void onStageCompleted(SparkListenerStageCompleted stageCompleted) {
stageInfoDeque.offer(info);
}

@Override
public void onTaskEnd(SparkListenerTaskEnd taskEnd) {
if (taskEnd == null || taskEnd.taskInfo() == null) {
return;
}
taskEndDeque.offer(taskEnd);
}

public long waitForExecutionId() {
try {
boolean found = latch.await(10, TimeUnit.SECONDS);
Expand Down
213 changes: 176 additions & 37 deletions src/main/java/com/google/cloud/gcs/IcebergBenchmark.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,12 @@
import java.sql.Timestamp;
import java.time.Instant;
import java.util.*;
import java.util.Deque;
import java.util.stream.Collectors;
import org.apache.spark.executor.TaskMetrics;
import org.apache.spark.scheduler.AccumulableInfo;
import org.apache.spark.scheduler.SparkListenerTaskEnd;
import org.apache.spark.scheduler.StageInfo;
import org.apache.spark.scheduler.TaskInfo;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
Expand Down Expand Up @@ -142,7 +143,9 @@ private StructType getResultsSchema() {
DataTypes.createStructField("analytics_core_enabled", DataTypes.BooleanType, true),
DataTypes.createStructField("client_type", DataTypes.StringType, true),
DataTypes.createStructField("total_batch_scan_time_ms", DataTypes.LongType, true),
DataTypes.createStructField("timestamp", DataTypes.TimestampType, false)
DataTypes.createStructField("timestamp", DataTypes.TimestampType, false),
DataTypes.createStructField("task_metrics_json", DataTypes.StringType, true),
DataTypes.createStructField("stage_metrics_json", DataTypes.StringType, true)
});
}

Expand Down Expand Up @@ -216,6 +219,7 @@ private void runBenchmark(
System.out.println("Waiting for 10 sec to synchronize SparkListener events");
Thread.sleep(10000);
processStageInfoFromDeque();
processTaskMetricsFromDeque();
} catch (IOException | InterruptedException e) {
System.err.println("Error listing SQL files: " + e.getMessage());
}
Expand Down Expand Up @@ -268,14 +272,17 @@ private void processStageInfoFromDeque() {

void updateQueryMetricFromStageInfo(Map<String, Object> queryMetric) {
if (!queryMetric.containsKey("stages")) {
queryMetric.put("stage_metrics_json", "[]");
return;
}
@SuppressWarnings("unchecked")
List<StageInfo> stages = (List<StageInfo>) queryMetric.get("stages");
if (stages.isEmpty()) {
queryMetric.put("stage_metrics_json", "[]");
return;
Comment on lines 274 to 282

@gemini-code-assist gemini-code-assist Bot Jan 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The checks for whether the 'stages' list exists and is empty can be combined for conciseness. This avoids having two separate return paths for the same outcome.

Suggested change
if (!queryMetric.containsKey("stages")) {
queryMetric.put("stage_metrics_json", "[]");
return;
}
@SuppressWarnings("unchecked")
List<StageInfo> stages = (List<StageInfo>) queryMetric.get("stages");
if (stages.isEmpty()) {
queryMetric.put("stage_metrics_json", "[]");
return;
@SuppressWarnings("unchecked")
List<StageInfo> stages = (List<StageInfo>) queryMetric.get("stages");
if (stages == null || stages.isEmpty()) {
queryMetric.put("stage_metrics_json", "[]");
return;
}

}
Map<String, String> metricJson = new HashMap<>();
Map<String, Object> metricJson = new HashMap<>();
List<Map<String, Object>> stageMetrics = new ArrayList<>();
long total_batch_scan_time_ms = 0;
long total_executor_run_time_ms = 0;
long total_executor_cpu_time_ms = 0;
Expand All @@ -284,17 +291,38 @@ void updateQueryMetricFromStageInfo(Map<String, Object> queryMetric) {
long total_batch_scan_node_cpu_time_ms = 0;
long total_batch_scan_node_gc_time_ms = 0;
for (StageInfo stageInfo : stages) {
Long submissionTime = null;
if (stageInfo.submissionTime().isDefined()) {
submissionTime = (Long) stageInfo.submissionTime().get();
}
Long completionTime = null;
if (stageInfo.completionTime().isDefined()) {
completionTime = (Long) stageInfo.completionTime().get();
}
TaskMetrics taskMetrics = stageInfo.taskMetrics();
long executorRunTime = 0;
long executorCpuTime = 0;
long jvmGCTime = 0;
if (taskMetrics != null) {
total_executor_run_time_ms += taskMetrics.executorRunTime();
total_executor_cpu_time_ms += taskMetrics.executorCpuTime();
total_executor_gc_time_ms += taskMetrics.jvmGCTime();
executorRunTime = taskMetrics.executorRunTime();
executorCpuTime = taskMetrics.executorCpuTime();
jvmGCTime = taskMetrics.jvmGCTime();
total_executor_run_time_ms += executorRunTime;
total_executor_cpu_time_ms += executorCpuTime;
total_executor_gc_time_ms += jvmGCTime;
}
Map<String, Object> stageMetric = new HashMap<>();
stageMetric.put("stageId", stageInfo.stageId());
stageMetric.put("submissionTime", submissionTime);
stageMetric.put("completionTime", completionTime);
stageMetric.put("custom_scan_time", 0);
if (stageInfo.accumulables() == null || stageInfo.accumulables().isEmpty()) {
stageMetrics.add(stageMetric);
continue;
}
Map<Object, AccumulableInfo> accumulables =
CollectionConverters.asJava(stageInfo.accumulables());
long beforeScanTime = total_batch_scan_time_ms;
for (AccumulableInfo accumInfo : accumulables.values()) {
if (accumInfo.name().isDefined() && accumInfo.value().isDefined()) {
String name = accumInfo.name().get();
Expand All @@ -311,7 +339,7 @@ void updateQueryMetricFromStageInfo(Map<String, Object> queryMetric) {
metricJson.put(
metricName,
String.valueOf(
Long.parseLong(metricJson.get(metricName)) + Long.parseLong(value.toString())));
Long.parseLong(metricJson.get(metricName).toString()) + Long.parseLong(value.toString())));
} else {
metricJson.put(metricName, value.toString());
}
Expand All @@ -322,36 +350,49 @@ void updateQueryMetricFromStageInfo(Map<String, Object> queryMetric) {
}
}
}
metricJson.put("total_executor_run_time_ms", String.valueOf(total_executor_run_time_ms));
metricJson.put("total_executor_cpu_time_ms", String.valueOf(total_executor_cpu_time_ms));
metricJson.put("total_executor_gc_time_ms", String.valueOf(total_executor_gc_time_ms));
metricJson.put(
"total_batch_scan_node_executor_run_time_ms",
String.valueOf(total_batch_scan_node_executor_run_time_ms));
metricJson.put(
"total_batch_scan_node_cpu_time_ms", String.valueOf(total_batch_scan_node_cpu_time_ms));
metricJson.put(
"total_batch_scan_node_gc_time_ms", String.valueOf(total_batch_scan_node_gc_time_ms));

metricJson.put(
"gcs.analytics-core.small-file.cache.threshold-bytes",
spark
.conf()
.get(
"spark.sql.catalog."
+ catalogName
+ ".gcs.analytics-core.small-file.cache.threshold-bytes",
"default"));
metricJson.put("execution_id", String.valueOf(queryMetric.get("execution_id")));
String json = "{}";
try {
json = mapper.writeValueAsString(metricJson);
} catch (Exception e) {
System.err.println("Error serializing metrics to JSON: " + e.getMessage());
}
queryMetric.put("metric_json", json);
queryMetric.put("total_batch_scan_time_ms", total_batch_scan_time_ms);
stageMetric.put("accumulable_custom_scan_time", total_batch_scan_time_ms - beforeScanTime);
stageMetrics.add(stageMetric);
}

metricJson.put("total_executor_run_time_ms", String.valueOf(total_executor_run_time_ms));
metricJson.put("total_executor_cpu_time_ms", String.valueOf(total_executor_cpu_time_ms));
metricJson.put("total_executor_gc_time_ms", String.valueOf(total_executor_gc_time_ms));
metricJson.put(
"total_batch_scan_node_executor_run_time_ms",
String.valueOf(total_batch_scan_node_executor_run_time_ms));
metricJson.put(
"total_batch_scan_node_cpu_time_ms", String.valueOf(total_batch_scan_node_cpu_time_ms));
metricJson.put(
"total_batch_scan_node_gc_time_ms", String.valueOf(total_batch_scan_node_gc_time_ms));

metricJson.put(
"gcs.analytics-core.small-file.cache.threshold-bytes",
spark
.conf()
.get(
"spark.sql.catalog."
+ catalogName
+ ".gcs.analytics-core.small-file.cache.threshold-bytes",
"default"));
metricJson.put("execution_id", String.valueOf(queryMetric.get("execution_id")));

String json = "{}";
try {
json = mapper.writeValueAsString(metricJson);
} catch (Exception e) {
System.err.println("Error serializing metrics to JSON: " + e.getMessage());
}
queryMetric.put("metric_json", json);

String stageMetricsJson = "[]";
try {
stageMetricsJson = mapper.writeValueAsString(stageMetrics);
} catch (Exception e) {
System.err.println("Error serializing stage metrics to JSON: " + e.getMessage());
}
queryMetric.put("stage_metrics_json", stageMetricsJson);
Comment on lines +379 to +393

@gemini-code-assist gemini-code-assist Bot Jan 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

This block contains duplicated logic for serializing objects to JSON and handling exceptions. This pattern is also repeated in updateQueryMetricFromTaskMetrics. To improve maintainability and reduce code duplication, you could extract this logic into a private helper method.

For example, you could add a helper method to the class:

private String toJson(Object value, String defaultValue) {
    try {
        return mapper.writeValueAsString(value);
    } catch (com.fasterxml.jackson.core.JsonProcessingException e) {
        System.err.println("Error serializing object to JSON: " + e.getMessage());
        return defaultValue;
    }
}

Using this helper would make the code here and in other places more concise and easier to maintain.

    queryMetric.put("metric_json", toJson(metricJson, "{}"));
    queryMetric.put("stage_metrics_json", toJson(stageMetrics, "[]"));


queryMetric.put("total_batch_scan_time_ms", total_batch_scan_time_ms);
}

private List<Row> createRowsFromBuffer() {
Expand All @@ -370,7 +411,10 @@ private List<Row> createRowsFromBuffer() {
map.get("analytics_core_enabled"),
map.get("client_type"),
map.get("total_batch_scan_time_ms"),
map.get("timestamp"));
map.get("timestamp"),
map.get("task_metrics_json"),
map.get("stage_metrics_json")
);
})
.collect(Collectors.toList());
}
Expand Down Expand Up @@ -401,4 +445,99 @@ private void flushResultsToCsv(String outputGcsPath) {
System.out.println(" -> Flushed " + resultsBuffer.size() + " results to " + finalOutputPath);
resultsBuffer.clear();
}

private void processTaskMetricsFromDeque() {
Deque<SparkListenerTaskEnd> taskEndDeque = CustomMetricListener.getTaskEndDeque();
System.out.println("Processing " + taskEndDeque.size() + " task metrics...");

while (!taskEndDeque.isEmpty()) {
SparkListenerTaskEnd taskEnd = taskEndDeque.poll();
int stageId = taskEnd.stageId();
Long executionId = listener.getStageToExecutionId().get(stageId);
if (executionId == null) {
continue;
}

Optional<Map<String, Object>> matchingResultOpt =
resultsBuffer.stream()
.filter(
result -> {
long queryExecutionId = ((Long) result.get("execution_id")).longValue();
return queryExecutionId == executionId;
})
Comment on lines +463 to +467

@gemini-code-assist gemini-code-assist Bot Jan 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The filter logic can be simplified to a more concise one-liner. Using equals() for comparing Long objects is also safer than == to avoid issues with object identity vs. value equality, especially for values outside the Long cache range.

              .filter(result -> executionId.equals(result.get("execution_id")))

.findFirst();
matchingResultOpt.ifPresentOrElse(
matchingResult -> {
@SuppressWarnings("unchecked")
List<SparkListenerTaskEnd> tasks =
(List<SparkListenerTaskEnd>)
matchingResult.computeIfAbsent("tasks", k -> new ArrayList<>());
tasks.add(taskEnd);
},
() -> {
System.out.println(
" -> Warning: Task from Stage ID "
+ stageId
+ " (execution_id "
+ executionId
+ ") not found in any query result buffer..");
});
}

for (Map<String, Object> result : resultsBuffer) {
updateQueryMetricFromTaskMetrics(result);
}
}

void updateQueryMetricFromTaskMetrics(Map<String, Object> queryMetric) {
if (!queryMetric.containsKey("tasks")) {
queryMetric.put("task_metrics_json", "[]");
return;
}
@SuppressWarnings("unchecked")
List<SparkListenerTaskEnd> tasks = (List<SparkListenerTaskEnd>) queryMetric.get("tasks");
if (tasks.isEmpty()) {
queryMetric.put("task_metrics_json", "[]");
return;
Comment on lines +493 to +501

@gemini-code-assist gemini-code-assist Bot Jan 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Similar to the 'stages' check, the checks for whether the 'tasks' list exists and is empty can be combined for better readability and conciseness.

    @SuppressWarnings("unchecked")
    List<SparkListenerTaskEnd> tasks = (List<SparkListenerTaskEnd>) queryMetric.get("tasks");
    if (tasks == null || tasks.isEmpty()) {
      queryMetric.put("task_metrics_json", "[]");
      return;
    }

}
List<Map<String, Object>> taskMaps = new ArrayList<>();
for (SparkListenerTaskEnd taskEnd : tasks) {
TaskInfo info = taskEnd.taskInfo();
if (info == null) {
continue;
}

// If a task has scantime it means, it is reading data from file system.
boolean doesTaskContributeToScanTime = false;
for (AccumulableInfo acc : CollectionConverters.asJava(info.accumulables())) {
if (acc.name().isDefined() && acc.name().get().equals("custom_scan_time")) {
doesTaskContributeToScanTime = true;
}
}
Comment on lines +511 to +516

@gemini-code-assist gemini-code-assist Bot Jan 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This for-loop can be expressed more concisely and efficiently using a Java Stream with anyMatch. The current implementation iterates through the entire list even after a match is found, whereas anyMatch will stop as soon as the condition is met.

      boolean doesTaskContributeToScanTime =
          CollectionConverters.asJava(info.accumulables()).stream()
              .anyMatch(
                  acc -> acc.name().isDefined() && "custom_scan_time".equals(acc.name().get()));


Map<String, Object> taskMap = new HashMap<>();
taskMap.put("stage_id", taskEnd.stageId());
taskMap.put("task_id", info.taskId());
taskMap.put("task_type", taskEnd.taskType());
taskMap.put("launch_time", info.launchTime());
taskMap.put("finish_time", info.finishTime());
taskMap.put("duration_ms", info.duration());
taskMap.put("host", info.host());
taskMap.put("status", info.status());
taskMap.put("failed", info.failed());
taskMap.put("doesTaskContributeToScanTime", doesTaskContributeToScanTime);

taskMaps.add(taskMap);
}

String json = "[]";
try {
json = mapper.writeValueAsString(taskMaps);
} catch (Exception e) {
System.err.println("Error serializing task metrics to JSON: " + e.getMessage());
}
queryMetric.put("task_metrics_json", json);
Comment on lines +533 to +539

@gemini-code-assist gemini-code-assist Bot Jan 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

This is another instance of duplicated JSON serialization logic. This can be replaced by the same helper method suggested for updateQueryMetricFromStageInfo to improve code reuse and maintainability.

    queryMetric.put("task_metrics_json", toJson(taskMaps, "[]"));

}


}