22from dataclasses import dataclass
33from functools import cached_property
44from mmap import ACCESS_READ , mmap
5- from typing import IO , override
5+ from operator import itemgetter
6+ from typing import IO , final , override
67
78import more_itertools as mit
89import numpy .typing as npt
910import torch
10- from jaxtyping import Shaped
1111from mcap .data_stream import ReadDataStream
1212from mcap .decoder import DecoderFactory
1313from mcap .opcode import Opcode
@@ -32,6 +32,7 @@ class MessageIndex:
3232 message_length : int
3333
3434
35+ @final
3536class McapTensorSource (TensorSource ):
3637 @validate_call (config = BaseModel .model_config )
3738 def __init__ (
@@ -47,8 +48,8 @@ def __init__(
4748 with bound_contextvars (
4849 path = path .as_posix (), topic = topic , message_decoder_factory = decoder_factory
4950 ):
50- self ._path : FilePath = path
51- self ._validate_crcs : bool = validate_crcs
51+ self ._path = path
52+ self ._validate_crcs = validate_crcs
5253
5354 summary = SeekingReader (
5455 stream = self ._file , validate_crcs = self ._validate_crcs
@@ -73,13 +74,14 @@ def __init__(
7374 logger .error (msg := "missing message decoder" )
7475 raise RuntimeError (msg )
7576
76- self ._message_decoder : Callable [[ bytes ], object ] = message_decoder
77- self ._chunk_indexes : tuple [ ChunkIndex , ...] = tuple (
77+ self ._message_decoder = message_decoder
78+ self ._chunk_indexes = tuple (
7879 chunk_index
7980 for chunk_index in summary .chunk_indexes
8081 if self ._channel .id in chunk_index .message_index_offsets
8182 )
82- self ._decoder : Callable [[bytes ], npt .ArrayLike ] = decoder
83+ self ._decoder = decoder
84+ self ._mmap = None
8385
8486 @property
8587 def _file (self ) -> IO [bytes ]:
@@ -89,42 +91,55 @@ def _file(self) -> IO[bytes]:
8991
9092 case None | mmap (closed = True ):
9193 with self ._path .open ("rb" ) as f :
92- self ._mmap : mmap = mmap (
93- fileno = f .fileno (), length = 0 , access = ACCESS_READ
94- )
94+ self ._mmap = mmap (fileno = f .fileno (), length = 0 , access = ACCESS_READ )
9595
9696 case _:
9797 raise RuntimeError
9898
9999 return self ._mmap # pyright: ignore[reportReturnType]
100100
101101 @override
102- def __getitem__ (self , indexes : Iterable [int ]) -> Shaped [Tensor , "b h w c" ]:
103- frames : Mapping [int , npt .ArrayLike ] = {}
104-
105- message_indexes_by_chunk_start_offset : Mapping [
106- int , Iterable [tuple [int , MessageIndex ]]
107- ] = mit .map_reduce (
108- zip (indexes , (self ._message_indexes [idx ] for idx in indexes ), strict = True ),
109- keyfunc = lambda x : x [1 ].chunk_start_offset ,
110- )
102+ def __getitem__ (self , indexes : int | Iterable [int ]) -> Tensor :
103+ match indexes :
104+ case Iterable ():
105+ arrays : Mapping [int , npt .ArrayLike ] = {}
106+ message_indexes = (self ._message_indexes [idx ] for idx in indexes )
107+ indexes_by_chunk_start_offset = mit .map_reduce (
108+ zip (indexes , message_indexes , strict = True ),
109+ keyfunc = lambda x : x [1 ].chunk_start_offset ,
110+ )
111+
112+ for chunk_start_offset , chunk_indexes in sorted (
113+ indexes_by_chunk_start_offset .items (), key = itemgetter (0 )
114+ ):
115+ _ = self ._file .seek (chunk_start_offset + 1 + 8 )
116+ chunk = Chunk .read (ReadDataStream (self ._file ))
117+ stream , _ = get_chunk_data_stream (
118+ chunk , validate_crc = self ._validate_crcs
119+ )
120+ for index , message_index in sorted (
121+ chunk_indexes , key = lambda x : x [1 ].message_start_offset
122+ ):
123+ stream .read (message_index .message_start_offset - stream .count ) # pyright: ignore[reportUnusedCallResult]
124+ message = Message .read (stream , message_index .message_length )
125+ decoded_message = self ._message_decoder (message .data )
126+ arrays [index ] = self ._decoder (decoded_message .data )
111127
112- for (
113- chunk_start_offset ,
114- chunk_message_indexes ,
115- ) in message_indexes_by_chunk_start_offset .items ():
116- self ._file .seek (chunk_start_offset + 1 + 8 ) # pyright: ignore[reportUnusedCallResult]
117- chunk = Chunk .read (ReadDataStream (self ._file ))
118- stream , _ = get_chunk_data_stream (chunk , validate_crc = self ._validate_crcs )
119- for frame_index , message_index in sorted (
120- chunk_message_indexes , key = lambda x : x [1 ].message_start_offset
121- ):
122- stream .read (message_index .message_start_offset - stream .count ) # pyright: ignore[reportUnusedCallResult]
123- message = Message .read (stream , message_index .message_length )
128+ tensors = [torch .from_numpy (arrays [idx ]) for idx in indexes ] # pyright: ignore[reportUnknownMemberType]
129+
130+ return torch .stack (tensors )
131+
132+ case _:
133+ message_index = self ._message_indexes [indexes ]
134+ _ = self ._file .seek (message_index .chunk_start_offset + 1 + 8 )
135+ chunk = Chunk .read (ReadDataStream (self ._file ))
136+ stream , _ = get_chunk_data_stream (chunk , self ._validate_crcs )
137+ _ = stream .read (message_index .message_start_offset - stream .count )
138+ message = Message .read (stream , length = message_index .message_length )
124139 decoded_message = self ._message_decoder (message .data )
125- frames [ frame_index ] = self ._decoder (decoded_message .data ) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
140+ array = self ._decoder (decoded_message .data )
126141
127- return torch .stack ([ torch . from_numpy (frames [ idx ]) for idx in indexes ] ) # pyright: ignore[reportUnknownMemberType]
142+ return torch .from_numpy (array ) # pyright: ignore[reportUnknownMemberType]
128143
129144 @override
130145 def __len__ (self ) -> int :
0 commit comments