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

---------
This commit is contained in:
Jason(Zhe-You) Liu authored and GitHub committed 2026-10-23 19:04:37 +08:00
1 parent af76df3863
commit 52237589d0
1 file changed
+16 -17
@@ -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())
}