Skip to content

Commit

Permalink
Refs #28477 -- Reduced complexity of aggregation over qualify queries.
Browse files Browse the repository at this point in the history
  • Loading branch information
charettes authored and felixxm committed Nov 11, 2022
1 parent 99b4f90 commit a9d2d8d
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 22 deletions.
34 changes: 20 additions & 14 deletions django/db/models/sql/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -447,13 +447,15 @@ def get_aggregation(self, using, added_aggregate_names):
if alias not in added_aggregate_names
}
# Existing usage of aggregation can be determined by the presence of
# selected aggregate and window annotations but also by filters against
# aliased aggregate and windows via HAVING / QUALIFY.
has_existing_aggregation = any(
getattr(annotation, "contains_aggregate", True)
or getattr(annotation, "contains_over_clause", True)
for annotation in existing_annotations.values()
) or any(self.where.split_having_qualify()[1:])
# selected aggregates but also by filters against aliased aggregates.
_, having, qualify = self.where.split_having_qualify()
has_existing_aggregation = (
any(
getattr(annotation, "contains_aggregate", True)
for annotation in existing_annotations.values()
)
or having
)
# Decide if we need to use a subquery.
#
# Existing aggregations would cause incorrect results as
Expand All @@ -468,6 +470,7 @@ def get_aggregation(self, using, added_aggregate_names):
isinstance(self.group_by, tuple)
or self.is_sliced
or has_existing_aggregation
or qualify
or self.distinct
or self.combinator
):
Expand All @@ -494,13 +497,16 @@ def get_aggregation(self, using, added_aggregate_names):
self.model._meta.pk.get_col(inner_query.get_initial_alias()),
)
inner_query.default_cols = False
# Mask existing annotations that are not referenced by
# aggregates to be pushed to the outer query.
annotation_mask = set()
for name in added_aggregate_names:
annotation_mask.add(name)
annotation_mask |= inner_query.annotations[name].get_refs()
inner_query.set_annotation_mask(annotation_mask)
if not qualify:
# Mask existing annotations that are not referenced by
# aggregates to be pushed to the outer query unless
# filtering against window functions is involved as it
# requires complex realising.
annotation_mask = set()
for name in added_aggregate_names:
annotation_mask.add(name)
annotation_mask |= inner_query.annotations[name].get_refs()
inner_query.set_annotation_mask(annotation_mask)

relabels = {t: "subquery" for t in inner_query.alias_map}
relabels[None] = "subquery"
Expand Down
22 changes: 14 additions & 8 deletions tests/expressions_window/tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
)
from django.db.models.lookups import Exact
from django.test import SimpleTestCase, TestCase, skipUnlessDBFeature
from django.test.utils import CaptureQueriesContext

from .models import Classification, Detail, Employee, PastEmployeeDepartment

Expand Down Expand Up @@ -1157,16 +1158,21 @@ def test_limited_filter(self):
)

def test_filter_count(self):
self.assertEqual(
Employee.objects.annotate(
department_salary_rank=Window(
Rank(), partition_by="department", order_by="-salary"
with CaptureQueriesContext(connection) as ctx:
self.assertEqual(
Employee.objects.annotate(
department_salary_rank=Window(
Rank(), partition_by="department", order_by="-salary"
)
)
.filter(department_salary_rank=1)
.count(),
5,
)
.filter(department_salary_rank=1)
.count(),
5,
)
self.assertEqual(len(ctx.captured_queries), 1)
sql = ctx.captured_queries[0]["sql"].lower()
self.assertEqual(sql.count("select"), 3)
self.assertNotIn("group by", sql)

@skipUnlessDBFeature("supports_frame_range_fixed_distance")
def test_range_n_preceding_and_following(self):
Expand Down

0 comments on commit a9d2d8d

Please sign in to comment.