support metadata in source

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

Author: BulkBeing

packages/pynumaflow/pynumaflow/_metadata.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from collections import defaultdict
+from typing import Optional
 from pynumaflow.proto.common import metadata_pb2
 
 """
@@ -46,13 +46,13 @@ def keys(self, group: str) -> list[str]:
         """
         Returns the list of keys for a given group.
         """
-        return list(self._data[group].keys())
+        return list(self._data.get(group, {}).keys())
 
-    def value(self, group: str, key: str) -> bytes:
+    def value(self, group: str, key: str) -> Optional[bytes]:
         """
         Returns the value for a given group and key.
         """
-        return self._data[group][key]
+        return self._data.get(group, {}).get(key)
 
 
 @dataclass
@@ -61,49 +61,95 @@ class UserMetadata:
     UserMetadata wraps the user-generated metadata groups per message. It is read-write to UDFs.
     """
 
-    _data: defaultdict[str, dict[str, bytes]] = field(default_factory=lambda: defaultdict(dict))
+    _data: dict[str, dict[str, bytes]] = field(default_factory=dict)
 
     def groups(self) -> list[str]:
         """
         Returns the list of group names for the user metadata.
         """
         return list(self._data.keys())
 
-    def keys(self, group: str) -> list[str]:
+    def keys(self, group: str) -> Optional[list[str]]:
         """
         Returns the list of keys for a given group.
         """
-        return list(self._data[group].keys())
+        keys = self._data.get(group)
+        if keys is None:
+            return None
+        return list(keys.keys())
 
-    def value(self, group: str, key: str) -> bytes:
+    def __contains__(self, group: str) -> bool:
         """
-        Returns the value for a given group and key.
+        Returns True if the group exists.
+        """
+        return group in self._data
+
+    def __getitem__(self, group: str) -> dict[str, bytes]:
+        """
+        Returns the data for a given group.
+        Raises KeyError if the group does not exist.
+        """
+        return self._data[group]
+
+    def __setitem__(self, group: str, data: dict[str, bytes]):
+        """
+        Sets the data for a given group.
+        """
+        self._data[group] = data
+
+    def __delitem__(self, group: str):
+        """
+        Removes the group and all its keys and values.
+        Raises KeyError if the group does not exist.
+        """
+        del self._data[group]
+
+    def __len__(self) -> int:
         """
-        return self._data[group][key]
+        Returns the number of groups.
+        """
+        return len(self._data)
+
+    def value(self, group: str, key: str) -> Optional[bytes]:
+        """
+        Returns the value for a given group and key. If the group or key does not exist, returns None.
+        """
+        value = self._data.get(group)
+        if value is None:
+            return None
+        return value.get(key)
 
     def add(self, group: str, key: str, value: bytes):
         """
         Adds the value for a given group and key.
         """
-        self._data[group][key] = value
+        self._data.setdefault(group, {})[key] = value
 
     def set_group(self, group: str, data: dict[str, bytes]):
         """
         Sets the data for a given group.
         """
         self._data[group] = data
 
-    def remove(self, group: str, key: str):
+    def remove(self, group: str, key: str) -> Optional[bytes]:
         """
-        Removes the key and its value for a given group.
+        Removes the key and its value for a given group and returns the value. If this key is the only key in the group, the group will be removed.
+        Returns None if the group or key does not exist.
         """
-        del self._data[group][key]
+        group_data = self._data.pop(group, None)
+        if group_data is None:
+            return None
+        value = group_data.pop(key, None)
+        if group_data:
+            self._data[group] = group_data
+        return value
 
-    def remove_group(self, group: str):
+    def remove_group(self, group: str) -> Optional[dict[str, bytes]]:
         """
-        Removes the group and all its keys and values.
+        Removes the group and all its keys and values and returns the data.
+        Returns None if the group does not exist.
         """
-        del self._data[group]
+        return self._data.pop(group, None)
 
     def clear(self):
         """

packages/pynumaflow/pynumaflow/shared/asynciter.py

@@ -1,23 +1,27 @@
 import asyncio
+from typing import Generic, TypeVar
+from collections.abc import AsyncIterator
 
 from pynumaflow._constants import STREAM_EOF
 
+T = TypeVar("T")
 
-class NonBlockingIterator:
+
+class NonBlockingIterator(Generic[T]):
     """An Async Interator backed by a queue"""
 
     __slots__ = "_queue"
 
-    def __init__(self, size=0):
-        self._queue = asyncio.Queue(maxsize=size)
+    def __init__(self, size: int = 0) -> None:
+        self._queue: asyncio.Queue[T] = asyncio.Queue(maxsize=size)
 
-    async def read_iterator(self):
+    async def read_iterator(self) -> AsyncIterator[T]:
         item = await self._queue.get()
         while True:
             if item == STREAM_EOF:
                 break
             yield item
             item = await self._queue.get()
 
-    async def put(self, item):
+    async def put(self, item: T) -> None:
         await self._queue.put(item)

packages/pynumaflow/pynumaflow/sinker/__init__.py

@@ -2,6 +2,16 @@
 
 from pynumaflow.sinker.server import SinkServer
 
+from pynumaflow._metadata import UserMetadata, SystemMetadata
 from pynumaflow.sinker._dtypes import Response, Responses, Datum, Sinker
 
-__all__ = ["Response", "Responses", "Datum", "Sinker", "SinkAsyncServer", "SinkServer"]
+__all__ = [
+    "Response",
+    "Responses",
+    "Datum",
+    "Sinker",
+    "SinkAsyncServer",
+    "SinkServer",
+    "UserMetadata",
+    "SystemMetadata",
+]

packages/pynumaflow/pynumaflow/sourcer/__init__.py

@@ -10,6 +10,7 @@
     Sourcer,
     SourceCallable,
 )
+from pynumaflow._metadata import UserMetadata
 from pynumaflow.sourcer.async_server import SourceAsyncServer
 
 __all__ = [
@@ -24,4 +25,5 @@
     "Sourcer",
     "SourceAsyncServer",
     "SourceCallable",
+    "UserMetadata",
 ]

packages/pynumaflow/pynumaflow/sourcer/_dtypes.py

@@ -5,6 +5,7 @@
 from datetime import datetime
 from typing import Callable, Optional
 
+from pynumaflow._metadata import UserMetadata
 from pynumaflow.shared.asynciter import NonBlockingIterator
 
 
@@ -56,15 +57,17 @@ class Message:
         event_time: event time of the message, usually extracted from the payload.
         keys: []string keys for vertex (optional)
         headers: dict of headers for the message (optional)
+        user_metadata: metadata for the message (optional)
     """
 
-    __slots__ = ("_payload", "_offset", "_event_time", "_keys", "_headers")
+    __slots__ = ("_payload", "_offset", "_event_time", "_keys", "_headers", "_user_metadata")
 
     _payload: bytes
     _offset: Offset
     _event_time: datetime
     _keys: list[str]
     _headers: dict[str, str]
+    _user_metadata: UserMetadata
 
     def __init__(
         self,
@@ -73,6 +76,7 @@ def __init__(
         event_time: datetime,
         keys: list[str] = None,
         headers: Optional[dict[str, str]] = None,
+        user_metadata: Optional[UserMetadata] = None,
     ):
         """
         Creates a Message object to send value to a vertex.
@@ -82,6 +86,7 @@ def __init__(
         self._event_time = event_time
         self._keys = keys or []
         self._headers = headers or {}
+        self._user_metadata = user_metadata or UserMetadata()
 
     @property
     def payload(self) -> bytes:
@@ -103,6 +108,11 @@ def event_time(self) -> datetime:
     def headers(self) -> dict[str, str]:
         return self._headers
 
+    @property
+    def user_metadata(self) -> UserMetadata:
+        """Returns the user metadata of the message."""
+        return self._user_metadata
+
 
 @dataclass(init=False)
 class ReadRequest:

packages/pynumaflow/pynumaflow/sourcer/servicer/async_servicer.py

@@ -1,5 +1,6 @@
 import asyncio
 from collections.abc import AsyncIterator
+from typing import Union
 
 from google.protobuf import timestamp_pb2 as _timestamp_pb2
 from google.protobuf import empty_pb2 as _empty_pb2
@@ -9,6 +10,7 @@
 from pynumaflow.sourcer import ReadRequest, Offset, NackRequest, AckRequest, SourceCallable
 from pynumaflow.proto.sourcer import source_pb2
 from pynumaflow.proto.sourcer import source_pb2_grpc
+from pynumaflow.sourcer._dtypes import Message
 from pynumaflow.types import NumaflowServicerContext
 from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_UDF_EXCEPTION_STRING
 
@@ -31,7 +33,7 @@ def _create_ack_handshake_response():
     )
 
 
-def _create_read_response(response):
+def _create_read_response(response: Message):
     """Create a read response from the handler result."""
     event_time_timestamp = _timestamp_pb2.Timestamp()
     event_time_timestamp.FromDatetime(dt=response.event_time)
@@ -41,6 +43,7 @@ def _create_read_response(response):
         offset=response.offset.as_dict,
         event_time=event_time_timestamp,
         headers=response.headers,
+        metadata=response.user_metadata._to_proto(),
     )
     status = source_pb2.ReadResponse.Status(eot=False, code=source_pb2.ReadResponse.Status.SUCCESS)
     return source_pb2.ReadResponse(result=result, status=status)
@@ -98,7 +101,7 @@ async def ReadFn(
             async for req in request_iterator:
                 # create an iterator to be provided to the user function where the responses will
                 # be streamed
-                niter = NonBlockingIterator()
+                niter: NonBlockingIterator[Union[Message, Exception]] = NonBlockingIterator()
                 riter = niter.read_iterator()
                 task = asyncio.create_task(self.__invoke_read(req, niter))
                 # Save a reference to the result of this function, to avoid a
@@ -121,7 +124,9 @@ async def ReadFn(
             _LOGGER.critical("User-Defined Source ReadFn error", exc_info=True)
             await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING)
 
-    async def __invoke_read(self, req, niter):
+    async def __invoke_read(
+        self, req: source_pb2.ReadRequest, niter: NonBlockingIterator[Union[Message, Exception]]
+    ):
         """Invoke the read handler and manage the iterator."""
         try:
             await self.__source_read_handler(

packages/pynumaflow/tests/source/test_async_source.py

@@ -1,13 +1,15 @@
 import asyncio
+from collections.abc import Iterator
 import logging
 import threading
 import unittest
 
 import grpc
 from google.protobuf import empty_pb2 as _empty_pb2
-from grpc.aio._server import Server
+from grpc.aio import Server
 
 from pynumaflow import setup_logging
+from pynumaflow._metadata import _user_and_system_metadata_from_proto
 from pynumaflow.proto.sourcer import source_pb2_grpc, source_pb2
 from pynumaflow.sourcer import (
     SourceAsyncServer,
@@ -100,13 +102,13 @@ def test_read_source(self) -> None:
             stub = source_pb2_grpc.SourceStub(channel)
 
             request = read_req_source_fn()
-            generator_response = None
             try:
-                generator_response = stub.ReadFn(
+                generator_response: Iterator[source_pb2.ReadResponse] = stub.ReadFn(
                     request_iterator=request_generator(1, request, "read")
                 )
             except grpc.RpcError as e:
                 logging.error(e)
+                raise
 
             counter = 0
             first = True
@@ -139,6 +141,20 @@ def test_read_source(self) -> None:
                     r.result.offset.partition_id,
                 )
 
+                print(r.result)
+                (user_metadata, sys_metadata) = _user_and_system_metadata_from_proto(
+                    r.result.metadata
+                )
+                print(user_metadata)
+
+                self.assertCountEqual(user_metadata.groups(), ["custom_info", "test_info"])
+                self.assertCountEqual(
+                    user_metadata.keys("custom_info"), ["custom_key", "custom_key2"]
+                )
+                self.assertIsNone(user_metadata.value("custom_info", "test_key"))
+                self.assertEqual(user_metadata.value("custom_info", "custom_key"), b"custom_value")
+                self.assertEqual(user_metadata.value("test_info", "test_key"), b"test_value")
+
             self.assertFalse(first)
             self.assertTrue(last)
 

packages/pynumaflow/tests/source/test_async_source_err.py

@@ -6,7 +6,7 @@
 
 import grpc
 
-from grpc.aio._server import Server
+from grpc.aio import Server
 
 from pynumaflow import setup_logging
 from pynumaflow.proto.sourcer import source_pb2_grpc