You are viewing a plain text version of this content. The canonical link for it is here.

Posted to commits@airflow.apache.org by "kaxil (via GitHub)" <gi...@apache.org> on 2026/01/22 17:18:07 UTC

[PR] Optimize RTIF delete_old_records query for DAGs with many mapped tasks [airflow]

kaxil opened a new pull request, #60951:
URL: https://github.com/apache/airflow/pull/60951

   Optimizes `RenderedTaskInstanceFields.delete_old_records()` to avoid scanning the RTIF table when determining which records to keep. This fixes scheduler crashes caused by slow RTIF cleanup queries on DAGs with many dynamically mapped tasks.
   
   **Problem**: For DAGs with 3k+ mapped task instances, the delete query was scanning ~100k RTIF records with a complex `NOT EXISTS` subquery containing a join, causing statement timeouts and high CPU utilization that prevented the scheduler from heartbeating.
   
   This PR simplifies the query to only use the `dag_run` table for finding recent run_ids, avoiding the expensive RTIF table scan entirely.
   
   ---
   
   ### Before (slow - scans RTIF table with join)
   
   ```sql
   DELETE FROM rendered_task_instance_fields 
   WHERE dag_id = :dag_id AND task_id = :task_id 
   AND NOT EXISTS (
     SELECT 1 FROM (
       SELECT DISTINCT rtif.dag_id, rtif.task_id, rtif.run_id, dr.logical_date 
       FROM rendered_task_instance_fields rtif
       JOIN dag_run dr ON rtif.dag_id = dr.dag_id AND rtif.run_id = dr.run_id 
       WHERE rtif.dag_id = :dag_id AND rtif.task_id = :task_id 
       ORDER BY dr.logical_date DESC LIMIT 30
     ) AS keep
     WHERE keep.dag_id = rtif.dag_id AND keep.task_id = rtif.task_id AND keep.run_id = rtif.run_id
   )
   ```
   
   ### After (fast - only queries dag_run table)
   
   ```sql
   DELETE FROM rendered_task_instance_fields 
   WHERE dag_id = :dag_id AND task_id = :task_id 
   AND run_id NOT IN (
     SELECT run_id FROM dag_run 
     WHERE dag_id = :dag_id 
     ORDER BY run_after DESC LIMIT 30
   )
   ```
   
   ## Benchmark Results (PostgreSQL)
   
   **Test Setup**: 100 DAG runs x 50 mapped TIs = 5,000 RTIF records, keeping 30 most recent runs
   
   | Metric | Old Query | New Query | Improvement |
   |--------|-----------|-----------|-------------|
   | Execution Time | 233 ms | 5.5 ms | **42x faster** |
   | Buffer Hits | 14,618 | 113 | 129x fewer |
   
   ### Query Plan Comparison
   
   **Old Query** - Nested loop with expensive heap fetches:
   ```
   Delete on rendered_task_instance_fields (actual time=233.12..233.12 rows=0 loops=1)
     Buffers: shared hit=14618
     ->  Nested Loop Anti Join (actual time=51.43..232.91 rows=3500 loops=1)
           Rows Removed by Join Filter: 126750
           ->  Index Scan using rendered_task_instance_fields_pkey ...
           ->  Subquery Scan on anon_1 (actual time=0.010..0.065 rows=26 loops=5000)
                 ->  Nested Loop (joins RTIF to dag_run)
                       Rows Removed by Join Filter: 495000
   Execution Time: 233.44 ms
   ```
   
   **New Query** - Simple index scan with hashed subplan:
   ```
   Delete on rendered_task_instance_fields (actual time=5.12..5.12 rows=0 loops=1)
     Buffers: shared hit=113
     ->  Index Scan using rendered_task_instance_fields_pkey ...
           Filter: (NOT (hashed SubPlan 1))
           SubPlan 1
             ->  Limit (cost=0.14..8.15 rows=1 width=524)
                   ->  Index Scan Backward using idx_dag_run_run_after on dag_run
   Execution Time: 5.55 ms
   ```
   
   ## Behavioral Change
   
   The semantics change slightly:
   
   | Scenario | Old Behavior | New Behavior |
   |----------|-------------|--------------|
   | Normal task (runs every dag run) | Keep last 30 runs | Keep last 30 runs |
   | Sparse task (runs occasionally) | Keep last 30 executions | Keep records from last 30 dag runs |
   
   This imo is acceptable because of the following, but looking forward for what others think too:
   1. Most tasks run every DAG run (not sparse)
   2. The purpose is to limit storage, not preserve specific historical records
   3. The performance improvement (42x) far outweighs this edge case
   
   ## Additional Changes
   
   - Uses `run_after` instead of `logical_date` for ordering (since `logical_date` can be NULL for manual runs)
   - Updated docstrings, config description, and documentation
   
   ## How to Reproduce
   
   To reproduce the benchmark, save this script and run inside breeze with PostgreSQL:
   
   ```bash
   breeze run --backend postgres python dev/rtif_benchmark.py
   ```
   
   <details>
   <summary>Benchmark Script</summary>
   
   ```python
   """
   Benchmark: RTIF delete_old_records query performance.
   Run: breeze run --backend postgres python dev/rtif_benchmark.py
   """
   from __future__ import annotations
   
   import time
   from datetime import datetime, timedelta, timezone
   
   from sqlalchemy import delete, exists, select, text
   
   from airflow.models.dagrun import DagRun
   from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF
   from airflow.utils.session import create_session
   from airflow.utils.state import DagRunState
   
   DAG_ID = "benchmark_dag"
   TASK_ID = "benchmark_task"
   NUM_RUNS = 100
   MAP_INDEX_COUNT = 50
   NUM_TO_KEEP = 30
   
   
   def setup_test_data():
       from airflow.utils.db import initdb
       initdb()
   
       with create_session() as session:
           session.execute(text(f"DELETE FROM rendered_task_instance_fields WHERE dag_id = '{DAG_ID}'"))
           session.execute(text(f"DELETE FROM dag_run WHERE dag_id = '{DAG_ID}'"))
           session.execute(text("ALTER TABLE rendered_task_instance_fields DROP CONSTRAINT IF EXISTS rtif_ti_fkey"))
           session.commit()
   
           base_date = datetime(2024, 1, 1, tzinfo=timezone.utc)
           for run_num in range(NUM_RUNS):
               run_id = f"run_{run_num:04d}"
               logical_date = base_date + timedelta(days=run_num)
               dr = DagRun(dag_id=DAG_ID, run_id=run_id, logical_date=logical_date,
                           data_interval=(logical_date, logical_date + timedelta(days=1)),
                           run_type="scheduled", state=DagRunState.SUCCESS)
               session.add(dr)
               session.flush()
   
               for map_idx in range(MAP_INDEX_COUNT):
                   session.execute(text("""
                       INSERT INTO rendered_task_instance_fields 
                       (dag_id, task_id, run_id, map_index, rendered_fields)
                       VALUES (:dag_id, :task_id, :run_id, :map_index, :fields)
                   """), {"dag_id": DAG_ID, "task_id": TASK_ID, "run_id": run_id, 
                          "map_index": map_idx, "fields": '{"cmd": "echo test"}'})
           session.commit()
       print(f"Created {NUM_RUNS} runs x {MAP_INDEX_COUNT} TIs = {NUM_RUNS * MAP_INDEX_COUNT} RTIF records")
   
   
   def old_query():
       """Old query with NOT EXISTS and JOIN."""
       tis_to_keep = (
           select(RTIF.dag_id, RTIF.task_id, RTIF.run_id, DagRun.logical_date)
           .where(RTIF.dag_id == DAG_ID, RTIF.task_id == TASK_ID)
           .join(RTIF.dag_run).distinct()
           .order_by(DagRun.logical_date.desc()).limit(NUM_TO_KEEP)
       ).subquery()
       return delete(RTIF).where(
           RTIF.dag_id == DAG_ID, RTIF.task_id == TASK_ID,
           ~exists(1).where(tis_to_keep.c.dag_id == RTIF.dag_id, 
                            tis_to_keep.c.task_id == RTIF.task_id, 
                            tis_to_keep.c.run_id == RTIF.run_id))
   
   
   def new_query():
       """New query with NOT IN on dag_run only."""
       run_ids = select(DagRun.run_id).where(DagRun.dag_id == DAG_ID).order_by(
           DagRun.run_after.desc()).limit(NUM_TO_KEEP)
       return delete(RTIF).where(
           RTIF.dag_id == DAG_ID, RTIF.task_id == TASK_ID,
           RTIF.run_id.not_in(run_ids.scalar_subquery()))
   
   
   def benchmark(session, stmt, label):
       start = time.perf_counter()
       result = session.execute(stmt)
       elapsed = time.perf_counter() - start
       print(f"{label}: {result.rowcount} rows in {elapsed*1000:.2f} ms")
       return elapsed
   
   
   def main():
       setup_test_data()
       with create_session() as session:
           old_time = benchmark(session, old_query(), "OLD")
           session.commit()
   
       setup_test_data()
       with create_session() as session:
           new_time = benchmark(session, new_query(), "NEW")
           session.commit()
   
       print(f"\nSpeedup: {old_time/new_time:.1f}x faster")
   
   
   if __name__ == "__main__":
       main()
   ```
   
   </details>


-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org