feat: introduce API to fail a message in udf/transformer (#370)

Signed-off-by: Vaibhav Tiwari <vaibhav.tiwari33@gmail.com>

Author: vaibhavtiwari33

packages/pynumaflow/pynumaflow/_constants.py

@@ -55,6 +55,7 @@
 DELIMITER = ":"
 DROP = "U+005C__DROP__"
 NACK = "U+005C__NACK__"
+FAIL = "U+005C__FAIL__"
 
 _PROCESS_COUNT = os.cpu_count()
 # Cap max value to 16

packages/pynumaflow/pynumaflow/batchmapper/__init__.py

@@ -1,4 +1,4 @@
-from pynumaflow._constants import DROP
+from pynumaflow._constants import DROP, FAIL
 
 from pynumaflow.batchmapper._dtypes import (
     Message,
@@ -14,6 +14,7 @@
     "Message",
     "Datum",
     "DROP",
+    "FAIL",
     "BatchMapAsyncServer",
     "BatchMapper",
     "BatchResponses",

packages/pynumaflow/pynumaflow/batchmapper/_dtypes.py

@@ -5,7 +5,7 @@
 from typing import TypeAlias, TypeVar
 from collections.abc import AsyncIterable, Callable
 
-from pynumaflow._constants import DROP, NACK
+from pynumaflow._constants import DROP, FAIL, NACK
 from pynumaflow._nack import NackOptions
 from pynumaflow._validate import _validate_message_fields
 
@@ -51,6 +51,10 @@ def to_nack(cls: type[M], opts: NackOptions | None = None) -> M:
         m._nack_options = opts
         return m
 
+    @classmethod
+    def to_fail(cls: type[M]) -> M:
+        return cls(b"", None, [FAIL])
+
     @property
     def nack_options(self) -> NackOptions | None:
         return self._nack_options

packages/pynumaflow/pynumaflow/mapper/__init__.py

@@ -5,12 +5,14 @@
 from pynumaflow.mapper._dtypes import Message, Messages, Datum, DROP, Mapper
 from pynumaflow._metadata import UserMetadata, SystemMetadata
 from pynumaflow._nack import NackOptions
+from pynumaflow._constants import FAIL
 
 __all__ = [
     "Message",
     "Messages",
     "Datum",
     "DROP",
+    "FAIL",
     "Mapper",
     "MapServer",
     "MapAsyncServer",

packages/pynumaflow/pynumaflow/mapper/_dtypes.py

@@ -6,7 +6,7 @@
 from collections.abc import Callable
 from warnings import warn
 
-from pynumaflow._constants import DROP, NACK
+from pynumaflow._constants import DROP, FAIL, NACK
 from pynumaflow._nack import NackOptions
 from pynumaflow._metadata import UserMetadata, SystemMetadata
 from pynumaflow._validate import _validate_message_fields
@@ -61,6 +61,10 @@ def to_nack(cls: type[M], opts: NackOptions | None = None) -> M:
         m._nack_options = opts
         return m
 
+    @classmethod
+    def to_fail(cls: type[M]) -> M:
+        return cls(b"", None, [FAIL])
+
     @property
     def nack_options(self) -> NackOptions | None:
         return self._nack_options

packages/pynumaflow/pynumaflow/mapstreamer/__init__.py

@@ -1,4 +1,4 @@
-from pynumaflow._constants import DROP
+from pynumaflow._constants import DROP, FAIL
 
 from pynumaflow.mapstreamer._dtypes import Message, Messages, Datum, MapStreamer
 from pynumaflow.mapstreamer.async_server import MapStreamAsyncServer
@@ -9,6 +9,7 @@
     "Messages",
     "Datum",
     "DROP",
+    "FAIL",
     "MapStreamAsyncServer",
     "MapStreamer",
     "NackOptions",

packages/pynumaflow/pynumaflow/mapstreamer/_dtypes.py

@@ -6,7 +6,7 @@
 from collections.abc import AsyncIterable, Callable
 from warnings import warn
 
-from pynumaflow._constants import DROP, NACK
+from pynumaflow._constants import DROP, FAIL, NACK
 from pynumaflow._nack import NackOptions
 from pynumaflow._validate import _validate_message_fields
 
@@ -51,6 +51,10 @@ def to_nack(cls: type[M], opts: NackOptions | None = None) -> M:
         m._nack_options = opts
         return m
 
+    @classmethod
+    def to_fail(cls: type[M]) -> M:
+        return cls(b"", None, [FAIL])
+
     @property
     def nack_options(self) -> NackOptions | None:
         return self._nack_options

packages/pynumaflow/pynumaflow/sourcetransformer/__init__.py

@@ -10,12 +10,14 @@
 from pynumaflow.sourcetransformer.async_server import SourceTransformAsyncServer
 from pynumaflow._metadata import UserMetadata, SystemMetadata
 from pynumaflow._nack import NackOptions
+from pynumaflow._constants import FAIL
 
 __all__ = [
     "Message",
     "Messages",
     "Datum",
     "DROP",
+    "FAIL",
     "SourceTransformServer",
     "SourceTransformer",
     "SourceTransformMultiProcServer",

packages/pynumaflow/pynumaflow/sourcetransformer/_dtypes.py

@@ -6,7 +6,7 @@
 from collections.abc import Awaitable, Callable
 from warnings import warn
 
-from pynumaflow._constants import DROP, NACK
+from pynumaflow._constants import DROP, FAIL, NACK
 from pynumaflow._nack import NackOptions
 from pynumaflow._metadata import UserMetadata, SystemMetadata
 from pynumaflow._validate import _validate_message_fields
@@ -70,6 +70,10 @@ def to_nack(
         m._nack_options = opts
         return m
 
+    @classmethod
+    def to_fail(cls: type[M], event_time: datetime) -> M:
+        return cls(b"", event_time, None, [FAIL])
+
     @property
     def nack_options(self) -> NackOptions | None:
         return self._nack_options

packages/pynumaflow/tests/batchmap/test_messages.py

@@ -1,7 +1,7 @@
 import pytest
 
 from pynumaflow.batchmapper import Message, DROP, BatchResponse, BatchResponses, NackOptions
-from pynumaflow._constants import NACK
+from pynumaflow._constants import FAIL, NACK
 from tests.batchmap.test_datatypes import TEST_ID
 from tests.testing_utils import mock_message
 
@@ -32,6 +32,14 @@ def test_message_default_nack_options():
     assert msg.nack_options is None
 
 
+def test_message_to_fail():
+    msg = Message.to_fail()
+    assert type(msg) is Message
+    assert msg.keys == []
+    assert msg.value == b""
+    assert msg.tags == [FAIL]
+
+
 def test_batch_responses_init():
     batch_responses = BatchResponses()
     batch_response1 = BatchResponse.from_id(TEST_ID)