fixes stream.count

Author: looselycoupled

btrdb/stream.py

@@ -20,13 +20,15 @@
 import uuid as uuidlib
 from copy import deepcopy
 from collections.abc import Sequence
+from datetime import timedelta
 
 from btrdb.utils.buffer import PointBuffer
 from btrdb.point import RawPoint, StatPoint
 from btrdb.transformers import StreamSetTransformer
 from btrdb.exceptions import BTrDBError, InvalidOperation
 from btrdb.utils.timez import currently_as_ns, to_nanoseconds
 from btrdb.utils.conversion import AnnotationEncoder, AnnotationDecoder
+from btrdb.utils.general import pointwidth as pw
 
 
 ##########################################################################
@@ -122,7 +124,7 @@ def exists(self):
                 return False
             raise bte
 
-    def count(self, start=MINIMUM_TIME, end=MAXIMUM_TIME, pointwidth=62, version=0):
+    def count(self, start=MINIMUM_TIME, end=MAXIMUM_TIME, pointwidth=62, precise=False, version=0):
         """
         Compute the total number of points in the stream
 
@@ -134,6 +136,9 @@ def count(self, start=MINIMUM_TIME, end=MAXIMUM_TIME, pointwidth=62, version=0):
         windows of time, you may also need to adjust the pointwidth to ensure
         that the count granularity is captured appropriately.
 
+        Alternatively you can set the precise argument to True which will
+        give an exact count to the nanosecond but may be slower to execute.
+
         Parameters
         ----------
         start : int or datetime like object, default: MINIMUM_TIME
@@ -145,7 +150,14 @@ def count(self, start=MINIMUM_TIME, end=MAXIMUM_TIME, pointwidth=62, version=0):
             :func:`btrdb.utils.timez.to_nanoseconds` for valid input types)
 
         pointwidth : int, default: 62
-            Specify the number of ns between data points (2**pointwidth)
+            Specify the number of ns between data points (2**pointwidth).  If the value
+            is too large for the time window than the next smallest, appropriate
+            pointwidth will be used.
+
+        precise : bool, default: False
+            Forces the use of a windows query instead of aligned_windows
+            to determine exact count down to the nanosecond.  This will
+            be some amount slower than the aligned_windows version.
 
         version : int, default: 0
             Version of the stream to query
@@ -155,8 +167,17 @@ def count(self, start=MINIMUM_TIME, end=MAXIMUM_TIME, pointwidth=62, version=0):
         int
             The total number of points in the stream for the specified window.
         """
-        points = self.aligned_windows(start, end, pointwidth, version)
-        return sum([point.count for point, _ in points])
+
+        if not precise:
+            pointwidth = min(pointwidth, pw.from_nanoseconds(to_nanoseconds(end) - to_nanoseconds(start))-1)
+            points = self.aligned_windows(start, end, pointwidth, version)
+            return  sum([point.count for point, _ in points])
+
+        depth = 0
+        width = to_nanoseconds(end) - to_nanoseconds(start)
+        points = self.windows(start, end, width, depth, version)
+        return  sum([point.count for point, _ in points])
+
 
     @property
     def btrdb(self):