Repository: apache/airflow
File: airflow-core/src/airflow/jobs/scheduler_job_runner.py (Python)

Code before:
                query = query.where(TI.pool.not_in(starved_pools))

            if starved_dags:
                query = query.where(TI.dag_id.not_in(starved_dags))

            if starved_tasks:
                query = query.where(tuple_(TI.dag_id, TI.task_id).not_in(starved_tasks))

            if starved_tasks_task_dagrun_concurrency:
                query = query.where(
                    tuple_(TI.dag_id, TI.run_id, TI.task_id).not_in(starved_tasks_task_dagrun_concurrency)
                )

            query = query.limit(max_tis)

            timer = Stats.timer("scheduler.critical_section_query_duration")
            timer.start()

            try:
                locked_query = with_row_locks(query, of=TI, session=session, skip_locked=True)
                task_instances_to_examine = session.scalars(locked_query).all()

                self.log.debug("Length of the tis to examine is %d", len(task_instances_to_examine))
                self.log.debug(
                    "TaskInstance selection is: %s",
                    dict(Counter(ti.dag_id for ti in task_instances_to_examine)),
                )

                timer.stop(send=True)
            except OperationalError as e:
                timer.stop(send=False)
                raise e

            # TODO[HA]: This was wrong before anyway, as it only looked at a sub-set of dags, not everything.
            # Stats.gauge('scheduler.tasks.pending', len(task_instances_to_examine))

            if not task_instances_to_examine:
                self.log.debug("No tasks to consider for execution.")
                break

            # Put one task instance on each line
            task_instance_str = "\n".join(f"\t{x!r}" for x in task_instances_to_examine)
            self.log.info("%s tasks up for execution:\n%s", len(task_instances_to_examine), task_instance_str)

            dag_id_to_team_name: dict[str, str | None] = {}
            if conf.getboolean("core", "multi_team"):
                # Batch query to resolve team names for all DAG IDs to optimize performance
                # Instead of individual queries in _try_to_load_executor(), resolve all team names upfront
                unique_dag_ids = {ti.dag_id for ti in task_instances_to_examine}
                dag_id_to_team_name = self._get_team_names_for_dag_ids(unique_dag_ids, session)
                self.log.debug(

Review comment (suggestion):
Not sure what happened to my earlier comment so re-adding: This new query is missing `.options(selectinload(TI.dag_model))` and `.with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")` which were on the original query. Every `ti.dag_model` access later will trigger a separate query - with 50 TIs that's 50+ extra queries per loop.

Code after:
                    DR.logical_date.label("logical_date_for_ordering"),
                    TI.map_index.label("map_index_for_ordering"),
                )
            ).subquery()

            # Select only rows where row_number <= max_active_tasks.
            query = (
                select(TI)
                .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
                .select_from(ranked_query)
                .join(
                    TI,
                    (TI.dag_id == ranked_query.c.dag_id)
                    & (TI.task_id == ranked_query.c.task_id)
                    & (TI.run_id == ranked_query.c.run_id)
                    & (TI.map_index == ranked_query.c.map_index),
                )
                .where(ranked_query.c.row_num <= ranked_query.c.dr_max_active_tasks)
                # Add the order_by columns from the ranked query for sqlite.
                .order_by(
                    -ranked_query.c.priority_weight_for_ordering,
                    ranked_query.c.logical_date_for_ordering,
                    ranked_query.c.map_index_for_ordering,
                )
                .options(selectinload(TI.dag_model))
            )

            query = query.limit(max_tis)

            timer = Stats.timer("scheduler.critical_section_query_duration")
            timer.start()

            try:
                locked_query = with_row_locks(query, of=TI, session=session, skip_locked=True)
                task_instances_to_examine = session.scalars(locked_query).all()

                if self.log.isEnabledFor(logging.DEBUG):
                    self.log.debug("Length of the tis to examine is %d", len(task_instances_to_examine))
                    self.log.debug(
                        "TaskInstance selection is: %s",
                        dict(Counter(ti.dag_id for ti in task_instances_to_examine)),
                    )

                timer.stop(send=True)
            except OperationalError as e:
                timer.stop(send=False)
                raise e

            # TODO[HA]: This was wrong before anyway, as it only looked at a sub-set of dags, not everything.
            # Stats.gauge('scheduler.tasks.pending', len(task_instances_to_examine))