Skip to content

Commit 2b36e56

Browse files
andystaplesCopilot
andauthored
Optimize deterministic set serialization without JSON parse round trips (#313)
* Optimize deterministic set serialization without JSON parse round trips Normalize set subtrees once, reuse normalized values for output, and skip singleton ordering keys while preserving legacy coercion defaults. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Preserve Unicode object-name normalization in set payloads Collapse equivalent UTF-16 surrogate-pair and non-BMP names before dictionary insertion, preserve lone surrogates, and consolidate test imports. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Align opt-in set discovery with tuple and abstract-set reconstruction Recognize hashable typed tuples recursively and reconstruct bare collections.abc.Set annotations as concrete sets only when enabled. Cover both opt-in names, discovery, activity execution, and unchanged legacy defaults. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Reduce dictionary-heavy set normalization overhead Compose only actual surrogate pairs, bypass recursive normalization for scalar dictionary fields, and avoid cycle-tracking writes for scalar-only records. Cover reviewer payload shapes and retain per-occurrence hooks and cycle detection. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent d2a3e7f commit 2b36e56

4 files changed

Lines changed: 393 additions & 22 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,15 @@ by one or more task hubs.
1414

1515
FIXED
1616

17+
- Optimized set and frozenset serialization, particularly for deeply nested
18+
frozensets and large singleton elements. JSON object names within sets are now
19+
ordered consistently after conversion to strings, including numeric names.
20+
Equivalent surrogate-pair and non-BMP object names retain the last value
21+
rather than emitting duplicate JSON properties.
22+
Set reconstruction remains opt-in.
23+
- With set reconstruction enabled, annotated sets of typed tuples are now
24+
recognized during input and activity-result discovery, and bare
25+
`collections.abc.Set` annotations reconstruct JSON arrays as concrete sets.
1726
- Implemented `__hash__` and `__repr__` for `EntityInstanceId`, making identity attributes immutable to guarantee hash stability when used in sets or dict keys.
1827
- Large-payload externalization now checks every eligible payload's maximum
1928
size before uploading or replacing any protobuf values, including carryover

‎durabletask/serialization.py‎

Lines changed: 77 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import inspect
3939
import json
4040
import logging
41+
import re
4142
import sys
4243
import types
4344
import typing
@@ -56,6 +57,8 @@
5657
# replay. Internal detail; re-exported from ``durabletask.internal.shared`` only
5758
# for backwards compatibility.
5859
_AUTO_SERIALIZED = "__durabletask_autoobject__"
60+
_SURROGATE_PAIR = re.compile("[\ud800-\udbff][\udc00-\udfff]")
61+
_JSON_SCALAR_TYPES = (int, str, float, type(None))
5962

6063

6164
class DataConverter(ABC):
@@ -137,6 +140,8 @@ class JsonDataConverter(DataConverter):
137140
own hooks. Sets and frozensets serialize to JSON arrays ordered
138141
deterministically by their canonical JSON representation (using lexical
139142
string order, so numeric sets like ``{2, 10}`` serialize as ``[10, 2]``).
143+
Within sets, object keys are normalized to strings and sorted, and custom
144+
hooks run once per occurrence. Singleton sets need no ordering key.
140145
141146
Deserialization (and value-level :meth:`coerce`) is **best-effort**: when a
142147
``target_type`` is supplied and the value cannot be coerced to it, the raw
@@ -247,6 +252,11 @@ def _is_supported_set_element(converter: DataConverter, elem_type: Any) -> bool:
247252
if elem_origin is frozenset:
248253
args = typing.get_args(elem_type)
249254
return not args or _is_supported_set_element(converter, args[0])
255+
if elem_origin is tuple:
256+
args = typing.get_args(elem_type)
257+
if len(args) == 2 and args[1] is Ellipsis:
258+
args = args[:1]
259+
return all(_is_supported_set_element(converter, arg) for arg in args)
250260
return False
251261
if not isinstance(elem_type, type):
252262
return False
@@ -368,25 +378,9 @@ def _encode_custom_object(o: Any) -> Any:
368378
if callable(to_json_hook):
369379
return to_json_hook(o)
370380
if isinstance(o, (set, frozenset, AbstractSet)):
371-
items: AbstractSet[Any] = cast("AbstractSet[Any]", o)
372-
# Sets and frozensets are unordered collections. To ensure deterministic
373-
# serialization across processes, worker replays, and PYTHONHASHSEED
374-
# variations, elements are individually encoded to canonical JSON
375-
# (with sort_keys=True), sorted lexically by their encoded JSON string,
376-
# and emitted as normalized parsed values.
377-
#
378-
# Emitting the normalized parsed values rather than original objects
379-
# guarantees identical order and eliminates leaks from dict insertion
380-
# order variations, NaN/non-finite comparison anomalies, or partial
381-
# ordering in frozensets. A consequence of sorting canonical JSON strings
382-
# is lexical ordering for numeric sets (e.g., {2, 10} serializes as
383-
# [10, 2]).
381+
items = cast(AbstractSet[Any], o)
384382
try:
385-
encoded_items = sorted(
386-
json.dumps(item, default=_encode_custom_object, sort_keys=True)
387-
for item in items
388-
)
389-
return [json.loads(encoded) for encoded in encoded_items]
383+
return _normalize_set_value(items, set())
390384
except Exception as e:
391385
raise TypeError(
392386
f"Cannot deterministically serialize set of type '{type(items).__name__}' "
@@ -415,6 +409,66 @@ def _encode_custom_object(o: Any) -> Any:
415409
raise TypeError(f"Object of type '{type(o).__name__}' is not JSON serializable")
416410

417411

412+
def _normalize_set_value(value: Any, active: set[int]) -> Any:
413+
"""Normalize only set subtrees; leave ordinary payload encoding unchanged."""
414+
identity = id(value)
415+
is_mapping = isinstance(value, dict)
416+
if not is_mapping and isinstance(value, _JSON_SCALAR_TYPES):
417+
return value
418+
if identity in active:
419+
raise ValueError("Circular reference detected")
420+
if is_mapping:
421+
# Match sort_keys=True, including rejection of incomparable keys.
422+
# JSON object names are strings; duplicate names keep the last value,
423+
# just as the previous dumps/loads normalization did.
424+
result: dict[str, Any] = {}
425+
string_names = True
426+
for key, item in sorted(cast(dict[Any, Any], value).items()):
427+
if isinstance(key, str):
428+
name = key
429+
if not str.isascii(name) and _SURROGATE_PAIR.search(name):
430+
# Ordinary Unicode names and lone surrogates need no codec
431+
# round trip; only a high/low pair changes after decoding.
432+
name = str.encode(name, "utf-16-le", "surrogatepass").decode(
433+
"utf-16-le", "surrogatepass")
434+
string_names = False
435+
else:
436+
if key is not None and not isinstance(key, (int, float)):
437+
raise TypeError(f"Unsupported JSON object key: {type(key).__name__}")
438+
name = json.dumps(key)
439+
string_names = False
440+
if isinstance(item, _JSON_SCALAR_TYPES):
441+
result[name] = item
442+
else:
443+
# Scalar-only records cannot cycle. Track the parent only
444+
# while descending into a child that could refer back to it.
445+
active.add(identity)
446+
try:
447+
result[name] = _normalize_set_value(item, active)
448+
finally:
449+
active.remove(identity)
450+
return result if string_names else dict(sorted(result.items()))
451+
active.add(identity)
452+
try:
453+
if isinstance(value, (list, tuple)):
454+
return [
455+
item if isinstance(item, _JSON_SCALAR_TYPES)
456+
else _normalize_set_value(item, active)
457+
for item in cast(Sequence[Any], value)
458+
]
459+
to_json_hook = getattr(cast(Any, type(value)), "to_json", None)
460+
if isinstance(value, AbstractSet) and not callable(to_json_hook):
461+
items = [_normalize_set_value(item, active) for item in cast(AbstractSet[Any], value)]
462+
if len(items) > 1:
463+
# list.sort evaluates each key once and discards the temporary
464+
# documents afterwards. Emit these same normalized values.
465+
items.sort(key=json.dumps)
466+
return items
467+
return _normalize_set_value(_encode_custom_object(value), active)
468+
finally:
469+
active.remove(identity)
470+
471+
418472
def _legacy_object_hook(d: dict[str, Any]) -> Any:
419473
# If the object carries the legacy marker, deserialize it as a SimpleNamespace.
420474
if d.pop(_AUTO_SERIALIZED, False):
@@ -458,6 +512,10 @@ def _coerce_to_type(value: Any, expected_type: Any, converter: DataConverter | N
458512
if isinstance(value, expected_type):
459513
return value
460514

515+
if (expected_type is AbstractSet and isinstance(value, list)
516+
and getattr(converter, "enable_generic_sets", False)):
517+
return _coerce_generic(value, expected_type, AbstractSet, converter)
518+
461519
from_json_hook = getattr(expected_type, "from_json", None)
462520
if callable(from_json_hook):
463521
return _invoke_from_json(from_json_hook, value, converter)
@@ -471,7 +529,7 @@ def _coerce_to_type(value: Any, expected_type: Any, converter: DataConverter | N
471529
except Exception as e:
472530
type_name = getattr(type_ctor, "__name__", None) or str(type_ctor)
473531
raise TypeError(
474-
f"Could not coerce value of type '{type(value).__name__}' to "
532+
f"Could not coerce value of type '{type(cast(object, value)).__name__}' to "
475533
f"'{type_name}'"
476534
) from e
477535

‎tests/durabletask/test_data_converter.py‎

Lines changed: 83 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,20 +6,24 @@
66
import inspect
77
import json
88
import logging
9+
from collections.abc import Set as AbstractSet
910
from dataclasses import dataclass, field
1011
from functools import wraps
11-
from typing import Any
12+
from typing import AbstractSet as TypingAbstractSet, Any
1213
from unittest.mock import patch
1314

1415
import pytest
16+
from durabletask import task
17+
from durabletask.internal import helpers
18+
from durabletask.internal import orchestrator_service_pb2 as pb
1519
from durabletask.internal.entity_state_shim import StateShim
1620
from durabletask.serialization import (
1721
DEFAULT_DATA_CONVERTER,
1822
DataConverter,
1923
JsonDataConverter,
2024
)
2125
from durabletask.internal import type_discovery
22-
from durabletask.worker import _ActivityExecutor, _Registry
26+
from durabletask.worker import _ActivityExecutor, _OrchestrationExecutor, _Registry
2327

2428

2529
@dataclass
@@ -423,6 +427,60 @@ def test_can_reconstruct_set_and_frozenset():
423427
assert opt_in_converter.can_reconstruct(set[dict[str, int]]) is False
424428

425429

430+
@pytest.mark.parametrize("option", ["enable_generic_sets", "enable_set_coercion"])
431+
@pytest.mark.parametrize("target_type,payload,expected", [
432+
(set[tuple[int, int]], "[[1, 2], [3, 4]]", {(1, 2), (3, 4)}),
433+
(frozenset[tuple[int, ...]], "[[1], [2, 3]]", frozenset({(1,), (2, 3)})),
434+
(set[tuple[int, frozenset[str]]], '[[1, ["b", "a"]]]', {(1, frozenset({"a", "b"}))}),
435+
(frozenset[frozenset[tuple[int, int]]], "[[[1, 2], [3, 4]]]",
436+
frozenset({frozenset({(1, 2), (3, 4)})})),
437+
(set[tuple[()]], "[[]]", {()}),
438+
(set[tuple[int, str | None]], '[[1, null], [2, "x"]]', {(1, None), (2, "x")}),
439+
(AbstractSet, "[2, 1, 2]", {1, 2}),
440+
(TypingAbstractSet, "[2, 1, 2]", {1, 2}),
441+
])
442+
def test_opt_in_set_discovery_matches_reconstruction(
443+
option: str, target_type: Any, payload: str, expected: Any,
444+
):
445+
converter = JsonDataConverter(**{option: True})
446+
447+
def activity(ctx: Any, values: Any) -> Any:
448+
assert isinstance(values, type(expected))
449+
assert values == expected
450+
return values
451+
452+
activity.__annotations__["values"] = target_type
453+
activity.__annotations__["return"] = target_type
454+
assert converter.can_reconstruct(target_type)
455+
assert converter.deserialize(payload, target_type) == expected
456+
assert converter.coerce(json.loads(payload), target_type) == expected
457+
assert type_discovery.activity_input_type(activity, converter) == target_type
458+
assert type_discovery.orchestrator_input_type(activity, converter) == target_type
459+
assert type_discovery.activity_output_type(activity, converter) == target_type
460+
registry = _Registry()
461+
name = registry.add_activity(activity)
462+
executor = _ActivityExecutor(registry, logging.getLogger("test_set_discovery"), converter)
463+
output = executor.execute("instance", name, 1, payload)
464+
assert converter.deserialize(output, target_type) == expected
465+
466+
default = JsonDataConverter()
467+
assert not default.can_reconstruct(target_type)
468+
assert type_discovery.activity_input_type(activity, default) is None
469+
assert type_discovery.orchestrator_input_type(activity, default) is None
470+
assert type_discovery.activity_output_type(activity, default) is None
471+
assert default.deserialize(payload, target_type) == json.loads(payload)
472+
473+
474+
@pytest.mark.parametrize("target_type", [
475+
set[tuple[list[int], int]],
476+
set[tuple[dict[str, int], ...]],
477+
set[tuple[frozenset[list[int]], int]],
478+
frozenset[tuple[set[int], ...]],
479+
])
480+
def test_set_tuple_discovery_rejects_unhashable_elements(target_type: Any):
481+
assert not JsonDataConverter(enable_generic_sets=True).can_reconstruct(target_type)
482+
483+
426484
def test_existing_instance_replay_compatibility_preserves_list():
427485
"""Preserve existing-instance replay behavior for recorded arrays under set annotations.
428486
@@ -453,6 +511,29 @@ def test_existing_instance_replay_compatibility_preserves_list():
453511
assert len(opt_in_restored.flags) == 2
454512

455513

514+
def test_default_converter_replays_recorded_arrays_under_set_annotations():
515+
def orchestrator(ctx: task.OrchestrationContext, data: TaggedOrder):
516+
yield ctx.call_activity(f"count-{len(data.tags)}")
517+
return {"tags": data.tags, "flags": data.flags}
518+
519+
registry = _Registry()
520+
name = registry.add_orchestrator(orchestrator)
521+
history = [
522+
helpers.new_orchestrator_started_event(),
523+
helpers.new_execution_started_event(
524+
name, "legacy-sets",
525+
encoded_input='{"item": "laptop", "tags": ["b", "a", "b"], "flags": [1, 2, 1]}',
526+
),
527+
helpers.new_task_scheduled_event(1, "count-3"),
528+
]
529+
executor = _OrchestrationExecutor(registry, logging.getLogger("test_set_replay"), JsonDataConverter())
530+
result = executor.execute("legacy-sets", history, [helpers.new_task_completed_event(1, "null")])
531+
assert len(result.actions) == 1
532+
completion = result.actions[0].completeOrchestration
533+
assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
534+
assert completion.result.value == '{"tags": ["b", "a", "b"], "flags": [1, 2, 1]}'
535+
536+
456537
def test_type_discovery_and_activity_execution_with_sets():
457538
"""Verify that type discovery recognizes set annotations when opted in,
458539
enabling activities to receive sets rather than raw lists."""

0 commit comments

Comments
 (0)