jason810496 commented on code in PR #74230:
URL: https://github.com/apache/airflow/pull/74230#discussion_r4209902795
##########
java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt:
##########
@@ -736,6 +933,13 @@ private val RESERVED_VIEW_NAMES =
"wait",
)
+/**
+ * What a group's view inherits from [Group] on top of [RESERVED_VIEW_NAMES].
+ * Only a group scope reserves them: a Dag's view extends [Deps], which has
+ * neither.
+ */
+private val RESERVED_GROUP_VIEW_NAMES = setOf("groupId", "nodes")
Review Comment:
Good catch on `before`. Confirmed: a `@Builder.Task long before(long rows)`
inside a group generates a non-varargs `before(Arg<? extends Long>)`, which
wins phase-1 overload resolution over the inherited `before(Flow...)`, so
`staging().before(extract())` registers a task and still compiles.
The set is now read from the interface rather than written by hand, so it
cannot drift:
```kotlin
private val RESERVED_GROUP_VIEW_NAMES: Set<String> =
Deps.TaskGroup::class.java
.methods
.filterNot { ReflectModifier.isStatic(it.modifiers) }
.map { it.name }
.toSet()
```
Reflecting on the built class gives exactly `[after, before, endpoints,
groupId, nodes]`. The static filter earns its keep: without it Kotlin's
synthetic `access$before$jd` bridges leak in. A processor test covers a group
task method named `before`.
Both `[Group]` links now read `Deps.TaskGroup`.
Fixed in b8ec1ff47ac.
##########
airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst:
##########
@@ -694,6 +694,66 @@ class that supplies only task bodies, for a Dag a Python
file declares, carries
for a run (see :ref:`java-sdk/arg-binding`), the binding at a parameter's
position is what the
task receives. Wired inputs are the fallback, which is what a native Java
Dag always uses.
+Task groups
+~~~~~~~~~~~
+
+A task group gathers tasks that the Airflow UI shows as one node, as Python's
``TaskGroup`` does.
+Everything declared in a group carries the group's ID as a prefix, so task
``stage`` in group
+``staging`` is the task ``staging.stage``. On the interface surface,
``taskGroup`` declares a group on
+the Dag or inside another group, and the group declares its tasks:
+
+.. code-block:: java
+
+ var staging = dag.taskGroup("staging");
+ var stage = staging.task("stage", Stage.class); //
"staging.stage"
+ staging.taskGroup("checks").task("nulls", Nulls.class).after(stage); //
"staging.checks.nulls"
+ extract.before(staging);
+
+With annotations, a ``@Builder.TaskGroup`` class holds the tasks of one group,
and nesting one in
+another nests the groups:
+
+.. code-block:: java
+
+ @Builder.TaskGroup // the group "Staging", after the
class
+ static class Staging {
+ @Builder.Task
+ public long stage(long rows) { ... } // the task "Staging.stage"
+
+ @Builder.TaskGroup(id = "checks")
+ static class Checks {
+ @Builder.Task
+ public void nulls(long staged) { ... } // "Staging.checks.nulls"
+ }
+ }
+
+ @Builder.Deps
+ static class Wiring implements EtlPipelineDeps {
+ void depends() {
+ var rows = extract();
+ var staged = staging().stage(rows);
+ staging().checks().nulls(staged);
+ extract().before(staging());
+ }
+ }
+
+The generated view nests the same way, so a group is both the namespace of
what it holds and a point
+in the flow: ``staging().checks().nulls(staged)`` reaches a task, and
``extract().before(staging())``
+orders the whole group. Task method names scope to their own group, so two
groups can each declare
+``run()``. A group class is ``static``, non-private, and needs a no-argument
constructor, because the
+generated task bodies instantiate it.
+
+A group stands at either end of ``before``, ``after`` and ``Flow.of``. As an
upstream it stands for
+its leaves, the tasks nothing else in the group runs after; as a downstream,
for its roots, the tasks
+that run after nothing else in the group. A group ID contains only ASCII
letters, digits,
+underscores, or dashes, and no task or other group in the Dag can share it.
+
+.. note::
+
+ A group's endpoints are read once, from everything the group holds by the
time the Dag is
Review Comment:
You are right, and the difference is reproducible. With `outer` holding `t`
and nested `inner` holding `i`, `inner.before(t); extract.before(outer);` gives
`outer.t` the upstreams {`outer.inner.i`}, while the swapped order gives
{`extract`, `outer.inner.i`}. Python is order-dependent in the same way, so
this is parity rather than a bug, but the note claimed more than the code
delivers.
The note and `TaskGroupRef`'s KDoc now say that what a group holds is read
once while edges resolve in the order they were drawn.
`shouldResolveGroupEndpointsInDrawingOrder` pins the swapped order against
{`extract`, `outer.inner.i`}.
Fixed in b8ec1ff47ac.
##########
airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst:
##########
@@ -694,6 +694,66 @@ class that supplies only task bodies, for a Dag a Python
file declares, carries
for a run (see :ref:`java-sdk/arg-binding`), the binding at a parameter's
position is what the
task receives. Wired inputs are the fallback, which is what a native Java
Dag always uses.
+Task groups
+~~~~~~~~~~~
+
+A task group gathers tasks that the Airflow UI shows as one node, as Python's
``TaskGroup`` does.
+Everything declared in a group carries the group's ID as a prefix, so task
``stage`` in group
+``staging`` is the task ``staging.stage``. On the interface surface,
``taskGroup`` declares a group on
+the Dag or inside another group, and the group declares its tasks:
+
+.. code-block:: java
+
+ var staging = dag.taskGroup("staging");
+ var stage = staging.task("stage", Stage.class); //
"staging.stage"
+ staging.taskGroup("checks").task("nulls", Nulls.class).after(stage); //
"staging.checks.nulls"
+ extract.before(staging);
+
+With annotations, a ``@Builder.TaskGroup`` class holds the tasks of one group,
and nesting one in
+another nests the groups:
+
+.. code-block:: java
+
+ @Builder.TaskGroup // the group "Staging", after the
class
+ static class Staging {
+ @Builder.Task
+ public long stage(long rows) { ... } // the task "Staging.stage"
+
+ @Builder.TaskGroup(id = "checks")
+ static class Checks {
+ @Builder.Task
+ public void nulls(long staged) { ... } // "Staging.checks.nulls"
+ }
+ }
+
+ @Builder.Deps
+ static class Wiring implements EtlPipelineDeps {
Review Comment:
Restored, so the example calls every task of `EtlPipelineDeps` and survives
a copy-paste:
```java
var rows = extract();
load(transform(rows, lit(0.9)));
rows.before(audit());
var staged = staging().stage(rows);
staging().checks().nulls(staged);
extract().before(staging());
```
Fixed in b8ec1ff47ac.
##########
java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/TaskGroupTest.kt:
##########
@@ -0,0 +1,250 @@
+/*
+ * 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.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 TaskGroupTest {
+ private fun upstreamIds(
+ dag: DagDef,
+ taskId: String,
+ ) = dag.expandGroupEdges().upstreamsOf(dag.tasks.getValue(taskId))
+
+ @Test
+ @DisplayName("Should prefix the IDs of tasks and groups declared in a group")
+ fun shouldPrefixIdsDeclaredInGroup() {
+ val dag = DagDef("d")
+ val staging = dag.taskGroup("staging")
+ val stage = staging.task<Unit>("stage", NoopTask::class.java)
+ val checks = staging.taskGroup("checks")
+ checks.task<Unit>("nulls", NoopTask::class.java)
+
+ assertEquals("staging.stage", stage.def.id)
+ assertEquals("staging.checks", checks.id)
+ assertEquals(listOf("staging.stage", "staging.checks.nulls"),
dag.tasks.keys.toList())
+ assertEquals(listOf("staging", "staging.checks"), dag.groups.keys.toList())
+ }
+
+ @Test
+ @DisplayName("Should reject a group ID that is not a plain identifier")
+ fun shouldRejectInvalidGroupId() {
+ val error = assertThrows(IllegalArgumentException::class.java) {
DagDef("d").taskGroup("a.b") }
+
+ assertEquals(
+ "Task group ID 'a.b' must contain only ASCII letters, digits,
underscores, or dashes",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject a group ID that a task already uses")
+ fun shouldRejectGroupIdTakenByTask() {
+ val dag = DagDef("d")
+ dag.task<Unit>("staging", NoopTask::class.java)
+
+ val error = assertThrows(IllegalArgumentException::class.java) {
dag.taskGroup("staging") }
+
+ assertEquals("Dag 'd' already has a task or task group with ID: staging",
error.message)
+ }
+
+ @Test
+ @DisplayName("Should reject a task ID that a group already uses")
+ fun shouldRejectTaskIdTakenByGroup() {
+ val dag = DagDef("d")
+ dag.taskGroup("staging")
+
+ val error =
+ assertThrows(IllegalArgumentException::class.java) {
dag.task<Unit>("staging", NoopTask::class.java) }
+
+ assertEquals("Dag 'd' already has a task group with ID: staging",
error.message)
+ }
+
+ @Test
+ @DisplayName("Should wire a group upstream from its leaves and downstream to
its roots on registration")
+ fun shouldExpandGroupEdgesOntoRootsAndLeaves() {
+ val dag = DagDef("d")
+ val extract = dag.task<Unit>("extract", NoopTask::class.java)
+ val staging = dag.taskGroup("staging")
+ val stage = staging.task<Unit>("stage", NoopTask::class.java)
+ staging.taskGroup("checks").task<Unit>("nulls",
NoopTask::class.java).after(stage)
+ val publish = dag.taskGroup("publish")
+ publish.task<Unit>("push", NoopTask::class.java)
+ val load = dag.task<Unit>("load", NoopTask::class.java)
+ extract.before(staging)
+ staging.before(publish)
+ publish.before(load)
+
+ Bundle().register(dag)
+
+ assertEquals(setOf("extract"), upstreamIds(dag, "staging.stage"))
+ assertEquals(setOf("staging.stage"), upstreamIds(dag,
"staging.checks.nulls"))
+ assertEquals(setOf("staging.checks.nulls"), upstreamIds(dag,
"publish.push"))
+ assertEquals(setOf("publish.push"), upstreamIds(dag, "load"))
+ }
+
+ @Test
+ @DisplayName("Should record each group's own edges the way Python's
TaskGroup does")
+ fun shouldRecordGroupEdgesOnGroups() {
+ val dag = DagDef("d")
+ val extract = dag.task<Unit>("extract", NoopTask::class.java)
+ val staging = dag.taskGroup("staging")
+ staging.task<Unit>("stage", NoopTask::class.java)
+ val publish = dag.taskGroup("publish")
+ publish.task<Unit>("push", NoopTask::class.java)
+ val load = dag.task<Unit>("load", NoopTask::class.java)
+ extract.before(staging)
+ staging.before(publish)
+ publish.before(load)
+
+ Bundle().register(dag)
+
+ val edges = dag.expandGroupEdges()
+ assertEquals(setOf("extract"), edges.edgesOf(staging.id).upstreamTaskIds)
+ assertEquals(setOf("publish"),
edges.edgesOf(staging.id).downstreamGroupIds)
+ assertEquals(emptySet<String>(),
edges.edgesOf(staging.id).downstreamTaskIds)
+ assertEquals(setOf("staging"), edges.edgesOf(publish.id).upstreamGroupIds)
+ assertEquals(setOf("staging.stage"),
edges.edgesOf(publish.id).upstreamTaskIds)
+ assertEquals(setOf("load"), edges.edgesOf(publish.id).downstreamTaskIds)
+ }
+
+ @Test
+ @DisplayName("Should step over a group with no tasks to the tasks beyond it")
+ fun shouldStepOverEmptyGroup() {
+ val dag = DagDef("d")
+ val extract = dag.task<Unit>("extract", NoopTask::class.java)
+ val empty = dag.taskGroup("empty")
+ val load = dag.task<Unit>("load", NoopTask::class.java)
+ extract.before(empty)
+ empty.before(load)
+
+ Bundle().register(dag)
+
+ assertEquals(setOf("extract"), upstreamIds(dag, "load"))
+ }
+
+ @Test
+ @DisplayName("Should step over an empty nested group to the tasks of the
group holding it")
+ fun shouldStepOverEmptyNestedGroup() {
Review Comment:
Correct, neither test could tell the orders apart. Added
`shouldPreferEarlierEdgeOverParentForEmptyGroup`, which draws
`extract.before(inner)` before `inner.before(load)` and asserts `load` waits on
{`extract`} rather than {`outer.t`}. It fails under a parent-first order.
Fixed in b8ec1ff47ac.
##########
java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt:
##########
@@ -403,15 +519,67 @@ class BuilderProcessor : AbstractProcessor() {
}
/**
- * Rejects a task method whose wiring-view twin would clash with a member
- * the view or the wiring class already has: `depends`, `lit`, or a method
- * of `Object`.
+ * Rejects two tasks of the Dag sharing an ID, a task and a task group
+ * sharing one, and two task methods whose generated classes would collide.
+ */
+ private fun checkIds(scope: Scope) {
+ val declarations = scope.allTasks()
+ val taskIds = mutableSetOf<String>()
+ declarations.forEach { decl ->
+ require(taskIds.add(decl.id)) { "Tasks in Dag have duplicate ID:
${decl.id}" }
+ }
+ scope.allGroups().forEach { group ->
Review Comment:
This one is already rejected at compile time, so I have not added a second
check. `collectScope` catches it while collecting, at `BuilderProcessor.kt:429`:
```kotlin
require(groups.none { it.id == localId }) {
"Class ${el.simpleName} declares more than one task group '$localId'"
}
```
Two group classes in one scope both carrying `@Builder.TaskGroup(id =
"checks")` fail the build there, not at Dag-parse time. Across nesting levels
the full IDs cannot collide either: `GROUP_ID` is `[A-Za-z0-9_-]+`, so a group
ID cannot contain the dot that joins the path, and two distinct paths cannot
produce one full ID. A guard in `checkIds` would be unreachable.
Nothing covered that message, so I added a test for it in b8ec1ff47ac.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]