Spark 3.4: Refactor JobGroupUtils (#8418)

apache · Aug 29, 2023 · 96bda46 · 96bda46
1 parent a2737a7
commit 96bda46
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 13 deletions.
diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/JobGroupInfo.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/JobGroupInfo.java
@@ -20,9 +20,13 @@
 
 /** Captures information about the current job which is used for displaying on the UI */
 public class JobGroupInfo {
-  private String groupId;
-  private String description;
-  private boolean interruptOnCancel;
+  private final String groupId;
+  private final String description;
+  private final boolean interruptOnCancel;
+
+  public JobGroupInfo(String groupId, String desc) {
+    this(groupId, desc, false);
+  }
 
   public JobGroupInfo(String groupId, String desc, boolean interruptOnCancel) {
     this.groupId = groupId;

diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/JobGroupUtils.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/JobGroupUtils.java
@@ -18,8 +18,10 @@
  */
 package org.apache.iceberg.spark;
 
+import java.util.function.Supplier;
 import org.apache.spark.SparkContext;
 import org.apache.spark.SparkContext$;
+import org.apache.spark.api.java.JavaSparkContext;
 
 public class JobGroupUtils {
 
@@ -43,4 +45,20 @@ public static void setJobGroupInfo(SparkContext sparkContext, JobGroupInfo info)
     sparkContext.setLocalProperty(
         JOB_INTERRUPT_ON_CANCEL, String.valueOf(info.interruptOnCancel()));
   }
+
+  public static <T> T withJobGroupInfo(
+      JavaSparkContext sparkContext, JobGroupInfo info, Supplier<T> supplier) {
+    return withJobGroupInfo(sparkContext.sc(), info, supplier);
+  }
+
+  public static <T> T withJobGroupInfo(
+      SparkContext sparkContext, JobGroupInfo info, Supplier<T> supplier) {
+    JobGroupInfo previousInfo = getJobGroupInfo(sparkContext);
+    try {
+      setJobGroupInfo(sparkContext, info);
+      return supplier.get();
+    } finally {
+      setJobGroupInfo(sparkContext, previousInfo);
+    }
+  }
 }
diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/actions/BaseSparkAction.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/actions/BaseSparkAction.java
@@ -67,7 +67,6 @@
 import org.apache.iceberg.spark.SparkTableUtil;
 import org.apache.iceberg.spark.source.SerializableTableWithSize;
 import org.apache.iceberg.util.Tasks;
-import org.apache.spark.SparkContext;
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.api.java.function.FlatMapFunction;
 import org.apache.spark.broadcast.Broadcast;
@@ -130,18 +129,11 @@ protected Map<String, String> options() {
   }
 
   protected <T> T withJobGroupInfo(JobGroupInfo info, Supplier<T> supplier) {
-    SparkContext context = spark().sparkContext();
-    JobGroupInfo previousInfo = JobGroupUtils.getJobGroupInfo(context);
-    try {
-      JobGroupUtils.setJobGroupInfo(context, info);
-      return supplier.get();
-    } finally {
-      JobGroupUtils.setJobGroupInfo(context, previousInfo);
-    }
+    return JobGroupUtils.withJobGroupInfo(sparkContext, info, supplier);
   }
 
   protected JobGroupInfo newJobGroupInfo(String groupId, String desc) {
-    return new JobGroupInfo(groupId + "-" + JOB_COUNTER.incrementAndGet(), desc, false);
+    return new JobGroupInfo(groupId + "-" + JOB_COUNTER.incrementAndGet(), desc);
   }
 
   protected Table newStaticTable(TableMetadata metadata, FileIO io) {