caoterry opened a new pull request, #73983:
URL: https://github.com/apache/airflow/pull/73983
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
Add two indexes on `asset_partition_dag_run`, matching the two queries the
partitioned-asset scheduling path runs on it:
- `idx_apdr_target_dag_id_partition_key_id (target_dag_id, partition_key,
id)` — `AssetManager._get_or_create_apdr` runs
`WHERE partition_key = ? AND target_dag_id = ? ORDER BY id DESC LIMIT 1`
for **every emitted partition key** (inside the
task-success request, under the asset row lock).
- `idx_apdr_created_dag_run_id_created_at_id (created_dag_run_id,
created_at, id)` —
`SchedulerJobRunner._create_dagruns_for_partitioned_asset_dags`
selects pending rows (`created_dag_run_id IS NULL`) ordered by
`created_at, id` on **every scheduler loop**.
Neither query had an index, so both were sequential scans of a table that
grows by one row per partition key per consumer and is
only trimmed by cascade when `dag_run` rows are cleaned.
**Why / measurements** (Airflow 3.3.2, Postgres 16, one asset, one
`PartitionedAssetTimetable` consumer with `IdentityMapper`,
100k keys emitted from 200 tasks of 500 keys each):
- The task-success request registering 500 keys took 3.6 s with an empty
table and 7.0 s at ~60k rows (`EXPLAIN`: 1,843 shared
buffers per key lookup), i.e. past the default 5 s `[workers]
execution_api_timeout`, so the Task SDK client started retrying.
Creating the composite index online brought the same request back to 3.4 s
(4 buffers per lookup) and the retries stopped.
- With the indexes in place (and two schedulers), 100k partition runs were
created and completed in 31 minutes on the same
machine; without them the run-completion rate was 3–10 runs/s (~3 h
projected).
Reproduction, harness and charts:
https://github.com/caoterry/airflow-100k-partitions (REPORT.md §4.2–4.3,
`docs/charts.md` §2).
**Notes for reviewers**
- Plain composite indexes rather than a partial index on `created_dag_run_id
IS NULL`, so the definition is identical on
Postgres, MySQL and SQLite (same approach as
`idx_asset_event_asset_id_partition_key`, migration 0127). Both indexed string
columns are `StringID` (250 chars), within MySQL's key-length limit.
- Index-only migration (`batch_alter_table`, no table rebuild);
`_REVISION_HEADS_MAP` and `migrations-ref.rst` updated.
Verified locally: `tests/unit/utils/test_db.py` (ORM vs. migrations) and
the migration-pattern tests, plus a SQLite
migrate → downgrade → migrate round trip. Newsfragment follows in a
separate commit once the PR number is known.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (Claude Code — Claude Fable 5.1; the change and this description
were drafted with it and verified locally by the author)
Generated-by: Claude Code following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)**
for more information. Note: commit author/co-author name and email in commits
become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal
([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
