Java SDK: Pick one of two tasks with a condition

.pre-commit-config.yaml

@@ -444,7 +444,7 @@ repos:
           ^java-sdk/scripts/ci/prek/check_serialization_conformance\.py$|
           ^java-sdk/sdk/schema/dag-schema\.json$|
           ^java-sdk/sdk/build\.gradle\.kts$|
-          ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Bundle|DagDef|Deps|Arg|TaskGroupRef)\.kt$|
+          ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef)\.kt$|
           ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde\.kt$|
           ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/(Refs|Fields)\.kt$|
           ^java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/.*$|

java-sdk/README.md

@@ -643,7 +643,7 @@ prek hook regenerate it.
 | capability: `task-args` | MUST † | ✓ | 3.4 |  |
 | capability: `dag-params` | MUST † | ✗ | – |  |
 | capability: `taskflow-dependencies` | MUST † | ✓ | 3.4 |  |
-| capability: `branching` | SHOULD † | ✗ | – |  |
+| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If / @Builder.If |
 | capability: `dag-test` | SHOULD † | ✗ | – |  |
 | capability: `task-group` | MAY † | ✓ | 3.4 |  |
 | capability: `dynamic-task-mapping` | MAY † | ✗ | – |  |

java-sdk/capabilities.yaml

@@ -107,7 +107,9 @@ capabilities:
     supported: true
     since: "3.4"
   branching:
-    supported: false
+    supported: true
+    since: "3.4"
+    note: "DagDef.If / @Builder.If"
   dag-test:
     supported: false
   task-group:

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

@@ -54,6 +54,18 @@ public void load(long transformed) {
     log.log(INFO, "Loaded {0}", transformed);
   }
 
+  @Builder.Task(id = "load_empty")
+  public void loadEmpty() {
+    log.log(INFO, "Nothing to load");
+  }
+
+  // A condition: its boolean picks one of the two loads, and the other is
+  // skipped.
+  @Builder.If(id = "has_rows")
+  public boolean hasRows(long transformed) {
+    return transformed > 0;
+  }
+
   // A task group: everything it declares is prefixed with its id, so this is
   // the task "checks.audit".
   @Builder.TaskGroup(id = "checks")
@@ -70,7 +82,8 @@ public void audit() {
   static class Wiring implements AnnotationExampleDeps {
     void depends() {
       var extracted = extract();
-      load(transform(extracted, lit(1.5)));
+      var transformed = transform(extracted, lit(1.5));
+      hasRows(transformed).Then(load(transformed)).Else(loadEmpty());
       // Ordering-only edge: the checks group runs after extract, with no data
       // flowing.
       extracted.before(checks());

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

@@ -55,6 +55,22 @@ public void execute(Context context, Client client) {
     }
   }
 
+  public static class LoadEmpty implements Task {
+    @Override
+    public void execute(Context context, Client client) {
+      log.log(INFO, "Nothing to load");
+    }
+  }
+
+  // A condition: its boolean picks one of the two loads, and the other is
+  // skipped.
+  public static class HasRows implements ConditionTask {
+    @Override
+    public boolean decide(Context context, Client client) {
+      return ((Number) client.getXCom("transform")).longValue() > 0;
+    }
+  }
+
   public static DagDef build() {
     var dag =
         new DagDef("java_native_interface_example")
@@ -69,8 +85,11 @@ public static DagDef build() {
             .config("doc_md", "Extracts a value and pushes it as an XCom.");
     var transform = dag.task("transform", Transform.class);
     var load = dag.task("load", Load.class);
+    var loadEmpty = dag.task("load_empty", LoadEmpty.class);
 
-    transform.after(extract).before(load);
+    transform.after(extract);
+    // With no task id given, the condition takes one from its class: "hasRows".
+    dag.If(HasRows.class).after(transform).Then(load).Else(loadEmpty);
     return dag;
   }
 }

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

@@ -96,6 +96,7 @@ import org.apache.airflow.sdk.internal.builderName as generatedBuilderName
 @SupportedAnnotationTypes(
   "org.apache.airflow.sdk.Builder.Dag",
   "org.apache.airflow.sdk.Builder.Task",
+  "org.apache.airflow.sdk.Builder.If",
   "org.apache.airflow.sdk.Builder.TaskGroup",
   "org.apache.airflow.sdk.Builder.TaskHandler",
   "org.apache.airflow.sdk.Builder.Deps",
@@ -356,29 +357,42 @@ class BuilderProcessor : AbstractProcessor() {
     builderName: ClassName,
     inGroup: Boolean,
   ): MethodSpec {
+    val condition = decl.kind == TaskKind.CONDITION
     val method =
       MethodSpec
         .methodBuilder(decl.method.simpleName.toString())
         .addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
-        .returns(ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible()))
+        .returns(
+          if (condition) {
+            CONDITION_REF_TYPE
+          } else {
+            ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible())
+          },
+        )
     decl.dataParams.forEach { method.addParameter(inType(it.type), it.name) }
     val def = taskDefCode(decl, CodeBlock.of($$"$T.$L", builderName, decl.className))
     // The view knows the group it belongs to, so the recorder is told where
     // the task goes instead of deriving it from the task's ID.
     val group = if (inGroup) CodeBlock.of("groupId()") else CodeBlock.of($$"$S", "")
+    val wrap = { call: CodeBlock -> if (condition) CodeBlock.of($$"$T.of($L)", CONDITION_REF_TYPE, call) else call }
     if (decl.dataParams.isEmpty()) {
-      method.addStatement($$"return $T.node($L, $L)", REFS_TYPE, group, def)
+      method.addStatement($$"return $L", wrap(CodeBlock.of($$"$T.node($L, $L)", REFS_TYPE, group, def)))
     } else {
       // The parameter names ride along so the serialized Dag can name each
       // argument, as the binding spec ADR-0007 defines requires.
       method.addStatement(
-        $$"return $T.call($L, $L, $T.of($L), $L)",
-        REFS_TYPE,
-        group,
-        def,
-        LIST_TYPE,
-        CodeBlock.join(decl.dataParams.map { CodeBlock.of($$"$S", it.name) }, ", "),
-        decl.dataParams.joinToString { it.name },
+        $$"return $L",
+        wrap(
+          CodeBlock.of(
+            $$"$T.call($L, $L, $T.of($L), $L)",
+            REFS_TYPE,
+            group,
+            def,
+            LIST_TYPE,
+            CodeBlock.join(decl.dataParams.map { CodeBlock.of($$"$S", it.name) }, ", "),
+            decl.dataParams.joinToString { it.name },
+          ),
+        ),
       )
     }
     return method.build()
@@ -396,7 +410,7 @@ class BuilderProcessor : AbstractProcessor() {
       CodeBlock
         .builder()
         .add($$"new $T($S, $L.class)", TASK_DEF_TYPE, decl.id, classRef)
-    explicitConfig(decl.method, TASK_ANNOTATION, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) ->
+    explicitConfig(decl.method, decl.kind.annotation, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) ->
       taskDef.add($$".config($S, $L)", key, value)
     }
     return taskDef.build()
@@ -421,16 +435,24 @@ class BuilderProcessor : AbstractProcessor() {
     val tasks = mutableListOf<TaskDeclaration>()
     for (inner in el.enclosedElements) {
       if (inner !is ExecutableElement) continue
-      val ann = inner.getAnnotation(Builder.Task::class.java) ?: continue
+      val task = inner.getAnnotation(Builder.Task::class.java)
+      val condition = inner.getAnnotation(Builder.If::class.java)
+      if (task == null && condition == null) continue
+      require(task == null || condition == null) {
+        "Method '${inner.simpleName}' carries both @Builder.Task and @Builder.If; a condition is a " +
+          "task, so it carries @Builder.If alone"
+      }
+      val kind = if (condition != null) TaskKind.CONDITION else TaskKind.TASK
       if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}")
-      val localId = ann.id.ifBlank { inner.simpleName.toString() }
+      checkDeciderReturn(kind, inner)
+      val localId = (condition?.id ?: task!!.id).ifBlank { inner.simpleName.toString() }
       require(tasks.none { it.method.simpleName.contentEquals(inner.simpleName) }) {
         "Class ${el.simpleName} overloads task method '${inner.simpleName}'; a method's name is the " +
           "name of its generated task class and of its wiring-view method, so rename one and keep its " +
-          "task id with @Builder.Task(id = \"$localId\")"
+          "task id with ${kind.spelling}(id = \"$localId\")"
       }
       tasks +=
-        TaskDeclaration(inner, (path + localId).joinToString("."), collectDataParams(inner), el, classPath)
+        TaskDeclaration(inner, (path + localId).joinToString("."), collectDataParams(inner), el, classPath, kind)
     }
 
     val groups = mutableListOf<GroupDeclaration>()
@@ -571,7 +593,7 @@ class BuilderProcessor : AbstractProcessor() {
       val name = decl.method.simpleName.toString()
       require(name !in reserved) {
         "Task method '$name' clashes with a member of the wiring view; rename the method and keep " +
-          "the task id with @Builder.Task(id = \"${decl.id.substringAfterLast('.')}\")"
+          "the task id with ${decl.kind.spelling}(id = \"${decl.id.substringAfterLast('.')}\")"
       }
     }
     val accessors = mutableMapOf<String, GroupDeclaration>()
@@ -682,13 +704,31 @@ class BuilderProcessor : AbstractProcessor() {
     }
   }
 
+  /**
+   * Checks what a decider hands back, which is what the SDK reads its decision
+   * from. A plain task may return anything, including nothing.
+   */
+  private fun checkDeciderReturn(
+    kind: TaskKind,
+    method: ExecutableElement,
+  ) {
+    if (kind != TaskKind.CONDITION) return
+    val returns = method.returnType
+    val boolean = returns.kind == TypeKind.BOOLEAN || with(processingEnv) { isType(returns, BOXED_BOOLEAN_TYPE) }
+    require(boolean) {
+      "@Builder.If method '${method.simpleName}' returns $returns, but a condition returns boolean: " +
+        "true runs the task named by Then, false the one named by Else"
+    }
+  }
+
   private fun buildTask(decl: TaskDeclaration): TypeSpec {
+    val condition = decl.kind == TaskKind.CONDITION
     val executeSpec =
       MethodSpec
-        .methodBuilder("execute")
+        .methodBuilder(if (condition) "decide" else "execute")
         .addAnnotation(Override::class.java)
         .addModifiers(Modifier.PUBLIC)
-        .returns(TypeName.VOID)
+        .returns(if (condition) TypeName.BOOLEAN else TypeName.VOID)
         .addParameter(CONTEXT_TYPE, "context")
         .addParameter(CLIENT_TYPE, "client")
         .addException(Exception::class.java)
@@ -734,10 +774,11 @@ class BuilderProcessor : AbstractProcessor() {
       }
     }
 
-    if (inner.returnType.kind == TypeKind.VOID) {
-      $$"new $T().$L($L)"
-    } else {
-      $$"client.setXCom(new $T().$L($L))"
+    when {
+      // The SDK pushes a condition's result itself, after it has skipped the side not taken.
+      condition -> $$"return new $T().$L($L)"
+      inner.returnType.kind == TypeKind.VOID -> $$"new $T().$L($L)"
+      else -> $$"client.setXCom(new $T().$L($L))"
     }.also {
       executeSpec.addStatement(
         it,
@@ -749,7 +790,7 @@ class BuilderProcessor : AbstractProcessor() {
 
     return TypeSpec
       .classBuilder(decl.className)
-      .addSuperinterface(Task::class.java)
+      .addSuperinterface(if (condition) CONDITION_TASK_TYPE else ClassName.get(Task::class.java))
       .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
       .addMethod(executeSpec.build())
       .build()
@@ -893,6 +934,7 @@ private class TaskDeclaration(
   val dataParams: List<DataParam>,
   val owner: TypeElement,
   val classPath: List<String> = emptyList(),
+  val kind: TaskKind = TaskKind.TASK,
 ) {
   val className: String =
     (classPath + method.simpleName.toString().replaceFirstChar(Char::uppercase)).joinToString("_")
@@ -924,12 +966,27 @@ private val DAG_SOURCE_TYPE = ClassName.get(DagSource::class.java)
 private val REFS_TYPE = ClassName.get(Refs::class.java)
 private val ARG_TYPE = ClassName.get(Arg::class.java)
 private val TASK_HANDLE_TYPE = ClassName.get(TaskRef::class.java)
+private val CONDITION_TASK_TYPE = ClassName.get(ConditionTask::class.java)
+private val CONDITION_REF_TYPE = ClassName.get(ConditionRef::class.java)
+private val BOXED_BOOLEAN_TYPE = ClassName.get("java.lang", "Boolean")
 private val DEPS_TYPE = ClassName.get(Deps::class.java)
 private val GROUP_TYPE = DEPS_TYPE.nestedClass("TaskGroup")
 private val LIST_TYPE = ClassName.get(List::class.java)
 
 private const val DAG_ANNOTATION = "org.apache.airflow.sdk.Builder.Dag"
-private const val TASK_ANNOTATION = "org.apache.airflow.sdk.Builder.Task"
+
+/**
+ * What a declared task is, and the annotation that declares it. A decider
+ * differs from a plain task in what it implements and in what its wiring-view
+ * method hands back, but carries the same configuration attributes.
+ */
+private enum class TaskKind(
+  val annotation: String,
+  val spelling: String,
+) {
+  TASK("org.apache.airflow.sdk.Builder.Task", "@Builder.Task"),
+  CONDITION("org.apache.airflow.sdk.Builder.If", "@Builder.If"),
+}
 
 private val RESERVED_VIEW_NAMES =
   setOf(