feat: source-transformer for pynumaflow-lite (#293)

Signed-off-by: Vigith Maurice <vigith@gmail.com>

Author: vigith

packages/pynumaflow-lite/Cargo.toml

@@ -59,3 +59,7 @@ path = "tests/bin/sink.rs"
 [[bin]]
 name = "test_source"
 path = "tests/bin/source.rs"
+
+[[bin]]
+name = "test_sourcetransform"
+path = "tests/bin/sourcetransform.rs"

packages/pynumaflow-lite/pynumaflow_lite/__init__.py

@@ -48,7 +48,13 @@
 except Exception:  # pragma: no cover
     sourcer = None
 
-# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, ReduceStreamer, Accumulator, Sinker, and Sourcer classes under the extension submodules for convenient access
+try:
+    sourcetransformer = _import_module(__name__ + ".sourcetransformer")
+except Exception:  # pragma: no cover
+    sourcetransformer = None
+
+# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, ReduceStreamer, Accumulator, Sinker,
+# Sourcer, and SourceTransformer classes under the extension submodules for convenient access
 from ._map_dtypes import Mapper
 from ._batchmapper_dtypes import BatchMapper
 from ._mapstream_dtypes import MapStreamer
@@ -58,6 +64,7 @@
 from ._accumulator_dtypes import Accumulator
 from ._sink_dtypes import Sinker
 from ._source_dtypes import Sourcer
+from ._sourcetransformer_dtypes import SourceTransformer
 
 if mapper is not None:
     try:
@@ -113,8 +120,15 @@
     except Exception:
         pass
 
+if sourcetransformer is not None:
+    try:
+        setattr(sourcetransformer, "SourceTransformer", SourceTransformer)
+    except Exception:
+        pass
+
 # Public API
-__all__ = ["mapper", "batchmapper", "mapstreamer", "reducer", "session_reducer", "reducestreamer", "accumulator", "sinker", "sourcer"]
+__all__ = ["mapper", "batchmapper", "mapstreamer", "reducer", "session_reducer", "reducestreamer", "accumulator",
+           "sinker", "sourcer", "sourcetransformer"]
 
 __doc__ = pynumaflow_lite.__doc__
 if hasattr(pynumaflow_lite, "__all__"):

packages/pynumaflow-lite/pynumaflow_lite/_sourcetransformer_dtypes.py

@@ -0,0 +1,35 @@
+from abc import ABCMeta, abstractmethod
+from pynumaflow_lite.sourcetransformer import Datum, Messages
+
+
+class SourceTransformer(metaclass=ABCMeta):
+    """
+    Provides an interface to write a SourceTransformer
+    which will be exposed over a gRPC server.
+    
+    A SourceTransformer is used for transforming and assigning event time
+    to input messages from a source.
+    """
+
+    def __call__(self, *args, **kwargs):
+        """
+        This allows to execute the handler function directly if
+        class instance is sent as a callable.
+        """
+        return self.handler(*args, **kwargs)
+
+    @abstractmethod
+    async def handler(self, keys: list[str], datum: Datum) -> Messages:
+        """
+        Implement this handler function which implements the SourceTransformer interface.
+        
+        Args:
+            keys: The keys associated with the message.
+            datum: The input datum containing value, event_time, watermark, and headers.
+            
+        Returns:
+            Messages: A collection of transformed messages with potentially modified
+                     event times and tags for conditional forwarding.
+        """
+        pass
+

packages/pynumaflow-lite/pynumaflow_lite/sourcetransformer.pyi

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from typing import Optional, List, Dict, Callable, Awaitable, Any
+import datetime as _dt
+
+# Re-export the Python ABC for user convenience and typing
+from ._sourcetransformer_dtypes import SourceTransformer as SourceTransformer
+
+
+class Messages:
+    def __init__(self) -> None: ...
+
+    def append(self, message: Message) -> None: ...
+
+    def __repr__(self) -> str: ...
+
+    def __str__(self) -> str: ...
+
+
+class Message:
+    keys: Optional[List[str]]
+    value: bytes
+    event_time: _dt.datetime
+    tags: Optional[List[str]]
+
+    def __init__(
+            self,
+            value: bytes,
+            event_time: _dt.datetime,
+            keys: Optional[List[str]] = ...,
+            tags: Optional[List[str]] = ...,
+    ) -> None: ...
+
+    @staticmethod
+    def message_to_drop(event_time: _dt.datetime) -> Message: ...
+
+
+class Datum:
+    # Read-only attributes provided by the extension
+    keys: List[str]
+    value: bytes
+    watermark: _dt.datetime
+    event_time: _dt.datetime
+    headers: Dict[str, str]
+
+    def __repr__(self) -> str: ...
+
+    def __str__(self) -> str: ...
+
+
+class SourceTransformAsyncServer:
+    def __init__(
+            self,
+            sock_file: str | None = ...,
+            info_file: str | None = ...,
+    ) -> None: ...
+
+    def start(self, py_func: Callable[..., Any]) -> Awaitable[None]: ...
+
+    def stop(self) -> None: ...
+
+
+__all__ = [
+    "Messages",
+    "Message",
+    "Datum",
+    "SourceTransformAsyncServer",
+    "SourceTransformer",
+]
+

packages/pynumaflow-lite/src/lib.rs

@@ -9,6 +9,7 @@ pub mod reducestream;
 pub mod session_reduce;
 pub mod sink;
 pub mod source;
+pub mod sourcetransform;
 
 use pyo3::prelude::*;
 
@@ -75,6 +76,13 @@ fn sourcer(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     Ok(())
 }
 
+/// Submodule: pynumaflow_lite.sourcetransformer
+#[pymodule]
+fn sourcetransformer(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
+    crate::sourcetransform::populate_py_module(m)?;
+    Ok(())
+}
+
 /// Top-level Python module `pynumaflow_lite` with submodules like `mapper`, `batchmapper`, and `mapstreamer`.
 #[pymodule]
 fn pynumaflow_lite(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
@@ -88,6 +96,7 @@ fn pynumaflow_lite(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(pyo3::wrap_pymodule!(accumulator))?;
     m.add_wrapped(pyo3::wrap_pymodule!(sinker))?;
     m.add_wrapped(pyo3::wrap_pymodule!(sourcer))?;
+    m.add_wrapped(pyo3::wrap_pymodule!(sourcetransformer))?;
 
     // Ensure it's importable as `pynumaflow_lite.mapper` as well as attribute access
     let binding = m.getattr("mapper")?;
@@ -170,5 +179,14 @@ fn pynumaflow_lite(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
         .getattr("modules")?
         .set_item(fullname, &sub)?;
 
+    // Ensure it's importable as `pynumaflow_lite.sourcetransformer` as well
+    let binding = m.getattr("sourcetransformer")?;
+    let sub = binding.cast::<PyModule>()?;
+    let fullname = "pynumaflow_lite.sourcetransformer";
+    sub.setattr("__name__", fullname)?;
+    py.import("sys")?
+        .getattr("modules")?
+        .set_item(fullname, &sub)?;
+
     Ok(())
 }