Adds count to stream and streamset (#57)

Author: bbengfort

btrdb/stream.py

@@ -128,6 +128,42 @@ def exists(self):
                 return False
             raise bte
 
+    def count(self, start=MINIMUM_TIME, end=MAXIMUM_TIME, pointwidth=62, version=0):
+        """
+        Compute the total number of points in the stream
+
+        Counts the number of points in the specified window and version. By
+        default returns the latest total count of points in the stream. This
+        helper method sums the counts of all StatPoints returned by
+        ``aligned_windows``. Because of this, note that the start and end
+        timestamps may be adjusted if they are not powers of 2. For smaller
+        windows of time, you may also need to adjust the pointwidth to ensure
+        that the count granularity is captured appropriately.
+
+        Parameters
+        ----------
+        start : int or datetime like object, default: MINIMUM_TIME
+            The start time in nanoseconds for the range to be queried. (see
+            :func:`btrdb.utils.timez.to_nanoseconds` for valid input types)
+
+        end : int or datetime like object, default: MAXIMUM_TIME
+            The end time in nanoseconds for the range to be queried. (see
+            :func:`btrdb.utils.timez.to_nanoseconds` for valid input types)
+
+        pointwidth : int, default: 62
+            Specify the number of ns between data points (2**pointwidth)
+
+        version : int, default: 0
+            Version of the stream to query
+
+        Returns
+        -------
+        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])
+
     @property
     def btrdb(self):
         """
@@ -772,6 +808,44 @@ def versions(self):
         """
         return self._pinned_versions if self._pinned_versions else self._latest_versions()
 
+    def count(self):
+        """
+        Compute the total number of points in the streams using filters.
+
+        Computes the total number of points across all streams using the
+        specified filters. By default, this returns the latest total count of
+        all points in the streams. The count is modified by start and end
+        filters or by pinning versions.
+
+        Note that this helper method sums the counts of all StatPoints returned
+        by ``aligned_windows``. Because of this the start and end timestamps
+        may be adjusted if they are not powers of 2. You can also set the
+        pointwidth property for smaller windows of time to ensure that the
+        count granularity is captured appropriately.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        int
+            The total number of points in all streams for the specified filters.
+        """
+        params = self._params_from_filters()
+        start = params.get("start", MINIMUM_TIME)
+        end = params.get("end", MAXIMUM_TIME)
+
+        pointwidth = self.pointwidth if self.pointwidth is not None else 62
+        versions = self._pinned_versions if self._pinned_versions else {}
+
+        count = 0
+        for s in self._streams:
+            version = versions.get(s.uuid, 0)
+            count += s.count(start, end, pointwidth, version)
+
+        return count
+
     def earliest(self):
         """
         Returns earliest points of data in streams using available filters.

tests/btrdb/test_stream.py

@@ -478,6 +478,28 @@ def test_aligned_windows(self):
         )
 
 
+    def test_count(self):
+        """
+        Test that stream count method uses aligned windows
+        """
+        uu = uuid.UUID('0d22a53b-e2ef-4e0a-ab89-b2d48fb2592a')
+        endpoint = Mock(Endpoint)
+        windows = [
+            [(StatPointProto(time=1,min=2,mean=3,max=4,count=5,stddev=6), StatPointProto(time=2,min=3,mean=4,max=5,count=6,stddev=7)), 42],
+            [(StatPointProto(time=3,min=4,mean=5,max=6,count=7,stddev=8), StatPointProto(time=4,min=5,mean=6,max=7,count=8,stddev=9)), 42],
+        ]
+        endpoint.alignedWindows = Mock(return_value=windows)
+        stream = Stream(btrdb=BTrDB(endpoint), uuid=uu)
+
+        assert stream.count() == 26
+        stream._btrdb.ep.alignedWindows.assert_called_once_with(
+            uu, MINIMUM_TIME, MAXIMUM_TIME, 62, 0
+        )
+
+        stream.count(10, 1000, 48, 1200)
+        stream._btrdb.ep.alignedWindows.assert_called_with(uu, 10, 1000, 48, 1200)
+
+
     ##########################################################################
     ## earliest/latest tests
     ##########################################################################
@@ -1083,6 +1105,50 @@ def test_currently_out_of_range(self, mocked):
         with pytest.raises(ValueError, match="current time is not included in filtered stream range"):
             streams.filter(start=0, end=10).current()
 
+    def test_count(self):
+        """
+        Test the stream set count method
+        """
+        uu1 = uuid.UUID('0d22a53b-e2ef-4e0a-ab89-b2d48fb2592a')
+        uu2 = uuid.UUID('4dadf38d-52a5-4b7a-ada9-a5d563f9538c')
+        endpoint = Mock(Endpoint)
+        windows = [
+            [(StatPointProto(time=1,min=2,mean=3,max=4,count=5,stddev=6), StatPointProto(time=2,min=3,mean=4,max=5,count=6,stddev=7)), 42],
+            [(StatPointProto(time=3,min=4,mean=5,max=6,count=7,stddev=8), StatPointProto(time=4,min=5,mean=6,max=7,count=8,stddev=9)), 42],
+        ]
+        endpoint.alignedWindows = Mock(return_value=windows)
+        streams = StreamSet([
+            Stream(btrdb=BTrDB(endpoint), uuid=uu1),
+            Stream(btrdb=BTrDB(endpoint), uuid=uu2),
+        ])
+
+        assert streams.count() == 52
+        endpoint.alignedWindows.assert_any_call(uu1, MINIMUM_TIME, MAXIMUM_TIME, 62, 0)
+        endpoint.alignedWindows.assert_any_call(uu2, MINIMUM_TIME, MAXIMUM_TIME, 62, 0)
+
+
+    def test_count_filtered(self):
+        """
+        Test the stream set count method with filters
+        """
+        uu1 = uuid.UUID('0d22a53b-e2ef-4e0a-ab89-b2d48fb2592a')
+        uu2 = uuid.UUID('4dadf38d-52a5-4b7a-ada9-a5d563f9538c')
+        endpoint = Mock(Endpoint)
+        endpoint.alignedWindows = Mock(return_value=[])
+        streams = StreamSet([
+            Stream(btrdb=BTrDB(endpoint), uuid=uu1),
+            Stream(btrdb=BTrDB(endpoint), uuid=uu2),
+        ])
+
+        streams = streams.filter(start=10, end=1000)
+        streams.pin_versions({uu1: 42, uu2: 99})
+        streams.pointwidth = 48
+
+        streams.count()
+        endpoint.alignedWindows.assert_any_call(uu1, 10, 1000, 48, 42)
+        endpoint.alignedWindows.assert_any_call(uu2, 10, 1000, 48, 99)
+
+
 
     ##########################################################################
     ## filter tests