Add fill_value option to ring buffer window

The `fill_value` option imputes missing values `NaN` or `None` in the
ring buffer with the corresponding value. If `None` is passed, no
imputation is done. By default, missing values are imputed with `NaN`.
This prevents returning outdated return values accidentally.

If missing values should be filled, the force_copy argument has to
be true to avoid overwriting the underlying data.

Signed-off-by: cwasicki <[email protected]>

Author: cwasicki

src/frequenz/sdk/timeseries/_ringbuffer/buffer.py

@@ -304,6 +304,7 @@ def window(
         end: datetime | int | None,
         *,
         force_copy: bool = True,
+        fill_value: float | None = np.nan,
     ) -> FloatArray:
         """Request a copy or view on the data between start timestamp and end timestamp.
 
@@ -326,13 +327,21 @@ def window(
             end: end time of the window.
             force_copy: optional, default True. If True, will always create a
                 copy of the data.
+            fill_value: If not None, will use this value to fill missing values.
+                If missing values should be filled, force_copy must be True.
+                Defaults to NaN to avoid returning outdated data unexpectedly.
 
         Raises:
             IndexError: When start and end are not both datetime or index.
+            ValueError: When fill_value is not None and force_copy is False.
 
         Returns:
             The requested window
         """
+        # We don't want to modify the original buffer
+        if fill_value is not None and not force_copy:
+            raise ValueError("fill_value only supported for force_copy=True")
+
         if self.count_covered() == 0:
             return np.array([]) if isinstance(self._buffer, np.ndarray) else []
 
@@ -359,7 +368,48 @@ def window(
         start_pos = self.to_internal_index(start)
         end_pos = self.to_internal_index(end)
 
-        return self._wrapped_buffer_window(self._buffer, start_pos, end_pos, force_copy)
+        window = self._wrapped_buffer_window(
+            self._buffer, start_pos, end_pos, force_copy
+        )
+
+        if fill_value is not None:
+            window = self._fill_gaps(window, fill_value, start, self.gaps)
+        return window
+
+    def _fill_gaps(
+        self,
+        data: FloatArray,
+        fill_value: float,
+        oldest_timestamp: datetime,
+        gaps: list[Gap],
+    ) -> FloatArray:
+        """Fill the gaps in the data with the given fill_value.
+
+        Args:
+            data: The data to fill.
+            fill_value: The value to fill the gaps with.
+            oldest_timestamp: The oldest timestamp in the data.
+            gaps: List of gaps to fill.
+
+        Returns:
+            The filled data.
+        """
+        assert isinstance(
+            data, (np.ndarray, list)
+        ), f"Unsupported data type {type(data)}"
+        for gap in gaps:
+            start_index = (gap.start - oldest_timestamp) // self._sampling_period
+            end_index = (gap.end - oldest_timestamp) // self._sampling_period
+            start_index = max(start_index, 0)
+            end_index = min(end_index, len(data))
+            if start_index < end_index:
+                if isinstance(data, np.ndarray):
+                    data[start_index:end_index] = fill_value
+                elif isinstance(data, list):
+                    data[start_index:end_index] = [fill_value] * (
+                        end_index - start_index
+                    )
+        return data
 
     @staticmethod
     def _wrapped_buffer_window(

tests/timeseries/test_ringbuffer.py

@@ -526,12 +526,12 @@ def get_orb(data: FloatArray) -> OrderedRingBuffer[FloatArray]:
 def test_window_datetime() -> None:
     """Test the window function with datetime."""
     buffer = get_orb(np.array([0, None, 2, 3, 4]))
-    win = buffer.window(dt(0), dt(3), force_copy=False)
+    win = buffer.window(dt(0), dt(3), force_copy=False, fill_value=None)
     assert [0, np.nan, 2] == list(win)
     buffer._buffer[1] = 1  # pylint: disable=protected-access
     # Test whether the window is a view or a copy
     assert [0, 1, 2] == list(win)
-    win = buffer.window(dt(0), dt(3), force_copy=False)
+    win = buffer.window(dt(0), dt(3), force_copy=False, fill_value=None)
     assert [0, 1, 2] == list(win)
     # Empty array
     assert 0 == buffer.window(dt(1), dt(1)).size
@@ -561,8 +561,45 @@ def test_window_index() -> None:
     assert [] == buffer.window(-3, 0)
 
 
+def test_window_index_fill_value() -> None:
+    """Test fill_value functionality of window function."""
+    # Init with dummy data of size 3 and fill with [0, nan, 2]
+    buffer = OrderedRingBuffer([4711.0] * 3, ONE_SECOND)
+    buffer.update(Sample(dt(0), Quantity(0)))
+    buffer.update(Sample(dt(1), None))
+    buffer.update(Sample(dt(2), Quantity(2)))
+
+    # Test fill_value for explicitly set gaps
+    assert [0.0, np.nan, 2.0] == buffer.window(0, None)
+    assert [0.0, np.nan, 2.0] == buffer.window(0, None, fill_value=None)
+    assert [0.0, 1.0, 2.0] == buffer.window(0, None, fill_value=1)
+
+    # initial nan is ignored (optional implementation decision)
+    buffer.update(Sample(dt(3), Quantity(3)))  # -> [nan, 2, 3]
+    assert [2.0, 3.0] == buffer.window(0, None)
+    assert [2.0, 3.0] == buffer.window(0, None, fill_value=1)
+
+    # Test fill_value for gaps implicitly created gaps by time jumps
+    buffer.update(Sample(dt(4), Quantity(4)))
+    buffer.update(Sample(dt(6), Quantity(6)))  # -> [4, ?, 6]
+    # Default is fill missing values with NaNs
+    assert [4.0, np.nan, 6.0] == buffer.window(0, None)
+    assert [4.0, np.nan, 6.0] == buffer.window(0, None, fill_value=np.nan)
+    assert [4.0, 5.0, 6.0] == buffer.window(0, None, fill_value=5)
+    # If missing values not filled, outdated values can be returned unexpectedly
+    assert [4.0, 2, 6.0] == buffer.window(0, None, fill_value=None)
+
+    # Some edge cases
+    buffer.update(Sample(dt(7), Quantity(7)))  # -> [?, 6, 7]
+    assert [6.0, 7.0] == buffer.window(0, None)
+    buffer.update(Sample(dt(8), Quantity(np.nan)))  # -> [6, 7, nan]
+    assert [6.0, 7.0, np.nan] == buffer.window(0, None)
+    buffer.update(Sample(dt(9), None))  # -> [7, nan, nan]
+    assert [7.0, np.nan, np.nan] == buffer.window(0, None)
+
+
 def test_window_fail() -> None:
-    """Test the window function with invalid indices."""
+    """Test the window function with invalid arguments."""
     buffer = get_orb([0.0, 1.0, 2.0, 3.0, 4.0])
     # Go crazy with the indices
     with pytest.raises(IndexError):
@@ -573,6 +610,9 @@ def test_window_fail() -> None:
         buffer.window(None, dt(2))
     with pytest.raises(IndexError):
         buffer.window(dt(2), None)
+    # Invalid argument combination
+    with pytest.raises(ValueError):
+        buffer.window(0, 1, force_copy=False, fill_value=0)
 
 
 def test_wrapped_buffer_window() -> None: