Fix N+1 query pattern in task instance states and count endpoints

[steveahnahn]

Problem

The previous implementation fetched all task instances from the database and then filtered by map_index in Python. For DAGs with mapped tasks containing large map indices, this caused unnecessary database load and memory usage.

Solution

Push the map_index filter to the SQL query, allowing the database to handle filtering efficiently:

• Move map_index filtering from Python to SQL in get_task_instance_states and get_task_instance_count endpoints
• Add map_index parameter to _get_group_tasks helper function to filter at the database level

[SameerMesiah97]

This looks good at first glance. But as the filter based on map index is a new addition (and likely not to be covered by existing tests), I think it would be worth adding a small test to lock in the expected behavior and guard against future regressions.

In particular, it would be a good idea to cover:

1. A mapped task within a task group that produces multiple task instances (i.e. multiple map indices for the same task ID).
2. The default behavior when map_index is not provided, ensuring all relevant task instances are returned (including unmapped tasks).

For (1), you could test both code paths by passing and omitting task_group_id. If feasible (and if doesn't make the test too bulky), you could cover all scenarios in a single parametrized test.

[steveahnahn]

Thanks for the review, the scenarios you mentioned are already covered in airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py:

• test_get_count_mix_of_task_and_task_group_dynamic_task_mapping
• test_get_task_states_mix_of_task_and_task_group_dynamic_task_mapping

The existing parametrized cases cover mapped tasks with multiple map indices, default behavior without map_index, and filtering with task_group_id. These tests will now exercise the new db-level filtering path.

I did end up adding one missed test combination, filtering by map_index without providing task_group_id

[github-actions]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

[steveahnahn]

@SameerMesiah97 it has been some time but tagging for an approval, thanks!

[jscheffl]

Looks good to me! LGTM.
As fixing performance I'd suggest to add this to 3.2.1 scope

[jscheffl]

Yandex-errors are also on main, is unrelated, merging therefore.

[github-actions]

Hi maintainer, this PR was merged without a milestone set.
We've automatically set the milestone to Airflow 3.2.1 based on: backport label targeting v3-2-test
If this milestone is not correct, please update it to the appropriate milestone.

This comment was generated by Milestone Tag Assistant.

[github-actions]

Backport successfully created: v3-2-test

Note: As of Merging PRs targeted for Airflow 3.X
the committer who merges the PR is responsible for backporting the PRs that are bug fixes (generally speaking) to the maintenance branches.

In matter of doubt please ask in #release-management Slack channel.

Status	Branch	Result
✅	v3-2-test