Support for using LIMIT N BY feature (simplify, additional testing, documentation)

CHANGELOG.md

@@ -4,6 +4,7 @@ Change Log
 Unreleased
 ----------
 - Support LowCardinality columns in ad-hoc queries
+- Support for LIMIT BY in querysets (utapyngo)
 
 v1.2.0
 ------

docs/class_reference.md

@@ -890,6 +890,14 @@ Adds a FINAL modifier to table, meaning data will be collapsed to final version.
 Can be used with `CollapsingMergeTree` engine only.
 
 
+#### limit_by(offset_limit, *fields)
+
+
+Adds a LIMIT BY clause to the query.
+- `offset_limit`: either an integer specifying the limit, or a tuple of integers (offset, limit).
+- `fields`: the field names to use in the clause.
+
+
 #### only(*field_names)
 
 
@@ -1013,6 +1021,14 @@ be names of grouping fields or calculated fields that this queryset was
 created with.
 
 
+#### limit_by(offset_limit, *fields)
+
+
+Adds a LIMIT BY clause to the query.
+- `offset_limit`: either an integer specifying the limit, or a tuple of integers (offset, limit).
+- `fields`: the field names to use in the clause.
+
+
 #### only(*field_names)
 
 

src/infi/clickhouse_orm/query.py

@@ -335,22 +335,21 @@ class QuerySet(object):
             return qs
 
     def limit_by(self, offset_limit, *fields):
+        """
+        Adds a LIMIT BY clause to the query.
+        - `offset_limit`: either an integer specifying the limit, or a tuple of integers (offset, limit).
+        - `fields`: the field names to use in the clause.
+        """
         if isinstance(offset_limit, six.integer_types):
             # Single limit
-            assert offset_limit >= 0, 'negative limits are not supported'
-            qs = copy(self)
-            qs._limit_by = (0, offset_limit)
-            qs._limit_by_fields = fields
-            return qs
-        else:
-            # Offset, limit
-            offset = offset_limit[0]
-            limit = offset_limit[1]
-            assert offset >= 0 and limit >= 0, 'negative limits are not supported'
-            qs = copy(self)
-            qs._limit_by = (offset, limit)
-            qs._limit_by_fields = fields
-            return qs
+            offset_limit = (0, offset_limit)
+        offset = offset_limit[0]
+        limit = offset_limit[1]
+        assert offset >= 0 and limit >= 0, 'negative limits are not supported'
+        qs = copy(self)
+        qs._limit_by = (offset, limit)
+        qs._limit_by_fields = fields
+        return qs
 
     def select_fields_as_sql(self):
         """
@@ -386,13 +385,13 @@ class QuerySet(object):
         if self._order_by:
             sql += '\nORDER BY ' + self.order_by_as_sql()
 
-        if self._limits:
-            sql += '\nLIMIT %d, %d' % self._limits
-
         if self._limit_by:
             sql += '\nLIMIT %d, %d' % self._limit_by
             sql += ' BY %s' % comma_join('`%s`' % field for field in self._limit_by_fields)
 
+        if self._limits:
+            sql += '\nLIMIT %d, %d' % self._limits
+
         return sql
 
     def order_by_as_sql(self):

tests/test_querysets.py

@@ -433,10 +433,21 @@ class AggregateTestCase(TestCaseWithData):
         self.assertEqual(qs.conditions_as_sql(), 'the__next__number > 1')
 
     def test_limit_by(self):
+        # Test without offset
         qs = Person.objects_in(self.database).aggregate('first_name', 'last_name', 'height', n='count()').\
             order_by('first_name', '-height').limit_by(1, 'first_name')
         self.assertEqual(qs.count(), 94)
         self.assertEqual(list(qs)[89].last_name, 'Bowen')
+        # Test with limit and offset, also mixing LIMIT with LIMIT BY
+        qs = Person.objects_in(self.database).filter(height__gt=1.67).order_by('height', 'first_name')
+        limited_qs = qs.limit_by((0, 3), 'height')
+        self.assertEquals([p.first_name for p in limited_qs[:3]], ['Amanda', 'Buffy', 'Dora'])
+        limited_qs = qs.limit_by((3, 3), 'height')
+        self.assertEquals([p.first_name for p in limited_qs[:3]], ['Elton', 'Josiah', 'Macaulay'])
+        limited_qs = qs.limit_by((6, 3), 'height')
+        self.assertEquals([p.first_name for p in limited_qs[:3]], ['Norman', 'Octavius', 'Oliver'])
+
+
 
 
 Color = Enum('Color', u'red blue green yellow brown white black')