From 52237589d0fbfd32117de4cd512eee44a74a46d1 Mon Sep 17 00:00:00 2001
From: "Jason(Zhe-You) Liu" <68415893+jason810496@users.noreply.github.com>
Date: Fri, 23 Oct 2026 19:04:37 +0800
Subject: [PATCH] Java SDK: register the Scala Spark example as task handlers
 (#75444)

* Java SDK: register the Scala Spark example as task handlers

scala_spark_examples.py owns the Dag and declares its three tasks with
@task.stub, but the bundle built a DagDef of the same ID, which is the Dag
definition path. Register one handler per task ID instead, so one Dag ID is
not declared on both sides.

* Java SDK: use the default XCom key overloads in the Scala Spark example

Client.getXCom and setXCom are @JvmOverloads, so Scala can call them with
only the task ID or value and get the return_value key and current Dag.

* Java SDK: shorten the Scala Spark bundle doc comment

---------
---
 .../airflow/example/ScalaSparkExample.scala   | 33 +++++++++----------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala
index f7f32bad8e2..d7d9d0924fd 100644
--- a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala
+++ b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala
@@ -19,7 +19,7 @@
 
 package org.apache.airflow.example
 
-import org.apache.airflow.sdk.{Bundle, BundleBuilder, Client, Context, DagDef, Server, Task}
+import org.apache.airflow.sdk.{Bundle, Client, Context, Server, Task}
 import org.apache.logging.log4j.{LogManager, Logger}
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.functions.sum
@@ -60,13 +60,11 @@ private object SparkEtl {
       .config("spark.driver.bindAddress", "127.0.0.1")
       .getOrCreate()
 
-  // 3-arg getXCom(key, dagId, taskId) reads the upstream return value; wire
-  // integers arrive boxed.
-  def readUpstreamLong(client: Client, context: Context, taskId: String): Long =
-    client.getXCom("return_value", context.dagRun.dagId, taskId).asInstanceOf[Number].longValue()
+  def readUpstreamLong(client: Client, taskId: String): Long =
+    client.getXCom(taskId).asInstanceOf[Number].longValue()
 
   def pushLong(client: Client, value: Long): Unit =
-    client.setXCom("return_value", java.lang.Long.valueOf(value))
+    client.setXCom(java.lang.Long.valueOf(value))
 }
 
 /** Builds the Spark DataFrame and reports how many records were extracted. */
@@ -95,7 +93,7 @@ class SparkTransform extends Task {
   override def execute(context: Context, client: Client): Unit = {
     // Read the upstream count only to demonstrate XCom passing between JVM
     // tasks; the aggregation below recomputes from the source dataset.
-    val extractedCount = SparkEtl.readUpstreamLong(client, context, SparkEtl.ExtractTaskId)
+    val extractedCount = SparkEtl.readUpstreamLong(client, SparkEtl.ExtractTaskId)
     log.info("Transform received {} records from extract", java.lang.Long.valueOf(extractedCount))
 
     val spark = SparkEtl.newSession("scala-spark-etl-transform")
@@ -116,7 +114,7 @@ class SparkLoad extends Task {
   private val log: Logger = LogManager.getLogger(classOf[SparkLoad])
 
   override def execute(context: Context, client: Client): Unit = {
-    val total = SparkEtl.readUpstreamLong(client, context, SparkEtl.TransformTaskId)
+    val total = SparkEtl.readUpstreamLong(client, SparkEtl.TransformTaskId)
     log.info("Loading aggregated revenue {}", java.lang.Long.valueOf(total))
 
     val spark = SparkEtl.newSession("scala-spark-etl-load")
@@ -132,18 +130,19 @@ class SparkLoad extends Task {
   }
 }
 
+/**
+ * Task handlers for the `scala_spark_example` Dag declared in `dags/scala_spark_examples.py`.
+ */
 object ScalaSparkExample {
-  def build(): DagDef =
-    new DagDef(SparkEtl.DagId)
-      .addTask(SparkEtl.ExtractTaskId, classOf[SparkExtract])
-      .addTask(SparkEtl.TransformTaskId, classOf[SparkTransform])
-      .addTask(SparkEtl.LoadTaskId, classOf[SparkLoad])
+  def build(): Bundle =
+    new Bundle()
+      .register(SparkEtl.DagId, SparkEtl.ExtractTaskId, classOf[SparkExtract])
+      .register(SparkEtl.DagId, SparkEtl.TransformTaskId, classOf[SparkTransform])
+      .register(SparkEtl.DagId, SparkEtl.LoadTaskId, classOf[SparkLoad])
 }
 
 /** Bundle entry point served to Airflow's Java coordinator. */
-object ScalaSparkBundleBuilder extends BundleBuilder {
-  override def getDags(): java.lang.Iterable[DagDef] = java.util.List.of(ScalaSparkExample.build())
-
+object ScalaSparkBundleBuilder {
   def main(args: Array[String]): Unit =
-    Server.create(args).serve(new Bundle(getDags()))
+    Server.create(args).serve(ScalaSparkExample.build())
 }