Adding stage level and task level metrics. by prudhvimaharishi · Pull Request #1 · ajayky-os/iceberg-benchmark-java

prudhvimaharishi · 2026-01-09T08:57:25Z

No description provided.

gemini-code-assist · 2026-01-09T08:57:46Z

Summary of Changes

Hello @prudhvimaharishi, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the observability of Spark jobs within the Iceberg benchmark by introducing granular metric collection at both the stage and task levels. By capturing and serializing detailed performance data, the changes provide a more comprehensive view of query execution, facilitating deeper analysis and optimization efforts.

Highlights

Task Event Capture: The CustomMetricListener has been enhanced to capture SparkListenerTaskEnd events, allowing for detailed task-level monitoring and data collection.
Stage Metrics Enhancement: Stage-level metrics now include more granular details such as submission and completion times, and individual stage metrics are serialized into a dedicated JSON field within the benchmark results.
Task Metrics Collection: A new mechanism has been implemented to collect and process detailed task metrics, including task ID, duration, host, status, and whether the task contributed to custom scan time.
Schema Expansion for Metrics: The benchmark results schema has been extended to store comprehensive stage and task metrics as JSON strings, enabling richer data analysis and deeper insights into query execution.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request successfully adds stage-level and task-level metrics collection to the benchmark. The changes in CustomMetricListener to capture SparkListenerTaskEnd events are correct. In IcebergBenchmark, the new logic to process and associate these metrics with query results is well-implemented. My review focuses on improving the maintainability and readability of the new code. I've identified several instances of code duplication, particularly in JSON serialization and pre-condition checks, which could be refactored into helper methods or more concise expressions. I've also suggested using Java Streams to simplify some data processing loops. These changes will make the new functionality more robust and easier to maintain in the future.

gemini-code-assist · 2026-01-09T09:00:20Z

+    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);


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, "[]"));

gemini-code-assist · 2026-01-09T09:00:21Z

+    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);


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, "[]"));

gemini-code-assist · 2026-01-09T09:00:21Z

    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;


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;

}

gemini-code-assist · 2026-01-09T09:00:21Z

+              .filter(
+                  result -> {
+                    long queryExecutionId = ((Long) result.get("execution_id")).longValue();
+                    return queryExecutionId == executionId;
+                  })


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")))

gemini-code-assist · 2026-01-09T09:00:21Z

+    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;


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; }

gemini-code-assist · 2026-01-09T09:00:21Z

+      boolean doesTaskContributeToScanTime = false;
+      for (AccumulableInfo acc : CollectionConverters.asJava(info.accumulables())) {
+        if (acc.name().isDefined() && acc.name().get().equals("custom_scan_time")) {
+          doesTaskContributeToScanTime = true;
+        }
+      }


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()));

Adding stage level and task level metrics.

cc87430

gemini-code-assist Bot reviewed Jan 9, 2026

View reviewed changes

Navigation Menu

Adding stage level and task level metrics. - #1

Adding stage level and task level metrics.#1
prudhvimaharishi wants to merge 1 commit into
ajayky-os:mainfrom
prudhvimaharishi:stage_task_metrics

prudhvimaharishi commented Jan 9, 2026

Uh oh!

gemini-code-assist Bot commented Jan 9, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

prudhvimaharishi commented Jan 9, 2026

Uh oh!

gemini-code-assist Bot commented Jan 9, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant