You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "henry3260 (via GitHub)" <gi...@apache.org> on 2026/03/22 05:47:45 UTC
[PR] avoid unnecessary COUNT queries in slice endpoint [airflow]
henry3260 opened a new pull request, #64054:
URL: https://github.com/apache/airflow/pull/64054
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] avoid unnecessary COUNT queries in slice endpoint [airflow]
Posted by "henry3260 (via GitHub)" <gi...@apache.org>.
henry3260 commented on PR #64054:
URL: https://github.com/apache/airflow/pull/64054#issuecomment-4299058025
Is someone can review my PR. Thanks!
--
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
Re: [PR] avoid unnecessary COUNT queries in slice endpoint [airflow]
Posted by "henry3260 (via GitHub)" <gi...@apache.org>.
henry3260 commented on PR #64054:
URL: https://github.com/apache/airflow/pull/64054#issuecomment-4318310413
I think I will close this PR. Initially, my goal was to reduce the count queries, but I realized that the resulting slicing logic becomes overly complicated and hard to read. The performance gain isn't worth the cost of maintainability. Thanks for the reviews!
--
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
Re: [PR] avoid unnecessary COUNT queries in slice endpoint [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #64054:
URL: https://github.com/apache/airflow/pull/64054#discussion_r3126579621
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py:
##########
@@ -323,6 +324,41 @@ def test_xcom_get_slice_accepts_include_prior_dates(
assert set(response.json()) == set(expected_xcoms)
+ @pytest.mark.parametrize(
+ ("start", "stop", "expected_count_calls"),
+ [
+ pytest.param(None, None, 0, id="no_slice"),
+ pytest.param(5, None, 0, id="positive_start_only"),
+ pytest.param(None, -5, 0, id="negative_stop_only"),
+ pytest.param(-5, None, 0, id="negative_start_only"),
+ pytest.param(5, -3, 1, id="positive_start_negative_stop"),
+ pytest.param(-5, 10, 1, id="negative_start_positive_stop"),
+ pytest.param(-5, -10, 0, id="both_negative_same_sign"),
+ ],
+ )
+ def test_xcom_slice_count_query_called_at_most_once(
+ self, client, dag_maker, session, start, stop, expected_count_calls
+ ):
+ """Test that get_query_count is called at most once (and only when necessary)
+ for slice requests with negative indices.
+
+ This validates the performance optimization in get_mapped_xcom_by_slice().
+ """
+ with dag_maker(dag_id="dag"):
+ EmptyOperator(task_id="task")
+ dag_maker.create_dagrun(run_id="runid")
+
+ with patch("airflow.api_fastapi.execution_api.routes.xcoms.get_query_count") as mock_count:
+ qs = {}
+ if start is not None:
+ qs["start"] = start
+ if stop is not None:
+ qs["stop"] = stop
+
+ client.get(f"/execution/xcoms/dag/runid/task/test_key/slice?{urllib.parse.urlencode(qs)}")
+
+ assert mock_count.call_count == expected_count_calls
Review Comment:
Two gaps that weaken this test:
1. `mock_count` has no `return_value` configured, so it returns a `MagicMock`. For the `positive_start_negative_stop` and `negative_start_positive_stop` cases, that propagates through `stop = total_count + stop` (or `stop - total_count`) and into `query.slice(...)`, where SQLAlchemy may raise depending on version.
2. `client.get(...)` is never inspected, so a 500 from the slice path still passes the assertion -- the test only protects the count-call optimization, not the code that actually uses `total_count`.
Suggest `mock_count.return_value = 0` to keep the arithmetic coherent, plus `assert response.status_code == 200` so this test catches regressions in the slice branch too.
--
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
Re: [PR] avoid unnecessary COUNT queries in slice endpoint [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #64054:
URL: https://github.com/apache/airflow/pull/64054#discussion_r3126577719
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -182,29 +201,29 @@ def get_mapped_xcom_by_slice(
query = query.offset(-stop)
elif start >= 0:
query = query.order_by(XComModel.map_index.asc())
- if (stop := params.stop) is None:
+ if stop is None:
if step >= 0:
query = query.offset(start)
else:
query = query.limit(start + 1)
else:
if stop < 0:
- stop += get_query_count(query, session=session)
+ stop = total_count + stop if total_count is not None else stop
Review Comment:
The `else stop` branch here is unreachable. At this point we're inside `elif start >= 0:` (so `start` is set), inside the `else:` at line 209 (so `stop` is set), and inside `if stop < 0:` -- which is exactly the `bool(start >= 0) != bool(stop >= 0)` condition at line 179 that sets `total_count`. So `total_count is not None` is always true here.
The defensive fallback just makes the read harder. You can drop it:
```python
stop += total_count
```
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -182,29 +201,29 @@ def get_mapped_xcom_by_slice(
query = query.offset(-stop)
elif start >= 0:
query = query.order_by(XComModel.map_index.asc())
- if (stop := params.stop) is None:
+ if stop is None:
if step >= 0:
query = query.offset(start)
else:
query = query.limit(start + 1)
else:
if stop < 0:
- stop += get_query_count(query, session=session)
+ stop = total_count + stop if total_count is not None else stop
if step >= 0:
query = query.slice(start, stop)
else:
query = query.slice(stop + 1, start + 1)
else:
query = query.order_by(XComModel.map_index.desc())
step = -step
- if (stop := params.stop) is None:
+ if stop is None:
if step > 0:
query = query.offset(-start - 1)
else:
query = query.limit(-start)
else:
if stop >= 0:
- stop -= get_query_count(query, session=session)
+ stop = stop - total_count if total_count is not None else stop
Review Comment:
Same as the symmetric branch at line 211: we're inside `else:` (start < 0), inside `else:` at line 224 (stop is set), and inside `if stop >= 0:`. That matches the `bool(start >= 0) != bool(stop >= 0)` guard at line 179, so `total_count is not None` is always true here and the `else stop` fallback is dead.
Simpler:
```python
stop -= total_count
```
--
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
Re: [PR] avoid unnecessary COUNT queries in slice endpoint [airflow]
Posted by "henry3260 (via GitHub)" <gi...@apache.org>.
henry3260 commented on code in PR #64054:
URL: https://github.com/apache/airflow/pull/64054#discussion_r3130758630
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -182,29 +201,29 @@ def get_mapped_xcom_by_slice(
query = query.offset(-stop)
elif start >= 0:
query = query.order_by(XComModel.map_index.asc())
- if (stop := params.stop) is None:
+ if stop is None:
if step >= 0:
query = query.offset(start)
else:
query = query.limit(start + 1)
else:
if stop < 0:
- stop += get_query_count(query, session=session)
+ stop = total_count + stop if total_count is not None else stop
Review Comment:
> The `else stop` branch here is unreachable. At this point we're inside `elif start >= 0:` (so `start` is set), inside the `else:` at line 209 (so `stop` is set), and inside `if stop < 0:` -- which is exactly the `bool(start >= 0) != bool(stop >= 0)` condition at line 179 that sets `total_count`. So `total_count is not None` is always true here.
>
> The defensive fallback just makes the read harder. You can drop it:
>
> ```python
> stop += total_count
> ```
Good Catch!
--
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
Re: [PR] avoid unnecessary COUNT queries in slice endpoint [airflow]
Posted by "henry3260 (via GitHub)" <gi...@apache.org>.
henry3260 closed pull request #64054: avoid unnecessary COUNT queries in slice endpoint
URL: https://github.com/apache/airflow/pull/64054
--
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
Re: [PR] avoid unnecessary COUNT queries in slice endpoint [airflow]
Posted by "henry3260 (via GitHub)" <gi...@apache.org>.
henry3260 commented on code in PR #64054:
URL: https://github.com/apache/airflow/pull/64054#discussion_r3130865768
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py:
##########
@@ -323,6 +324,41 @@ def test_xcom_get_slice_accepts_include_prior_dates(
assert set(response.json()) == set(expected_xcoms)
+ @pytest.mark.parametrize(
+ ("start", "stop", "expected_count_calls"),
+ [
+ pytest.param(None, None, 0, id="no_slice"),
+ pytest.param(5, None, 0, id="positive_start_only"),
+ pytest.param(None, -5, 0, id="negative_stop_only"),
+ pytest.param(-5, None, 0, id="negative_start_only"),
+ pytest.param(5, -3, 1, id="positive_start_negative_stop"),
+ pytest.param(-5, 10, 1, id="negative_start_positive_stop"),
+ pytest.param(-5, -10, 0, id="both_negative_same_sign"),
+ ],
+ )
+ def test_xcom_slice_count_query_called_at_most_once(
+ self, client, dag_maker, session, start, stop, expected_count_calls
+ ):
+ """Test that get_query_count is called at most once (and only when necessary)
+ for slice requests with negative indices.
+
+ This validates the performance optimization in get_mapped_xcom_by_slice().
+ """
+ with dag_maker(dag_id="dag"):
+ EmptyOperator(task_id="task")
+ dag_maker.create_dagrun(run_id="runid")
+
+ with patch("airflow.api_fastapi.execution_api.routes.xcoms.get_query_count") as mock_count:
+ qs = {}
+ if start is not None:
+ qs["start"] = start
+ if stop is not None:
+ qs["stop"] = stop
+
+ client.get(f"/execution/xcoms/dag/runid/task/test_key/slice?{urllib.parse.urlencode(qs)}")
+
+ assert mock_count.call_count == expected_count_calls
Review Comment:
> Two gaps that weaken this test:
>
> 1. `mock_count` has no `return_value` configured, so it returns a `MagicMock`. For the `positive_start_negative_stop` and `negative_start_positive_stop` cases, that propagates through `stop = total_count + stop` (or `stop - total_count`) and into `query.slice(...)`, where SQLAlchemy may raise depending on version.
> 2. `client.get(...)` is never inspected, so a 500 from the slice path still passes the assertion -- the test only protects the count-call optimization, not the code that actually uses `total_count`.
>
> Suggest `mock_count.return_value = 0` to keep the arithmetic coherent, plus `assert response.status_code == 200` so this test catches regressions in the slice branch too.
Yes, you are right. I didn't notice that. Applied!
--
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
Re: [PR] avoid unnecessary COUNT queries in slice endpoint [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #64054:
URL: https://github.com/apache/airflow/pull/64054#discussion_r2971481370
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py:
##########
Review Comment:
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]`
2) Add this line right above the assertions:
`expected = [v for v in xcom_values if v is not None][key]`
3) 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).
--
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