Added options to improve reader/decoder parsing efficiency.

python/fusion_engine_client/messages/defs.py

@@ -250,6 +250,11 @@ class MessageHeader:
     _FORMAT = '<BBHIBBHIII'
     _SIZE: int = struct.calcsize(_FORMAT)
 
+    # The header layout as a numpy structured type, matching _FORMAT, for reading many headers at once.
+    _NUMPY_DTYPE = np.dtype([('sync0', 'u1'), ('sync1', 'u1'), ('reserved', '<u2'), ('crc', '<u4'),
+                             ('protocol_version', 'u1'), ('message_version', 'u1'), ('message_type', '<u2'),
+                             ('sequence_number', '<u4'), ('payload_size_bytes', '<u4'), ('source_identifier', '<u4')])
+
     # The offset of the payload_size_bytes field within the header.
     PAYLOAD_SIZE_OFFSET_BYTES = 16
 
@@ -596,6 +601,21 @@ def unpack(self, buffer: bytes, offset: int = 0, message_version: int = _UNSPECI
     def calcsize(self) -> int:
         raise NotImplementedError('calcsize() not implemented.')
 
+    @classmethod
+    def unpack_array(cls, batch: 'MessageBatch') -> Dict[str, np.ndarray]:
+        """!
+        @brief Deserialize the fields of every message in a batch at once.
+
+        This is a vectorized counterpart to @ref unpack(), for reading many messages quickly. The fields available
+        depend on the class. Classes that do not support it raise `NotImplementedError`. Call it through
+        @ref MessageBatch.unpack_array(), which checks the message types first.
+
+        @param batch A @ref MessageBatch containing only messages of this type.
+
+        @return A `dict` mapping field names to numpy arrays, one entry per message.
+        """
+        raise NotImplementedError('unpack_array() not implemented.')
+
     def get_p1_time(self) -> Timestamp:
         measurement_details = getattr(self, 'details', None)
         if isinstance(measurement_details, MeasurementDetails):

python/fusion_engine_client/messages/measurements.py

@@ -1,6 +1,6 @@
 import math
 import struct
-from typing import Dict, Sequence, Tuple
+from typing import Dict, Sequence, Tuple, Union
 
 from construct import Array, Struct, Padding, Float32l, Int16sl, Int32sl
 import numpy as np
@@ -1743,6 +1743,11 @@ class InputDataWrapperMessage(MessagePayload):
 
     _STRUCT = struct.Struct('<IB x H')
 
+    # The fixed fields at the start of the payload as a numpy structured type, matching _STRUCT, for reading many
+    # messages at once.
+    _NUMPY_DTYPE = np.dtype([('system_time_cs_low', '<u4'), ('system_time_cs_high', 'u1'), ('reserved', 'u1'),
+                             ('data_type', '<u2')])
+
     # The location of the data_type field within the message.
     DATA_TYPE_OFFSET_BYTES = 6
 
@@ -1797,6 +1802,35 @@ def unpack(self, buffer: bytes, offset: int = 0, message_version: int = MessageP
     def calcsize(self) -> int:
         return self._STRUCT.size + len(self.data)
 
+    @classmethod
+    def unpack_array(cls, batch) -> Dict[str, np.ndarray]:
+        """!
+        @brief Deserialize the fixed fields of every message in a batch at once.
+
+        See @ref MessagePayload.unpack_array(). Use @ref get_data() to get the wrapped data of individual messages.
+
+        @param batch A @ref MessageBatch containing only messages of this type.
+
+        @return A `dict` containing `system_time_ns` (`int64`) and `data_type` (`uint16`) arrays.
+        """
+        fields = batch.get_payload_fields(cls._NUMPY_DTYPE)
+        system_time_cs = (fields['system_time_cs_high'].astype(np.int64) << 32) | fields['system_time_cs_low']
+        return {
+            'system_time_ns': system_time_cs * cls.CENTI_NANO_SCALE_FACTOR,
+            'data_type': fields['data_type'],
+        }
+
+    @classmethod
+    def get_data(cls, payload: Union[bytes, memoryview]) -> Union[bytes, memoryview]:
+        """!
+        @brief Get the wrapped data from a serialized payload without deserializing the rest of it.
+
+        @param payload The serialized message payload, e.g., from @ref MessageBatch.get_payload().
+
+        @return The wrapped data, of the same type as `payload`.
+        """
+        return payload[cls._STRUCT.size:]
+
     def get_fe_content_header(self) -> Optional[MessageHeader]:
         if self._fe_content_header is None and self.data_type == InputDataType.M_TYPE_FUSION_ENGINE_MESSAGE:
             try:

python/fusion_engine_client/parsers/__init__.py

@@ -1,4 +1,5 @@
 from .encoder import FusionEngineEncoder
 from .decoder import FusionEngineDecoder
 from .file_index import FileIndex
+from .message_batch import MessageBatch
 from .mixed_log_reader import MixedLogReader

python/fusion_engine_client/parsers/decoder.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import List, Dict, Callable, Optional, Tuple, Union
+from typing import Iterable, List, Dict, Callable, Optional, Tuple, Union
 
 from ..messages import MessageHeader, MessageType, MessagePayload, message_type_to_class
 from ..utils import trace as logging
@@ -34,7 +34,8 @@ class WarnOnError(IntEnum):
     def __init__(self, max_payload_len_bytes: int = MessageHeader._MAX_EXPECTED_SIZE_BYTES,
                  warn_on_unrecognized: bool = False, warn_on_gap: bool = False,
                  warn_on_error: Union[bool, str, WarnOnError] = WarnOnError.LIKELY,
-                 return_bytes: bool = False, return_offset: bool = False):
+                 return_bytes: bool = False, return_offset: bool = False,
+                 message_types: Optional[Iterable[Union[MessageType, type]]] = None, return_payload: bool = True):
         """!
         @brief Construct a new decoder instance.
 
@@ -53,6 +54,12 @@ def __init__(self, max_payload_len_bytes: int = MessageHeader._MAX_EXPECTED_SIZE
         @param return_bytes If `True`, return a `bytes` object with the raw data (header + payload) in addition to the
                decoded message object.
         @param return_offset If `True`, the byte offset into the data stream at which the message was found.
+        @param message_types If set, only return the listed message types (@ref MessageType values or @ref
+               MessagePayload classes). Other messages are still validated, but are skipped without copying or
+               deserializing their payloads, and do not trigger callbacks.
+        @param return_payload If `False`, do not deserialize message payloads, and omit them from the returned
+               results and from the callback arguments. This is much faster when only the header, bytes, or offset are
+               needed.
         """
         self._warn_on_unrecognized = warn_on_unrecognized
         self._warn_on_seq_skip = warn_on_gap
@@ -71,6 +78,12 @@ def __init__(self, max_payload_len_bytes: int = MessageHeader._MAX_EXPECTED_SIZE
 
         self._return_bytes = return_bytes
         self._return_offset = return_offset
+        self._return_payload = return_payload
+
+        if message_types is None:
+            self._message_types = None
+        else:
+            self._message_types = set(t.get_type() if MessagePayload.is_subclass(t) else t for t in message_types)
 
         self._max_payload_len_bytes = max_payload_len_bytes
         self._buffer = bytearray()
@@ -97,6 +110,11 @@ def add_callback(self, type: Optional[MessageType], callback: Union[Callback, Ca
                the uninterpreted payload contents. If `return_bytes == True`, the function will be called with a third
                argument containing the uninterpreted message contents (including the header).
         """
+        # Callbacks are only called for message types allowed by the message type filter.
+        if type is not None and self._message_types is not None and type not in self._message_types:
+            _logger.warning('Callback registered for %s, which is excluded by the message type filter. It will not '
+                            'be called.' % MessageType.get_type_string(type))
+
         # `self._callbacks` is a `defaultdict(list)` so no need to initialize list on first entry.
         self._callbacks[type].append(callback)
 
@@ -284,36 +302,47 @@ def on_data(self, data: Union[bytes, int]) -> List[Union[MessageTuple, MessageWi
                                         expected_sequence_number, self._header.sequence_number)
             self._last_sequence_number = self._header.sequence_number
 
-            # Get the class for the received message type and deserialize the message payload. If cls is not None, it is
-            # a child of @ref MessagePayload that maps to the received @ref MessageType.
-            cls = message_type_to_class.get(self._header.message_type, None)
-            payload = bytes(buffer[offset + header_size:offset + self._msg_len])
+            # Skip messages the user did not request.
+            if self._message_types is not None and self._header.message_type not in self._message_types:
+                self._bytes_processed += self._msg_len
+                offset += self._msg_len
+                self._header = None
+                self._msg_len = 0
+                continue
+
+            # Get the class for the received message type and deserialize the message payload, if requested. If cls is
+            # not None, it is a child of @ref MessagePayload that maps to the received @ref MessageType.
             contents = None
-            if cls is not None:
-                contents = cls()
-                try:
-                    contents.unpack(buffer=payload, offset=0,
-                                    message_version=self._header.message_version)
-                    if debug_enabled:
-                        _logger.debug('Decoded FusionEngine message %s.', repr(contents))
-                except NotImplementedError as e:
+            if self._return_payload:
+                cls = message_type_to_class.get(self._header.message_type, None)
+                payload = bytes(buffer[offset + header_size:offset + self._msg_len])
+                if cls is not None:
+                    contents = cls()
+                    try:
+                        contents.unpack(buffer=payload, offset=0,
+                                        message_version=self._header.message_version)
+                        if debug_enabled:
+                            _logger.debug('Decoded FusionEngine message %s.', repr(contents))
+                    except NotImplementedError as e:
+                        print_func = _logger.warning if self._warn_on_unrecognized else _logger.debug
+                        print_func('Error deserializing message %s payload: %s', self._header.get_type_string(), e)
+                    except Exception as e:
+                        # unpack() may fail if the payload length in the header differs from the length expected by
+                        # the class, the payload contains an illegal value, etc.
+                        _logger.error('Error deserializing message %s payload: %s', self._header.get_type_string(), e)
+                # If cls is None, we don't have a class for the message type. Return a copy of the payload bytes.
+                else:
                     print_func = _logger.warning if self._warn_on_unrecognized else _logger.debug
-                    print_func('Error deserializing message %s payload: %s', self._header.get_type_string(), e)
-                except Exception as e:
-                    # unpack() may fail if the payload length in the header differs from the length expected by the
-                    # class, the payload contains an illegal value, etc.
-                    _logger.error('Error deserializing message %s payload: %s', self._header.get_type_string(), e)
-            # If cls is None, we don't have a class for the message type. Return a copy of the payload bytes.
-            else:
-                print_func = _logger.warning if self._warn_on_unrecognized else _logger.debug
-                print_func('Decoded unknown FusionEngine message. [type=%d, payload_size=%d B]',
-                           self._header.message_type, self._header.payload_size_bytes)
+                    print_func('Decoded unknown FusionEngine message. [type=%d, payload_size=%d B]',
+                               self._header.message_type, self._header.payload_size_bytes)
 
-            if contents is None:
-                contents = payload
+                if contents is None:
+                    contents = payload
 
             # Store the result.
-            result = [self._header, contents]
+            result = [self._header]
+            if self._return_payload:
+                result.append(contents)
             if self._return_bytes:
                 result.append(buffer[offset:offset + self._msg_len])
             if self._return_offset:

python/fusion_engine_client/parsers/message_batch.py

@@ -0,0 +1,158 @@
+from typing import Dict, Type, Union
+
+import numpy as np
+
+from ..messages import MessageHeader, MessagePayload
+
+
+class MessageBatch(object):
+    """!
+    @brief A set of FusionEngine messages from a file, read all at once into numpy arrays.
+
+    A batch is returned by @ref MixedLogReader.read_batch(). It holds the header of every message as numpy arrays, one
+    entry per message in file order. Payloads are read from the file only on request: one message at a time with
+    @ref get_payload(), or the fields of every message at once with @ref unpack_array().
+
+    Use numpy indexing to select messages, e.g., `batch[batch.message_type == MessageType.POSE]`.
+    """
+
+    # The number of messages to gather at once. Gathering builds a temporary array of one 8-byte offset per byte read,
+    # so this bounds its size.
+    _GATHER_CHUNK_SIZE = 1 << 18
+
+    def __init__(self, buffer: np.ndarray, offset: np.ndarray, headers: np.ndarray, p1_time: np.ndarray,
+                 message_index: np.ndarray):
+        """!
+        @brief Construct a batch from data already read by @ref MixedLogReader.read_batch().
+
+        @param buffer The contents of the data file, as a `uint8` array.
+        @param offset The file offset of each message (in bytes).
+        @param headers The header of each message, as an array of @ref MessageHeader._NUMPY_DTYPE.
+        @param p1_time The P1 time of each message (in seconds), or `NaN` if the message has none.
+        @param message_index The 0-based index of each message within the file.
+        """
+        self._buffer = buffer
+        self._headers = headers
+        self.offset = offset
+        self.p1_time = p1_time
+        self.message_index = message_index
+
+    @classmethod
+    def _read_headers(cls, buffer: np.ndarray, offset: np.ndarray) -> np.ndarray:
+        """!
+        @brief Read the header of every message at the specified offsets.
+
+        @param buffer The contents of the data file, as a `uint8` array.
+        @param offset The file offset of each message (in bytes).
+
+        @return An array of @ref MessageHeader._NUMPY_DTYPE, one entry per message.
+        """
+        return cls._gather(buffer, offset, MessageHeader._NUMPY_DTYPE)
+
+    @property
+    def message_type(self) -> np.ndarray:
+        return self._headers['message_type']
+
+    @property
+    def message_version(self) -> np.ndarray:
+        return self._headers['message_version']
+
+    @property
+    def sequence_number(self) -> np.ndarray:
+        return self._headers['sequence_number']
+
+    @property
+    def source_identifier(self) -> np.ndarray:
+        return self._headers['source_identifier']
+
+    @property
+    def payload_size_bytes(self) -> np.ndarray:
+        return self._headers['payload_size_bytes']
+
+    @property
+    def size_bytes(self) -> np.ndarray:
+        """!
+        @brief Get the size of each complete message, header included (in bytes).
+        """
+        return MessageHeader.calcsize() + self.payload_size_bytes.astype(np.int64)
+
+    def get_bytes(self, i: int) -> memoryview:
+        """!
+        @brief Get one complete serialized message, header included, without copying it.
+
+        @param i The index of the message within this batch.
+
+        @return A `memoryview` of the message.
+        """
+        start = int(self.offset[i])
+        return memoryview(self._buffer[start:start + int(self.size_bytes[i])])
+
+    def get_payload(self, i: int) -> memoryview:
+        """!
+        @brief Get the serialized payload of one message without copying it.
+
+        @param i The index of the message within this batch.
+
+        @return A `memoryview` of the payload.
+        """
+        start = int(self.offset[i]) + MessageHeader.calcsize()
+        return memoryview(self._buffer[start:start + int(self.payload_size_bytes[i])])
+
+    def get_payload_fields(self, dtype: Union[np.dtype, list]) -> np.ndarray:
+        """!
+        @brief Read fixed-size fields from the start of every message's payload at once.
+
+        This is intended for use by @ref MessagePayload.unpack_array() implementations.
+
+        @param dtype A numpy structured type describing the fields at the start of the payload.
+
+        @return An array of `dtype`, one entry per message.
+        """
+        dtype = np.dtype(dtype)
+        if np.any(self.payload_size_bytes < dtype.itemsize):
+            raise ValueError('Payload too small for the requested fields. [size=%d B, min=%d B]' %
+                             (np.min(self.payload_size_bytes), dtype.itemsize))
+        return self._gather(self._buffer, self.offset + MessageHeader.calcsize(), dtype)
+
+    def unpack_array(self, cls: Type[MessagePayload]) -> Dict[str, np.ndarray]:
+        """!
+        @brief Deserialize the fields of every message in the batch at once.
+
+        See @ref MessagePayload.unpack_array().
+
+        @param cls The @ref MessagePayload class of the messages. Every message in the batch must be of this type.
+
+        @return A `dict` mapping field names to numpy arrays, one entry per message.
+        """
+        if np.any(self.message_type != int(cls.get_type())):
+            raise ValueError('Batch contains messages other than %s. Select them first, e.g., '
+                             'batch[batch.message_type == %s].' % (cls.__name__, cls.get_type()))
+        return cls.unpack_array(self)
+
+    def __len__(self):
+        return len(self.offset)
+
+    def __getitem__(self, key) -> 'MessageBatch':
+        """!
+        @brief Select a subset of the messages using numpy indexing (a boolean mask, indices, or a slice).
+
+        @param key The numpy index.
+
+        @return A new @ref MessageBatch containing the selected messages.
+        """
+        # A scalar index would select a single numpy element instead of an array.
+        if isinstance(key, (int, np.integer)):
+            raise TypeError('Batch index must select an array of messages. Use batch[i:i + 1] to select one '
+                            'message, or get_payload(i) to get its payload.')
+
+        return MessageBatch(buffer=self._buffer, offset=self.offset[key], headers=self._headers[key],
+                            p1_time=self.p1_time[key], message_index=self.message_index[key])
+
+    @classmethod
+    def _gather(cls, buffer: np.ndarray, offset: np.ndarray, dtype: np.dtype) -> np.ndarray:
+        result = np.empty(len(offset), dtype=dtype)
+        byte_offsets = np.arange(dtype.itemsize)
+        for start in range(0, len(offset), cls._GATHER_CHUNK_SIZE):
+            chunk = offset[start:start + cls._GATHER_CHUNK_SIZE].astype(np.int64)
+            result[start:start + len(chunk)] = buffer[chunk[:, None] + byte_offsets].view(dtype).reshape(-1)
+        return result