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

Code before:
                )
            )

            if starved_pools:
                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

Review comment (suggestion):
The window function limits selection to `max_active_tasks` per dag_run without considering currently executing tasks. For example, if a dag_run has 3 tasks already executing and `max_active_tasks=4`, this query would select 4 scheduled tasks (since `row_num` goes 1-4), but only 1 can actually be queued. This means up to 3 extra tasks are retrieved and then rejected at line 701.

Consider incorporating the current task count (`dr_count` from line 505) into the row number filter to improve efficiency:
1. Add `dr_task_concurrency_subquery.c.dr_count` to the columns at line 533
2. Change line 559 to: `where(ranked_query.c.row_num <= (ranked_query.c.dr_max_active_tasks - func.coalesce(ranked_query.c.dr_count, 0)))`

This would reduce unnecessary database retrievals and subsequent filtering in Python code.

Code after:
                    .over(
                        partition_by=[TI.dag_id, TI.run_id],
                        order_by=[-TI.priority_weight, DR.logical_date, TI.map_index],
                    )
                    .label("row_num"),
                    DM.max_active_tasks.label("dr_max_active_tasks"),
                    # Create columns for the order_by checks here for sqlite.
                    TI.priority_weight.label("priority_weight_for_ordering"),
                    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)),
                    )