ColtenOuO opened a new pull request, #71185:
URL: https://github.com/apache/airflow/pull/71185
## Summary
`delete_dag` forces SQLAlchemy's `"fetch"` synchronization strategy on every
bulk
delete it issues. That strategy reads the primary key of every deleted row
back from
the database so it can mark matching in-memory objects as deleted — via
`RETURNING`
on backends that support it, or a second full `SELECT` beforehand on those
that don't.
The session holds nothing but the Dag's own `DagModel` row, so there is no
identity
map to synchronize. The keys are read back, matched against an effectively
empty map,
and discarded — at a cost proportional to the Dag's history, once per table
with a
`dag_id` column.
Removing the override restores SQLAlchemy's default `"auto"` strategy, which
evaluates
the criteria in Python against the objects already loaded. Synchronization
still
happens; it just no longer needs a round-trip whose size is set by the row
count
instead of the identity map.
### Before / after
200,000 rows in a table shaped like `task_instance` (its eight declared
indexes plus
the `dag_run` FK), with one ORM object in the session:
| backend | `synchronize_session` | time | peak Python heap | statements |
|---|---|---|---|---|
| PostgreSQL 16 | `"fetch"` (before) | 0.625s | 42.1 MiB | 1 (`DELETE …
RETURNING id`) |
| PostgreSQL 16 | default (after) | **0.055s** | **0.0 MiB** | 1 (`DELETE`) |
| MySQL 8.4 | `"fetch"` (before) | 24.341s | 43.6 MiB | 2 (`SELECT id …` +
`DELETE`) |
| MySQL 8.4 | default (after) | 26.572s | **0.0 MiB** | 1 (`DELETE`) |
The memory saving is the point, and it scales with the Dag: ~211 bytes per
deleted
row, so ~1 GiB of transient heap in the API server for a Dag with five
million task
instances. The time saving is real on PostgreSQL only — InnoDB's own delete
cost
dwarfs the extra `SELECT`, so the MySQL times are noise in both directions.
<details>
<summary>Benchmark script</summary>
```python
"""
Measure what delete_dag's synchronize_session strategy costs.
Usage:
pip install 'sqlalchemy>=2.0.50' psycopg2-binary pymysql
python bench.py postgresql+psycopg2://postgres:[email protected]:5432/test
python bench.py mysql+pymysql://root:[email protected]:3306/test
"""
from __future__ import annotations
import sys
import time
import tracemalloc
from sqlalchemy import (
Column, DateTime, ForeignKeyConstraint, Index, Integer, String,
create_engine, delete, event, select,
)
from sqlalchemy.orm import Session, declarative_base
URL = sys.argv[1] if len(sys.argv) > 1 else
"postgresql+psycopg2://postgres:[email protected]:5432/test"
ROWS = int(sys.argv[2]) if len(sys.argv) > 2 else 200_000
CHUNK = 10_000
Base = declarative_base()
class DagRun(Base):
__tablename__ = "bench_dag_run"
dag_id = Column(String(190), primary_key=True)
run_id = Column(String(190), primary_key=True)
class TaskInstance(Base):
__tablename__ = "bench_task_instance"
id = Column(Integer, primary_key=True)
dag_id = Column(String(190))
task_id = Column(String(190))
run_id = Column(String(190))
map_index = Column(Integer, default=-1)
state = Column(String(20))
pool = Column(String(190))
priority_weight = Column(Integer)
trigger_id = Column(Integer)
last_heartbeat_at = Column(DateTime)
dag_version_id = Column(Integer)
__table_args__ = (
Index("bench_ti_dag_state", "dag_id", "state"),
Index("bench_ti_dag_run", "dag_id", "run_id"),
Index("bench_ti_state", "state"),
Index("bench_ti_state_lkp", "dag_id", "task_id", "run_id", "state"),
Index("bench_ti_pool", "pool", "state", "priority_weight"),
Index("bench_ti_trigger_id", "trigger_id"),
Index("bench_ti_heartbeat", "last_heartbeat_at"),
Index("bench_ti_dag_version_id", "dag_version_id"),
ForeignKeyConstraint(["dag_id", "run_id"], ["bench_dag_run.dag_id",
"bench_dag_run.run_id"]),
)
engine = create_engine(URL)
def seed() -> None:
Base.metadata.drop_all(engine)
Base.metadata.create_all(engine)
with Session(engine) as s:
s.add_all([DagRun(dag_id="d", run_id=f"r{i}") for i in range(200)])
s.commit()
rows = [
{
"id": i, "dag_id": "d", "task_id": f"t{i % 50}", "run_id":
f"r{i % 200}",
"state": "success", "pool": "default_pool",
"priority_weight": 1,
"dag_version_id": 1,
}
for i in range(1, ROWS + 1)
]
for start in range(0, len(rows), CHUNK):
s.bulk_insert_mappings(TaskInstance, rows[start : start + CHUNK])
s.commit()
def measure(strategy):
seed()
statements: list[str] = []
def capture(_conn, _cursor, statement, _parameters, _context,
_executemany):
statements.append(" ".join(statement.split()))
with Session(engine) as s:
s.begin()
s.scalars(select(TaskInstance).limit(1)).all() # the one object
delete_dag loads
stmt = delete(TaskInstance).where(TaskInstance.dag_id == "d")
if strategy is not None:
stmt = stmt.execution_options(synchronize_session=strategy)
event.listen(engine, "before_cursor_execute", capture)
tracemalloc.start()
started = time.monotonic()
s.execute(stmt)
elapsed = time.monotonic() - started
peak = tracemalloc.get_traced_memory()[1]
tracemalloc.stop()
event.remove(engine, "before_cursor_execute", capture)
s.rollback()
return elapsed, peak / 1024 / 1024, statements
print(f"{URL.split('://')[0]}, deleting {ROWS:,} rows\n")
print(f"{'synchronize_session':<22} {'time':>9} {'peak heap':>12} SQL")
for label, strategy in (('"fetch" (before)', "fetch"), ("default (after)",
None)):
elapsed, peak, statements = measure(strategy)
print(f"{label:<22} {elapsed:>8.3f}s {peak:>10.1f} MiB
{len(statements)} statement(s)")
for statement in statements:
print(f"{'':<22} {'':>9} {'':>12} {statement[:88]}")
```
</details>
The regression test asserts on the emitted SQL rather than the query count,
because on
backends with `RETURNING` the count is identical either way. It also pins a
property
worth keeping: the delete loop runs over every mapped model with a `dag_id`,
and
`"auto"` silently falls back to `"fetch"` for criteria it cannot evaluate in
Python, so
the assertions catch a future model that would quietly reintroduce the
round-trip.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes — Claude Code (Opus 5) for writting the regression test and
Benchmark script
--
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]
