Add TaskMap and TaskInstance.map_id #20286

uranusjr · 2021-12-14T10:25:07Z

TaskMap is populated by a TaskInstance when it pushes its return value to XCom.

Mostly for feedback for now. The pushing part needs some additional logic from #19965 to discover downstreams that need mapping.

airflow/models/taskmap.py

tests/providers/jira/operators/test_jira.py

airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py

uranusjr · 2022-01-10T05:06:10Z

I think this is ready-ish. Probably needs more eyes on the migration code.

ashb

Migration looks good.

ashb · 2022-01-10T11:47:16Z

airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py

Not that it matters, but for consistency should we have the same default here as we do on TI's map_index column?

The default on TI.map_index is mainly added so we don’t need to change too much existing code (otherwise we’d need to add a ton of map_index=-1), but this being a new class, I think I prefer being explicit instead.

airflow/models/baseoperator.py

airflow/models/taskinstance.py

ashb

We need more tests of _record_task_map_for_downstreams behaviour though please.

airflow/models/taskinstance.py

uranusjr · 2022-01-12T14:49:33Z

Alright, modified the code a bit and added a few tests. We can’t put mapped task in a DAG yet (bag_dag kept throwing AttributeError at me about a ton of things expected on BaseOperator but not yet present on MappedOperator), so a bit of mocking is needed. But otherwise things seem to work as expected.

ashb · 2022-01-12T15:16:39Z

so a bit of mocking is needed. But otherwise things seem to work as expected.

Yeah, I guess those'll be fixed by #20743

airflow/models/taskinstance.py

TaskMap is populated by a TaskInstance when it pushes its return value to XCom.

On MySQL you must not drop an index "needed by a primary key", so we need to drop the primary key first. On both MySQL and MSSQL, a column used by a primary key cannot be nullable (which we already accounted for, but I forgot to add nullable=False in the migration file).

This call actually makes more sense in TI._execute_task(), which calls xcom_push() for the return value, since we get to remove a somewhat hackish if check.

This makes it easier to keep the logic in sync with BaseOperator.

An XCom value is arbitrary and can potentially be a giant blob of data, which would make the log unreadable. Instead, only log the *type* of the XCom value instead; since the offending TaskInstance's key is also logged, the user can always look up the actual value in the XCom storage if needed. The exception still contains a reference to the actual value so the catching frame can access it. Since that value is already in memory anyway, keeping it a bit longer shouldn't be too big a problem.

We need to traverse the entire DAG to find mapped decendants because task groups are not currently recorded in the upstream-downstream lists, only operators. This should change in the future, but until then we need to make do.

airflow/models/taskmap.py

ashb

Doc change but LGTM

github-actions · 2022-01-13T17:33:05Z

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

Co-authored-by: Ash Berlin-Taylor <ash_github@firemirror.com>

ashb · 2022-01-17T09:44:55Z

airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py

+    with op.batch_alter_table("task_instance") as batch_op:
+        # I think we always use this name for TaskInstance after 7b2661a43ba3?
+        batch_op.drop_constraint("task_instance_pkey", type_="primary")
+        batch_op.add_column(Column("map_index", Integer, nullable=False, default=-1))


@uranusjr Did this need server_default=-1 instead? - re #20876

I’m investigating if it works as expected on all engines (if not I’ll just do an UPDATE SET map_index = -1 after this line)

Doing it with a server default is much preferred -- as otherwise an UPDATE requires us to re-write/touch every row, but a server default (at least on Postgres) doesn't.

boring-cyborg bot added the kind:documentation label Dec 14, 2021

uranusjr force-pushed the task-map-table branch 2 times, most recently from 637d31c to 13105b0 Compare December 15, 2021 12:48

ashb reviewed Dec 15, 2021

View reviewed changes

airflow/models/taskmap.py Outdated Show resolved Hide resolved

ashb reviewed Dec 15, 2021

View reviewed changes

tests/providers/jira/operators/test_jira.py Outdated Show resolved Hide resolved

uranusjr added area:dynamic-task-mapping AIP-42 and removed kind:documentation labels Jan 4, 2022

uranusjr force-pushed the task-map-table branch from 13105b0 to aded37b Compare January 7, 2022 06:18

ashb reviewed Jan 7, 2022

View reviewed changes

airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py Outdated Show resolved Hide resolved

uranusjr force-pushed the task-map-table branch 4 times, most recently from cc39c58 to 76cb924 Compare January 10, 2022 05:03

uranusjr marked this pull request as ready for review January 10, 2022 05:06

uranusjr requested review from XD-DENG, kaxil and potiuk as code owners January 10, 2022 05:06

uranusjr requested a review from ashb January 10, 2022 05:07

ashb reviewed Jan 10, 2022

View reviewed changes

airflow/models/taskinstance.py Outdated Show resolved Hide resolved

uranusjr force-pushed the task-map-table branch from f59e703 to 9a41876 Compare January 12, 2022 14:48

ashb reviewed Jan 12, 2022

View reviewed changes

airflow/models/taskinstance.py Outdated Show resolved Hide resolved

uranusjr force-pushed the task-map-table branch from 9a41876 to 572a6d5 Compare January 13, 2022 06:14

uranusjr requested review from mik-laj and vikramkoka as code owners January 13, 2022 09:24

Add TaskMap and TaskInstance.map_id

c8ba1ec

TaskMap is populated by a TaskInstance when it pushes its return value to XCom.

uranusjr added 18 commits January 13, 2022 17:24

Only push task map when needed as dependencies

8a66a1b

Fix task_reschedule migration

d7e7450

Fix TaskMap table name

d97c086

Must re-import TaskGrounp at runtime

74dd354

Use client-side default

fff3a1e

Use StringID in migration

dcfd72d

Add map_index as TI and TR init arg

be4212a

Record map_index to refresh_from_db test

f726cf6

Spellchecker doesn't think "downstreams" is a noun

f1bd8d3

Move TaskMap creation logic to TaskMap init

0ab3226

Fail ti if unmappable value is pushed for mapping

d05053d

Move task map recording call a stack higher

3fe7283

This call actually makes more sense in TI._execute_task(), which calls xcom_push() for the return value, since we get to remove a somewhat hackish if check.

Some tests (with a bit mocking)

3c0941f

Move decorator-mapping logic to MappedOperator

50c2cd7

This makes it easier to keep the logic in sync with BaseOperator.

Note on has_mapped_dependants implementation

338a457

We need to traverse the entire DAG to find mapped decendants because task groups are not currently recorded in the upstream-downstream lists, only operators. This should change in the future, but until then we need to make do.

Add 'unmappable' as a valid word

b7038e9

uranusjr force-pushed the task-map-table branch from dea0f77 to b7038e9 Compare January 13, 2022 09:24

ashb reviewed Jan 13, 2022

View reviewed changes

airflow/models/taskmap.py Outdated Show resolved Hide resolved

ashb approved these changes Jan 13, 2022

View reviewed changes

github-actions bot added the full tests needed We need to run full set of tests for this PR to merge label Jan 13, 2022

Clarify TaskMap costring

02b6bb4

Co-authored-by: Ash Berlin-Taylor <ash_github@firemirror.com>

uranusjr merged commit d48a3a3 into apache:main Jan 14, 2022

uranusjr deleted the task-map-table branch January 14, 2022 05:17

dcardinha mentioned this pull request Jan 14, 2022

Airflow database upgrade fails with "psycopg2.errors.NotNullViolation: column "map_index" of relation "task_instance" contains null value"s #20876

Closed

2 tasks

ashb reviewed Jan 17, 2022

View reviewed changes

ephraimbuddy added the changelog:skip Changes that should be skipped from the changelog (CI, tests, etc..) label Apr 11, 2022

jedcunningham added this to the Airflow 2.3.0 milestone Apr 26, 2022

Add TaskMap and TaskInstance.map_id #20286

Add TaskMap and TaskInstance.map_id #20286

Uh oh!

Conversation

uranusjr commented Dec 14, 2021

Uh oh!

Uh oh!

Uh oh!

Uh oh!

uranusjr commented Jan 10, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ashb left a comment

Choose a reason for hiding this comment

Uh oh!

ashb Jan 10, 2022

Choose a reason for hiding this comment

Uh oh!

uranusjr Jan 12, 2022

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ashb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

uranusjr commented Jan 12, 2022

Uh oh!

ashb commented Jan 12, 2022

Uh oh!

Uh oh!

Uh oh!

ashb left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jan 13, 2022

Uh oh!

ashb Jan 17, 2022

Choose a reason for hiding this comment

Uh oh!

uranusjr Jan 17, 2022

Choose a reason for hiding this comment

Uh oh!

ashb Jan 17, 2022

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

uranusjr commented Jan 10, 2022 •

edited

Loading