[HUDI-8596] Fix Hudi Spark SQL cancelling issue (#12358)

apache · Jan 3, 2025 · d6d4e88 · d6d4e88
1 parent 4af4454
commit d6d4e88
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/...di-spark-client/src/main/java/org/apache/hudi/client/common/HoodieSparkEngineContext.java b/...di-spark-client/src/main/java/org/apache/hudi/client/common/HoodieSparkEngineContext.java
@@ -197,7 +197,7 @@ public Option<String> getProperty(EngineProperty key) {
 
   @Override
   public void setJobStatus(String activeModule, String activityDescription) {
-    javaSparkContext.setJobGroup(activeModule, activityDescription);
+    javaSparkContext.setJobDescription(String.format("%s:%s", activeModule, activityDescription));
   }
 
   @Override

diff --git a/...park-client/src/test/java/org/apache/hudi/client/common/TestHoodieSparkEngineContext.java b/...park-client/src/test/java/org/apache/hudi/client/common/TestHoodieSparkEngineContext.java
@@ -48,4 +48,23 @@ void testAddRemoveCachedDataIds() {
     assertEquals(Arrays.asList(1, 2, 3), context.removeCachedDataIds(HoodieDataCacheKey.of(basePath, instantTime)));
     assertTrue(context.getCachedDataIds(HoodieDataCacheKey.of(basePath, instantTime)).isEmpty());
   }
+
+  @Test
+  void testSetJobStatus() {
+    // Test data
+    String jobGroupId = "jobGroupId";
+    String initialJobDescription = "initialJobDescription";
+    String activeModule = "TestModule";
+    String activityDescription = "Running test operation";
+    String expectedJobDescription = String.format("%s:%s", activeModule, activityDescription);
+    context.getJavaSparkContext().setJobGroup(jobGroupId, initialJobDescription);
+    assertEquals(jobGroupId, context.getJavaSparkContext().getLocalProperty("spark.jobGroup.id"));
+    assertEquals(initialJobDescription, context.getJavaSparkContext().getLocalProperty("spark.job.description"));
+
+    // Set the job status
+    context.setJobStatus(activeModule, activityDescription);
+    assertEquals(expectedJobDescription, context.getJavaSparkContext().getLocalProperty("spark.job.description"));
+    // Assert jobGroupId does not change
+    assertEquals(jobGroupId, context.getJavaSparkContext().getLocalProperty("spark.jobGroup.id"));
+  }
 }