FrankYang0529 commented on code in PR #75190:
URL: https://github.com/apache/airflow/pull/75190#discussion_r4344221024
##########
ts-sdk/src/coordinator/serde.ts:
##########
@@ -87,15 +87,31 @@ const TASK_MODULE = "airflow.sdk.coordinators.node";
*/
const TASK_LANGUAGE = "typescript";
-/** Python's `DAG.edge_info`, keying a group end by its join node as Python
does. */
-function serializeEdgeInfo(dag: Dag): SerializedValue {
+/**
+ * Python's `DAG.edge_info`, keying a group end by its join node as Python
does.
+ *
+ * This writes the labels in the order they were last set. Python writes a
label when it draws the
+ * edge. So when two labels land on one entry, the label set last is the one
kept.
+ */
+function serializeEdgeInfo(dag: Dag, graph: DagGraph): SerializedValue {
const groups = getDagTaskGroups(dag);
+ const grouped = new Set([...groups.values()].flatMap((group) =>
group.taskIds));
const info: Record<string, Record<string, SerializedValue>> = {};
- for (const { upstream, downstream, label } of getDagOrderEdges(dag)) {
+ const put = (from: string, to: string, label: string): void => {
+ (info[from] ??= {})[to] = { label };
+ };
+ for (const { upstream, downstream, label } of getDagLabelledEdges(dag)) {
if (label === undefined) continue;
const from = groups.has(upstream) ? upstream + DOWNSTREAM_JOIN_SUFFIX :
upstream;
const to = groups.has(downstream) ? downstream + UPSTREAM_JOIN_SUFFIX :
downstream;
- (info[from] ??= {})[to] = { label };
+ put(from, to, label);
+ // Python's `extract >> Label("x") >> group` also labels the edge from
extract to each root of
+ // the group. Like the Java SDK, this code labels the edges to the roots
for an edge drawn with
+ // `before` or `after`, but only when no group holds extract.
+ if (grouped.has(upstream)) continue;
+ for (const root of graph.groupRootsReached.get(upstream)?.get(downstream)
?? []) {
+ put(upstream, root, label);
Review Comment:
Each label now records which end called `before` or `after`. The root or
leaf edges get the label only when that end is a task that no group holds. So
`staging.after(Label.of("rows").on(extracted))` labels only the join entry now.
--
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]
