Skip to content

Java SDK: Choose one of several tasks with a branch - #74439

Merged
jason810496 merged 7 commits into
jason/java-sdk/native-dag-iffrom
jason/java-sdk/native-dag-branch
Oct 9, 2026
Merged

jason810496 merged 7 commits into
jason/java-sdk/native-dag-iffrom
jason/java-sdk/native-dag-branch

Conversation

@jason810496

@jason810496 jason810496 commented Oct 8, 2026 •

Copy link
Copy Markdown
Member

Merge order:

  1. Java SDK: Pack each native Dag's source file into the bundle JAR #74096 — Pack each native Dag's source file into the bundle JAR
  2. Java SDK: Pick one of two tasks with a condition #74438 — Pick one of two tasks with a condition
  3. Java SDK: Choose one of several tasks with a branch #74439 — Choose one of several tasks with a branch (current one)
  4. Java SDK: Trigger another Dag's run from a native Java Dag #74440 — Trigger another Dag's run from a native Java Dag

Why

A condition picks between two tasks, but a Dag that routes work by size or by region needs more than two, and nesting conditions to get there makes the graph unreadable.

What changes

DagDef.Branch and @Builder.Branch declare a task that chooses one of the tasks listed with option(...); every other one is skipped. Neither style hard-codes a task id string.

// Interface style: the case is its own class, so javac checks the choice.
dag.Branch(PickPath.class).option(handleLong).option(handleShort);

public class PickPath implements BranchTask {
  @Override
  public Class<? extends Task> choose(Context context, Client client) {
    return ((Number) client.getXCom("extract")).longValue() > 1000
        ? HandleLong.class
        : HandleShort.class;
  }
}
// Annotation style: a task method's implementation is generated, so there is no class
// to name. The method returns a TaskId from the generated TaskIds instead.
@Builder.Branch(id = "pick_report")
public TaskId pickReport(long transformed) {
  return transformed > 100 ? EtlBuilder.TaskIds.REPORT_LONG : EtlBuilder.TaskIds.REPORT_SHORT;
}

@Builder.Deps
static class Wiring implements EtlDeps {
  void depends() {
    pickReport(transformed).option(reportLong()).option(reportShort());
  }
}

EtlBuilder.TaskIds holds one constant per task of the Dag and is generated only for a Dag that declares a branch.

  • The chosen task's ID becomes the branch's return value, as it does for Python's branch operators.
  • A branch chooses exactly one case, so a task that runs after several of them needs a trigger rule such as none_failed_min_one_success.
  • Two cases that share a Task class are rejected where the second is named: a BranchTask could not tell them apart.
  • The serialization conformance check now covers the shared switch Dag.

Was generative AI tooling used to co-author this PR?

@jason810496
jason810496 added this pull request to stack #74441 October 8, 2026 05:01
@jason810496
jason810496 removed this pull request from stack #74441 October 8, 2026 05:06
@jason810496
jason810496 added this pull request to stack #74442 October 8, 2026 05:06
@jason810496
jason810496 force-pushed the jason/java-sdk/native-dag-branch branch from 80beb7a to bc81537 Compare October 8, 2026 05:33
Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt Outdated
Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt Outdated
Comment thread java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt Outdated
Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt Outdated
Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt Outdated
Comment thread java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt Outdated
Comment thread java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BranchTest.kt Outdated
A condition picks between two tasks, but a Dag that routes work by size or
by region needs more than two, and splitting that into nested conditions
makes the graph unreadable.

`DagDef.Branch` and `@Builder.Branch` declare a task that chooses one of
the tasks listed with `option(...)`; every other one is skipped. The chosen
task's ID becomes the branch's return value, as it does for Python's branch
operators. A branch chooses exactly one case, so a task that runs after
several of them needs a trigger rule such as
`none_failed_min_one_success`.

Each authoring style names the case in the way that style can type-check. A
`BranchTask` returns the case's own class, so javac checks it; two cases
that share a class are rejected where the second is named, since the branch
could not tell them apart. A `@Builder.Branch` method has no such class to
name -- a task method's implementation is generated -- so it returns a
`TaskId` from the generated `<Dag>Builder.TaskIds`, which holds one
constant per task of the Dag. Either way nothing hard-codes a task id
string, and the generated holder is emitted only for a Dag that declares a
branch.
@jason810496
jason810496 force-pushed the jason/java-sdk/native-dag-branch branch from bc81537 to e643f06 Compare October 8, 2026 11:46
@jason810496
jason810496 marked this pull request as ready for review October 8, 2026 13:00
@jason810496
jason810496 requested a review from uranusjr as a code owner October 8, 2026 13:00

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the review.

Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt Outdated
Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt Outdated
Comment thread java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt Outdated
Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt Outdated
Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt Outdated
Comment thread java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt Outdated
Comment thread java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BranchTest.kt Outdated
@jason810496
jason810496 merged commit c3043c2 into main Oct 9, 2026
192 checks passed
@jason810496
jason810496 deleted the jason/java-sdk/native-dag-branch branch October 9, 2026 06:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants