This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 52237589d0f Java SDK: register the Scala Spark example as task
handlers (#75444)
52237589d0f is described below
commit 52237589d0fbfd32117de4cd512eee44a74a46d1
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Fri Oct 23 19:04:37 2026 +0800
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
---------
---
.../apache/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())
}
