feat: Improve Writing Performance Further

src/pybag/encoding/__init__.py

@@ -88,7 +88,7 @@ def encode(self, type_str: str, value: Any) -> None:
         ...  # pragma: no cover
 
     @abstractmethod
-    def save(self) -> bytes:
+    def save(self) -> Any:
         ...  # pragma: no cover
 
     # Primitive encoders -------------------------------------------------

src/pybag/encoding/cdr.py

@@ -1,13 +1,32 @@
 import logging
 import struct
+from dataclasses import dataclass
 from typing import Any
 
 from pybag.encoding import MessageDecoder, MessageEncoder
 from pybag.io.raw_reader import BytesReader
-from pybag.io.raw_writer import BytesWriter
+from pybag.io.raw_writer import BaseWriter, BytesWriter
 
 logger = logging.getLogger(__name__)
 
+@dataclass(slots=True)
+class SerializedMessage:
+    header: bytes
+    payload: memoryview
+
+    def __len__(self) -> int:
+        return len(self.header) + len(self.payload)
+
+    def write(self, writer: BaseWriter) -> None:
+        writer.write(self.header)
+        writer.write(self.payload)
+
+    def to_bytes(self) -> bytes:
+        return self.header + self.payload.tobytes()
+
+    def __bytes__(self) -> bytes:  # pragma: no cover - convenience
+        return self.to_bytes()
+
 
 class CdrDecoder(MessageDecoder):
     def __init__(self, data: bytes):
@@ -149,13 +168,16 @@ def __init__(self, *, little_endian: bool = True) -> None:
     def encoding(cls) -> str:
         return "cdr"
 
+    def reset(self) -> None:
+        self._payload.clear()
+
     def encode(self, type_str: str, value: Any) -> None:
         """Encode ``value`` based on ``type_str``."""
         getattr(self, type_str)(value)
 
-    def save(self) -> bytes:
+    def save(self) -> SerializedMessage:
         """Return the encoded byte stream."""
-        return self._header + self._payload.as_bytes()
+        return SerializedMessage(self._header, self._payload.as_memoryview())
 
     # Primitive encoders -------------------------------------------------
 

src/pybag/io/raw_writer.py

@@ -66,6 +66,9 @@ def size(self) -> int:
     def as_bytes(self) -> bytes:
         return bytes(self._buffer)
 
+    def as_memoryview(self) -> memoryview:
+        return memoryview(self._buffer)
+
     def clear(self) -> None:
         self._buffer.clear()
 

src/pybag/mcap/record_writer.py

@@ -1,6 +1,7 @@
 import struct
 from typing import Any, Callable
 
+from pybag.encoding.cdr import SerializedMessage
 from pybag.io.raw_writer import BaseWriter
 from pybag.mcap.records import (
     AttachmentIndexRecord,
@@ -137,14 +138,22 @@ def write_channel(cls, writer: BaseWriter, record: ChannelRecord) -> None:
 
     @classmethod
     def write_message(cls, writer: BaseWriter, record: MessageRecord) -> None:
-        payload = (
-              cls._encode_uint16(record.channel_id)
-            + cls._encode_uint32(record.sequence)
-            + cls._encode_timestamp(record.log_time)
-            + cls._encode_timestamp(record.publish_time)
-            + record.data
-        )
-        cls._write_record(writer, RecordType.MESSAGE, payload)
+        # Avoid creating a large temporary payload buffer which would copy the
+        # message data by constructing the fixed-size header in-place and
+        # writing it before the pre-serialized message bytes.
+        payload_length = 2 + 4 + 8 + 8 + len(record.data)
+        header = bytearray(1 + 8 + 2 + 4 + 8 + 8)
+        struct.pack_into("<B", header, 0, int(RecordType.MESSAGE))
+        struct.pack_into("<Q", header, 1, payload_length)
+        struct.pack_into("<H", header, 9, record.channel_id)
+        struct.pack_into("<I", header, 11, record.sequence)
+        struct.pack_into("<Q", header, 15, record.log_time)
+        struct.pack_into("<Q", header, 23, record.publish_time)
+        writer.write(header)
+        if isinstance(record.data, SerializedMessage):
+            record.data.write(writer)
+        else:
+            writer.write(record.data)
 
     @classmethod
     def write_chunk(cls, writer: BaseWriter, record: ChunkRecord) -> None:

src/pybag/mcap/records.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from enum import IntEnum
 
+from pybag.encoding.cdr import SerializedMessage
+
 
 @dataclass
 class HeaderRecord:
@@ -38,7 +40,7 @@ class MessageRecord:
     sequence: int
     log_time: int
     publish_time: int
-    data: bytes
+    data: bytes | SerializedMessage
 
 
 @dataclass

src/pybag/mcap_writer.py

@@ -26,6 +26,8 @@
 
 logger = logging.getLogger(__name__)
 
+DEFAULT_CHUNK_SIZE = 4 * 1024 * 1024
+
 
 class McapFileWriter:
     """High level writer for producing MCAP files."""
@@ -35,7 +37,7 @@ def __init__(
         writer: BaseWriter,
         *,
         profile: str = "ros2",
-        chunk_size: int | None = None,
+        chunk_size: int | None = DEFAULT_CHUNK_SIZE,
         chunk_compression: Literal["lz4", "zstd"] | None = None,
     ) -> None:
         self._writer = CrcWriter(writer)
@@ -86,16 +88,17 @@ def open(
         file_path: str | Path,
         *,
         profile: str = "ros2",
-        chunk_size: int | None = None,
-        chunk_compression: Literal["lz4", "zstd"] | None = "lz4",
+        chunk_size: int | None = DEFAULT_CHUNK_SIZE,
+        chunk_compression: Literal["lz4", "zstd"] | None = None,
     ) -> "McapFileWriter":
         """Create a writer backed by a file on disk.
 
         Args:
             file_path: The path to the file to write to.
             profile: The profile to use for the MCAP file.
-            chunk_size: The size of the chunk to write to in bytes.
-            chunk_compression: The compression to use for the chunk.
+            chunk_size: The target chunk size in bytes. Pass ``None`` to disable
+                chunking.
+            chunk_compression: Optional compression to apply to each chunk.
 
         Returns:
             A writer backed by a file on disk.

src/pybag/serialize.py

@@ -2,7 +2,7 @@
 from typing import Callable
 
 from pybag.encoding import MessageEncoder
-from pybag.encoding.cdr import CdrEncoder
+from pybag.encoding.cdr import CdrEncoder, SerializedMessage
 from pybag.mcap.records import ChannelRecord, SchemaRecord
 from pybag.schema import SchemaEncoder
 from pybag.schema.compiler import compile_serializer
@@ -21,6 +21,7 @@ def __init__(
         self._schema_encoder = schema_encoder
         self._message_encoder = message_encoder
         self._compiled: dict[type[Message], Callable[[MessageEncoder, Message], None]] = {}
+        self._encoders: dict[tuple[type[Message], bool], CdrEncoder] = {}
 
     @property
     def schema_encoding(self) -> str:
@@ -31,6 +32,15 @@ def message_encoding(self) -> str:
         return self._message_encoder.encoding()
 
     def serialize_message(self, message: Message, *, little_endian: bool = True) -> bytes:
+        serialized = self.serialize_message_view(message, little_endian=little_endian)
+        return serialized.to_bytes()
+
+    def serialize_message_view(
+        self,
+        message: Message,
+        *,
+        little_endian: bool = True,
+    ) -> SerializedMessage:
         if not is_dataclass(message):  # pragma: no cover - defensive programming
             raise TypeError("Expected a dataclass instance")
 
@@ -40,7 +50,15 @@ def serialize_message(self, message: Message, *, little_endian: bool = True) ->
             serializer = compile_serializer(schema, sub_schemas)
             self._compiled[message_type] = serializer
 
-        encoder = self._message_encoder(little_endian=little_endian)
+        key = (message_type, little_endian)
+        encoder = self._encoders.get(key)
+        if encoder is None:
+            encoder = self._message_encoder(little_endian=little_endian)  # type: ignore[call-arg]
+            self._encoders[key] = encoder
+            # Cache is keyed by endianness, so the header stays valid.
+        else:
+            encoder.reset()
+
         serializer(encoder, message)
         return encoder.save()
 

tests/test_mcap_writer.py

@@ -91,7 +91,7 @@ class Example:
 
     with tempfile.TemporaryDirectory() as tmpdir:
         file_path = Path(tmpdir) / "test.mcap"
-        with McapFileWriter.open(file_path, profile="ros2") as mcap:
+        with McapFileWriter.open(file_path, profile="ros2", chunk_size=None) as mcap:
             channel_id = mcap.add_channel("/example", Example)
             mcap.write_message("/example", 1, Example(5))
         reader = CrcReader(BytesReader(file_path.read_bytes()))