Repository: apache/airflow
File: airflow-core/src/airflow/models/dagrun.py (Python)

Code before:
                # If we get here, it could be that we are moving from non-mapped to mapped
                # after task instance clearing or this ti is not yet expanded. Safe to clear
                # the db references.
                ti.clear_db_references(session=session)
            try:
                expanded_tis, _ = TaskMap.expand_mapped_task(ti.task, self.run_id, session=session)
            except NotMapped:  # Not a mapped task, nothing needed.
                return None
            if expanded_tis:
                return expanded_tis
            return ()

        # Check dependencies.
        expansion_happened = False
        # Set of task ids for which was already done _revise_map_indexes_if_mapped
        revised_map_index_task_ids: set[str] = set()
        for schedulable in itertools.chain(schedulable_tis, additional_tis):
            if TYPE_CHECKING:
                assert isinstance(schedulable.task, Operator)
            old_state = schedulable.state
            if not schedulable.are_dependencies_met(
                session=session,
                dep_context=(
                    dep_context
                    if old_state != TaskInstanceState.UP_FOR_RESCHEDULE
                    else dep_context_with_reschedule
                ),
            ):
                old_states[schedulable.key] = old_state
                continue
            # If schedulable is not yet expanded, try doing it now. This is
            # called in two places: First and ideally in the mini scheduler at
            # the end of LocalTaskJob, and then as an "expansion of last resort"
            # in the scheduler to ensure that the mapped task is correctly
            # expanded before executed. Also see _revise_map_indexes_if_mapped
            # docstring for additional information.
            new_tis = None
            if schedulable.map_index < 0:
                new_tis = _expand_mapped_task_if_needed(schedulable)
                if new_tis is not None:
                    additional_tis.extend(new_tis)
                    expansion_happened = True
            if new_tis is None and schedulable.state in SCHEDULEABLE_STATES:
                # It's enough to revise map index once per task id,
                # checking the map index for each mapped task significantly slows down scheduling
                if schedulable.task.task_id not in revised_map_index_task_ids:
                    ready_tis.extend(
                        self._revise_map_indexes_if_mapped(
                            schedulable.task, dag_version_id=schedulable.dag_version_id, session=session
                        )
                    )

Review comment (suggestion):
I wonder if this conditional logic can be built into `are_dependencies_met` since that function knows the old state.

We can have a dep class similar to `ReadyToRescheduleDep` to contain this logic. The dep class knows about the ti when it evaluates and can handle `ti.state` in `_get_dep_statuses`.

Code after:

        # Check dependencies.
        expansion_happened = False
        # Set of task ids for which was already done _revise_map_indexes_if_mapped
        revised_map_index_task_ids: set[str] = set()
        for schedulable in itertools.chain(schedulable_tis, additional_tis):
            if TYPE_CHECKING:
                assert isinstance(schedulable.task, Operator)
            old_state = schedulable.state
            if not schedulable.are_dependencies_met(session=session, dep_context=dep_context):
                old_states[schedulable.key] = old_state
                continue
            # If schedulable is not yet expanded, try doing it now. This is
            # called in two places: First and ideally in the mini scheduler at
            # the end of LocalTaskJob, and then as an "expansion of last resort"
            # in the scheduler to ensure that the mapped task is correctly
            # expanded before executed. Also see _revise_map_indexes_if_mapped
            # docstring for additional information.
            new_tis = None
            if schedulable.map_index < 0:
                new_tis = _expand_mapped_task_if_needed(schedulable)
                if new_tis is not None:
                    additional_tis.extend(new_tis)
                    expansion_happened = True
            if new_tis is None and schedulable.state in SCHEDULEABLE_STATES:
                # It's enough to revise map index once per task id,
                # checking the map index for each mapped task significantly slows down scheduling
                if schedulable.task.task_id not in revised_map_index_task_ids:
                    ready_tis.extend(
                        self._revise_map_indexes_if_mapped(
                            schedulable.task, dag_version_id=schedulable.dag_version_id, session=session
                        )
                    )
                    revised_map_index_task_ids.add(schedulable.task.task_id)

                # _revise_map_indexes_if_mapped might mark the current task as REMOVED
                # after calculating mapped task length, so we need to re-check
                # the task state to ensure it's still schedulable
                if schedulable.state in SCHEDULEABLE_STATES:
                    ready_tis.append(schedulable)

        # Check if any ti changed state
        tis_filter = TI.filter_for_tis(old_states)
        if tis_filter is not None:
            fresh_tis = session.scalars(select(TI).where(tis_filter)).all()
            changed_tis = any(ti.state != old_states[ti.key] for ti in fresh_tis)

        return ready_tis, changed_tis, expansion_happened

    def _are_premature_tis(
        self,