feat: source using rust sdk (#283)

Signed-off-by: Vigith Maurice <[email protected]>

Author: vigith

packages/pynumaflow-lite/Cargo.toml

@@ -51,3 +51,7 @@ path = "tests/bin/accumulator.rs"
 [[bin]]
 name = "test_sink"
 path = "tests/bin/sink.rs"
+
+[[bin]]
+name = "test_source"
+path = "tests/bin/source.rs"

packages/pynumaflow-lite/pynumaflow_lite/__init__.py

@@ -33,13 +33,25 @@
 except Exception:  # pragma: no cover
     accumulator = None
 
-# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, and Accumulator classes under the extension submodules for convenient access
+try:
+    sinker = _import_module(__name__ + ".sinker")
+except Exception:  # pragma: no cover
+    sinker = None
+
+try:
+    sourcer = _import_module(__name__ + ".sourcer")
+except Exception:  # pragma: no cover
+    sourcer = None
+
+# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, Accumulator, Sinker, and Sourcer classes under the extension submodules for convenient access
 from ._map_dtypes import Mapper
 from ._batchmapper_dtypes import BatchMapper
 from ._mapstream_dtypes import MapStreamer
 from ._reduce_dtypes import Reducer
 from ._session_reduce_dtypes import SessionReducer
 from ._accumulator_dtypes import Accumulator
+from ._sink_dtypes import Sinker
+from ._source_dtypes import Sourcer
 
 if mapper is not None:
     try:
@@ -77,8 +89,20 @@
     except Exception:
         pass
 
+if sinker is not None:
+    try:
+        setattr(sinker, "Sinker", Sinker)
+    except Exception:
+        pass
+
+if sourcer is not None:
+    try:
+        setattr(sourcer, "Sourcer", Sourcer)
+    except Exception:
+        pass
+
 # Public API
-__all__ = ["mapper", "batchmapper", "mapstreamer", "reducer", "session_reducer", "accumulator"]
+__all__ = ["mapper", "batchmapper", "mapstreamer", "reducer", "session_reducer", "accumulator", "sinker", "sourcer"]
 
 __doc__ = pynumaflow_lite.__doc__
 if hasattr(pynumaflow_lite, "__all__"):

packages/pynumaflow-lite/pynumaflow_lite/__init__.pyi

@@ -9,5 +9,6 @@ from . import reducer as reducer
 from . import session_reducer as session_reducer
 from . import accumulator as accumulator
 from . import sinker as sinker
+from . import sourcer as sourcer
 
-__all__ = ['mapper', 'batchmapper', 'mapstreamer', 'reducer', 'session_reducer', 'accumulator', 'sinker']
+__all__ = ['mapper', 'batchmapper', 'mapstreamer', 'reducer', 'session_reducer', 'accumulator', 'sinker', 'sourcer']

packages/pynumaflow-lite/pynumaflow_lite/_source_dtypes.py

@@ -0,0 +1,125 @@
+from abc import ABCMeta, abstractmethod
+from collections.abc import AsyncIterator
+from pynumaflow_lite.sourcer import (
+    Message,
+    Offset,
+    ReadRequest,
+    AckRequest,
+    NackRequest,
+    PendingResponse,
+    PartitionsResponse,
+)
+
+
+class Sourcer(metaclass=ABCMeta):
+    """
+    Provides an interface to write a User Defined Source.
+    
+    A Sourcer must implement the following handlers:
+    - read_handler: Read messages from the source
+    - ack_handler: Acknowledge processed messages
+    - pending_handler: Return the number of pending messages
+    - partitions_handler: Return the partitions this source handles
+    
+    Optionally, you can implement:
+    - nack_handler: Negatively acknowledge messages (default: no-op)
+    """
+
+    def __call__(self, *args, **kwargs):
+        """
+        This allows to execute the handler function directly if
+        class instance is sent as a callable.
+        """
+        return self.read_handler(*args, **kwargs)
+
+    @abstractmethod
+    async def read_handler(self, request: ReadRequest) -> AsyncIterator[Message]:
+        """
+        Read messages from the source.
+        
+        Args:
+            request: ReadRequest containing num_records and timeout
+            
+        Yields:
+            Message: Messages to be sent to the next vertex
+            
+        Example:
+            async def read_handler(self, request: ReadRequest) -> AsyncIterator[Message]:
+                for i in range(request.num_records):
+                    yield Message(
+                        payload=f"message-{i}".encode(),
+                        offset=Offset(str(i).encode(), partition_id=0),
+                        event_time=datetime.now(),
+                        keys=["key1"],
+                        headers={"x-txn-id": str(uuid.uuid4())}
+                    )
+        """
+        pass
+
+    @abstractmethod
+    async def ack_handler(self, request: AckRequest) -> None:
+        """
+        Acknowledge that messages have been processed.
+        
+        Args:
+            request: AckRequest containing the list of offsets to acknowledge
+            
+        Example:
+            async def ack_handler(self, request: AckRequest) -> None:
+                for offset in request.offsets:
+                    # Remove from pending set, mark as processed, etc.
+                    self.pending_offsets.remove(offset.offset)
+        """
+        pass
+
+    @abstractmethod
+    async def pending_handler(self) -> PendingResponse:
+        """
+        Return the number of pending messages yet to be processed.
+        
+        Returns:
+            PendingResponse: Response containing the count of pending messages.
+                            Return count=-1 if the source doesn't support detecting backlog.
+            
+        Example:
+            async def pending_handler(self) -> PendingResponse:
+                return PendingResponse(count=len(self.pending_offsets))
+        """
+        pass
+
+    @abstractmethod
+    async def partitions_handler(self) -> PartitionsResponse:
+        """
+        Return the partitions associated with this source.
+        
+        This is used by the platform to determine the partitions to which
+        the watermark should be published. If your source doesn't have the
+        concept of partitions, return the replica ID.
+        
+        Returns:
+            PartitionsResponse: Response containing the list of partition IDs
+            
+        Example:
+            async def partitions_handler(self) -> PartitionsResponse:
+                return PartitionsResponse(partitions=[self.partition_id])
+        """
+        pass
+
+    async def nack_handler(self, request: NackRequest) -> None:
+        """
+        Negatively acknowledge messages (optional).
+        
+        This is called when messages could not be processed and should be
+        retried or handled differently. Default implementation is a no-op.
+        
+        Args:
+            request: NackRequest containing the list of offsets to nack
+            
+        Example:
+            async def nack_handler(self, request: NackRequest) -> None:
+                for offset in request.offsets:
+                    # Add back to pending, mark for retry, etc.
+                    self.nacked_offsets.add(offset.offset)
+        """
+        pass
+

packages/pynumaflow-lite/pynumaflow_lite/sourcer.pyi

@@ -0,0 +1,135 @@
+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 ._source_dtypes import Sourcer as Sourcer
+
+
+class Message:
+    """A message to be sent from the source."""
+    payload: bytes
+    offset: Offset
+    event_time: _dt.datetime
+    keys: List[str]
+    headers: Dict[str, str]
+
+    def __init__(
+        self,
+        payload: bytes,
+        offset: Offset,
+        event_time: _dt.datetime,
+        keys: Optional[List[str]] = ...,
+        headers: Optional[Dict[str, str]] = ...,
+    ) -> None: ...
+
+    def __repr__(self) -> str: ...
+
+    def __str__(self) -> str: ...
+
+
+class Offset:
+    """The offset of a message."""
+    offset: bytes
+    partition_id: int
+
+    def __init__(
+        self,
+        offset: bytes,
+        partition_id: int = ...,
+    ) -> None: ...
+
+    def __repr__(self) -> str: ...
+
+    def __str__(self) -> str: ...
+
+
+class ReadRequest:
+    """A request to read messages from the source."""
+    num_records: int
+    timeout_ms: int
+
+    def __init__(
+        self,
+        num_records: int,
+        timeout_ms: int = ...,
+    ) -> None: ...
+
+    def __repr__(self) -> str: ...
+
+
+class AckRequest:
+    """A request to acknowledge messages."""
+    offsets: List[Offset]
+
+    def __init__(
+        self,
+        offsets: List[Offset],
+    ) -> None: ...
+
+    def __repr__(self) -> str: ...
+
+
+class NackRequest:
+    """A request to negatively acknowledge messages."""
+    offsets: List[Offset]
+
+    def __init__(
+        self,
+        offsets: List[Offset],
+    ) -> None: ...
+
+    def __repr__(self) -> str: ...
+
+
+class PendingResponse:
+    """Response for pending messages count."""
+    count: int
+
+    def __init__(
+        self,
+        count: int = ...,
+    ) -> None: ...
+
+    def __repr__(self) -> str: ...
+
+
+class PartitionsResponse:
+    """Response for partitions."""
+    partitions: List[int]
+
+    def __init__(
+        self,
+        partitions: List[int],
+    ) -> None: ...
+
+    def __repr__(self) -> str: ...
+
+
+class SourceAsyncServer:
+    """Async Source Server that can be started from Python code."""
+
+    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__ = [
+    "Message",
+    "Offset",
+    "ReadRequest",
+    "AckRequest",
+    "NackRequest",
+    "PendingResponse",
+    "PartitionsResponse",
+    "SourceAsyncServer",
+    "Sourcer",
+]
+

packages/pynumaflow-lite/src/lib.rs

@@ -7,6 +7,7 @@ pub mod pyrs;
 pub mod reduce;
 pub mod session_reduce;
 pub mod sink;
+pub mod source;
 
 use pyo3::prelude::*;
 
@@ -59,6 +60,13 @@ fn sinker(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     Ok(())
 }
 
+/// Submodule: pynumaflow_lite.sourcer
+#[pymodule]
+fn sourcer(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
+    crate::source::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<()> {
@@ -70,6 +78,7 @@ fn pynumaflow_lite(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(pyo3::wrap_pymodule!(session_reducer))?;
     m.add_wrapped(pyo3::wrap_pymodule!(accumulator))?;
     m.add_wrapped(pyo3::wrap_pymodule!(sinker))?;
+    m.add_wrapped(pyo3::wrap_pymodule!(sourcer))?;
 
     // Ensure it's importable as `pynumaflow_lite.mapper` as well as attribute access
     let binding = m.getattr("mapper")?;
@@ -134,5 +143,14 @@ fn pynumaflow_lite(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
         .getattr("modules")?
         .set_item(fullname, &sub)?;
 
+    // Ensure it's importable as `pynumaflow_lite.sourcer` as well
+    let binding = m.getattr("sourcer")?;
+    let sub = binding.downcast::<PyModule>()?;
+    let fullname = "pynumaflow_lite.sourcer";
+    sub.setattr("__name__", fullname)?;
+    py.import("sys")?
+        .getattr("modules")?
+        .set_item(fullname, &sub)?;
+
     Ok(())
 }

packages/pynumaflow-lite/src/sink/mod.rs

@@ -398,4 +398,3 @@ pub(crate) fn populate_py_module(m: &Bound<PyModule>) -> PyResult<()> {
 
     Ok(())
 }
-