Stream tabular data (#1201)

* Streaming table writes works

* Test streaming appends

* array-ref needs a shape for the first update

* Fix typing on tuples of ints

* Test update that does not specify patch.

* If no patch is given, slice to the current shape.

Author: danielballan

tiled/_tests/test_subscription.py

@@ -4,8 +4,11 @@
 import uuid
 
 import numpy as np
+import pandas as pd
+import pyarrow
 import pytest
 import tifffile
+from pandas.testing import assert_frame_equal
 from starlette.testclient import WebSocketDenialResponse
 
 from ..client import from_context
@@ -304,7 +307,8 @@ def callback(sub):
         assert event.wait(timeout=5.0), "Timeout waiting for messages"
 
 
-def test_subscribe_to_array_registered(tiled_websocket_context, tmp_path):
+def test_subscribe_to_array_registered_with_patch(tiled_websocket_context, tmp_path):
+    "Writer specifies the region of the update (patch)."
     context = tiled_websocket_context
     client = from_context(context)
     container_sub = client.subscribe()
@@ -356,7 +360,7 @@ def on_child_created(update):
             data_sources=[data_source],
             metadata={},
             specs=[],
-            key="test_subscribe_to_array_registered",
+            key="test_subscribe_to_array_registered_with_patch",
         )
         actual = x.read()  # smoke test
         np.testing.assert_array_equal(actual, arr[:2])
@@ -394,3 +398,162 @@ def on_child_created(update):
     assert update.patch.extend
     actual_streamed = update.data()
     np.testing.assert_array_equal(actual_streamed, arr[2:])
+
+
+def test_subscribe_to_array_registered_without_patch(tiled_websocket_context, tmp_path):
+    "Writer does not specify the region of the update (patch)."
+    context = tiled_websocket_context
+    client = from_context(context)
+    container_sub = client.subscribe()
+
+    updates = []
+    event = threading.Event()
+
+    def on_array_updated(update):
+        updates.append(update)
+        event.set()
+
+    def on_child_created(update):
+        array_sub = update.child().subscribe()
+        array_sub.new_data.add_callback(on_array_updated)
+        array_sub.start_in_thread(1)
+
+    container_sub.child_created.add_callback(on_child_created)
+
+    arr = np.random.random((3, 7, 13))
+    tifffile.imwrite(tmp_path / "image1.tiff", arr[0])
+    tifffile.imwrite(tmp_path / "image2.tiff", arr[1])
+
+    # Register just the first two images.
+    structure = ArrayStructure.from_array(arr[:2])
+    data_source = DataSource(
+        management=Management.external,
+        mimetype="multipart/related;type=image/tiff",
+        structure_family=StructureFamily.array,
+        structure=structure,
+        assets=[
+            Asset(
+                data_uri=f"file://{tmp_path}/image1.tiff",
+                is_directory=False,
+                parameter="data_uris",
+                num=1,
+            ),
+            Asset(
+                data_uri=f"file://{tmp_path}/image2.tiff",
+                is_directory=False,
+                parameter="data_uris",
+                num=2,
+            ),
+        ],
+    )
+
+    with container_sub.start_in_thread(1):
+        x = client.new(
+            structure_family=StructureFamily.array,
+            data_sources=[data_source],
+            metadata={},
+            specs=[],
+            key="test_subscribe_to_array_registered_without_patch",
+        )
+        actual = x.read()  # smoke test
+        np.testing.assert_array_equal(actual, arr[:2])
+        # Add the third image.
+        tifffile.imwrite(tmp_path / "image3.tiff", arr[2])
+        updated_structure = ArrayStructure.from_array(arr[:])
+        updated_data_source = copy.deepcopy(x.data_sources()[0])
+        updated_data_source.structure = updated_structure
+        updated_data_source.assets.append(
+            Asset(
+                data_uri=f"file://{tmp_path}/image3.tiff",
+                is_directory=False,
+                parameter="data_uris",
+                num=3,
+            ),
+        )
+        x.context.http_client.put(
+            x.uri.replace("/metadata/", "/data_source/", 1),
+            content=safe_json_dump(
+                {
+                    "data_source": updated_data_source,
+                }
+            ),
+        ).raise_for_status()
+        assert event.wait(timeout=5.0), "Timeout waiting for messages"
+        x.close_stream()
+        client.close_stream()
+        x.refresh()
+        actual_updated = x.read()
+        np.testing.assert_array_equal(actual_updated, arr[:])
+    (update,) = updates
+    assert update.patch is None
+    actual_streamed = update.data()
+    np.testing.assert_array_equal(actual_streamed, arr[:])
+
+
+def test_streaming_table_write(tiled_websocket_context):
+    context = tiled_websocket_context
+    client = from_context(context)
+    updates = []
+    event = threading.Event()
+    key = "test_streaming_table_write"
+
+    def collect(update):
+        updates.append(update)
+        event.set()
+
+    df1 = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+    df2 = pd.DataFrame({"a": [7, 8, 9], "b": [10, 11, 12]})
+    x = client.write_table(df1, key=key)
+
+    sub = client[key].subscribe()
+    sub.new_data.add_callback(collect)
+    with sub.start_in_thread(1):
+        assert event.wait(timeout=5.0), "Timeout waiting for messages"
+        actual = updates[0].data()
+        assert_frame_equal(actual, df1)
+        event.clear()
+        x.write(df2)
+        assert event.wait(timeout=5.0), "Timeout waiting for messages"
+        assert not updates[1].append
+        actual_updated = updates[1].data()
+        assert_frame_equal(actual_updated, df2)
+
+
+def test_streaming_table_appends(tiled_websocket_context):
+    context = tiled_websocket_context
+    client = from_context(context)
+    updates = []
+    event = threading.Event()
+    key = "test_streaming_table_append"
+
+    def collect(update):
+        updates.append(update)
+        event.set()
+
+    df1 = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+    df2 = pd.DataFrame({"a": [7, 8, 9], "b": [10, 11, 12]})
+    table1 = pyarrow.Table.from_pandas(df1, preserve_index=False)
+    table2 = pyarrow.Table.from_pandas(df2, preserve_index=False)
+    x = client.create_appendable_table(table1.schema, key=key)
+
+    sub = client[key].subscribe()
+    sub.new_data.add_callback(collect)
+    with sub.start_in_thread(1):
+        x.append_partition(0, table1)
+        assert event.wait(timeout=5.0), "Timeout waiting for messages"
+        assert updates[0].append
+        streamed1 = updates[0].data()
+        streamed1_pyarrow = pyarrow.Table.from_pandas(streamed1, preserve_index=False)
+        assert streamed1_pyarrow == table1
+        event.clear()
+        x.append_partition(0, table2)
+        assert event.wait(timeout=5.0), "Timeout waiting for messages"
+        assert updates[1].append
+        streamed2 = updates[1].data()
+        streamed2_pyarrow = pyarrow.Table.from_pandas(streamed2, preserve_index=False)
+        assert streamed2_pyarrow == table2
+        streaming_combined = pyarrow.concat_tables(
+            [streamed1_pyarrow, streamed2_pyarrow]
+        )
+        expected_combined = pyarrow.concat_tables([table1, table2])
+        assert streaming_combined == expected_combined

tiled/catalog/adapter.py

@@ -879,6 +879,7 @@ async def put_data_source(self, data_source, patch):
                 "timestamp": datetime.now().isoformat(),
                 "data_source": data_source.model_dump(),
                 "patch": patch.model_dump() if patch else None,
+                "shape": structure["shape"],
             }
 
             # Cache data in Redis with a TTL, and publish
@@ -1279,28 +1280,59 @@ class CatalogSparseAdapter(CatalogArrayAdapter):
 
 
 class CatalogTableAdapter(CatalogNodeAdapter):
+    def make_ws_schema(self):
+        return {
+            "type": "table-schema",
+            "version": 1,
+            "arrow_schema": self.structure().arrow_schema,
+        }
+
+    async def _stream(self, media_type, entry, body, partition, append):
+        sequence = await self.context.streaming_cache.incr_seq(self.node.id)
+        metadata = {
+            "type": "table-data",
+            "sequence": sequence,
+            "timestamp": datetime.now().isoformat(),
+            "mimetype": media_type,
+            "partition": partition,
+            "append": append,
+        }
+
+        await self.context.streaming_cache.set(
+            self.node.id, sequence, metadata, payload=body
+        )
+
     async def get(self, *args, **kwargs):
         return (await self.get_adapter()).get(*args, **kwargs)
 
     async def read(self, *args, **kwargs):
         return await ensure_awaitable((await self.get_adapter()).read, *args, **kwargs)
 
-    async def write(self, *args, **kwargs):
-        return await ensure_awaitable((await self.get_adapter()).write, *args, **kwargs)
+    async def write(self, media_type, deserializer, entry, body):
+        if self.context.streaming_cache:
+            await self._stream(media_type, entry, body, None, False)
+        data = await ensure_awaitable(deserializer, body)
+        return await ensure_awaitable((await self.get_adapter()).write, data)
 
     async def read_partition(self, *args, **kwargs):
         return await ensure_awaitable(
             (await self.get_adapter()).read_partition, *args, **kwargs
         )
 
-    async def write_partition(self, *args, **kwargs):
+    async def write_partition(self, media_type, deserializer, entry, body, partition):
+        if self.context.streaming_cache:
+            await self._stream(media_type, entry, body, partition, False)
+        data = await ensure_awaitable(deserializer, body)
         return await ensure_awaitable(
-            (await self.get_adapter()).write_partition, *args, **kwargs
+            (await self.get_adapter()).write_partition, partition, data
         )
 
-    async def append_partition(self, *args, **kwargs):
+    async def append_partition(self, media_type, deserializer, entry, body, partition):
+        if self.context.streaming_cache:
+            await self._stream(media_type, entry, body, partition, True)
+        data = await ensure_awaitable(deserializer, body)
         return await ensure_awaitable(
-            (await self.get_adapter()).append_partition, *args, **kwargs
+            (await self.get_adapter()).append_partition, partition, data
         )
 
 

tiled/client/dataframe.py

@@ -1,5 +1,7 @@
+import concurrent.futures
 import functools
 import warnings
+from typing import TYPE_CHECKING, Optional
 from urllib.parse import parse_qs, urlparse
 
 import dask
@@ -18,6 +20,9 @@
     retry_context,
 )
 
+if TYPE_CHECKING:
+    from .stream import TableSubscription
+
 _EXTRA_CHARS_PER_ITEM = len("&column=")
 
 
@@ -302,6 +307,25 @@ def export(self, filepath, columns=None, *, format=None):
             params=params,
         )
 
+    def subscribe(
+        self,
+        executor: Optional[concurrent.futures.Executor] = None,
+    ) -> "TableSubscription":
+        """
+        Subscribe to streaming updates about this table.
+
+        Returns
+        -------
+        subscription : Subscription
+        executor : concurrent.futures.Executor, optional
+            Launches tasks asynchronously, in response to updates. By default,
+            a concurrent.futures.ThreadPoolExecutor is used.
+        """
+        # Keep this import here to defer the websockets import until/unless needed.
+        from .stream import TableSubscription
+
+        return TableSubscription(self.context, self.path_parts, executor)
+
 
 # Subclass with a public class that adds the dask-specific methods.
 

tiled/client/stream.py

@@ -30,6 +30,8 @@
     ChildMetadataUpdated,
     ContainerSchema,
     Schema,
+    TableData,
+    TableSchema,
     Update,
 )
 from ..structures.core import STRUCTURE_TYPES, StructureFamily
@@ -497,6 +499,36 @@ def process(self, update: Update):
         self.new_data.process(update)
 
 
+class TableSubscription(Subscription):
+    """
+    Subscribe to streaming updates from an table.
+
+    Parameters
+    ----------
+    context : tiled.client.Context
+        Provides connection to Tiled server
+    segments : list[str]
+        Path to node of interest, given as a list of path segments
+    executor : concurrent.futures.Executor, optional
+        Launches tasks asynchronously, in response to updates. By default,
+        a concurrent.futures.ThreadPoolExecutor is used.
+    """
+
+    def __init__(
+        self,
+        context: Context,
+        segments: List[str] = None,
+        executor: Optional[concurrent.futures.Executor] = None,
+    ):
+        super().__init__(context, segments, executor)
+        self.new_data: CallbackRegistry["LiveTableData"] = CallbackRegistry(
+            self.executor
+        )
+
+    def process(self, update: Update):
+        self.new_data.process(update)
+
+
 class UnparseableMessage(RuntimeError):
     "Message can be decoded but cannot be interpreted by the application"
     pass
@@ -609,16 +641,38 @@ def data(self):
                 ).read()
         # Decode payload (bytes) into array.
         numpy_dtype = self.data_type.to_numpy_dtype()
-        return numpy.frombuffer(content, dtype=numpy_dtype).reshape(self.patch.shape)
+        if self.patch:
+            shape = self.patch.shape
+        else:
+            shape = self.shape
+        return numpy.frombuffer(content, dtype=numpy_dtype).reshape(shape)
+
+
+class LiveTableData(TableData):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    subscription: TableSubscription
+
+    def data(self):
+        "Get table"
+        # Registration occurs on import. Ensure this is imported.
+        from ..serialization import array
+
+        del array
+
+        # Decode payload (bytes) into array.
+        deserializer = default_deserialization_registry.dispatch("table", self.mimetype)
+        return deserializer(self.payload)
 
 
 SCHEMA_MESSAGE_TYPES = {
     "array-schema": ArraySchema,
     "container-schema": ContainerSchema,
+    "table-schema": TableSchema,
 }
 UPDATE_MESSAGE_TYPES = {
     "container-child-created": LiveChildCreated,
     "container-child-metadata-updated": LiveChildMetadataUpdated,
     "array-data": LiveArrayData,
     "array-ref": LiveArrayRef,
+    "table-data": LiveTableData,
 }