feat: Batch Result serialization

- Adds serialization for batch result in the serdes module.

Unfortunately we need to do an adhoc import as we are dealing
with cyclical dependencies.

Author: FullyTyped

src/aws_durable_execution_sdk_python/serdes.py

@@ -62,6 +62,7 @@ class TypeTag(StrEnum):
     TUPLE = "t"
     LIST = "l"
     DICT = "m"
+    BATCH_RESULT = "br"
 
 
 @dataclass(frozen=True)
@@ -206,7 +207,17 @@ def dispatcher(self):
 
     def encode(self, obj: Any) -> EncodedValue:
         """Encode container using dispatcher for recursive elements."""
+        # Import here to avoid circular dependency
+        # concurrency -> child_handler -> serdes -> concurrency
+        from aws_durable_execution_sdk_python.concurrency import BatchResult
+
         match obj:
+            case BatchResult():
+                # Encode BatchResult as dict with special tag
+                return EncodedValue(
+                    TypeTag.BATCH_RESULT,
+                    self._wrap(obj.to_dict(), self.dispatcher).value,
+                )
             case list():
                 return EncodedValue(
                     TypeTag.LIST, [self._wrap(v, self.dispatcher) for v in obj]
@@ -230,7 +241,15 @@ def encode(self, obj: Any) -> EncodedValue:
 
     def decode(self, tag: TypeTag, value: Any) -> Any:
         """Decode container using dispatcher for recursive elements."""
+        # Import here to avoid circular dependency
+        from aws_durable_execution_sdk_python.concurrency import BatchResult
+
         match tag:
+            case TypeTag.BATCH_RESULT:
+                # Decode BatchResult from dict - value is already the dict structure
+                # First decode it as a dict to unwrap all nested EncodedValues
+                decoded_dict = self.decode(TypeTag.DICT, value)
+                return BatchResult.from_dict(decoded_dict)
             case TypeTag.LIST:
                 if not isinstance(value, list):
                     msg = f"Expected list, got {type(value)}"
@@ -295,6 +314,11 @@ def encode(self, obj: Any) -> EncodedValue:
             case list() | tuple() | dict():
                 return self.container_codec.encode(obj)
             case _:
+                # Check if it's a BatchResult (handled by container_codec)
+                from aws_durable_execution_sdk_python.concurrency import BatchResult
+
+                if isinstance(obj, BatchResult):
+                    return self.container_codec.encode(obj)
                 msg = f"Unsupported type: {type(obj)}"
                 raise SerDesError(msg)
 
@@ -316,7 +340,7 @@ def decode(self, tag: TypeTag, value: Any) -> Any:
                 return self.decimal_codec.decode(tag, value)
             case TypeTag.DATETIME | TypeTag.DATE:
                 return self.datetime_codec.decode(tag, value)
-            case TypeTag.LIST | TypeTag.TUPLE | TypeTag.DICT:
+            case TypeTag.LIST | TypeTag.TUPLE | TypeTag.DICT | TypeTag.BATCH_RESULT:
                 return self.container_codec.decode(tag, value)
             case _:
                 msg = f"Unknown type tag: {tag}"

tests/concurrency_test.py

@@ -1,5 +1,6 @@
 """Tests for the concurrency module."""
 
+import json
 import random
 import threading
 import time
@@ -101,29 +102,6 @@ def test_batch_item_from_dict():
     assert item.result == "success_result"
     assert item.error is None
 
-
-def test_batch_item_from_dict_with_error():
-    """Test BatchItem from_dict with error object."""
-    error_data = {
-        "message": "Test error",
-        "type": "TestError",
-        "data": None,
-        "stackTrace": None,
-    }
-    data = {
-        "index": 1,
-        "status": "FAILED",
-        "result": None,
-        "error": error_data,
-    }
-
-    item = BatchItem.from_dict(data)
-    assert item.index == 1
-    assert item.status == BatchItemStatus.FAILED
-    assert item.result is None
-    assert item.error is not None
-
-
 def test_batch_result_creation():
     """Test BatchResult creation."""
     items = [
@@ -2676,3 +2654,79 @@ def mock_get_checkpoint_result(operation_id):
         assert len(result.all) == 1
         assert result.all[0].status == BatchItemStatus.SUCCEEDED
         assert result.all[0].result == "re_executed_result"
+
+
+def test_batch_item_from_dict_with_error():
+    """Test BatchItem.from_dict() with error."""
+    data = {
+        "index": 3,
+        "status": "FAILED",
+        "result": None,
+        "error": {
+            "ErrorType": "ValueError",
+            "ErrorMessage": "bad value",
+            "StackTrace": [],
+        },
+    }
+
+    item = BatchItem.from_dict(data)
+
+    assert item.index == 3
+    assert item.status == BatchItemStatus.FAILED
+    assert item.error.type == "ValueError"
+    assert item.error.message == "bad value"
+
+
+def test_batch_result_with_mixed_statuses():
+    """Test BatchResult serialization with mixed item statuses."""
+    result = BatchResult(
+        all=[
+            BatchItem(0, BatchItemStatus.SUCCEEDED, result="success"),
+            BatchItem(
+                1,
+                BatchItemStatus.FAILED,
+                error=ErrorObject(message="msg", type="E", data=None, stack_trace=[]),
+            ),
+            BatchItem(2, BatchItemStatus.STARTED),
+        ],
+        completion_reason=CompletionReason.FAILURE_TOLERANCE_EXCEEDED,
+    )
+
+    serialized = json.dumps(result.to_dict())
+    deserialized = BatchResult.from_dict(json.loads(serialized))
+
+    assert len(deserialized.all) == 3
+    assert deserialized.all[0].status == BatchItemStatus.SUCCEEDED
+    assert deserialized.all[1].status == BatchItemStatus.FAILED
+    assert deserialized.all[2].status == BatchItemStatus.STARTED
+    assert deserialized.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED
+
+
+def test_batch_result_empty_list():
+    """Test BatchResult serialization with empty items list."""
+    result = BatchResult(all=[], completion_reason=CompletionReason.ALL_COMPLETED)
+
+    serialized = json.dumps(result.to_dict())
+    deserialized = BatchResult.from_dict(json.loads(serialized))
+
+    assert len(deserialized.all) == 0
+    assert deserialized.completion_reason == CompletionReason.ALL_COMPLETED
+
+
+def test_batch_result_complex_nested_data():
+    """Test BatchResult with complex nested data structures."""
+    complex_result = {
+        "users": [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}],
+        "metadata": {"count": 2, "timestamp": "2025-10-31"},
+    }
+
+    result = BatchResult(
+        all=[BatchItem(0, BatchItemStatus.SUCCEEDED, result=complex_result)],
+        completion_reason=CompletionReason.ALL_COMPLETED,
+    )
+
+    serialized = json.dumps(result.to_dict())
+    deserialized = BatchResult.from_dict(json.loads(serialized))
+
+    assert deserialized.all[0].result == complex_result
+    assert deserialized.all[0].result["users"][0]["name"] == "Alice"

tests/operation/map_test.py

@@ -1,5 +1,6 @@
 """Tests for map operation."""
 
+import json
 from unittest.mock import Mock, patch
 
 # Mock the executor.execute method
@@ -750,3 +751,46 @@ def get_checkpoint_result(self, operation_id):
         # Verify replay was called, execute was not
         mock_replay.assert_called_once()
         mock_execute.assert_not_called()
+
+
+
+def test_map_result_serialization_roundtrip():
+    """Test that map operation BatchResult can be serialized and deserialized."""
+
+    items = ["a", "b", "c"]
+
+    def func(ctx, item, idx, items):
+        return {"item": item.upper(), "index": idx}
+
+    class MockExecutionState:
+        durable_execution_arn = "arn:test"
+
+        def get_checkpoint_result(self, operation_id):
+            mock_result = Mock()
+            mock_result.is_succeeded.return_value = False
+            return mock_result
+
+    execution_state = MockExecutionState()
+    map_context = Mock()
+    map_context._create_step_id_for_logical_step = Mock(side_effect=["1", "2", "3"]) # noqa SLF001
+    map_context.create_child_context = Mock(return_value=Mock())
+    operation_identifier = OperationIdentifier("test_op", "parent", "test_map")
+
+    # Execute map
+    result = map_handler(
+        items, func, MapConfig(), execution_state, map_context, operation_identifier
+    )
+
+    # Serialize the BatchResult
+    serialized = json.dumps(result.to_dict())
+
+    # Deserialize
+    deserialized = BatchResult.from_dict(json.loads(serialized))
+
+    # Verify all data preserved
+    assert len(deserialized.all) == 3
+    assert deserialized.all[0].result == {"item": "A", "index": 0}
+    assert deserialized.all[1].result == {"item": "B", "index": 1}
+    assert deserialized.all[2].result == {"item": "C", "index": 2}
+    assert deserialized.completion_reason == result.completion_reason
+    assert all(item.status == BatchItemStatus.SUCCEEDED for item in deserialized.all)

tests/operation/parallel_test.py

@@ -1,5 +1,6 @@
 """Tests for the parallel operation module."""
 
+import json
 from unittest.mock import Mock, patch
 
 import pytest
@@ -734,3 +735,58 @@ def get_checkpoint_result(self, operation_id):
         # Verify replay was called, execute was not
         mock_replay.assert_called_once()
         mock_execute.assert_not_called()
+
+
+
+def test_parallel_result_serialization_roundtrip():
+    """Test that parallel operation BatchResult can be serialized and deserialized."""
+
+    def func1(ctx):
+        return [1, 2, 3]
+
+    def func2(ctx):
+        return {"status": "complete", "count": 42}
+
+    def func3(ctx):
+        return "simple string"
+
+    callables = [func1, func2, func3]
+
+    class MockExecutionState:
+        durable_execution_arn = "arn:test"
+
+        def get_checkpoint_result(self, operation_id):
+            mock_result = Mock()
+            mock_result.is_succeeded.return_value = False
+            return mock_result
+
+    execution_state = MockExecutionState()
+    parallel_context = Mock()
+    parallel_context._create_step_id_for_logical_step = Mock( # noqa SLF001
+        side_effect=["1", "2", "3"]
+    )
+    parallel_context.create_child_context = Mock(return_value=Mock())
+    operation_identifier = OperationIdentifier("test_op", "parent", "test_parallel")
+
+    # Execute parallel
+    result = parallel_handler(
+        callables,
+        ParallelConfig(),
+        execution_state,
+        parallel_context,
+        operation_identifier,
+    )
+
+    # Serialize the BatchResult
+    serialized = json.dumps(result.to_dict())
+
+    # Deserialize
+    deserialized = BatchResult.from_dict(json.loads(serialized))
+
+    # Verify all data preserved
+    assert len(deserialized.all) == 3
+    assert deserialized.all[0].result == [1, 2, 3]
+    assert deserialized.all[1].result == {"status": "complete", "count": 42}
+    assert deserialized.all[2].result == "simple string"
+    assert deserialized.completion_reason == result.completion_reason
+    assert all(item.status == BatchItemStatus.SUCCEEDED for item in deserialized.all)