Skip to content

Fix N+1 XCom fetches when iterating a mapped task's lazy sequence - #40

Open
nancy-wang-dd wants to merge 1 commit into
DataDog:mainfrom
nancy-wang-dd:nancy-wang-dd/batch-lazy-xcom-slice-fetch
Open

Fix N+1 XCom fetches when iterating a mapped task's lazy sequence#40
nancy-wang-dd wants to merge 1 commit into
DataDog:mainfrom
nancy-wang-dd:nancy-wang-dd/batch-lazy-xcom-slice-fetch

Conversation

@nancy-wang-dd

@nancy-wang-dd nancy-wang-dd commented Jul 16, 2026

Copy link
Copy Markdown

Forward iteration over LazyXComSequence (task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py) issued one Execution API HTTP request per item. A task iterating a mapped upstream task's full XCom output — directly, or via zip()/concat() on XComArgs — fired a sequential HTTP call per item, which Datadog APM's "Sequential API Calls" detector flagged under airflow-worker's execute_workload span.

Jira Ticket

What

LazyXComIterator.next now fetches items in chunks (50 at a time) using the existing GetXComSequenceSlice / get_sequence_slice endpoint, instead of one GetXComSequenceItem request per item. This mirrors the batching already used by BaseXCom.get_all().

How it behaves

  • Forward iteration (dir=1, the common case): buffers a chunk via one slice request, refilling as the buffer is exhausted.
  • Backward iteration (dir=-1, rare): unchanged — still one per-item request, since it wasn't worth batching. - Single-item (seq[i]) and slice (seq[a:b]) a
  • No API contract changes — reuses the slice endpoint that already existed and is used elsewhere.

Test plan

  • Rewrote test_iter to assert the new batched-slice call pattern (previously asserted the old per-item calls). - Added test_iter_backward_still_uses_per_ite had no prior coverage.
  • Added test_iter_preserves_falsy_values — guards the buffer-emptiness check against falsy-but-real items (0, "", False, None).
  • Ran the full task-sdk test suite; 19 pre-existing failures confirmed unrelated by reproducing them identically on a git stash'd clean tree.
  • Ran an independent adversarial review pass (buffer boundary math, slice-vs-item deserialization equivalence,
    concurrency/shared-state, error-handling divessues found.

Was generative AI tooling used to co-author this PR?

  • Yes — Claude Code (Sonnet 5)

Generated-by: Claude Code (Sonnet 5) following the guidelines (https://github.com/apache/airflow/blob/main/requests.rst#gen-ai-assisted-contributions)

Forward iteration over LazyXComSequence previously issued one HTTP
request per item via the Execution API. A task iterating a mapped
upstream task's full output (e.g. via zip()/concat() on XComArgs)
would fire a sequential HTTP call per item instead of fetching in
bulk. Forward iteration now fetches items in chunks using the
existing slice endpoint, matching the batching already used by
BaseXCom.get_all(). Backward iteration and single-item/slice access
are unchanged.
@nancy-wang-dd
nancy-wang-dd marked this pull request as ready for review July 16, 2026 20:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant