feat: introduce API to fail a message in udf/transformer

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)

packages/pynumaflow/tests/map/test_async_mapper.py

@@ -16,8 +16,13 @@
 from pynumaflow.proto.common import metadata_pb2
 from pynumaflow.proto.mapper import map_pb2, map_pb2_grpc
 from tests.conftest import create_async_loop, start_async_server, teardown_async_server
-from tests.map.utils import get_test_datums, async_nack_map_handler, NACK_TEST_OPTIONS
-from pynumaflow._constants import NACK
+from tests.map.utils import (
+    get_test_datums,
+    async_nack_map_handler,
+    async_fail_map_handler,
+    NACK_TEST_OPTIONS,
+)
+from pynumaflow._constants import FAIL, NACK
 
 pytestmark = pytest.mark.integration
 
@@ -28,6 +33,7 @@
 
 SOCK_PATH = "unix:///tmp/async_map.sock"
 NACK_SOCK_PATH = "unix:///tmp/async_map_nack.sock"
+FAIL_SOCK_PATH = "unix:///tmp/async_map_fail.sock"
 
 
 def request_generator(req):
@@ -156,6 +162,51 @@ def test_map_nack(async_nack_map_server):
         assert result.nack_options.reason == NACK_TEST_OPTIONS.reason
 
 
+async def _start_fail_server(udfs):
+    _server_options = [
+        ("grpc.max_send_message_length", MAX_MESSAGE_SIZE),
+        ("grpc.max_receive_message_length", MAX_MESSAGE_SIZE),
+    ]
+    server = grpc.aio.server(options=_server_options)
+    map_pb2_grpc.add_MapServicer_to_server(udfs, server)
+    server.add_insecure_port(FAIL_SOCK_PATH)
+    logging.info("Starting fail server on %s", FAIL_SOCK_PATH)
+    await server.start()
+    return server, FAIL_SOCK_PATH
+
+
+@pytest.fixture(scope="module")
+def async_fail_map_server():
+    """Module-scoped fixture: async gRPC map server whose handler fails every message."""
+    loop = create_async_loop()
+
+    server_obj = MapAsyncServer(mapper_instance=async_fail_map_handler)
+    udfs = server_obj.servicer
+    server = start_async_server(loop, _start_fail_server(udfs))
+
+    yield loop
+
+    teardown_async_server(loop, server)
+
+
+def test_map_fail(async_fail_map_server):
+    with grpc.insecure_channel(FAIL_SOCK_PATH) as channel:
+        stub = map_pb2_grpc.MapStub(channel)
+        request = get_test_datums()
+        generator_response = stub.MapFn(request_iterator=request_generator(request))
+
+        responses = list(generator_response)
+
+    # 1 handshake + 3 data responses
+    assert len(responses) == 4
+    assert responses[0].handshake.sot
+
+    for resp in responses[1:]:
+        assert len(resp.results) == 1
+        result = resp.results[0]
+        assert FAIL in result.tags
+
+
 def test_map(map_stub):
     request = get_test_datums()
     generator_response: Iterator[map_pb2.MapResponse] = map_stub.MapFn(

packages/pynumaflow/tests/map/test_messages.py

@@ -1,7 +1,7 @@
 import pytest
 
 from pynumaflow.mapper import Messages, Message, DROP, Mapper, Datum, NackOptions
-from pynumaflow._constants import NACK
+from pynumaflow._constants import FAIL, NACK
 from tests.testing_utils import mock_message
 
 
@@ -52,6 +52,14 @@ def test_message_to_nack_default_opts():
     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_message_default_nack_options():
     msg = Message(mock_message())
     assert msg.nack_options is None

packages/pynumaflow/tests/map/test_sync_mapper.py

@@ -12,8 +12,9 @@
     get_test_datums,
     nack_map_handler,
     NACK_TEST_OPTIONS,
+    fail_map_handler,
 )
-from pynumaflow._constants import NACK
+from pynumaflow._constants import FAIL, NACK
 from tests.conftest import collect_responses, drain_responses, send_test_requests
 
 
@@ -131,6 +132,28 @@ def test_map_nack():
     assert code == StatusCode.OK
 
 
+def test_map_fail():
+    my_server = MapServer(mapper_instance=fail_map_handler)
+    services = {map_pb2.DESCRIPTOR.services_by_name["Map"]: my_server.servicer}
+    test_server = server_from_dictionary(services, strict_real_time())
+
+    test_datums = get_test_datums(handshake=True)
+    method = _invoke_map_fn(test_server)
+    send_test_requests(method, test_datums)
+    responses = collect_responses(method)
+
+    metadata, code, details = method.termination()
+    # 1 handshake + 3 data responses
+    assert len(responses) == 4
+    assert responses[0].handshake.sot
+
+    for resp in responses[1:]:
+        assert len(resp.results) == 1
+        result = resp.results[0]
+        assert FAIL in result.tags
+    assert code == StatusCode.OK
+
+
 def test_invalid_input():
     with pytest.raises(TypeError):
         MapServer()

packages/pynumaflow/tests/map/utils.py

@@ -16,6 +16,14 @@ async def async_nack_map_handler(keys: list[str], datum: Datum) -> Messages:
     return Messages(Message.to_nack(NACK_TEST_OPTIONS))
 
 
+def fail_map_handler(keys: list[str], datum: Datum) -> Messages:
+    return Messages(Message.to_fail())
+
+
+async def async_fail_map_handler(keys: list[str], datum: Datum) -> Messages:
+    return Messages(Message.to_fail())
+
+
 async def async_map_error_fn(keys: list[str], datum: Datum) -> Messages:
     raise ValueError("error invoking map")
 

packages/pynumaflow/tests/mapstream/test_messages.py

@@ -1,7 +1,7 @@
 import pytest
 
 from pynumaflow.mapstreamer import Messages, Message, DROP, NackOptions
-from pynumaflow._constants import NACK
+from pynumaflow._constants import FAIL, NACK
 from tests.testing_utils import mock_message
 
 
@@ -26,6 +26,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_message_key():
     mock_obj = {"Keys": ["test-key"], "Value": mock_message()}
     msg = Message(value=mock_obj["Value"], keys=mock_obj["Keys"])

packages/pynumaflow/tests/sourcetransform/test_messages.py

@@ -11,7 +11,7 @@
     SystemMetadata,
     NackOptions,
 )
-from pynumaflow._constants import NACK
+from pynumaflow._constants import FAIL, NACK
 from tests.testing_utils import mock_new_event_time
 
 
@@ -83,6 +83,15 @@ def test_message_to_nack_default_opts():
     assert msgt.nack_options is None
 
 
+def test_message_to_fail():
+    msgt = Message.to_fail(mock_event_time())
+    assert isinstance(msgt, Message)
+    assert msgt.keys == []
+    assert msgt.value == b""
+    assert msgt.tags == [FAIL]
+    assert msgt.event_time == mock_event_time()
+
+
 def test_message_default_nack_options():
     msgt = Message(mock_message_t(), mock_event_time())
     assert msgt.nack_options is None