Skip to content

avoid unnecessary COUNT queries in slice endpoint#64054

Open
henry3260 wants to merge 1 commit intoapache:mainfrom
henry3260:avoid-count-get-mapped-xcom-by-slice
Open

avoid unnecessary COUNT queries in slice endpoint#64054
henry3260 wants to merge 1 commit intoapache:mainfrom
henry3260:avoid-count-get-mapped-xcom-by-slice

Conversation

@henry3260
Copy link
Contributor

@henry3260 henry3260 commented Mar 22, 2026

Why

The get_mapped_xcom_by_slice endpoint previously issued multiple COUNT queries in several negative slicing scenarios start or stop were negative. Each get_query_count() call performs an extra database query, which is unnecessary and hurts performance for common slice operations.

This change moves the total count calculation to the top of the function and only executes it when both start and stop are explicitly provided and they have different signs (one positive, one negative). In all other cases involving negative indices we can correctly compute the slice bounds without knowing the total length, avoiding the COUNT query entirely.

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

  • Read the 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) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@boring-cyborg boring-cyborg bot added area:API Airflow's REST/HTTP API area:task-sdk labels Mar 22, 2026
@henry3260 henry3260 force-pushed the avoid-count-get-mapped-xcom-by-slice branch from 63db595 to d008c44 Compare March 22, 2026 07:29
@henry3260 henry3260 marked this pull request as ready for review March 22, 2026 08:54
@henry3260 henry3260 force-pushed the avoid-count-get-mapped-xcom-by-slice branch from d008c44 to 0eae1d5 Compare March 22, 2026 10:01
Copy link
Contributor

@SameerMesiah97 SameerMesiah97 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Solid performance improvement. But this needs more coverage for E2E behavior.

pytest.param(slice(-1, -3, -1), id="-1:-3:-1"),
],
)
def test_xcom_get_with_slice(self, client, dag_maker, session, key):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right now, your new test only covers the performance impact of your changes (i.e. when COUNT is called). It doesn’t validate the actual slicing behavior.

We already have E2E coverage in test_xcom_get_with_slice, but the dataset there is very small (effectively length 3), so some of the more complex cases introduced here don’t really get exercised (e.g. mixed-sign slicing that requires COUNT normalization, or early return paths).

I think a simple approach would be:

  1. Modify this test to includde a slightly larger (but still sparse) dataset, for example:

xcom_values = [0, None, 2, None, 4, 5, None, 7, 8, 9]

  1. Add this line right above the assertions:

expected = [v for v in xcom_values if v is not None][key]

  1. Assert against that instead of a hardcoded list.

I would also add a few more targeted cases to cover the different scenarios (same-sign, mixed-sign, early return, etc).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:API Airflow's REST/HTTP API area:task-sdk

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants