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 5a64dfcc78e Java SDK: Add edge labels to native Dags (#74523)
5a64dfcc78e is described below
commit 5a64dfcc78e2a8d1f3f77a61c23237ef010943d6
Author: PoAn Yang <[email protected]>
AuthorDate: Tue Oct 20 09:52:36 2026 +0800
Java SDK: Add edge labels to native Dags (#74523)
* Java SDK: Add edge labels to native Dags
Signed-off-by: PoAn Yang <[email protected]>
* Apply edge labels in the order they were set
A task edge's own label always won over a group edge's label, and a
relabeled group edge kept its first position, so edge_info differed from
Python's, where the label set last wins. Also note that Python's Label
rewrites an edge when one end is in no task group.
* Java SDK: Write the edge label text before the task
Signed-off-by: PoAn Yang <[email protected]>
---------
Signed-off-by: PoAn Yang <[email protected]>
Co-authored-by: ZHE YOU LIU <[email protected]>
---
.../language-sdks/java.rst | 54 +++-
.../example/nativedag/AnnotationExample.java | 5 +-
.../example/nativedag/InterfaceExample.java | 6 +-
.../org/apache/airflow/sdk/AirflowProcessorTest.kt | 79 +++++
.../ci/prek/check_serialization_conformance.py | 2 +-
.../main/kotlin/org/apache/airflow/sdk/DagDef.kt | 46 ++-
.../src/main/kotlin/org/apache/airflow/sdk/Deps.kt | 69 ++++-
.../main/kotlin/org/apache/airflow/sdk/Label.kt | 93 ++++++
.../org/apache/airflow/sdk/execution/Serde.kt | 52 +++-
.../kotlin/org/apache/airflow/sdk/EdgeLabelTest.kt | 323 +++++++++++++++++++++
.../airflow/sdk/conformance/SerializeJava.kt | 6 +-
11 files changed, 719 insertions(+), 16 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
index e8e63811c96..93178c747cd 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
@@ -682,9 +682,10 @@ returns the same handle every time, so it names one node
wherever it appears; on
arguments is called once, and the wiring fails if it is called again with
arguments, so hold its
handle in a local and reuse that.
-``before``, ``after`` and ``Flow.of`` work here exactly as they do on the
interface surface; inside
-the wiring class ``Flow`` is inherited by simple name, so it needs no import
and never collides with
-``java.util.concurrent.Flow``.
+``before``, ``after``, ``Flow.of`` and ``Label`` work here exactly as they do
on the interface
+surface. Inside the wiring class, ``Flow`` is inherited by simple name, so it
needs no import and
+never collides with ``java.util.concurrent.Flow``. ``Label`` is not
inherited, so the wiring class
+has to import it with ``import org.apache.airflow.sdk.Label;``.
Every ``@Airflow.Dag`` class declares a wiring class, because the graph is
what the Dag owns. A
class that supplies only task bodies, for a Dag a Python file declares, carries
@@ -759,6 +760,53 @@ underscores, or dashes, and no task or other group in the
Dag can share it.
drawn, as Python resolves them, so drawing an inner edge before or after
an outer one gives
different upstreams.
+Edge labels
+~~~~~~~~~~~
+
+An edge can carry a label. The Airflow UI shows the label on the edge in the
graph, as it does for
+Python's ``Label``. ``Label`` is a class in the ``org.apache.airflow.sdk``
package. ``Label.of`` takes
+the label text, and ``on`` takes the task or group at one end of the edge you
want to label. Pass the
+result to ``before`` or ``after`` like any other point in the flow. The text
then sits between the two
+tasks in the code, as the label sits between them in the graph:
+
+.. code-block:: java
+
+ loaded.before(Label.of("when empty").on(reportEmpty)); // load >>
Label("when empty") >> report_empty
+ cleaned.after(Label.of("always").on(loaded)); // load >>
Label("always") >> cleanup
+
+A label goes on edges, not on tasks. It covers only the edges between the
receiver of ``before`` or
+``after`` and the tasks and groups passed to the label's ``on``. So each edge
of a fan-out can carry
+its own label. ``on`` also takes several tasks or groups. In that case, the
label goes on the edge
+to or from each of them:
+
+.. code-block:: java
+
+ checked.before(Label.of("rows found").on(processed), Label.of("no
rows").on(reportEmpty));
+ // load >> Label("forks") >> [report_empty, report_suspicious]
+ loaded.before(Label.of("forks").on(reportEmpty, reportSuspicious));
+
+When an edge already exists, declaring it again with a label only adds the
label. That is how the
+sides of a condition get their labels:
+
+.. code-block:: java
+
+ var hasRows = dag.If(HasRows.class).Then(load).Else(loadEmpty);
+ hasRows.before(Label.of("rows found").on(load), Label.of("no
rows").on(loadEmpty));
+
+When an edge already has a label, a new label replaces it. Use the value that
``on`` returns only in
+the arguments of ``before`` and ``after``. Calling ``before`` or ``after`` on
that value fails. When
+an edge goes to or from a task group, the UI draws it to or from the group's
own node and shows the
+label there.
+
+.. note::
+
+ A label never changes which tasks an edge connects. Python's ``Label``
can change those tasks.
+ It does so when the end written to the left of the ``Label`` is in a task
group and the other end
+ is not in that same group. Python then replaces the left end with the
group that contains it.
+ For example, if ``v`` is in group ``g`` and ``n`` is in no group, ``v >>
Label("x") >> n`` draws
+ edges from the leaves of ``g`` instead of from ``v``. In Java, the end
that matches this
+ left end is the receiver of ``before`` or ``after``.
+
Configuration attributes
~~~~~~~~~~~~~~~~~~~~~~~~
diff --git
a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
index 01639dd6f7e..6e751dd724a 100644
---
a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
+++
b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
@@ -110,7 +110,10 @@ public class AnnotationExample {
var extracted = extract();
var transformed = transform(extracted, lit(1.5));
var loaded = load(transformed);
- hasRows(transformed).Then(loaded).Else(loadEmpty());
+ var gate = hasRows(transformed).Then(loaded).Else(loadEmpty());
+ // Then and Else already declared these edges. This call only adds
+ // labels to the edges.
+ gate.before(Label.of("rows found").on(loaded), Label.of("no
rows").on(loadEmpty()));
pickReport(transformed).Case(reportLong()).Case(reportShort());
loaded.before(triggerDownstream());
// Ordering-only edge: the checks group runs after extract, with no data
diff --git
a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
index fcf095ee72e..47a8f6786af 100644
---
a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
+++
b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
@@ -123,8 +123,12 @@ public class InterfaceExample {
transform.after(extract);
// With no task id given, a decider takes one from its class: "hasRows" and
// "pickReport".
- dag.If(HasRows.class).after(transform).Then(load).Else(loadEmpty);
+ var hasRows =
dag.If(HasRows.class).after(transform).Then(load).Else(loadEmpty);
dag.Switch(PickReport.class).after(transform).Case(reportLong).Case(reportShort);
+ // The Airflow UI shows these labels on the edges from hasRows. Then and
+ // Else already declared both of those edges. Declaring the edges again
+ // here only adds the labels.
+ hasRows.before(Label.of("rows found").on(load), Label.of("no
rows").on(loadEmpty));
load.before(trigger);
return dag;
}
diff --git
a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/AirflowProcessorTest.kt
b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/AirflowProcessorTest.kt
index 368b9a58532..c42d52413bd 100644
---
a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/AirflowProcessorTest.kt
+++
b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/AirflowProcessorTest.kt
@@ -1379,6 +1379,85 @@ class AirflowProcessorTest {
.contains("public static final class First_Run implements Task {")
}
+ @Test
+ @DisplayName("let a wiring class label an edge with Label.of(...).on(...)")
+ fun compileWiringThatLabelsAnEdge() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+ import org.apache.airflow.sdk.Airflow;
+ import org.apache.airflow.sdk.Label;
+ @Airflow.Dag
+ public class TestExample {
+ @Airflow.Task public void extract() {}
+ @Airflow.Task public void load() {}
+
+ @Airflow.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() { extract().before(Label.of("rows").on(load())); }
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).succeeded()
+ }
+
+ @Test
+ @DisplayName("fail to compile a label passed to before without on")
+ fun failOnLabelWithoutOn() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+ import org.apache.airflow.sdk.Airflow;
+ import org.apache.airflow.sdk.Label;
+ @Airflow.Dag
+ public class TestExample {
+ @Airflow.Task public void extract() {}
+
+ @Airflow.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() { extract().before(Label.of("rows")); }
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).failed()
+ assertThat(compilation).hadErrorContaining(
+ "org.apache.airflow.sdk.Label cannot be converted to
org.apache.airflow.sdk.Deps.Flow",
+ )
+ }
+
+ @Test
+ @DisplayName("fail to compile Label.on called with no task")
+ fun failOnLabelOnWithoutTask() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+ import org.apache.airflow.sdk.Airflow;
+ import org.apache.airflow.sdk.Label;
+ @Airflow.Dag
+ public class TestExample {
+ @Airflow.Task public void extract() {}
+
+ @Airflow.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() { extract().before(Label.of("rows").on()); }
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).failed()
+ assertThat(compilation).hadErrorContaining(
+ "method on in class org.apache.airflow.sdk.Label cannot be applied to
given types",
+ )
+ }
+
@Test
@DisplayName("reject a task group ID that is not a plain identifier")
fun rejectInvalidTaskGroupId() {
diff --git a/java-sdk/scripts/ci/prek/check_serialization_conformance.py
b/java-sdk/scripts/ci/prek/check_serialization_conformance.py
index 677ffcdb217..d6f454ec3e7 100755
--- a/java-sdk/scripts/ci/prek/check_serialization_conformance.py
+++ b/java-sdk/scripts/ci/prek/check_serialization_conformance.py
@@ -56,7 +56,7 @@ if __name__ == "__main__":
"--sdk",
"java",
"--supports",
- "literal_inputs,branch,switch,trigger_dag_run",
+ "literal_inputs,branch,switch,trigger_dag_run,edge_labels",
"--",
*serializer,
]
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
index 06630cc89f4..a4150e8a9b6 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
@@ -25,6 +25,7 @@ import org.apache.airflow.sdk.internal.SchemaFields
import org.apache.airflow.sdk.internal.checkConfigValue
import org.apache.airflow.sdk.internal.deriveTaskId
import org.apache.airflow.sdk.internal.validateTaskInput
+import java.util.concurrent.atomic.AtomicLong
import kotlin.Throws
/**
@@ -63,6 +64,9 @@ class DagDef(
/** Edges with a task group at either end, in the order drawn. */
internal val groupEdges = linkedSetOf<Pair<Endpoint, Endpoint>>()
+ /** The labels that [Label.on] put on [groupEdges], keyed by the edge. */
+ internal val groupEdgeLabels = mutableMapOf<Pair<Endpoint, Endpoint>,
EdgeLabel>()
+
/**
* Whether a [Bundle] has taken this Dag. A decider checked when the Dag was
* registered cannot be changed afterwards, because nothing would check the
@@ -317,6 +321,8 @@ class DagDef(
internal fun expandGroupEdges(): GroupExpansion {
val upstreams = mutableMapOf<String, MutableSet<String>>()
val edges = mutableMapOf<String, MutableGroupEdges>()
+ val labels = mutableMapOf<Pair<String, String>, EdgeLabel>()
+ val grouped = groups.values.flatMapTo(mutableSetOf()) { it.taskIds }
fun edgesOf(groupId: String) = edges.getOrPut(groupId) {
MutableGroupEdges() }
@@ -359,7 +365,21 @@ class DagDef(
for ((upstream, downstream) in groupEdges) {
val from = leavesOf(upstream).map { it.id }
- rootsOf(downstream).forEach { task -> upstreams.getOrPut(task.id) {
linkedSetOf() } += from }
+ val roots = rootsOf(downstream)
+ roots.forEach { task -> upstreams.getOrPut(task.id) { linkedSetOf() } +=
from }
+ // Python's `extract >> Label("x") >> group` also writes the label on
+ // the edge from extract to each root of the group. The code below does
+ // the same when no group holds extract. It does not copy the rarer
+ // cases where Python labels such edges. The UI never draws the edges
+ // from extract to the roots, because it draws an edge to the group's
+ // node instead.
+ val label = groupEdgeLabels[upstream to downstream]
+ if (label != null && upstream is TaskDef && upstream.id !in grouped &&
downstream is TaskGroupRef) {
+ roots.forEach { root ->
+ val ends = upstream.id to root.id
+ if (labels[ends].let { it == null || it.order < label.order })
labels[ends] = label
+ }
+ }
if (downstream is TaskGroupRef) {
edgesOf(downstream.id).upstreamTaskIds += from
if (upstream is TaskGroupRef) edgesOf(downstream.id).upstreamGroupIds
+= upstream.id
@@ -373,17 +393,36 @@ class DagDef(
edgesOf(upstream.id).downstreamTaskIds += downstream.id
}
}
- return GroupExpansion(upstreams, edges)
+ return GroupExpansion(upstreams, edges, labels)
+ }
+}
+
+/**
+ * A label that [Label.on] put on an edge.
+ * [order] grows with each label set, so when two labels land on the same
edge, the one set later wins, as in Python.
+ */
+internal class EdgeLabel(
+ val text: String,
+) {
+ val order: Long = nextOrder.getAndIncrement()
+
+ private companion object {
+ val nextOrder = AtomicLong()
}
}
/**
* The task edges a Dag's task-group edges stand for, and the edges each group
* records for itself, as [DagDef.expandGroupEdges] worked them out.
+ *
+ * @property labels The label that each of those task edges takes from its
+ * group edge. Each key is the pair of task IDs at the two ends of the
+ * task edge, upstream first.
*/
internal class GroupExpansion(
private val upstreams: Map<String, Set<String>>,
private val edges: Map<String, GroupEdges>,
+ val labels: Map<Pair<String, String>, EdgeLabel>,
) {
/** Every task [def] runs after: the edges it carries, plus the ones a group
edge implies. */
fun upstreamsOf(def: TaskDef): Set<String> =
def.upstreams.mapTo(linkedSetOf()) { it.id } + upstreams[def.id].orEmpty()
@@ -458,6 +497,9 @@ class TaskDef(
/** Name of the task parameter each of [inputs] feeds, in the same order. */
internal val inputNames = mutableListOf<String>()
internal val upstreams = linkedSetOf<TaskDef>()
+
+ /** The labels that [Label.on] put on edges from [upstreams], keyed by the
upstream task. */
+ internal val upstreamLabels = mutableMapOf<TaskDef, EdgeLabel>()
internal var owner: DagDef? = null
/** What this task decides to run, for a condition or a switch; null
otherwise. */
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt
index 1cf90491841..34d6071e69f 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt
@@ -27,6 +27,7 @@ import org.apache.airflow.sdk.internal.Refs
*
* An [Airflow.Deps] class inherits [lit] and [Flow] by simple name, so it
needs
* no import and `Flow` does not collide with `java.util.concurrent.Flow`.
+ * [Label] is not inherited, so the class has to import it.
*/
interface Deps {
/**
@@ -56,12 +57,16 @@ interface Deps {
*
* Variadic, so one call fans out, and it returns its own receiver: a
* fan-out has no single next task to hand back. Declaring an edge that
- * already exists changes nothing.
+ * already exists changes nothing. The only exception is a label put on
+ * the edge with [Label.on], which replaces the label the edge had.
*
* @param next Tasks that run after the ones here.
* @return This point in the flow.
+ * @throws IllegalArgumentException if [Label.on] returned this point in
+ * the flow, or returned a point that this one holds.
*/
fun before(vararg next: Flow): Flow {
+ requireUnlabeled(this, "before")
next.forEach { link(this, it) }
return this
}
@@ -75,8 +80,11 @@ interface Deps {
*
* @param previous Tasks that run before the ones here.
* @return This point in the flow.
+ * @throws IllegalArgumentException if [Label.on] returned this point in
+ * the flow, or returned a point that this one holds.
*/
fun after(vararg previous: Flow): Flow {
+ requireUnlabeled(this, "after")
previous.forEach { link(it, this) }
return this
}
@@ -143,6 +151,47 @@ internal class FlowSet(
override fun endpoints(): List<Endpoint> = flows.flatMap { it.endpoints() }
}
+/** A point in the flow whose edges [Label.on] labels with [text]. */
+internal class LabeledFlow(
+ internal val flow: Deps.Flow,
+ internal val text: String,
+) : Deps.Flow {
+ override fun nodes(): List<TaskDef> = flow.nodes()
+
+ override fun endpoints(): List<Endpoint> = flow.endpoints()
+}
+
+/**
+ * Each endpoint of this flow, paired with the label that [Label.on] put on
+ * it, or with null when no label is on it. When a labeled flow is labeled
+ * again, the outer label replaces the inner label.
+ */
+private val Deps.Flow.labeledEndpoints: List<Pair<Endpoint, String?>>
+ get() =
+ when (this) {
+ is LabeledFlow -> flow.endpoints().map { it to text }
+ is FlowSet -> flows.flatMap { it.labeledEndpoints }
+ else -> endpoints().map { it to null }
+ }
+
+/**
+ * Rejects a call of [verb] on a flow that [Label.on] returned, or on a flow
+ * that holds such a flow. A label goes on the edges between the receiver of
+ * [verb] and the flows passed to [Label.on]. A label on the receiver would
+ * therefore label no edge.
+ */
+private fun requireUnlabeled(
+ receiver: Deps.Flow,
+ verb: String,
+) {
+ val labeled = receiver.labeledEndpoints.filter { (_, text) -> text != null }
+ require(labeled.isEmpty()) {
+ "Cannot call $verb on ${labeled.joinToString { (end, text) ->
"${end.diagnosticName} labeled \"$text\"" }}: " +
+ "a label goes on the edges between the receiver of before or after and
the flows it is put on, so pass " +
+ "Label.of(...).on(...) to before or after as an argument instead"
+ }
+}
+
/**
* Draws an ordering edge from each endpoint of [upstream] to each of
* [downstream]. An edge between two tasks is recorded on the downstream task;
@@ -153,18 +202,26 @@ private fun link(
upstream: Deps.Flow,
downstream: Deps.Flow,
) {
- for (up in upstream.endpoints()) {
- for (down in downstream.endpoints()) {
+ for ((up, upLabel) in upstream.labeledEndpoints) {
+ for ((down, downLabel) in downstream.labeledEndpoints) {
+ // `before` and `after` reject a labeled receiver. So a label always
+ // comes from the flow passed as an argument, whichever end of the edge
+ // that flow is.
+ val label = upLabel ?: downLabel
if (up is TaskDef && down is TaskDef) {
down.dependsOn(up)
+ label?.let { down.upstreamLabels[up] = EdgeLabel(it) }
} else {
val upDag = up.owningDag
val downDag = down.owningDag
require(upDag == null || downDag == null || upDag === downDag) {
- "Cannot order ${up.label} of Dag '${upDag?.id}' before ${down.label}
of " +
+ "Cannot order ${up.diagnosticName} of Dag '${upDag?.id}' before
${down.diagnosticName} of " +
"Dag '${downDag?.id}'; an edge stays inside one Dag"
}
- (upDag ?: downDag)?.let { it.groupEdges += up to down }
+ (upDag ?: downDag)?.let { dag ->
+ dag.groupEdges += up to down
+ label?.let { dag.groupEdgeLabels[up to down] = EdgeLabel(it) }
+ }
}
}
}
@@ -179,7 +236,7 @@ internal val Endpoint.owningDag: DagDef?
}
/** How an endpoint is named in a diagnostic. */
-internal val Endpoint.label: String
+internal val Endpoint.diagnosticName: String
get() =
when (this) {
is TaskDef -> "task '$id'"
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Label.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Label.kt
new file mode 100644
index 00000000000..e8758f308d4
--- /dev/null
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Label.kt
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk
+
+/**
+ * The text of an edge label. The Airflow UI shows the label on the edge in
+ * the graph, as it does for Python's `Label`. The text comes before the task
+ * that the label is put on. That way the text sits between the two tasks of
+ * the edge in the code, just as the label sits between them in the graph:
+ *
+ * ```java
+ * loaded.before(Label.of("when empty").on(reportEmpty)); // load >>
Label("when empty") >> report_empty
+ * cleaned.after(Label.of("always").on(loaded)); // load >>
Label("always") >> cleanup
+ * ```
+ */
+class Label private constructor(
+ private val text: String,
+) {
+ /**
+ * Labels each edge that a call to [Deps.Flow.before] or [Deps.Flow.after]
+ * draws between the call's receiver and [first] or one of [rest].
+ *
+ * Each edge of a fan-out can therefore carry its own label, and one label
+ * can go on several edges:
+ *
+ * ```java
+ * checked.before(Label.of("rows found").on(processed), Label.of("no
rows").on(reportEmpty));
+ * // load >> Label("forks") >> [report_empty, report_suspicious]
+ * loaded.before(Label.of("forks").on(reportEmpty, reportSuspicious));
+ * ```
+ *
+ * Passing a task's handle into another task's arguments declares an
+ * edge. Calling `Then`, `Else` or `Case` also declares an edge. To label
+ * one of those edges, declare it again with a label. Declaring an edge
+ * again only applies the label. A new label replaces the label the edge
+ * had, as Python's `DAG.set_edge_info` does. Declaring the edge again
+ * without a label keeps its label.
+ *
+ * When an edge goes to or from a task group, the UI draws it to or from
+ * the group's own node and shows the label there. A label never changes
+ * which tasks an edge connects. Python's `Label` can change those tasks.
+ * It does so when the end written to the left of the `Label` is in a task
+ * group and the other end is not in that same group. Python then replaces
+ * the left end with the group that contains it. For example, if `v` is in
+ * group `g` and `n` is in no group, `v >> Label("x") >> n` draws edges
+ * from the last tasks in `g` instead of from `v`. In Java, the end that
+ * matches this left end is the receiver of [Deps.Flow.before] or
+ * [Deps.Flow.after].
+ *
+ * @param first The task, task group, or set of tasks and groups at one end
+ * of each edge that gets the label.
+ * @param rest More tasks, task groups, or sets, labeled the same way as
+ * [first].
+ * @return A point in the flow to pass to [Deps.Flow.before] or
+ * [Deps.Flow.after]. Calling [Deps.Flow.before] or [Deps.Flow.after] on
+ * this point fails, because the label goes on the edges between the
+ * receiver of that call and the flows passed to this method.
+ */
+ fun on(
+ first: Deps.Flow,
+ vararg rest: Deps.Flow,
+ ): Deps.Flow = LabeledFlow(FlowSet(listOf(first, *rest)), text)
+
+ companion object {
+ /**
+ * Makes a label with [text], to put on edges with [on].
+ *
+ * @throws IllegalArgumentException if [text] is blank.
+ */
+ @JvmStatic
+ fun of(text: String): Label {
+ require(text.isNotBlank()) { "An edge label cannot be blank; pass the
text to show on the edge" }
+ return Label(text)
+ }
+ }
+}
diff --git
a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt
index e5ad4d6eacf..0482fc96ad8 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt
@@ -22,6 +22,8 @@ package org.apache.airflow.sdk.execution
import com.fasterxml.jackson.databind.ObjectMapper
import org.apache.airflow.sdk.Bundle
import org.apache.airflow.sdk.DagDef
+import org.apache.airflow.sdk.EdgeLabel
+import org.apache.airflow.sdk.Endpoint
import org.apache.airflow.sdk.GroupEdges
import org.apache.airflow.sdk.GroupExpansion
import org.apache.airflow.sdk.LiteralArg
@@ -110,7 +112,7 @@ internal fun serializeDag(
"tasks" to dag.tasks.map { (taskId, def) -> serializeTask(taskId, def,
downstream[taskId]) },
"dag_dependencies" to serializeDagDependencies(dag),
"task_group" to serializeTaskGroups(dag, expansion),
- "edge_info" to emptyMap<String, Any?>(),
+ "edge_info" to serializeEdgeInfo(dag, expansion),
"params" to emptyList<Any?>(),
"deadline" to null,
"allowed_run_types" to null,
@@ -505,6 +507,54 @@ private fun taskGroupObject(
/** The edges this group records for itself; the root group records none. */
private fun TaskGroupRef?.edges(expansion: GroupExpansion): GroupEdges =
this?.let { expansion.edgesOf(it.id) } ?: GroupEdges()
+/**
+ * The label of every labeled edge, keyed first by the ID of the edge's
+ * upstream end and then by the ID of its downstream end. Python's
+ * `DAG.edge_info` holds the labels the same way.
+ *
+ * The Airflow UI draws an edge to or from a task group as an edge to or from
+ * a join node of the group. So the label of such an edge goes under the ID of
+ * the join node. Python puts the label there too.
+ */
+private fun serializeEdgeInfo(
+ dag: DagDef,
+ expansion: GroupExpansion,
+): Map<String, Map<String, Map<String, String>>> {
+ val info = linkedMapOf<String, MutableMap<String, Map<String, String>>>()
+
+ fun put(
+ upstream: String,
+ downstream: String,
+ label: String,
+ ) {
+ info.getOrPut(upstream) { linkedMapOf() }[downstream] = mapOf("label" to
label)
+ }
+
+ // Python writes a label when its edge is drawn, so a later label replaces
an earlier one on the same edge.
+ // The labels go in in the order they were set to give the same result.
+ val labels = mutableListOf<Triple<String, String, EdgeLabel>>()
+ expansion.labels.forEach { (ends, label) -> labels += Triple(ends.first,
ends.second, label) }
+ dag.tasks.values.forEach { def ->
+ def.upstreamLabels.forEach { (upstream, label) -> labels +=
Triple(upstream.id, def.id, label) }
+ }
+ dag.groupEdgeLabels.forEach { (ends, label) ->
+ labels += Triple(ends.first.toEdgeInfoId(".downstream_join_id"),
ends.second.toEdgeInfoId(".upstream_join_id"), label)
+ }
+ labels.sortedBy { it.third.order }.forEach { (upstream, downstream, label)
-> put(upstream, downstream, label.text) }
+ return info
+}
+
+/**
+ * The ID that `edge_info` uses for this end of an edge. For a task, it is the
+ * task's own ID. For a group, it is the group's ID followed by [groupSuffix],
+ * which names one of the group's join nodes.
+ */
+private fun Endpoint.toEdgeInfoId(groupSuffix: String): String =
+ when (this) {
+ is TaskDef -> id
+ is TaskGroupRef -> id + groupSuffix
+ }
+
/**
* Recursively serializes a value with Airflow's type/var encoding, matching
* Python's `BaseSerialization.serialize()` output: primitives pass through,
diff --git
a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/EdgeLabelTest.kt
b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/EdgeLabelTest.kt
new file mode 100644
index 00000000000..7596097158f
--- /dev/null
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/EdgeLabelTest.kt
@@ -0,0 +1,323 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk
+
+import org.apache.airflow.sdk.Deps.Flow
+import org.apache.airflow.sdk.execution.serializeDag
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertThrows
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+
+internal class EdgeLabelTest {
+ private fun edgeInfo(dag: DagDef): Any? = serializeDag(dag, "",
".")["edge_info"]
+
+ private fun labels(vararg edges: Triple<String, String, String>):
Map<String, Map<String, Map<String, String>>> =
+ edges
+ .groupBy { it.first }
+ .mapValues { (_, fromOne) -> fromOne.associate { (_, downstream, label)
-> downstream to mapOf("label" to label) } }
+
+ private fun upstreamIds(
+ dag: DagDef,
+ taskId: String,
+ ) = dag.expandGroupEdges().upstreamsOf(dag.tasks.getValue(taskId))
+
+ private fun DagDef.noop(id: String) = task<Unit>(id, NoopTask::class.java)
+
+ private fun TaskGroupRef.noop(id: String) = task<Unit>(id,
NoopTask::class.java)
+
+ @Test
+ @DisplayName("Should label the edge that before draws to the flow the label
is on")
+ fun shouldLabelEdgeDrawnByBefore() {
+ val dag = DagDef("d")
+ val load = dag.noop("load")
+ val reportEmpty = dag.noop("report_empty")
+
+ load.before(Label.of("when empty").on(reportEmpty))
+
+ assertEquals(setOf("load"), upstreamIds(dag, "report_empty"))
+ assertEquals(labels(Triple("load", "report_empty", "when empty")),
edgeInfo(dag))
+ }
+
+ @Test
+ @DisplayName("Should label the edge that after draws from the flow the label
is on")
+ fun shouldLabelEdgeDrawnByAfter() {
+ val dag = DagDef("d")
+ val load = dag.noop("load")
+ val cleanup = dag.noop("cleanup")
+
+ cleanup.after(Label.of("always").on(load))
+
+ assertEquals(setOf("load"), upstreamIds(dag, "cleanup"))
+ assertEquals(labels(Triple("load", "cleanup", "always")), edgeInfo(dag))
+ }
+
+ @Test
+ @DisplayName("Should give each edge of a fan-out its own label, and none to
an edge without one")
+ fun shouldLabelEachEdgeOfFanOut() {
+ val dag = DagDef("d")
+ val check = dag.noop("check")
+ val process = dag.noop("process")
+ val reportEmpty = dag.noop("report_empty")
+ val audit = dag.noop("audit")
+
+ check.before(Label.of("rows found").on(process), Label.of("no
rows").on(reportEmpty), audit)
+
+ assertEquals(setOf("check"), upstreamIds(dag, "audit"))
+ assertEquals(
+ labels(Triple("check", "process", "rows found"), Triple("check",
"report_empty", "no rows")),
+ edgeInfo(dag),
+ )
+ }
+
+ @Test
+ @DisplayName("Should label the edge from every flow the label is on, with
the outer label of a labeled flow winning")
+ fun shouldLabelEveryEdgeToASet() {
+ val dag = DagDef("d")
+ val left = dag.noop("left")
+ val right = dag.noop("right")
+ val staging = dag.taskGroup("staging")
+ staging.noop("stage")
+ val join = dag.noop("join")
+
+ join.after(Label.of("done").on(Label.of("inner").on(left),
Label.of("inner").on(staging), right))
+
+ assertEquals(
+ labels(
+ Triple("left", "join", "done"),
+ Triple("right", "join", "done"),
+ Triple("staging.downstream_join_id", "join", "done"),
+ ),
+ edgeInfo(dag),
+ )
+ }
+
+ @Test
+ @DisplayName("Should only add the label to an edge that a condition already
declared")
+ fun shouldLabelEdgeDeclaredAlready() {
+ val dag = DagDef("d")
+ val load = dag.noop("load")
+ val reportEmpty = dag.noop("report_empty")
+ val gate = dag.If(HasRows::class.java).Then(load).Else(reportEmpty)
+
+ gate.before(Label.of("rows found").on(load), Label.of("no
rows").on(reportEmpty))
+
+ assertEquals(listOf(dag.tasks.getValue("hasRows")),
load.def.upstreams.toList())
+ assertEquals(
+ labels(Triple("hasRows", "load", "rows found"), Triple("hasRows",
"report_empty", "no rows")),
+ edgeInfo(dag),
+ )
+ }
+
+ @Test
+ @DisplayName("Should replace an earlier label, and keep it when the edge is
declared again without one")
+ fun shouldReplaceEarlierLabel() {
+ val dag = DagDef("d")
+ val load = dag.noop("load")
+ val cleanup = dag.noop("cleanup")
+ val publish = dag.taskGroup("publish")
+ publish.noop("push")
+
+ load.before(Label.of("always").on(cleanup, publish))
+ load.before(Label.of("when empty").on(cleanup, publish))
+ load.before(cleanup, publish)
+ cleanup.after(load)
+
+ assertEquals(
+ labels(
+ Triple("load", "cleanup", "when empty"),
+ Triple("load", "publish.push", "when empty"),
+ Triple("load", "publish.upstream_join_id", "when empty"),
+ ),
+ edgeInfo(dag),
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject before or after called on a labeled flow, and
draw no edge")
+ fun shouldRejectLabeledReceiver() {
+ val dag = DagDef("d")
+ val load = dag.noop("load")
+ val cleanup = dag.noop("cleanup")
+ val audit = dag.noop("audit")
+
+ val calls =
+ mapOf<String, () -> Unit>(
+ "before" to { Label.of("always").on(load).before(cleanup) },
+ "after" to { Label.of("always").on(load).after(cleanup) },
+ "before on a set" to { Flow.of(audit,
Label.of("always").on(load)).before(cleanup) },
+ )
+
+ calls.forEach { (call, draw) ->
+ val error = assertThrows(IllegalArgumentException::class.java, { draw()
}, call)
+ assertEquals(
+ "Cannot call ${call.substringBefore(' ')} on task 'load' labeled
\"always\": a label goes on the " +
+ "edges between the receiver of before or after and the flows it is
put on, so pass " +
+ "Label.of(...).on(...) to before or after as an argument instead",
+ error.message,
+ call,
+ )
+ }
+ assertEquals(emptySet<String>(), upstreamIds(dag, "cleanup"))
+ assertEquals(emptySet<String>(), upstreamIds(dag, "load"))
+ }
+
+ @Test
+ @DisplayName("Should reject an empty or blank label")
+ fun shouldRejectBlankLabel() {
+ listOf("", " ").forEach { text ->
+ val error = assertThrows(IllegalArgumentException::class.java, {
Label.of(text) }, "\"$text\"")
+ assertEquals("An edge label cannot be blank; pass the text to show on
the edge", error.message)
+ }
+ }
+
+ @Test
+ @DisplayName("Should put the label of an edge to or from a group under the
group's join node")
+ fun shouldLabelGroupEdgesAtJoinNodes() {
+ val dag = DagDef("d")
+ val extract = dag.noop("extract")
+ val transform = dag.taskGroup("transform")
+ transform.noop("clean")
+ val publish = dag.taskGroup("publish")
+ publish.noop("push")
+ val notify = dag.noop("notify")
+
+ extract.before(Label.of("to transform").on(transform))
+ transform.before(Label.of("to publish").on(publish))
+ publish.before(Label.of("to notify").on(notify))
+
+ // Python also labels the edge from extract to the group's first task,
when no group holds extract.
+ assertEquals(
+ labels(
+ Triple("extract", "transform.clean", "to transform"),
+ Triple("extract", "transform.upstream_join_id", "to transform"),
+ Triple("transform.downstream_join_id", "publish.upstream_join_id", "to
publish"),
+ Triple("publish.downstream_join_id", "notify", "to notify"),
+ ),
+ edgeInfo(dag),
+ )
+ }
+
+ @Test
+ @DisplayName("Should label the task edges of a group edge from a task in no
group, whichever verb draws it")
+ fun shouldLabelTaskEdgesOfGroupEdgeEitherWay() {
+ val dag = DagDef("d")
+ val extract = dag.noop("extract")
+ val transform = dag.taskGroup("transform")
+ transform.noop("clean")
+
+ transform.after(Label.of("rows").on(extract))
+
+ assertEquals(
+ labels(Triple("extract", "transform.clean", "rows"), Triple("extract",
"transform.upstream_join_id", "rows")),
+ edgeInfo(dag),
+ )
+ }
+
+ @Test
+ @DisplayName("Should leave the task edges of a group edge from a task in a
group unlabeled")
+ fun shouldNotLabelTaskEdgesOfGroupEdgeFromGroupedTask() {
+ val dag = DagDef("d")
+ val transform = dag.taskGroup("transform")
+ val clean = transform.noop("clean")
+ val nulls = transform.taskGroup("checks").noop("nulls")
+ val publish = dag.taskGroup("publish")
+ publish.noop("push")
+
+ Flow.of(clean, nulls).before(Label.of("rows").on(publish))
+
+ assertEquals(setOf("transform.clean", "transform.checks.nulls"),
upstreamIds(dag, "publish.push"))
+ assertEquals(
+ labels(
+ Triple("transform.clean", "publish.upstream_join_id", "rows"),
+ Triple("transform.checks.nulls", "publish.upstream_join_id", "rows"),
+ ),
+ edgeInfo(dag),
+ )
+ }
+
+ @Test
+ @DisplayName("Should give a task edge the label set last, whether by the
edge itself or by its group edge")
+ fun shouldGiveTaskEdgeTheLabelSetLast() {
+ val dag = DagDef("d")
+ val extract = dag.noop("extract")
+ val transform = dag.taskGroup("transform")
+ val clean = transform.noop("clean")
+ val validate = transform.noop("validate")
+
+ extract.before(Label.of("rows").on(clean), validate)
+ extract.before(Label.of("to transform").on(transform))
+
+ // Python's later Label writes over the earlier one, so clean's edge takes
the group edge's label.
+ assertEquals(
+ labels(
+ Triple("extract", "transform.clean", "to transform"),
+ Triple("extract", "transform.validate", "to transform"),
+ Triple("extract", "transform.upstream_join_id", "to transform"),
+ ),
+ edgeInfo(dag),
+ )
+
+ extract.before(Label.of("rows").on(clean))
+
+ assertEquals(
+ labels(
+ Triple("extract", "transform.clean", "rows"),
+ Triple("extract", "transform.validate", "to transform"),
+ Triple("extract", "transform.upstream_join_id", "to transform"),
+ ),
+ edgeInfo(dag),
+ )
+ }
+
+ @Test
+ @DisplayName("Should let a later group edge relabel a task edge that an
earlier one labeled")
+ fun shouldLetLaterGroupEdgeRelabelTaskEdge() {
+ val dag = DagDef("d")
+ val extract = dag.noop("extract")
+ val outer = dag.taskGroup("outer")
+ val inner = outer.taskGroup("inner")
+ inner.noop("i")
+
+ extract.before(Label.of("a").on(outer))
+ extract.before(Label.of("b").on(inner))
+
+ assertEquals(
+ labels(
+ Triple("extract", "outer.inner.i", "b"),
+ Triple("extract", "outer.upstream_join_id", "a"),
+ Triple("extract", "outer.inner.upstream_join_id", "b"),
+ ),
+ edgeInfo(dag),
+ )
+
+ // Relabeling the outer group edge is the last write, as it is in Python.
+ extract.before(Label.of("c").on(outer))
+
+ assertEquals(
+ labels(
+ Triple("extract", "outer.inner.i", "c"),
+ Triple("extract", "outer.upstream_join_id", "c"),
+ Triple("extract", "outer.inner.upstream_join_id", "b"),
+ ),
+ edgeInfo(dag),
+ )
+ }
+}
diff --git
a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
index f7d11aa4284..a535ff2c9e0 100644
---
a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
+++
b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
@@ -39,6 +39,7 @@ import org.apache.airflow.sdk.ConditionTask
import org.apache.airflow.sdk.Context
import org.apache.airflow.sdk.DagDef
import org.apache.airflow.sdk.Deps
+import org.apache.airflow.sdk.Label
import org.apache.airflow.sdk.SwitchRef
import org.apache.airflow.sdk.SwitchTask
import org.apache.airflow.sdk.Task
@@ -253,7 +254,10 @@ private fun buildDag(case: JsonNode): DagDef {
case.path("order_edges").forEach { edge ->
val node = { id: String -> groups[id] as Deps.Flow? ?: tasks.getValue(id) }
- node(edge[0].asText()).before(node(edge[1].asText()))
+ val downstream = node(edge[1].asText())
+ // A third item is the label of the edge.
+ val labeled = if (edge.size() > 2)
Label.of(edge[2].asText()).on(downstream) else downstream
+ node(edge[0].asText()).before(labeled)
}
return dag
}
