Skip to content

Commit d4013cc

Browse files
authored
Java SDK: Pick one of two tasks with a condition (#74438)
* Java SDK: Pick one of two tasks with a condition A native Java Dag could only run every task it declared: it had no way to let a task's result decide what runs next, so a Dag that loads only when there are rows had to be split in two or fall back to a Python stub. `DagDef.If` and `@Builder.If` declare a task whose boolean names the side that runs; the other is skipped. The deciding task carries `_can_skip_downstream`, so Airflow skips a cleared side again rather than running it, and the SDK writes the `skipmixin_key` XCom before the skip, as Python's SkipMixin does. It skips only the side not taken, so a task after both sides needs a trigger rule such as `none_failed_min_one_success`. `If` is capitalized because `if` is a Java keyword and `@Builder.If` reads the same way. With no task id given, one is derived from the condition class's simple name, so `HasRows.class` becomes `hasRows`. The runtime reaches `SkipDownstreamTasks` through the transport client only, never through the task-facing `Client`, so a task body cannot skip tasks the Dag did not declare it may. * Java SDK: Let a wiring view name a condition's sides in more than one call * Java SDK: Reject a decider class registered as an ordinary task * Java SDK: Name a condition's sides Then and Else * Java SDK: Pin the order a condition skips and records what it skipped * Java SDK: Drop the unused switch in the condition test's decider
1 parent 2bda192 commit d4013cc

27 files changed

Lines changed: 1274 additions & 36 deletions

File tree

‎.pre-commit-config.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -461,7 +461,7 @@ repos:
461461
^java-sdk/scripts/ci/prek/check_serialization_conformance\.py$|
462462
^java-sdk/sdk/schema/dag-schema\.json$|
463463
^java-sdk/sdk/build\.gradle\.kts$|
464-
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Bundle|DagDef|Deps|Arg|TaskGroupRef)\.kt$|
464+
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef)\.kt$|
465465
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde\.kt$|
466466
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/(Refs|Fields)\.kt$|
467467
^java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/.*$|

‎java-sdk/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -643,7 +643,7 @@ prek hook regenerate it.
643643
| capability: `task-args` | MUST † | ✓ | 3.4 | |
644644
| capability: `dag-params` | MUST † | ✗ | – | |
645645
| capability: `taskflow-dependencies` | MUST † | ✓ | 3.4 | |
646-
| capability: `branching` | SHOULD † | ✗ | – | |
646+
| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If / @Builder.If |
647647
| capability: `dag-test` | SHOULD † | ✗ | – | |
648648
| capability: `task-group` | MAY † | ✓ | 3.4 | |
649649
| capability: `dynamic-task-mapping` | MAY † | ✗ | – | |

‎java-sdk/capabilities.yaml‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,9 @@ capabilities:
107107
supported: true
108108
since: "3.4"
109109
branching:
110-
supported: false
110+
supported: true
111+
since: "3.4"
112+
note: "DagDef.If / @Builder.If"
111113
dag-test:
112114
supported: false
113115
task-group:

‎java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,18 @@ public void load(long transformed) {
5454
log.log(INFO, "Loaded {0}", transformed);
5555
}
5656

57+
@Builder.Task(id = "load_empty")
58+
public void loadEmpty() {
59+
log.log(INFO, "Nothing to load");
60+
}
61+
62+
// A condition: its boolean picks one of the two loads, and the other is
63+
// skipped.
64+
@Builder.If(id = "has_rows")
65+
public boolean hasRows(long transformed) {
66+
return transformed > 0;
67+
}
68+
5769
// A task group: everything it declares is prefixed with its id, so this is
5870
// the task "checks.audit".
5971
@Builder.TaskGroup(id = "checks")
@@ -70,7 +82,8 @@ public void audit() {
7082
static class Wiring implements AnnotationExampleDeps {
7183
void depends() {
7284
var extracted = extract();
73-
load(transform(extracted, lit(1.5)));
85+
var transformed = transform(extracted, lit(1.5));
86+
hasRows(transformed).Then(load(transformed)).Else(loadEmpty());
7487
// Ordering-only edge: the checks group runs after extract, with no data
7588
// flowing.
7689
extracted.before(checks());

‎java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,22 @@ public void execute(Context context, Client client) {
5555
}
5656
}
5757

58+
public static class LoadEmpty implements Task {
59+
@Override
60+
public void execute(Context context, Client client) {
61+
log.log(INFO, "Nothing to load");
62+
}
63+
}
64+
65+
// A condition: its boolean picks one of the two loads, and the other is
66+
// skipped.
67+
public static class HasRows implements ConditionTask {
68+
@Override
69+
public boolean decide(Context context, Client client) {
70+
return ((Number) client.getXCom("transform")).longValue() > 0;
71+
}
72+
}
73+
5874
public static DagDef build() {
5975
var dag =
6076
new DagDef("java_native_interface_example")
@@ -69,8 +85,11 @@ public static DagDef build() {
6985
.config("doc_md", "Extracts a value and pushes it as an XCom.");
7086
var transform = dag.task("transform", Transform.class);
7187
var load = dag.task("load", Load.class);
88+
var loadEmpty = dag.task("load_empty", LoadEmpty.class);
7289

73-
transform.after(extract).before(load);
90+
transform.after(extract);
91+
// With no task id given, the condition takes one from its class: "hasRows".
92+
dag.If(HasRows.class).after(transform).Then(load).Else(loadEmpty);
7493
return dag;
7594
}
7695
}

‎java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt‎

Lines changed: 80 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ import org.apache.airflow.sdk.internal.builderName as generatedBuilderName
9696
@SupportedAnnotationTypes(
9797
"org.apache.airflow.sdk.Builder.Dag",
9898
"org.apache.airflow.sdk.Builder.Task",
99+
"org.apache.airflow.sdk.Builder.If",
99100
"org.apache.airflow.sdk.Builder.TaskGroup",
100101
"org.apache.airflow.sdk.Builder.TaskHandler",
101102
"org.apache.airflow.sdk.Builder.Deps",
@@ -356,29 +357,42 @@ class BuilderProcessor : AbstractProcessor() {
356357
builderName: ClassName,
357358
inGroup: Boolean,
358359
): MethodSpec {
360+
val condition = decl.kind == TaskKind.CONDITION
359361
val method =
360362
MethodSpec
361363
.methodBuilder(decl.method.simpleName.toString())
362364
.addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
363-
.returns(ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible()))
365+
.returns(
366+
if (condition) {
367+
CONDITION_REF_TYPE
368+
} else {
369+
ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible())
370+
},
371+
)
364372
decl.dataParams.forEach { method.addParameter(inType(it.type), it.name) }
365373
val def = taskDefCode(decl, CodeBlock.of($$"$T.$L", builderName, decl.className))
366374
// The view knows the group it belongs to, so the recorder is told where
367375
// the task goes instead of deriving it from the task's ID.
368376
val group = if (inGroup) CodeBlock.of("groupId()") else CodeBlock.of($$"$S", "")
377+
val wrap = { call: CodeBlock -> if (condition) CodeBlock.of($$"$T.of($L)", CONDITION_REF_TYPE, call) else call }
369378
if (decl.dataParams.isEmpty()) {
370-
method.addStatement($$"return $T.node($L, $L)", REFS_TYPE, group, def)
379+
method.addStatement($$"return $L", wrap(CodeBlock.of($$"$T.node($L, $L)", REFS_TYPE, group, def)))
371380
} else {
372381
// The parameter names ride along so the serialized Dag can name each
373382
// argument, as the binding spec ADR-0007 defines requires.
374383
method.addStatement(
375-
$$"return $T.call($L, $L, $T.of($L), $L)",
376-
REFS_TYPE,
377-
group,
378-
def,
379-
LIST_TYPE,
380-
CodeBlock.join(decl.dataParams.map { CodeBlock.of($$"$S", it.name) }, ", "),
381-
decl.dataParams.joinToString { it.name },
384+
$$"return $L",
385+
wrap(
386+
CodeBlock.of(
387+
$$"$T.call($L, $L, $T.of($L), $L)",
388+
REFS_TYPE,
389+
group,
390+
def,
391+
LIST_TYPE,
392+
CodeBlock.join(decl.dataParams.map { CodeBlock.of($$"$S", it.name) }, ", "),
393+
decl.dataParams.joinToString { it.name },
394+
),
395+
),
382396
)
383397
}
384398
return method.build()
@@ -396,7 +410,7 @@ class BuilderProcessor : AbstractProcessor() {
396410
CodeBlock
397411
.builder()
398412
.add($$"new $T($S, $L.class)", TASK_DEF_TYPE, decl.id, classRef)
399-
explicitConfig(decl.method, TASK_ANNOTATION, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) ->
413+
explicitConfig(decl.method, decl.kind.annotation, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) ->
400414
taskDef.add($$".config($S, $L)", key, value)
401415
}
402416
return taskDef.build()
@@ -421,16 +435,24 @@ class BuilderProcessor : AbstractProcessor() {
421435
val tasks = mutableListOf<TaskDeclaration>()
422436
for (inner in el.enclosedElements) {
423437
if (inner !is ExecutableElement) continue
424-
val ann = inner.getAnnotation(Builder.Task::class.java) ?: continue
438+
val task = inner.getAnnotation(Builder.Task::class.java)
439+
val condition = inner.getAnnotation(Builder.If::class.java)
440+
if (task == null && condition == null) continue
441+
require(task == null || condition == null) {
442+
"Method '${inner.simpleName}' carries both @Builder.Task and @Builder.If; a condition is a " +
443+
"task, so it carries @Builder.If alone"
444+
}
445+
val kind = if (condition != null) TaskKind.CONDITION else TaskKind.TASK
425446
if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}")
426-
val localId = ann.id.ifBlank { inner.simpleName.toString() }
447+
checkDeciderReturn(kind, inner)
448+
val localId = (condition?.id ?: task!!.id).ifBlank { inner.simpleName.toString() }
427449
require(tasks.none { it.method.simpleName.contentEquals(inner.simpleName) }) {
428450
"Class ${el.simpleName} overloads task method '${inner.simpleName}'; a method's name is the " +
429451
"name of its generated task class and of its wiring-view method, so rename one and keep its " +
430-
"task id with @Builder.Task(id = \"$localId\")"
452+
"task id with ${kind.spelling}(id = \"$localId\")"
431453
}
432454
tasks +=
433-
TaskDeclaration(inner, (path + localId).joinToString("."), collectDataParams(inner), el, classPath)
455+
TaskDeclaration(inner, (path + localId).joinToString("."), collectDataParams(inner), el, classPath, kind)
434456
}
435457

436458
val groups = mutableListOf<GroupDeclaration>()
@@ -571,7 +593,7 @@ class BuilderProcessor : AbstractProcessor() {
571593
val name = decl.method.simpleName.toString()
572594
require(name !in reserved) {
573595
"Task method '$name' clashes with a member of the wiring view; rename the method and keep " +
574-
"the task id with @Builder.Task(id = \"${decl.id.substringAfterLast('.')}\")"
596+
"the task id with ${decl.kind.spelling}(id = \"${decl.id.substringAfterLast('.')}\")"
575597
}
576598
}
577599
val accessors = mutableMapOf<String, GroupDeclaration>()
@@ -682,13 +704,31 @@ class BuilderProcessor : AbstractProcessor() {
682704
}
683705
}
684706

707+
/**
708+
* Checks what a decider hands back, which is what the SDK reads its decision
709+
* from. A plain task may return anything, including nothing.
710+
*/
711+
private fun checkDeciderReturn(
712+
kind: TaskKind,
713+
method: ExecutableElement,
714+
) {
715+
if (kind != TaskKind.CONDITION) return
716+
val returns = method.returnType
717+
val boolean = returns.kind == TypeKind.BOOLEAN || with(processingEnv) { isType(returns, BOXED_BOOLEAN_TYPE) }
718+
require(boolean) {
719+
"@Builder.If method '${method.simpleName}' returns $returns, but a condition returns boolean: " +
720+
"true runs the task named by Then, false the one named by Else"
721+
}
722+
}
723+
685724
private fun buildTask(decl: TaskDeclaration): TypeSpec {
725+
val condition = decl.kind == TaskKind.CONDITION
686726
val executeSpec =
687727
MethodSpec
688-
.methodBuilder("execute")
728+
.methodBuilder(if (condition) "decide" else "execute")
689729
.addAnnotation(Override::class.java)
690730
.addModifiers(Modifier.PUBLIC)
691-
.returns(TypeName.VOID)
731+
.returns(if (condition) TypeName.BOOLEAN else TypeName.VOID)
692732
.addParameter(CONTEXT_TYPE, "context")
693733
.addParameter(CLIENT_TYPE, "client")
694734
.addException(Exception::class.java)
@@ -734,10 +774,11 @@ class BuilderProcessor : AbstractProcessor() {
734774
}
735775
}
736776

737-
if (inner.returnType.kind == TypeKind.VOID) {
738-
$$"new $T().$L($L)"
739-
} else {
740-
$$"client.setXCom(new $T().$L($L))"
777+
when {
778+
// The SDK pushes a condition's result itself, after it has skipped the side not taken.
779+
condition -> $$"return new $T().$L($L)"
780+
inner.returnType.kind == TypeKind.VOID -> $$"new $T().$L($L)"
781+
else -> $$"client.setXCom(new $T().$L($L))"
741782
}.also {
742783
executeSpec.addStatement(
743784
it,
@@ -749,7 +790,7 @@ class BuilderProcessor : AbstractProcessor() {
749790

750791
return TypeSpec
751792
.classBuilder(decl.className)
752-
.addSuperinterface(Task::class.java)
793+
.addSuperinterface(if (condition) CONDITION_TASK_TYPE else ClassName.get(Task::class.java))
753794
.addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
754795
.addMethod(executeSpec.build())
755796
.build()
@@ -893,6 +934,7 @@ private class TaskDeclaration(
893934
val dataParams: List<DataParam>,
894935
val owner: TypeElement,
895936
val classPath: List<String> = emptyList(),
937+
val kind: TaskKind = TaskKind.TASK,
896938
) {
897939
val className: String =
898940
(classPath + method.simpleName.toString().replaceFirstChar(Char::uppercase)).joinToString("_")
@@ -924,12 +966,27 @@ private val DAG_SOURCE_TYPE = ClassName.get(DagSource::class.java)
924966
private val REFS_TYPE = ClassName.get(Refs::class.java)
925967
private val ARG_TYPE = ClassName.get(Arg::class.java)
926968
private val TASK_HANDLE_TYPE = ClassName.get(TaskRef::class.java)
969+
private val CONDITION_TASK_TYPE = ClassName.get(ConditionTask::class.java)
970+
private val CONDITION_REF_TYPE = ClassName.get(ConditionRef::class.java)
971+
private val BOXED_BOOLEAN_TYPE = ClassName.get("java.lang", "Boolean")
927972
private val DEPS_TYPE = ClassName.get(Deps::class.java)
928973
private val GROUP_TYPE = DEPS_TYPE.nestedClass("TaskGroup")
929974
private val LIST_TYPE = ClassName.get(List::class.java)
930975

931976
private const val DAG_ANNOTATION = "org.apache.airflow.sdk.Builder.Dag"
932-
private const val TASK_ANNOTATION = "org.apache.airflow.sdk.Builder.Task"
977+
978+
/**
979+
* What a declared task is, and the annotation that declares it. A decider
980+
* differs from a plain task in what it implements and in what its wiring-view
981+
* method hands back, but carries the same configuration attributes.
982+
*/
983+
private enum class TaskKind(
984+
val annotation: String,
985+
val spelling: String,
986+
) {
987+
TASK("org.apache.airflow.sdk.Builder.Task", "@Builder.Task"),
988+
CONDITION("org.apache.airflow.sdk.Builder.If", "@Builder.If"),
989+
}
933990

934991
private val RESERVED_VIEW_NAMES =
935992
setOf(

0 commit comments

Comments
 (0)