perf: improve count query performance for simple queries (#516)

For simple queries (queries that map directly to single Datastore
queries, ie not multiqueries, or queries with post filters),
`Query.count` can be otimized by setting a high offset and counting the
number of entities skipped.

Note that this still requires Datastore to assemble and iterate over a
result set and Datastore will still only "skip" a certain number of
entities at a time. (1000 at the time of this writing.) So this doesn't
dramatically impact the amount of work that has to be done on the
Datastore side, nor does it reduce the number gRPC calls necessary to
count a large result set. It does reduce significantly, however, the amount of
I/O required, which has sped up some large counts in testing by a factor
of around 4x.

Author: Chris Rossi

google/cloud/ndb/_datastore_query.py

@@ -35,7 +35,8 @@
 log = logging.getLogger(__name__)
 
 MoreResultsType = query_pb2.QueryResultBatch.MoreResultsType
-MORE_RESULTS_TYPE_NOT_FINISHED = MoreResultsType.Value("NOT_FINISHED")
+NO_MORE_RESULTS = MoreResultsType.Value("NO_MORE_RESULTS")
+NOT_FINISHED = MoreResultsType.Value("NOT_FINISHED")
 MORE_RESULTS_AFTER_LIMIT = MoreResultsType.Value("MORE_RESULTS_AFTER_LIMIT")
 
 ResultType = query_pb2.EntityResult.ResultType
@@ -112,6 +113,70 @@ def fetch(query):
     raise tasklets.Return(entities)
 
 
+def count(query):
+    """Count query results.
+
+    Args:
+        query (query.QueryOptions): The query spec.
+
+    Returns:
+        tasklets.Future: Results is int: Number of results that would be
+            returned by the query.
+    """
+    filters = query.filters
+    if filters:
+        if filters._multiquery or filters._post_filters():
+            return _count_brute_force(query)
+
+    return _count_by_skipping(query)
+
+
+@tasklets.tasklet
+def _count_brute_force(query):
+    query = query.copy(projection=["__key__"], order_by=None)
+    results = iterate(query, raw=True)
+    count = 0
+    limit = query.limit
+    while (yield results.has_next_async()):
+        count += 1
+        if limit and count == limit:
+            break
+
+        results.next()
+
+    raise tasklets.Return(count)
+
+
+@tasklets.tasklet
+def _count_by_skipping(query):
+    limit = query.limit
+    query = query.copy(projection=["__key__"], order_by=None, limit=1)
+    count = 0
+    more_results = NOT_FINISHED
+    cursor = None
+
+    while more_results != NO_MORE_RESULTS:
+        if limit:
+            offset = limit - count - 1
+        else:
+            offset = 10000
+
+        query = query.copy(offset=offset, start_cursor=cursor)
+        response = yield _datastore_run_query(query)
+        batch = response.batch
+
+        more_results = batch.more_results
+        count += batch.skipped_results
+        count += len(batch.entity_results)
+
+        if limit and count >= limit:
+            break
+
+        cursor = Cursor(batch.end_cursor)
+
+    raise tasklets.Return(count)
+
+
 def iterate(query, raw=False):
     """Get iterator for query results.
 
@@ -307,9 +372,7 @@ def _next_batch(self):
             for result_pb in response.batch.entity_results
         ]
 
-        self._has_next_batch = more_results = (
-            batch.more_results == MORE_RESULTS_TYPE_NOT_FINISHED
-        )
+        self._has_next_batch = more_results = batch.more_results == NOT_FINISHED
 
         self._more_results_after_limit = batch.more_results == MORE_RESULTS_AFTER_LIMIT
 
@@ -935,3 +998,19 @@ def to_websafe_string(self):
     def urlsafe(self):
         # Documented in official Legacy NDB docs
         return base64.urlsafe_b64encode(self.cursor)
+
+    def __eq__(self, other):
+        if isinstance(other, Cursor):
+            return self.cursor == other.cursor
+
+        return NotImplemented
+
+    def __ne__(self, other):
+        # required for Python 2.7 compatibility
+        result = self.__eq__(other)
+        if result is NotImplemented:
+            result = False
+        return not result
+
+    def __hash__(self):
+        return hash(self.cursor)

google/cloud/ndb/query.py

@@ -2172,7 +2172,6 @@ def count(self, limit=None, **kwargs):
         """
         return self.count_async(_options=kwargs["_options"]).result()
 
-    @tasklets.tasklet
     @_query_options
     @utils.keyword_only(
         offset=None,
@@ -2201,19 +2200,7 @@ def count_async(self, limit=None, **kwargs):
         # Avoid circular import in Python 2.7
         from google.cloud.ndb import _datastore_query
 
-        _options = kwargs["_options"]
-        options = _options.copy(projection=["__key__"], order_by=None)
-        results = _datastore_query.iterate(options, raw=True)
-        count = 0
-        limit = options.limit
-        while (yield results.has_next_async()):
-            count += 1
-            if limit and count == limit:
-                break
-
-            results.next()
-
-        raise tasklets.Return(count)
+        return _datastore_query.count(kwargs["_options"])
 
     @_query_options
     @utils.keyword_only(

tests/unit/test__datastore_query.py

@@ -81,6 +81,172 @@ def test_fetch(iterate):
         iterate.assert_called_once_with("foo")
 
 
+class Test_count:
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @mock.patch("google.cloud.ndb._datastore_query.iterate")
+    def test_count_brute_force(iterate):
+        class DummyQueryIterator:
+            def __init__(self, items):
+                self.items = list(items)
+
+            def has_next_async(self):
+                return utils.future_result(bool(self.items))
+
+            def next(self):
+                return self.items.pop()
+
+        iterate.return_value = DummyQueryIterator(range(5))
+        query = query_module.QueryOptions(
+            filters=mock.Mock(_multiquery=True, spec=("_multiquery",))
+        )
+
+        future = _datastore_query.count(query)
+        assert future.result() == 5
+        iterate.assert_called_once_with(
+            query_module.QueryOptions(filters=query.filters, projection=["__key__"]),
+            raw=True,
+        )
+
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @mock.patch("google.cloud.ndb._datastore_query.iterate")
+    def test_count_brute_force_with_limit(iterate):
+        class DummyQueryIterator:
+            def __init__(self, items):
+                self.items = list(items)
+
+            def has_next_async(self):
+                return utils.future_result(bool(self.items))
+
+            def next(self):
+                return self.items.pop()
+
+        iterate.return_value = DummyQueryIterator(range(5))
+        query = query_module.QueryOptions(
+            filters=mock.Mock(
+                _multiquery=False,
+                _post_filters=mock.Mock(return_value=True),
+                spec=("_multiquery", "_post_filters"),
+            ),
+            limit=3,
+        )
+
+        future = _datastore_query.count(query)
+        assert future.result() == 3
+        iterate.assert_called_once_with(
+            query_module.QueryOptions(
+                filters=query.filters, projection=["__key__"], limit=3
+            ),
+            raw=True,
+        )
+
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")
+    def test_count_by_skipping(run_query):
+        run_query.side_effect = utils.future_results(
+            mock.Mock(
+                batch=mock.Mock(
+                    more_results=_datastore_query.NOT_FINISHED,
+                    skipped_results=1000,
+                    entity_results=[],
+                    end_cursor=b"himom",
+                    spec=(
+                        "more_results",
+                        "skipped_results",
+                        "entity_results",
+                        "end_cursor",
+                    ),
+                ),
+                spec=("batch",),
+            ),
+            mock.Mock(
+                batch=mock.Mock(
+                    more_results=_datastore_query.NO_MORE_RESULTS,
+                    skipped_results=100,
+                    entity_results=[],
+                    end_cursor=b"hellodad",
+                    spec=(
+                        "more_results",
+                        "skipped_results",
+                        "entity_results",
+                        "end_cursor",
+                    ),
+                ),
+                spec=("batch",),
+            ),
+        )
+
+        query = query_module.QueryOptions()
+        future = _datastore_query.count(query)
+        assert future.result() == 1100
+
+        expected = [
+            mock.call(
+                query_module.QueryOptions(
+                    limit=1,
+                    offset=10000,
+                    projection=["__key__"],
+                )
+            ),
+            (
+                (
+                    query_module.QueryOptions(
+                        limit=1,
+                        offset=10000,
+                        projection=["__key__"],
+                        start_cursor=_datastore_query.Cursor(b"himom"),
+                    ),
+                ),
+                {},
+            ),
+        ]
+        assert run_query.call_args_list == expected
+
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")
+    def test_count_by_skipping_with_limit(run_query):
+        run_query.return_value = utils.future_result(
+            mock.Mock(
+                batch=mock.Mock(
+                    more_results=_datastore_query.MORE_RESULTS_AFTER_LIMIT,
+                    skipped_results=99,
+                    entity_results=[object()],
+                    end_cursor=b"himom",
+                    spec=(
+                        "more_results",
+                        "skipped_results",
+                        "entity_results",
+                        "end_cursor",
+                    ),
+                ),
+                spec=("batch",),
+            )
+        )
+
+        query = query_module.QueryOptions(
+            filters=mock.Mock(
+                _multiquery=False,
+                _post_filters=mock.Mock(return_value=None),
+                spec=("_multiquery", "_post_filters"),
+            ),
+            limit=100,
+        )
+        future = _datastore_query.count(query)
+        assert future.result() == 100
+
+        run_query.assert_called_once_with(
+            query_module.QueryOptions(
+                limit=1,
+                offset=99,
+                projection=["__key__"],
+                filters=query.filters,
+            )
+        )
+
+
 class Test_iterate:
     @staticmethod
     @mock.patch("google.cloud.ndb._datastore_query._QueryIteratorImpl")
@@ -1448,3 +1614,22 @@ def test_urlsafe():
         urlsafe = base64.urlsafe_b64encode(b"123")
         cursor = _datastore_query.Cursor(b"123")
         assert cursor.urlsafe() == urlsafe
+
+    @staticmethod
+    def test__eq__same():
+        assert _datastore_query.Cursor(b"123") == _datastore_query.Cursor(b"123")
+        assert not _datastore_query.Cursor(b"123") != _datastore_query.Cursor(b"123")
+
+    @staticmethod
+    def test__eq__different():
+        assert _datastore_query.Cursor(b"123") != _datastore_query.Cursor(b"234")
+        assert not _datastore_query.Cursor(b"123") == _datastore_query.Cursor(b"234")
+
+    @staticmethod
+    def test__eq__different_type():
+        assert _datastore_query.Cursor(b"123") != b"234"
+        assert not _datastore_query.Cursor(b"123") == b"234"
+
+    @staticmethod
+    def test__hash__():
+        assert hash(_datastore_query.Cursor(b"123")) == hash(b"123")