Merge pull request smartfile#4 from klem4/feature

поддержка aggregate(Sum)

django_sphinx_db/backend/models.py

@@ -8,7 +8,7 @@
 from django.db.models.query import QuerySet
 from django.utils.log import getLogger
 from django_sphinx_db.backend.sphinx.compiler import SphinxWhereNode, SphinxExtraWhere, SphinxQLCompiler
-
+from django_sphinx_db.backend.sphinx import aggregates as sphinx_aggregates
 
 def sphinx_escape(value):
     if type(value) not in (str, unicode):
@@ -42,6 +42,8 @@ class SphinxQuery(Query):
     _clonable = ('options', 'match', 'group_limit', 'group_order_by',
                  'with_meta')
 
+    aggregates_module = sphinx_aggregates
+
     def __init__(self, *args, **kwargs):
         kwargs.setdefault('where', SphinxWhereNode)
         super(SphinxQuery, self).__init__(*args, **kwargs)

django_sphinx_db/backend/sphinx/aggregates.py

@@ -0,0 +1,84 @@
+# coding: utf-8
+from django.db.models.fields import IntegerField, FloatField
+
+ordinal_aggregate_field = IntegerField()
+computed_aggregate_field = FloatField()
+
+from django.db.models.sql.aggregates import Aggregate as base_aggregate
+
+
+class Aggregate(base_aggregate):
+    def as_sql(self, qn, connection):
+        if hasattr(self.col, 'as_sql'):
+            field_name = self.col.as_sql(qn, connection)
+        elif isinstance(self.col, (list, tuple)):
+            # единственное отличие от базового класса
+            # если прилетает tuple с алиасом имени таблицы
+            # не учитываем его, а берем только имя филда
+
+            # базовый метод полностью:
+            # def as_sql(self, qn, connection):
+            #     "Return the aggregate, rendered as SQL."
+            #
+            #     if hasattr(self.col, 'as_sql'):
+            #         field_name = self.col.as_sql(qn, connection)
+            #     elif isinstance(self.col, (list, tuple)):
+            #         field_name = '.'.join([qn(c) for c in self.col])
+            #     else:
+            #         field_name = self.col
+            #
+            #     params = {
+            #         'function': self.sql_function,
+            #         'field': field_name
+            #     }
+            #     params.update(self.extra)
+            #
+            #     return self.sql_template % params
+
+            field_name = self.col[-1]
+        else:
+            field_name = self.col
+
+        params = {
+            'function': self.sql_function,
+            'field': field_name
+        }
+        params.update(self.extra)
+
+        return self.sql_template % params
+
+
+class Avg(Aggregate):
+    is_computed = True
+    sql_function = 'AVG'
+
+class Count(Aggregate):
+    is_ordinal = True
+    sql_function = 'COUNT'
+    sql_template = '%(function)s(%(distinct)s%(field)s)'
+
+    def __init__(self, col, distinct=False, **extra):
+        super(Count, self).__init__(col, distinct=distinct and 'DISTINCT ' or '', **extra)
+
+class Max(Aggregate):
+    sql_function = 'MAX'
+
+class Min(Aggregate):
+    sql_function = 'MIN'
+
+class StdDev(Aggregate):
+    is_computed = True
+
+    def __init__(self, col, sample=False, **extra):
+        super(StdDev, self).__init__(col, **extra)
+        self.sql_function = sample and 'STDDEV_SAMP' or 'STDDEV_POP'
+
+class Sum(Aggregate):
+    sql_function = 'SUM'
+
+class Variance(Aggregate):
+    is_computed = True
+
+    def __init__(self, col, sample=False, **extra):
+        super(Variance, self).__init__(col, **extra)
+        self.sql_function = sample and 'VAR_SAMP' or 'VAR_POP'

django_sphinx_db/tests.py

@@ -2,6 +2,7 @@
 
 from django.test import TestCase
 from django.db import models
+from django.db.models import Sum
 from backend.models import SphinxModel, sphinx_escape
 
 class TagsIndex(SphinxModel):
@@ -148,3 +149,11 @@ def testCastIntToChar(self):
             pass
         except AttributeError:
             self.fail('Fail while casting int to char')
+
+    def testAggregateSum(self):
+        qs = TagsIndex.objects.match(u'"ТВ"'.encode('utf-8'))
+
+        iter_sum = sum(qs.values_list('id', flat=True))
+        aggr_sum = qs.aggregate(Sum('id'))['id__sum']
+
+        self.assertEqual(iter_sum, aggr_sum)