perf(ws): collapse snapshot apply into a single dict walk (#263)

kalshi/ws/models/orderbook_delta.py

@@ -1,30 +1,77 @@
 """Orderbook delta and snapshot message models."""
+
 from __future__ import annotations
 
-from typing import Literal
+from decimal import Decimal
+from typing import Annotated, Any, Literal
+
+from pydantic import AliasChoices, BaseModel, BeforeValidator, Field
+
+from kalshi.types import DollarDecimal, FixedPointCount, _coerce_decimal
+
+
+def _levels_to_dict(value: Any) -> Any:
+    """Collapse the snapshot wire payload into a single dict in one walk.
+
+    The wire format for ``yes_dollars_fp`` / ``no_dollars_fp`` is a list of
+    ``[price, count]`` string pairs. Previously the model validated that into
+    ``list[tuple[Decimal, Decimal]]`` and ``OrderbookManager`` then called
+    ``dict(...)`` on it — a second per-level Python walk that dominated
+    snapshot apply cost for deep books (#263).
 
-from pydantic import AliasChoices, BaseModel, Field
+    Coercing each price/count to ``Decimal`` inline and returning a fully
+    typed ``dict[Decimal, Decimal]`` means (a) pydantic-core just bounces
+    the already-validated map back without a second walk, and (b)
+    ``_apply_snapshot_inplace`` adopts the dict directly with no rebuild.
 
-from kalshi.types import DollarDecimal, FixedPointCount
+    Accepts ``None`` (→ empty), an already-built dict (direct constructor /
+    re-validation — keys/values still coerced so callers can pass strings),
+    or any iterable of 2-element ``(price, count)`` pairs (wire shape).
+    """
+    if value is None:
+        return {}
+    coerce = _coerce_decimal
+    if isinstance(value, dict):
+        # Fast path when keys/values are already Decimal (avoids hashing the
+        # same Decimal twice on dict rebuild).
+        if all(type(k) is Decimal and type(v) is Decimal for k, v in value.items()):
+            return value
+        return {coerce(k): coerce(v) for k, v in value.items()}
+    out: dict[Decimal, Decimal] = {}
+    for row in value:
+        price, count = row
+        out[coerce(price)] = coerce(count)
+    return out
+
+
+# Field type is ``dict[Decimal, Decimal]`` (not ``dict[DollarDecimal, FixedPointCount]``)
+# so pydantic-core does not re-run the per-item ``BeforeValidator(_coerce_decimal)``
+# walk on values our ``BeforeValidator`` has already coerced. The on-the-wire
+# semantics are unchanged — prices are dollar-decimals, counts are fixed-point
+# counts — they are simply both ``Decimal`` once parsed.
+PriceCountMap = Annotated[dict[Decimal, Decimal], BeforeValidator(_levels_to_dict)]
 
 
 class OrderbookSnapshotPayload(BaseModel):
     """Payload for orderbook_snapshot messages.
 
     Wire format per AsyncAPI spec: ``yes_dollars_fp`` and ``no_dollars_fp`` are
     arrays of ``[price_in_dollars, contract_count_fp]`` string pairs. Each row
-    is two JSON strings, e.g. ``["0.5500", "100.00"]``. Consumers should parse
-    both elements as :class:`decimal.Decimal`.
+    is two JSON strings, e.g. ``["0.5500", "100.00"]``.
+
+    These are validated directly into ``dict[Decimal, Decimal]`` (price ->
+    count) in a single walk so ``OrderbookManager._apply_snapshot_inplace``
+    can adopt the map with no second iteration (#263).
     """
 
     market_ticker: str
     market_id: str
-    yes: list[tuple[DollarDecimal, FixedPointCount]] = Field(
-        default=[],
+    yes: PriceCountMap = Field(
+        default_factory=dict,
         validation_alias=AliasChoices("yes_dollars_fp", "yes"),
     )
-    no: list[tuple[DollarDecimal, FixedPointCount]] = Field(
-        default=[],
+    no: PriceCountMap = Field(
+        default_factory=dict,
         validation_alias=AliasChoices("no_dollars_fp", "no"),
     )
     model_config = {"extra": "allow", "populate_by_name": True}
@@ -58,6 +105,7 @@ class OrderbookDeltaPayload(BaseModel):
 
 class OrderbookSnapshotMessage(BaseModel):
     """Full orderbook snapshot, sent on initial subscribe."""
+
     type: str = "orderbook_snapshot"
     sid: int
     seq: int
@@ -67,6 +115,7 @@ class OrderbookSnapshotMessage(BaseModel):
 
 class OrderbookDeltaMessage(BaseModel):
     """Incremental orderbook update."""
+
     type: str = "orderbook_delta"
     sid: int
     seq: int

kalshi/ws/orderbook.py

@@ -1,4 +1,5 @@
 """Local orderbook manager from WebSocket snapshots and deltas."""
+
 from __future__ import annotations
 
 import logging
@@ -62,12 +63,10 @@ def to_orderbook(self) -> Orderbook:
         if self._cached is not None:
             return self._cached
         yes_levels = [
-            OrderbookLevel(price=price, quantity=qty)
-            for price, qty in sorted(self.yes.items())
+            OrderbookLevel(price=price, quantity=qty) for price, qty in sorted(self.yes.items())
         ]
         no_levels = [
-            OrderbookLevel(price=price, quantity=qty)
-            for price, qty in sorted(self.no.items())
+            OrderbookLevel(price=price, quantity=qty) for price, qty in sorted(self.no.items())
         ]
         ob = Orderbook(ticker=self.ticker, yes=yes_levels, no=no_levels)
         self._cached = ob
@@ -191,9 +190,7 @@ def _apply_delta_inplace(self, msg: OrderbookDeltaMessage) -> bool:
 
         return True
 
-    def apply_snapshot(
-        self, msg: OrderbookSnapshotMessage, sid: int | None = None
-    ) -> Orderbook:
+    def apply_snapshot(self, msg: OrderbookSnapshotMessage, sid: int | None = None) -> Orderbook:
         """Initialize (or reset) a book from a full snapshot.
 
         Public wrapper that mutates in place then materializes the

scripts/bench_orderbook_delta.py

@@ -28,28 +28,41 @@ def _build_snapshot(ticker: str, depth: int) -> OrderbookSnapshotMessage:
     yes = [(Decimal("0." + f"{i:04d}"), Decimal("100.00")) for i in range(1, depth + 1)]
     no = [(Decimal("0." + f"{i:04d}"), Decimal("100.00")) for i in range(1, depth + 1)]
     return OrderbookSnapshotMessage(
-        sid=1, seq=1,
+        sid=1,
+        seq=1,
         msg=OrderbookSnapshotPayload(
-            market_ticker=ticker, market_id=ticker, yes=yes, no=no,
+            market_ticker=ticker,
+            market_id=ticker,
+            yes=yes,
+            no=no,
         ),
     )
 
 
 def _build_delta(
-    ticker: str, seq: int, price: Decimal, delta: Decimal, side: str,
+    ticker: str,
+    seq: int,
+    price: Decimal,
+    delta: Decimal,
+    side: str,
 ) -> OrderbookDeltaMessage:
     return OrderbookDeltaMessage(
-        sid=1, seq=seq,
+        sid=1,
+        seq=seq,
         msg=OrderbookDeltaPayload(
-            market_ticker=ticker, market_id=ticker,
-            price=price, delta=delta, side=side,  # type: ignore[arg-type]
+            market_ticker=ticker,
+            market_id=ticker,
+            price=price,
+            delta=delta,
+            side=side,  # type: ignore[arg-type]
         ),
     )
 
 
 def main() -> None:
-    p = argparse.ArgumentParser(description=__doc__,
-        formatter_class=argparse.RawDescriptionHelpFormatter)
+    p = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
     p.add_argument("--depth", type=int, default=100)
     p.add_argument("--updates", type=int, default=100_000)
     args = p.parse_args()
@@ -60,12 +73,16 @@ def main() -> None:
 
     # Pre-build deltas so the loop only measures _apply_delta_inplace.
     prices = [Decimal("0." + f"{(i % args.depth) + 1:04d}") for i in range(args.updates)]
-    deltas = [_build_delta(
-        ticker, seq=i + 2,
-        price=prices[i],
-        delta=Decimal("1.00") if i % 2 == 0 else Decimal("-1.00"),
-        side="yes" if i % 2 == 0 else "no",
-    ) for i in range(args.updates)]
+    deltas = [
+        _build_delta(
+            ticker,
+            seq=i + 2,
+            price=prices[i],
+            delta=Decimal("1.00") if i % 2 == 0 else Decimal("-1.00"),
+            side="yes" if i % 2 == 0 else "no",
+        )
+        for i in range(args.updates)
+    ]
 
     gc.collect()
     t0 = time.perf_counter()

tests/ws/test_models.py

@@ -99,7 +99,10 @@ def test_parse_snapshot(self) -> None:
         assert msg.seq == 1
         assert msg.msg.market_ticker == "ECON-GDP-25Q1"
         assert len(msg.msg.yes) == 2
-        assert msg.msg.yes[0] == (Decimal("0.50"), Decimal("100.00"))
+        assert msg.msg.yes == {
+            Decimal("0.50"): Decimal("100.00"),
+            Decimal("0.55"): Decimal("200.00"),
+        }
 
     def test_parse_delta(self) -> None:
         raw = {
@@ -129,8 +132,8 @@ def test_snapshot_empty_book(self) -> None:
             "msg": {"market_ticker": "T", "market_id": "x", "yes": [], "no": []},
         }
         msg = OrderbookSnapshotMessage.model_validate(raw)
-        assert msg.msg.yes == []
-        assert msg.msg.no == []
+        assert msg.msg.yes == {}
+        assert msg.msg.no == {}
 
     def test_delta_with_optional_fields(self) -> None:
         raw = {
@@ -980,9 +983,7 @@ def test_market_positions_position_volume_parse_as_decimal(self) -> None:
             {
                 "type": "market_positions",
                 "sid": 1,
-                "msg": market_positions_payload_dict(
-                    position_fp="100", volume_fp="200"
-                ),
+                "msg": market_positions_payload_dict(position_fp="100", volume_fp="200"),
             }
         )
         assert isinstance(msg.msg.position, Decimal)
@@ -1004,11 +1005,15 @@ def test_orderbook_snapshot_yes_no_parse_as_decimal_tuples(self) -> None:
                 },
             }
         )
-        assert msg.msg.yes[0] == (Decimal("0.50"), Decimal("100.00"))
-        assert isinstance(msg.msg.yes[0][0], Decimal)
-        assert isinstance(msg.msg.yes[0][1], Decimal)
-        assert isinstance(msg.msg.no[0][0], Decimal)
-        assert isinstance(msg.msg.no[0][1], Decimal)
+        assert msg.msg.yes == {
+            Decimal("0.50"): Decimal("100.00"),
+            Decimal("0.55"): Decimal("200.00"),
+        }
+        assert msg.msg.no == {Decimal("0.45"): Decimal("150.00")}
+        for d in (msg.msg.yes, msg.msg.no):
+            for key, value in d.items():
+                assert isinstance(key, Decimal)
+                assert isinstance(value, Decimal)
 
     def test_rfq_created_contracts_parses_as_decimal(self) -> None:
         payload = RfqCreatedPayload.model_validate(
@@ -1189,6 +1194,8 @@ def test_orderbook_delta_payload_side_rejects_trailing_whitespace(self) -> None:
                     "side": "yes ",
                 }
             )
+
+
 # ---------- P2#3: populate_by_name across every WS payload model ----------
 
 _SDK_NAME_PAYLOADS: list[tuple[type, dict[str, object]]] = [

tests/ws/test_orderbook.py

@@ -1,4 +1,5 @@
 """Tests for OrderbookManager."""
+
 from __future__ import annotations
 
 from decimal import Decimal
@@ -21,17 +22,19 @@ def make_snapshot(
     # that ``OrderbookSnapshotPayload.yes/no: list[tuple[str, str]]`` expects.
     # Callers pass list-of-list literals for readability; direct constructor
     # arguments would trip mypy strict on the list vs tuple mismatch.
-    return OrderbookSnapshotMessage.model_validate({
-        "type": "orderbook_snapshot",
-        "sid": 1,
-        "seq": seq,
-        "msg": {
-            "market_ticker": ticker,
-            "market_id": "id",
-            "yes": yes or [],
-            "no": no or [],
-        },
-    })
+    return OrderbookSnapshotMessage.model_validate(
+        {
+            "type": "orderbook_snapshot",
+            "sid": 1,
+            "seq": seq,
+            "msg": {
+                "market_ticker": ticker,
+                "market_id": "id",
+                "yes": yes or [],
+                "no": no or [],
+            },
+        }
+    )
 
 
 def make_delta(
@@ -242,16 +245,10 @@ def __contains__(self, key: object) -> bool:
         # If O(1), they're identical.
         def measure(n: int) -> int:
             mgr = OrderbookManager()
-            mgr.apply_snapshot(
-                make_snapshot(yes=[[f"0.{i:04d}", "100"] for i in range(1, n + 1)])
-            )
+            mgr.apply_snapshot(make_snapshot(yes=[[f"0.{i:04d}", "100"] for i in range(1, n + 1)]))
             counting = CountingDict(mgr._books["T"].yes)
-            mgr._books["T"] = _BookState(
-                ticker="T", yes=counting, no=mgr._books["T"].no
-            )
-            mgr.apply_delta(
-                make_delta(price=f"0.{n // 2:04d}", delta="1", side="yes")
-            )
+            mgr._books["T"] = _BookState(ticker="T", yes=counting, no=mgr._books["T"].no)
+            mgr.apply_delta(make_delta(price=f"0.{n // 2:04d}", delta="1", side="yes"))
             return counting.accesses
 
         small = measure(10)
@@ -271,9 +268,7 @@ def test_fractional_delta_matches_wire_format(self) -> None:
         """
         mgr = OrderbookManager()
         mgr.apply_snapshot(make_snapshot(yes=[["0.0100", "0.00"]]))
-        book = mgr.apply_delta(
-            make_delta(price="0.0100", delta="1.00", side="yes")
-        )
+        book = mgr.apply_delta(make_delta(price="0.0100", delta="1.00", side="yes"))
         assert book is not None
         # Decimal("0.00") + Decimal("1.00") == Decimal("1.00"), not Decimal("1")
         assert book.yes[0].quantity == Decimal("1.00")
@@ -309,23 +304,20 @@ def test_apply_snapshot_inplace_records_sid(self) -> None:
     def test_apply_delta_inplace_returns_true_on_known_ticker(self) -> None:
         mgr = OrderbookManager()
         mgr._apply_snapshot_inplace(make_snapshot(yes=[["0.50", "10.00"]]))
-        ok = mgr._apply_delta_inplace(
-            make_delta(price="0.50", delta="5.00", side="yes")
-        )
+        ok = mgr._apply_delta_inplace(make_delta(price="0.50", delta="5.00", side="yes"))
         assert ok is True
         book = mgr.get("T")
         assert book is not None
         assert book.yes[0].quantity == Decimal("15.00")
 
     def test_apply_delta_inplace_returns_false_for_unknown_ticker(self) -> None:
         mgr = OrderbookManager()
-        ok = mgr._apply_delta_inplace(
-            make_delta(ticker="OTHER", price="0.50", delta="5.00")
-        )
+        ok = mgr._apply_delta_inplace(make_delta(ticker="OTHER", price="0.50", delta="5.00"))
         assert ok is False
 
     def test_apply_delta_inplace_does_not_call_to_orderbook(
-        self, monkeypatch: object,
+        self,
+        monkeypatch: object,
     ) -> None:
         """Hot-path invariant: the in-place variant MUST NOT materialize
         an Orderbook. Spy on ``_BookState.to_orderbook`` and assert zero
@@ -345,9 +337,7 @@ def spy(self):  # type: ignore[no-untyped-def]
         mgr = OrderbookManager()
         mgr._apply_snapshot_inplace(make_snapshot(yes=[["0.50", "10.00"]]))
         for _ in range(5):
-            mgr._apply_delta_inplace(
-                make_delta(price="0.50", delta="1.00", side="yes")
-            )
+            mgr._apply_delta_inplace(make_delta(price="0.50", delta="1.00", side="yes"))
         assert calls == []
 
     def test_public_apply_delta_still_returns_orderbook(self) -> None:
@@ -367,9 +357,7 @@ def test_public_apply_delta_returns_none_for_unknown_ticker(self) -> None:
 
     def test_public_apply_snapshot_returns_orderbook(self) -> None:
         mgr = OrderbookManager()
-        book = mgr.apply_snapshot(
-            make_snapshot(yes=[["0.10", "5.00"]], no=[["0.90", "5.00"]])
-        )
+        book = mgr.apply_snapshot(make_snapshot(yes=[["0.10", "5.00"]], no=[["0.90", "5.00"]]))
         assert book is not None
         assert book.ticker == "T"
         assert len(book.yes) == 1