Fixed #27849 -- Added filtering support to aggregates.

4 files changed, 75 insertions, 14 deletions

diff --git a/django/contrib/postgres/aggregates/statistics.py b/django/contrib/postgres/aggregates/statistics.py
index b9a8ba07c5..19f26ec53c 100644
--- a/django/contrib/postgres/aggregates/statistics.py
+++ b/django/contrib/postgres/aggregates/statistics.py
@@ -8,10 +8,10 @@ __all__ = [
 
 
 class StatAggregate(Aggregate):
-    def __init__(self, y, x, output_field=FloatField()):
+    def __init__(self, y, x, output_field=FloatField(), filter=None):
         if not x or not y:
             raise ValueError('Both y and x must be provided.')
-        super().__init__(y, x, output_field=output_field)
+        super().__init__(y, x, output_field=output_field, filter=filter)
 
     def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
         return super().resolve_expression(query, allow_joins, reuse, summarize)
@@ -22,9 +22,9 @@ class Corr(StatAggregate):
 
 
 class CovarPop(StatAggregate):
-    def __init__(self, y, x, sample=False):
+    def __init__(self, y, x, sample=False, filter=None):
         self.function = 'COVAR_SAMP' if sample else 'COVAR_POP'
-        super().__init__(y, x)
+        super().__init__(y, x, filter=filter)
 
 
 class RegrAvgX(StatAggregate):
@@ -38,8 +38,8 @@ class RegrAvgY(StatAggregate):
 class RegrCount(StatAggregate):
     function = 'REGR_COUNT'
 
-    def __init__(self, y, x):
-        super().__init__(y=y, x=x, output_field=IntegerField())
+    def __init__(self, y, x, filter=None):
+        super().__init__(y=y, x=x, output_field=IntegerField(), filter=filter)
 
     def convert_value(self, value, expression, connection):
         if value is None:
diff --git a/django/db/backends/base/features.py b/django/db/backends/base/features.py
index 7626595741..22c2990e77 100644
--- a/django/db/backends/base/features.py
+++ b/django/db/backends/base/features.py
@@ -229,6 +229,10 @@ class BaseDatabaseFeatures:
     supports_select_difference = True
     supports_slicing_ordering_in_compound = False
 
+    # Does the database support SQL 2003 FILTER (WHERE ...) in aggregate
+    # expressions?
+    supports_aggregate_filter_clause = False
+
     # Does the backend support indexing a TextField?
     supports_index_on_text_field = True
 
diff --git a/django/db/backends/postgresql/features.py b/django/db/backends/postgresql/features.py
index 647fb9dc7f..3c7a7af80f 100644
--- a/django/db/backends/postgresql/features.py
+++ b/django/db/backends/postgresql/features.py
@@ -51,6 +51,10 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     $$ LANGUAGE plpgsql;"""
 
     @cached_property
+    def supports_aggregate_filter_clause(self):
+        return self.connection.pg_version >= 90400
+
+    @cached_property
     def has_select_for_update_skip_locked(self):
         return self.connection.pg_version >= 90500
 
diff --git a/django/db/models/aggregates.py b/django/db/models/aggregates.py
index 2472a24663..b70ea03838 100644
--- a/django/db/models/aggregates.py
+++ b/django/db/models/aggregates.py
@@ -2,8 +2,9 @@
 Classes to represent the definitions of aggregate functions.
 """
 from django.core.exceptions import FieldError
-from django.db.models.expressions import Func, Star
+from django.db.models.expressions import Case, Func, Star, When
 from django.db.models.fields import DecimalField, FloatField, IntegerField
+from django.db.models.query_utils import Q
 
 __all__ = [
     'Aggregate', 'Avg', 'Count', 'Max', 'Min', 'StdDev', 'Sum', 'Variance',
@@ -13,12 +14,36 @@ __all__ = [
 class Aggregate(Func):
     contains_aggregate = True
     name = None
+    filter_template = '%s FILTER (WHERE %%(filter)s)'
+
+    def __init__(self, *args, filter=None, **kwargs):
+        self.filter = filter
+        super().__init__(*args, **kwargs)
+
+    def get_source_fields(self):
+        # Don't return the filter expression since it's not a source field.
+        return [e._output_field_or_none for e in super().get_source_expressions()]
+
+    def get_source_expressions(self):
+        source_expressions = super().get_source_expressions()
+        if self.filter:
+            source_expressions += [self.filter]
+        return source_expressions
+
+    def set_source_expressions(self, exprs):
+        if self.filter:
+            self.filter = exprs.pop()
+        return super().set_source_expressions(exprs)
 
     def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
         # Aggregates are not allowed in UPDATE queries, so ignore for_save
         c = super().resolve_expression(query, allow_joins, reuse, summarize)
+        if c.filter:
+            c.filter = c.filter.resolve_expression(query, allow_joins, reuse, summarize)
         if not summarize:
-            expressions = c.get_source_expressions()
+            # Call Aggregate.get_source_expressions() to avoid
+            # returning self.filter and including that in this loop.
+            expressions = super(Aggregate, c).get_source_expressions()
             for index, expr in enumerate(expressions):
                 if expr.contains_aggregate:
                     before_resolved = self.get_source_expressions()[index]
@@ -36,6 +61,29 @@ class Aggregate(Func):
     def get_group_by_cols(self):
         return []
 
+    def as_sql(self, compiler, connection, **extra_context):
+        if self.filter:
+            if connection.features.supports_aggregate_filter_clause:
+                filter_sql, filter_params = self.filter.as_sql(compiler, connection)
+                template = self.filter_template % extra_context.get('template', self.template)
+                sql, params = super().as_sql(compiler, connection, template=template, filter=filter_sql)
+                return sql, params + filter_params
+            else:
+                copy = self.copy()
+                copy.filter = None
+                condition = When(Q())
+                source_expressions = copy.get_source_expressions()
+                condition.set_source_expressions([self.filter, source_expressions[0]])
+                copy.set_source_expressions([Case(condition)] + source_expressions[1:])
+                return super(Aggregate, copy).as_sql(compiler, connection, **extra_context)
+        return super().as_sql(compiler, connection, **extra_context)
+
+    def _get_repr_options(self):
+        options = super()._get_repr_options()
+        if self.filter:
+            options.update({'filter': self.filter})
+        return options
+
 
 class Avg(Aggregate):
     function = 'AVG'
@@ -52,7 +100,7 @@ class Avg(Aggregate):
             expression = self.get_source_expressions()[0]
             from django.db.backends.oracle.functions import IntervalToSeconds, SecondsToInterval
             return compiler.compile(
-                SecondsToInterval(Avg(IntervalToSeconds(expression)))
+                SecondsToInterval(Avg(IntervalToSeconds(expression), filter=self.filter))
             )
         return super().as_sql(compiler, connection)
 
@@ -62,16 +110,19 @@ class Count(Aggregate):
     name = 'Count'
     template = '%(function)s(%(distinct)s%(expressions)s)'
 
-    def __init__(self, expression, distinct=False, **extra):
+    def __init__(self, expression, distinct=False, filter=None, **extra):
         if expression == '*':
             expression = Star()
+        if isinstance(expression, Star) and filter is not None:
+            raise ValueError('Star cannot be used with filter. Please specify a field.')
         super().__init__(
             expression, distinct='DISTINCT ' if distinct else '',
-            output_field=IntegerField(), **extra
+            output_field=IntegerField(), filter=filter, **extra
         )
 
     def _get_repr_options(self):
-        return {'distinct': self.extra['distinct'] != ''}
+        options = super()._get_repr_options()
+        return dict(options, distinct=self.extra['distinct'] != '')
 
     def convert_value(self, value, expression, connection):
         if value is None:
@@ -97,7 +148,8 @@ class StdDev(Aggregate):
         super().__init__(expression, output_field=FloatField(), **extra)
 
     def _get_repr_options(self):
-        return {'sample': self.function == 'STDDEV_SAMP'}
+        options = super()._get_repr_options()
+        return dict(options, sample=self.function == 'STDDEV_SAMP')
 
     def convert_value(self, value, expression, connection):
         if value is None:
@@ -127,7 +179,8 @@ class Variance(Aggregate):
         super().__init__(expression, output_field=FloatField(), **extra)
 
     def _get_repr_options(self):
-        return {'sample': self.function == 'VAR_SAMP'}
+        options = super()._get_repr_options()
+        return dict(options, sample=self.function == 'VAR_SAMP')
 
     def convert_value(self, value, expression, connection):
         if value is None: