From be4110ad92ce4061b8b09c75efe1ffc5c72cb445 Mon Sep 17 00:00:00 2001
From: James Ong <zhiyong@hudson-trading.com>
Date: Tue, 18 Oct 2022 11:46:32 +0800
Subject: [PATCH 1/9] optimize task filtering

---
 airflow/models/taskinstance.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 54ce2a2d2e030..aaf97baccdf23 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -2439,6 +2439,34 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
                 TaskInstance.task_id == first_task_id,
             )
 
+        # group by dag_id and map_index while creating multiple groups for the other 2 params
+        if all(t.dag_id == dag_id and t.map_index == map_index for t in tis):
+
+            # naively create 2 groups, 1 grouped by run_id and the other by task_id then use the smaller group
+            run_id_groups = defaultdict(list)
+            task_id_groups = defaultdict(list)
+            for t in tis:
+                run_id_groups[t.run_id].append(t.task_id)
+                task_id_groups[t.task_id].append(t.run_id)
+            filter_condition = []
+            if len(run_id_groups) <= len(task_id_groups):
+                for run_id, task_ids in run_id_groups.items():
+                    filter_condition.append(
+                        and_(TaskInstance.run_id == run_id, TaskInstance.task_id.in_(task_ids)
+                    )
+                )
+            else:
+                for task_id, run_ids in task_id_groups.items():
+                    filter_condition.append(
+                        and_(TaskInstance.task_id == task_id, TaskInstance.run_id.in_(run_ids)
+                    )
+                )
+            return and_(
+                TaskInstance.dag_id == dag_id,
+                TaskInstance.map_index == map_index,
+                or_(*filter_condition)
+            )   
+
         return tuple_in_condition(
             (TaskInstance.dag_id, TaskInstance.task_id, TaskInstance.run_id, TaskInstance.map_index),
             (ti.key.primary for ti in tis),

From 51a0fb9e0daf11512bd2f8ef0ff577ffbcab8798 Mon Sep 17 00:00:00 2001
From: James Ong <zhiyong@hudson-trading.com>
Date: Tue, 18 Oct 2022 17:21:11 +0800
Subject: [PATCH 2/9] generalize filtering

---
 airflow/models/taskinstance.py | 92 ++++++++++++++++++++--------------
 1 file changed, 55 insertions(+), 37 deletions(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index aaf97baccdf23..f412edfa5ac8e 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -2417,60 +2417,78 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
         first_task_id = first.task_id
         # Common path optimisations: when all TIs are for the same dag_id and run_id, or same dag_id
         # and task_id -- this can be over 150x faster for huge numbers of TIs (20k+)
-        if all(t.dag_id == dag_id and t.run_id == run_id and t.map_index == map_index for t in tis):
+
+        # pre-compute the set of dag_id, run_id, map_indices and task_ids
+        dag_ids, run_ids, map_indices, task_ids = set(), set(), set(), set()
+        for t in tis:
+            dag_ids.add(t.dag_id)
+            run_ids.add(t.run_id)
+            map_indices.add(t.map_index)
+            task_ids.add(t.task_id)
+
+        if dag_ids == {dag_id} and run_ids == {run_id} and map_indices == {map_index}:
             return and_(
                 TaskInstance.dag_id == dag_id,
                 TaskInstance.run_id == run_id,
                 TaskInstance.map_index == map_index,
-                TaskInstance.task_id.in_(t.task_id for t in tis),
+                TaskInstance.task_id.in_(task_ids),
             )
-        if all(t.dag_id == dag_id and t.task_id == first_task_id and t.map_index == map_index for t in tis):
+        if dag_ids == {dag_id} and task_ids == {first_task_id} and map_indices == {map_index}:
             return and_(
                 TaskInstance.dag_id == dag_id,
-                TaskInstance.run_id.in_(t.run_id for t in tis),
+                TaskInstance.run_id.in_(run_ids),
                 TaskInstance.map_index == map_index,
                 TaskInstance.task_id == first_task_id,
             )
-        if all(t.dag_id == dag_id and t.run_id == run_id and t.task_id == first_task_id for t in tis):
+        if dag_ids == {dag_id} and run_ids == {run_id} and task_ids == {task_id}:
             return and_(
                 TaskInstance.dag_id == dag_id,
                 TaskInstance.run_id == run_id,
-                TaskInstance.map_index.in_(t.map_index for t in tis),
+                TaskInstance.map_index.in_(map_indices),
                 TaskInstance.task_id == first_task_id,
             )
 
-        # group by dag_id and map_index while creating multiple groups for the other 2 params
-        if all(t.dag_id == dag_id and t.map_index == map_index for t in tis):
-
-            # naively create 2 groups, 1 grouped by run_id and the other by task_id then use the smaller group
-            run_id_groups = defaultdict(list)
-            task_id_groups = defaultdict(list)
-            for t in tis:
-                run_id_groups[t.run_id].append(t.task_id)
-                task_id_groups[t.task_id].append(t.run_id)
-            filter_condition = []
-            if len(run_id_groups) <= len(task_id_groups):
-                for run_id, task_ids in run_id_groups.items():
-                    filter_condition.append(
-                        and_(TaskInstance.run_id == run_id, TaskInstance.task_id.in_(task_ids)
-                    )
-                )
-            else:
-                for task_id, run_ids in task_id_groups.items():
-                    filter_condition.append(
-                        and_(TaskInstance.task_id == task_id, TaskInstance.run_id.in_(run_ids)
-                    )
-                )
-            return and_(
-                TaskInstance.dag_id == dag_id,
-                TaskInstance.map_index == map_index,
-                or_(*filter_condition)
-            )   
+        filter_condition = []
+        # create 2 groups, 1 grouped by task_id the other by map_index and use the smaller group
+        task_id_groups, map_index_groups = defaultdict(lambda: defaultdict(list)), defaultdict(lambda: defaultdict(list))
+        for t in tis:
+            task_id_groups[(t.dag_id, t.run_id)][t.task_id].append(t.map_index)
+            map_index_groups[(t.dag_id, t.run_id)][t.map_index].append(t.task_id)
+
+        # this assumes that most dags have dag_id as the largest grouping, followed by run_id. even if its not, this is still
+        # a significant optimization over querying for every single tuple key
+        for cur_dag_id in dag_ids:
+            for cur_run_id in run_ids:
+                # we compare the group size between task_id and map_index and use the smaller group
+                dag_task_id_groups = task_id_groups[(cur_dag_id, cur_run_id)]
+                dag_map_index_groups = task_id_groups[(cur_dag_id, cur_run_id)]
+
+                if len(dag_task_id_groups) <= len(dag_map_index_groups):
+                    for cur_task_id, cur_map_indices in dag_task_id_groups.items():
+                        filter_condition.append(
+                            and_(
+                                TaskInstance.dag_id == cur_dag_id,
+                                TaskInstance.run_id == cur_run_id, 
+                                TaskInstance.task_id == cur_task_id,
+                                TaskInstance.map_index.in_(cur_map_indices)
+                            )
+                        )
+                else:
+                    for cur_map_index, cur_task_ids in dag_map_index_groups.items():
+                        filter_condition.append(
+                            and_(
+                                TaskInstance.dag_id == cur_dag_id,
+                                TaskInstance.run_id == cur_run_id,
+                                TaskInstance.task_id.in_(cur_task_ids),
+                                TaskInstance.map_index == cur_map_index,
+                            )
+                        )       
+                
+        return and_(
+            TaskInstance.map_index == map_index,
+            or_(*filter_condition)
+        )   
 
-        return tuple_in_condition(
-            (TaskInstance.dag_id, TaskInstance.task_id, TaskInstance.run_id, TaskInstance.map_index),
-            (ti.key.primary for ti in tis),
-        )
 
     @classmethod
     def ti_selector_condition(cls, vals: Collection[str | tuple[str, int]]) -> ColumnOperators:

From 76b80e4774d4945df011169f17e3b708a14b6cb1 Mon Sep 17 00:00:00 2001
From: James Ong <zhiyong@hudson-trading.com>
Date: Tue, 18 Oct 2022 17:21:55 +0800
Subject: [PATCH 3/9] fix bug on incorect vairable naming for task_id

---
 airflow/models/taskinstance.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index f412edfa5ac8e..cbaa2eabf4155 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -2440,7 +2440,7 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
                 TaskInstance.map_index == map_index,
                 TaskInstance.task_id == first_task_id,
             )
-        if dag_ids == {dag_id} and run_ids == {run_id} and task_ids == {task_id}:
+        if dag_ids == {dag_id} and run_ids == {run_id} and task_ids == {first_task_id}:
             return and_(
                 TaskInstance.dag_id == dag_id,
                 TaskInstance.run_id == run_id,
@@ -2484,10 +2484,7 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
                             )
                         )       
                 
-        return and_(
-            TaskInstance.map_index == map_index,
-            or_(*filter_condition)
-        )   
+        return or_(*filter_condition)
 
 
     @classmethod

From 0309336d8f91c58170770e00fa42005bded8b876 Mon Sep 17 00:00:00 2001
From: James Ong <zhiyong@hudson-trading.com>
Date: Wed, 19 Oct 2022 13:35:23 +0800
Subject: [PATCH 4/9] add task mapping to existing test

---
 airflow/models/taskinstance.py             |  4 ++--
 tests/sensors/test_external_task_sensor.py | 21 ++++++++++++++++-----
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index cbaa2eabf4155..5be0307ccdb80 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -2415,8 +2415,6 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
         run_id = first.run_id
         map_index = first.map_index
         first_task_id = first.task_id
-        # Common path optimisations: when all TIs are for the same dag_id and run_id, or same dag_id
-        # and task_id -- this can be over 150x faster for huge numbers of TIs (20k+)
 
         # pre-compute the set of dag_id, run_id, map_indices and task_ids
         dag_ids, run_ids, map_indices, task_ids = set(), set(), set(), set()
@@ -2426,6 +2424,8 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
             map_indices.add(t.map_index)
             task_ids.add(t.task_id)
 
+        # Common path optimisations: when all TIs are for the same dag_id and run_id, or same dag_id
+        # and task_id -- this can be over 150x faster for huge numbers of TIs (20k+)
         if dag_ids == {dag_id} and run_ids == {run_id} and map_indices == {map_index}:
             return and_(
                 TaskInstance.dag_id == dag_id,
diff --git a/tests/sensors/test_external_task_sensor.py b/tests/sensors/test_external_task_sensor.py
index 80f538e8680ec..8aad10c67e845 100644
--- a/tests/sensors/test_external_task_sensor.py
+++ b/tests/sensors/test_external_task_sensor.py
@@ -24,6 +24,7 @@
 import unittest
 import zipfile
 from datetime import time, timedelta
+from airflow.decorators import task
 
 import pytest
 
@@ -1116,6 +1117,10 @@ def dag_bag_head_tail():
     dag_bag = DagBag(dag_folder=DEV_NULL, include_examples=False)
 
     with DAG("head_tail", start_date=DEFAULT_DATE, schedule="@daily") as dag:
+        @task
+        def dummy_task(x: int):
+            return x
+
         head = ExternalTaskSensor(
             task_id="head",
             external_dag_id=dag.dag_id,
@@ -1123,7 +1128,8 @@ def dag_bag_head_tail():
             execution_delta=timedelta(days=1),
             mode="reschedule",
         )
-        body = EmptyOperator(task_id="body")
+        
+        body = dummy_task.expand(x=[i for i in range(5)])
         tail = ExternalTaskMarker(
             task_id="tail",
             external_dag_id=dag.dag_id,
@@ -1153,15 +1159,20 @@ def test_clear_overlapping_external_task_marker(dag_bag_head_tail, session):
         )
         session.add(dagrun)
         for task in dag.tasks:
-            ti = TaskInstance(task=task)
-            dagrun.task_instances.append(ti)
+            if task.task_id == "dummy_task":
+                for map_index in range(5):
+                    ti = TaskInstance(task=task, map_index=map_index)
+                    dagrun.task_instances.append(ti)
+            else:
+                ti = TaskInstance(task=task)
+                dagrun.task_instances.append(ti)
             ti.state = TaskInstanceState.SUCCESS
     session.flush()
 
     # The next two lines are doing the same thing. Clearing the first "head" with "Future"
     # selected is the same as not selecting "Future". They should take similar amount of
     # time too because dag.clear() uses visited_external_tis to keep track of visited ExternalTaskMarker.
-    assert dag.clear(start_date=DEFAULT_DATE, dag_bag=dag_bag_head_tail, session=session) == 30
+    assert dag.clear(start_date=DEFAULT_DATE, dag_bag=dag_bag_head_tail, session=session) == 70
     assert (
         dag.clear(
             start_date=DEFAULT_DATE,
@@ -1169,7 +1180,7 @@ def test_clear_overlapping_external_task_marker(dag_bag_head_tail, session):
             dag_bag=dag_bag_head_tail,
             session=session,
         )
-        == 30
+        == 70
     )
 
 

From 9ac4b31b64f3953872638099635b9a85d0bd5037 Mon Sep 17 00:00:00 2001
From: James Ong <zhiyong@hudson-trading.com>
Date: Tue, 1 Nov 2022 15:50:44 +0800
Subject: [PATCH 5/9] fix typo

---
 airflow/models/taskinstance.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 5be0307ccdb80..ad2f6190022a3 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -2449,7 +2449,8 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
             )
 
         filter_condition = []
-        # create 2 groups, 1 grouped by task_id the other by map_index and use the smaller group
+        # create 2 nested groups, both primarily grouped by dag_id and run_id, 
+        # and in the nested group 1 grouped by task_id the other by map_index
         task_id_groups, map_index_groups = defaultdict(lambda: defaultdict(list)), defaultdict(lambda: defaultdict(list))
         for t in tis:
             task_id_groups[(t.dag_id, t.run_id)][t.task_id].append(t.map_index)
@@ -2461,7 +2462,7 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
             for cur_run_id in run_ids:
                 # we compare the group size between task_id and map_index and use the smaller group
                 dag_task_id_groups = task_id_groups[(cur_dag_id, cur_run_id)]
-                dag_map_index_groups = task_id_groups[(cur_dag_id, cur_run_id)]
+                dag_map_index_groups = map_index_groups[(cur_dag_id, cur_run_id)]
 
                 if len(dag_task_id_groups) <= len(dag_map_index_groups):
                     for cur_task_id, cur_map_indices in dag_task_id_groups.items():

From b57f9b0f62a1098611a24c9258433054db4c7102 Mon Sep 17 00:00:00 2001
From: James Ong <zhiyong@hudson-trading.com>
Date: Wed, 2 Nov 2022 21:55:13 +0800
Subject: [PATCH 6/9] remove test since recursively clearing mapped tasks does
 not work now

---
 tests/sensors/test_external_task_sensor.py | 21 +++++----------------
 1 file changed, 5 insertions(+), 16 deletions(-)

diff --git a/tests/sensors/test_external_task_sensor.py b/tests/sensors/test_external_task_sensor.py
index 8aad10c67e845..80f538e8680ec 100644
--- a/tests/sensors/test_external_task_sensor.py
+++ b/tests/sensors/test_external_task_sensor.py
@@ -24,7 +24,6 @@
 import unittest
 import zipfile
 from datetime import time, timedelta
-from airflow.decorators import task
 
 import pytest
 
@@ -1117,10 +1116,6 @@ def dag_bag_head_tail():
     dag_bag = DagBag(dag_folder=DEV_NULL, include_examples=False)
 
     with DAG("head_tail", start_date=DEFAULT_DATE, schedule="@daily") as dag:
-        @task
-        def dummy_task(x: int):
-            return x
-
         head = ExternalTaskSensor(
             task_id="head",
             external_dag_id=dag.dag_id,
@@ -1128,8 +1123,7 @@ def dummy_task(x: int):
             execution_delta=timedelta(days=1),
             mode="reschedule",
         )
-        
-        body = dummy_task.expand(x=[i for i in range(5)])
+        body = EmptyOperator(task_id="body")
         tail = ExternalTaskMarker(
             task_id="tail",
             external_dag_id=dag.dag_id,
@@ -1159,20 +1153,15 @@ def test_clear_overlapping_external_task_marker(dag_bag_head_tail, session):
         )
         session.add(dagrun)
         for task in dag.tasks:
-            if task.task_id == "dummy_task":
-                for map_index in range(5):
-                    ti = TaskInstance(task=task, map_index=map_index)
-                    dagrun.task_instances.append(ti)
-            else:
-                ti = TaskInstance(task=task)
-                dagrun.task_instances.append(ti)
+            ti = TaskInstance(task=task)
+            dagrun.task_instances.append(ti)
             ti.state = TaskInstanceState.SUCCESS
     session.flush()
 
     # The next two lines are doing the same thing. Clearing the first "head" with "Future"
     # selected is the same as not selecting "Future". They should take similar amount of
     # time too because dag.clear() uses visited_external_tis to keep track of visited ExternalTaskMarker.
-    assert dag.clear(start_date=DEFAULT_DATE, dag_bag=dag_bag_head_tail, session=session) == 70
+    assert dag.clear(start_date=DEFAULT_DATE, dag_bag=dag_bag_head_tail, session=session) == 30
     assert (
         dag.clear(
             start_date=DEFAULT_DATE,
@@ -1180,7 +1169,7 @@ def test_clear_overlapping_external_task_marker(dag_bag_head_tail, session):
             dag_bag=dag_bag_head_tail,
             session=session,
         )
-        == 70
+        == 30
     )
 
 

From 891677aa979b289b111c1524cd72535f78aa53bf Mon Sep 17 00:00:00 2001
From: James Ong <zhiyong@hudson-trading.com>
Date: Wed, 16 Nov 2022 14:58:04 -0500
Subject: [PATCH 7/9] fix typing

---
 airflow/models/taskinstance.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index ad2f6190022a3..164e9f2fd3382 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -11,7 +11,7 @@
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF dict[dict[int, list[int]]]
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
@@ -2451,7 +2451,8 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
         filter_condition = []
         # create 2 nested groups, both primarily grouped by dag_id and run_id, 
         # and in the nested group 1 grouped by task_id the other by map_index
-        task_id_groups, map_index_groups = defaultdict(lambda: defaultdict(list)), defaultdict(lambda: defaultdict(list))
+        task_id_groups: dict[tuple, dict[Any, list[Any]]] = defaultdict(lambda: defaultdict(list))
+        map_index_groups: dict[tuple, dict[Any, list[Any]]] = defaultdict(lambda: defaultdict(list))
         for t in tis:
             task_id_groups[(t.dag_id, t.run_id)][t.task_id].append(t.map_index)
             map_index_groups[(t.dag_id, t.run_id)][t.map_index].append(t.task_id)

From 59a355d484a6423cd1e0189792a882ca057d3ced Mon Sep 17 00:00:00 2001
From: James Ong <zhiyong@hudson-trading.com>
Date: Wed, 16 Nov 2022 15:32:19 -0500
Subject: [PATCH 8/9] remove typo

---
 airflow/models/taskinstance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 164e9f2fd3382..0adc510ebfe75 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -11,7 +11,7 @@
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF dict[dict[int, list[int]]]
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.

From 5bddbce5c507c7158764f369c6328bf15aa76c88 Mon Sep 17 00:00:00 2001
From: James Ong <zhiyong@hudson-trading.com>
Date: Wed, 16 Nov 2022 17:10:09 -0500
Subject: [PATCH 9/9] run pre commit

---
 airflow/models/taskinstance.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 0adc510ebfe75..9ec2854be7ded 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -2449,16 +2449,16 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
             )
 
         filter_condition = []
-        # create 2 nested groups, both primarily grouped by dag_id and run_id, 
-        # and in the nested group 1 grouped by task_id the other by map_index
+        # create 2 nested groups, both primarily grouped by dag_id and run_id,
+        # and in the nested group 1 grouped by task_id the other by map_index.
         task_id_groups: dict[tuple, dict[Any, list[Any]]] = defaultdict(lambda: defaultdict(list))
         map_index_groups: dict[tuple, dict[Any, list[Any]]] = defaultdict(lambda: defaultdict(list))
         for t in tis:
             task_id_groups[(t.dag_id, t.run_id)][t.task_id].append(t.map_index)
             map_index_groups[(t.dag_id, t.run_id)][t.map_index].append(t.task_id)
 
-        # this assumes that most dags have dag_id as the largest grouping, followed by run_id. even if its not, this is still
-        # a significant optimization over querying for every single tuple key
+        # this assumes that most dags have dag_id as the largest grouping, followed by run_id. even
+        # if its not, this is still  a significant optimization over querying for every single tuple key
         for cur_dag_id in dag_ids:
             for cur_run_id in run_ids:
                 # we compare the group size between task_id and map_index and use the smaller group
@@ -2470,9 +2470,9 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
                         filter_condition.append(
                             and_(
                                 TaskInstance.dag_id == cur_dag_id,
-                                TaskInstance.run_id == cur_run_id, 
+                                TaskInstance.run_id == cur_run_id,
                                 TaskInstance.task_id == cur_task_id,
-                                TaskInstance.map_index.in_(cur_map_indices)
+                                TaskInstance.map_index.in_(cur_map_indices),
                             )
                         )
                 else:
@@ -2484,10 +2484,9 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
                                 TaskInstance.task_id.in_(cur_task_ids),
                                 TaskInstance.map_index == cur_map_index,
                             )
-                        )       
-                
-        return or_(*filter_condition)
+                        )
 
+        return or_(*filter_condition)
 
     @classmethod
     def ti_selector_condition(cls, vals: Collection[str | tuple[str, int]]) -> ColumnOperators:
