Merge pull request #533 from scipp/history-buffer-service

Buffer history of results in `DataService`

Author: SimonHeybrock

src/ess/livedata/config/workflows.py

@@ -76,7 +76,10 @@ def is_empty(self) -> bool:
     def _get_value(self) -> sc.DataArray:
         if self._to_nxlog is None:
             raise ValueError("No data accumulated")
-        return self._to_nxlog.get()
+        # Return latest value. Will be aggregated into a timeseries in frontend (if a
+        # plot requests it). This accumulator may be fully replaced once it is clear how
+        # we want to handle obtaining the full history (e.g., after frontend restarts).
+        return self._to_nxlog.get()[-1]
 
     def _do_push(self, value: sc.DataArray) -> None:
         if self._to_nxlog is None:

src/ess/livedata/core/job.py

@@ -184,7 +184,11 @@ def add(self, data: JobData) -> JobReply:
                     remapped_aux_data[field_name] = value
 
             # Pass data to workflow with field names (not stream names)
-            self._processor.accumulate({**data.primary_data, **remapped_aux_data})
+            self._processor.accumulate(
+                {**data.primary_data, **remapped_aux_data},
+                start_time=data.start_time,
+                end_time=data.end_time,
+            )
             if data.is_active():
                 if self._start_time is None:
                     self._start_time = data.start_time

src/ess/livedata/dashboard/correlation_histogram.py

@@ -392,11 +392,25 @@ def add_correlation_processor(
         items: dict[ResultKey, sc.DataArray],
     ) -> None:
         """Add a correlation histogram processor with DataService subscription."""
+        from .extractors import FullHistoryExtractor
+
         self._processors.append(processor)
 
-        # Create subscriber that merges data and sends to processor
+        # Create subscriber that merges data and sends to processor.
+        # Use FullHistoryExtractor to get complete timeseries history needed for
+        # correlation histogram computation.
+        # TODO We should update the plotter to operate more efficiently by simply
+        # subscribing to the changes. This will likely require a new extractor type as
+        # well as changes in the plotter, so we defer this for now.
         assembler = MergingStreamAssembler(set(items))
-        subscriber = DataSubscriber(assembler, processor)
+        extractors = {key: FullHistoryExtractor() for key in items}
+
+        # Create factory that sends initial data to processor and returns it
+        def processor_pipe_factory(data: dict[ResultKey, sc.DataArray]):
+            processor.send(data)
+            return processor
+
+        subscriber = DataSubscriber(assembler, processor_pipe_factory, extractors)
         self._data_service.register_subscriber(subscriber)
 
     def get_timeseries(self) -> list[ResultKey]:
@@ -412,7 +426,23 @@ def create_2d_config(self) -> CorrelationHistogramConfigurationAdapter:
 
 
 def _is_timeseries(da: sc.DataArray) -> bool:
-    return da.dims == ('time',) and 'time' in da.coords
+    """Check if data represents a timeseries.
+
+    When DataService uses LatestValueExtractor (default), it returns the latest value
+    from a timeseries buffer as a 0D scalar with a time coordinate. This function
+    identifies such values as originating from a timeseries.
+
+    Parameters
+    ----------
+    da:
+        DataArray to check.
+
+    Returns
+    -------
+    :
+        True if the data is a 0D scalar with a time coordinate.
+    """
+    return da.ndim == 0 and 'time' in da.coords
 
 
 class CorrelationHistogramProcessor:

src/ess/livedata/dashboard/data_service.py

@@ -2,32 +2,75 @@
 # Copyright (c) 2025 Scipp contributors (https://github.com/scipp)
 from __future__ import annotations
 
-from collections import UserDict
-from collections.abc import Callable, Hashable
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Hashable, Iterator, Mapping, MutableMapping
 from contextlib import contextmanager
-from typing import TypeVar
+from typing import Any, Generic, TypeVar
 
-from .data_subscriber import DataSubscriber
+from .extractors import LatestValueExtractor, UpdateExtractor
+from .temporal_buffer_manager import TemporalBufferManager
 
 K = TypeVar('K', bound=Hashable)
 V = TypeVar('V')
 
 
-class DataService(UserDict[K, V]):
+class DataServiceSubscriber(ABC, Generic[K]):
+    """Base class for data service subscribers with cached keys and extractors."""
+
+    def __init__(self) -> None:
+        """Initialize subscriber and cache keys from extractors."""
+        # Cache keys from extractors to avoid repeated computation
+        self._keys = set(self.extractors.keys())
+
+    @property
+    def keys(self) -> set[K]:
+        """Return the set of data keys this subscriber depends on."""
+        return self._keys
+
+    @property
+    @abstractmethod
+    def extractors(self) -> Mapping[K, UpdateExtractor]:
+        """
+        Return extractors for obtaining data views.
+
+        Returns a mapping from key to the extractor to use for that key.
+        """
+
+    @abstractmethod
+    def trigger(self, store: dict[K, Any]) -> None:
+        """Trigger the subscriber with updated data."""
+
+
+class DataService(MutableMapping[K, V]):
     """
     A service for managing and retrieving data and derived data.
 
     New data is set from upstream Kafka topics. Subscribers are typically plots that
     provide a live view of the data.
+
+    Uses buffers internally for storage, but presents a dict-like interface
+    that returns the latest value for each key.
     """
 
-    def __init__(self) -> None:
-        super().__init__()
-        self._subscribers: list[DataSubscriber[K]] = []
-        self._key_change_subscribers: list[Callable[[set[K], set[K]], None]] = []
+    def __init__(
+        self,
+        buffer_manager: TemporalBufferManager | None = None,
+    ) -> None:
+        """
+        Initialize DataService.
+
+        Parameters
+        ----------
+        buffer_manager:
+            Manager for buffer sizing. If None, creates a new TemporalBufferManager.
+        """
+        if buffer_manager is None:
+            buffer_manager = TemporalBufferManager()
+        self._buffer_manager = buffer_manager
+        self._default_extractor = LatestValueExtractor()
+        self._subscribers: list[DataServiceSubscriber[K]] = []
+        self._update_callbacks: list[Callable[[set[K]], None]] = []
         self._pending_updates: set[K] = set()
-        self._pending_key_additions: set[K] = set()
-        self._pending_key_removals: set[K] = set()
         self._transaction_depth = 0
 
     @contextmanager
@@ -48,30 +91,94 @@ def transaction(self):
     def _in_transaction(self) -> bool:
         return self._transaction_depth > 0
 
-    def register_subscriber(self, subscriber: DataSubscriber[K]) -> None:
+    def _get_extractors(self, key: K) -> list[UpdateExtractor]:
         """
-        Register a subscriber for updates.
+        Collect extractors for a key from all subscribers.
+
+        Examines all subscribers that need this key.
+
+        Parameters
+        ----------
+        key:
+            The key to collect extractors for.
+
+        Returns
+        -------
+        :
+            List of extractors from all subscribers for this key.
+        """
+        extractors = []
+
+        for subscriber in self._subscribers:
+            subscriber_extractors = subscriber.extractors
+            if key in subscriber_extractors:
+                extractor = subscriber_extractors[key]
+                extractors.append(extractor)
+
+        return extractors
+
+    def _build_subscriber_data(
+        self, subscriber: DataServiceSubscriber[K]
+    ) -> dict[K, Any]:
+        """
+        Extract data for a subscriber based on its extractors.
 
         Parameters
         ----------
         subscriber:
-            The subscriber to register. Must implement the DataSubscriber interface.
+            The subscriber to extract data for.
+
+        Returns
+        -------
+        :
+            Dictionary mapping keys to extracted data (None values filtered out).
         """
-        self._subscribers.append(subscriber)
+        subscriber_data = {}
 
-    def subscribe_to_changed_keys(
-        self, subscriber: Callable[[set[K], set[K]], None]
-    ) -> None:
+        for key, extractor in subscriber.extractors.items():
+            buffered_data = self._buffer_manager.get_buffered_data(key)
+            if buffered_data is not None:
+                subscriber_data[key] = extractor.extract(buffered_data)
+
+        return subscriber_data
+
+    def register_subscriber(self, subscriber: DataServiceSubscriber[K]) -> None:
         """
-        Register a subscriber for key change updates (additions/removals).
+        Register a subscriber for updates with extractor-based data access.
+
+        Triggers the subscriber immediately with existing data using its extractors.
 
         Parameters
         ----------
         subscriber:
-            A callable that accepts two sets: added_keys and removed_keys.
+            The subscriber to register.
         """
-        self._key_change_subscribers.append(subscriber)
-        subscriber(set(self.data.keys()), set())
+        self._subscribers.append(subscriber)
+
+        # Add extractors for keys this subscriber needs
+        for key in subscriber.keys:
+            if key in self._buffer_manager:
+                extractor = subscriber.extractors[key]
+                self._buffer_manager.add_extractor(key, extractor)
+
+        # Trigger immediately with existing data using subscriber's extractors
+        existing_data = self._build_subscriber_data(subscriber)
+        subscriber.trigger(existing_data)
+
+    def register_update_callback(self, callback: Callable[[set[K]], None]) -> None:
+        """
+        Register a callback for key update notifications.
+
+        Callback receives only the set of updated key names, not the data.
+        Use this for infrastructure that needs to know what changed but will
+        query data itself.
+
+        Parameters
+        ----------
+        callback:
+            Callable that accepts a set of updated keys.
+        """
+        self._update_callbacks.append(callback)
 
     def _notify_subscribers(self, updated_keys: set[K]) -> None:
         """
@@ -82,40 +189,46 @@ def _notify_subscribers(self, updated_keys: set[K]) -> None:
         updated_keys
             The set of data keys that were updated.
         """
+        # Notify extractor-based subscribers
         for subscriber in self._subscribers:
-            if not isinstance(subscriber, DataSubscriber):
-                subscriber(updated_keys)
-                continue
             if updated_keys & subscriber.keys:
-                # Pass only the data that the subscriber is interested in
-                subscriber_data = {
-                    key: self.data[key] for key in subscriber.keys if key in self.data
-                }
+                subscriber_data = self._build_subscriber_data(subscriber)
                 subscriber.trigger(subscriber_data)
 
-    def _notify_key_change_subscribers(self) -> None:
-        """Notify subscribers about key changes (additions/removals)."""
-        if not self._pending_key_additions and not self._pending_key_removals:
-            return
+        # Notify update callbacks with just key names
+        for callback in self._update_callbacks:
+            callback(updated_keys)
 
-        for subscriber in self._key_change_subscribers:
-            subscriber(
-                self._pending_key_additions.copy(), self._pending_key_removals.copy()
-            )
+    def __getitem__(self, key: K) -> V:
+        """Get the latest value for a key."""
+        buffered_data = self._buffer_manager.get_buffered_data(key)
+        if buffered_data is None:
+            raise KeyError(key)
+        return self._default_extractor.extract(buffered_data)
 
     def __setitem__(self, key: K, value: V) -> None:
-        if key not in self.data:
-            self._pending_key_additions.add(key)
-        super().__setitem__(key, value)
+        """Set a value, storing it in a buffer."""
+        if key not in self._buffer_manager:
+            extractors = self._get_extractors(key)
+            self._buffer_manager.create_buffer(key, extractors)
+        self._buffer_manager.update_buffer(key, value)
         self._pending_updates.add(key)
         self._notify_if_not_in_transaction()
 
     def __delitem__(self, key: K) -> None:
-        self._pending_key_removals.add(key)
-        super().__delitem__(key)
+        """Delete a key and its buffer."""
+        self._buffer_manager.delete_buffer(key)
         self._pending_updates.add(key)
         self._notify_if_not_in_transaction()
 
+    def __iter__(self) -> Iterator[K]:
+        """Iterate over keys."""
+        return iter(self._buffer_manager)
+
+    def __len__(self) -> int:
+        """Return the number of keys."""
+        return len(self._buffer_manager)
+
     def _notify_if_not_in_transaction(self) -> None:
         """Notify subscribers if not in a transaction."""
         if not self._in_transaction:
@@ -127,6 +240,3 @@ def _notify(self) -> None:
             pending = set(self._pending_updates)
             self._pending_updates.clear()
             self._notify_subscribers(pending)
-        self._notify_key_change_subscribers()
-        self._pending_key_additions.clear()
-        self._pending_key_removals.clear()