diff --git a/python/fusion_engine_client/messages/defs.py b/python/fusion_engine_client/messages/defs.py index 10e5d762..7ae621f6 100644 --- a/python/fusion_engine_client/messages/defs.py +++ b/python/fusion_engine_client/messages/defs.py @@ -250,6 +250,11 @@ class MessageHeader: _FORMAT = ' 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): diff --git a/python/fusion_engine_client/messages/measurements.py b/python/fusion_engine_client/messages/measurements.py index 7f0f0e57..721de3c9 100644 --- a/python/fusion_engine_client/messages/measurements.py +++ b/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(' 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: diff --git a/python/fusion_engine_client/parsers/__init__.py b/python/fusion_engine_client/parsers/__init__.py index 8c563d58..e7cedbc6 100644 --- a/python/fusion_engine_client/parsers/__init__.py +++ b/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 diff --git a/python/fusion_engine_client/parsers/decoder.py b/python/fusion_engine_client/parsers/decoder.py index c51ec251..188cd1b8 100644 --- a/python/fusion_engine_client/parsers/decoder.py +++ b/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: diff --git a/python/fusion_engine_client/parsers/message_batch.py b/python/fusion_engine_client/parsers/message_batch.py new file mode 100644 index 00000000..76e756d5 --- /dev/null +++ b/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 diff --git a/python/fusion_engine_client/parsers/mixed_log_reader.py b/python/fusion_engine_client/parsers/mixed_log_reader.py index 2c038cc1..a57e35d7 100644 --- a/python/fusion_engine_client/parsers/mixed_log_reader.py +++ b/python/fusion_engine_client/parsers/mixed_log_reader.py @@ -9,6 +9,7 @@ import numpy as np from . import fast_indexer, file_index +from .message_batch import MessageBatch from ..messages import MessageType, MessageHeader, MessagePayload, Timestamp, message_type_to_class from ..utils import trace as logging from ..utils.time_range import TimeRange @@ -212,6 +213,35 @@ def set_max_bytes(self, max_bytes): def get_bytes_read(self): return self.total_bytes_read + def read_batch(self) -> MessageBatch: + """! + @brief Read every message that meets the current filter criteria at once. + + This is much faster than iterating over the messages when reading many of them, since it reads all of their + headers in one vectorized operation instead of one Python call per message. It selects the same messages the + iterator would return from the start of the file, honoring the message type, time range, source identifier, + and `max_bytes` criteria, and it does not change the iterator's position. + + Payloads are not deserialized up front. See @ref MessageBatch for reading them. + + @return A @ref MessageBatch. + """ + # Read the headers of every message in the filtered index. + offset = self.index.offset.astype(np.int64) + if len(offset) == 0: + buffer = np.zeros(0, dtype=np.uint8) + else: + buffer = np.memmap(self.input_path, dtype=np.uint8, mode='r') + headers = MessageBatch._read_headers(buffer, offset) + + # Apply the filters the index does not cover. + keep = offset + MessageHeader.calcsize() + headers['payload_size_bytes'].astype(np.int64) <= self.max_bytes + if self.requested_source_ids is not None and len(self.requested_source_ids) > 0: + keep &= np.isin(headers['source_identifier'], list(self.requested_source_ids)) + + return MessageBatch(buffer=buffer, offset=offset[keep], headers=headers[keep], p1_time=self.index.time[keep], + message_index=self.index.message_index[keep]) + def next(self): return self.read_next() @@ -297,8 +327,9 @@ def _read_next(self, require_p1_time=False, require_system_time=False, force_eof # In either case, we will skip past this header and see if there any further candidate headers later in # the file. # - # If the CRC fails, either because we found an invalid header or because a valid message got corrupted, - # validate_crc() will raise a ValueError and we will skip forward in the same manner. + # Without an index, if the CRC fails, either because we found an invalid header or because a valid + # message got corrupted, validate_crc() will raise a ValueError and we will skip forward in the same + # manner. payload_bytes = self.input_file.read(header.payload_size_bytes) read_len += len(payload_bytes) self.total_bytes_read += len(payload_bytes) @@ -307,7 +338,11 @@ def _read_next(self, require_p1_time=False, require_system_time=False, force_eof header.get_type_string()) data += payload_bytes - header.validate_crc(data) + + # The indexer validated the CRC of every message it indexed, and an index is regenerated if the data + # file has changed since, so only validate the CRC when reading without an index. + if self.index is None: + header.validate_crc(data) # Verify that source ID is correct. if (self.requested_source_ids is not None and len(self.requested_source_ids) > 0 and @@ -350,6 +385,7 @@ def _read_next(self, require_p1_time=False, require_system_time=False, force_eof self.time_range is not None or \ require_p1_time or require_system_time + payload = None if need_payload: cls = message_type_to_class.get(header.message_type, None) if cls is not None: diff --git a/python/tests/test_decoder.py b/python/tests/test_decoder.py index f7556788..6503c773 100644 --- a/python/tests/test_decoder.py +++ b/python/tests/test_decoder.py @@ -108,6 +108,41 @@ def func4(header, payload): return func_helper(header, payload, 3) assert counters[3] == 5 +# Check that a message type filter returns only the requested types, and skips their callbacks. +def test_message_types(): + test_bytes = P1_POSE_MESSAGE1 + P1_POSE_AUX_MESSAGE3 + P1_POSE_MESSAGE2 + decoder = FusionEngineDecoder(message_types=[PoseAuxMessage], return_offset=True) + callback_types = [] + decoder.add_callback(None, lambda header, payload, offset: callback_types.append(header.message_type)) + ret = decoder.on_data(test_bytes) + assert [r[0].message_type for r in ret] == [MessageType.POSE_AUX] + assert isinstance(ret[0][1], PoseAuxMessage) + assert ret[0][2] == len(P1_POSE_MESSAGE1) + assert callback_types == [MessageType.POSE_AUX] + + +# Check that registering a callback for a type excluded by the message type filter prints a warning. +def test_message_types_callback_warning(caplog): + decoder = FusionEngineDecoder(message_types=[PoseAuxMessage]) + decoder.add_callback(MessageType.POSE_AUX, lambda header, payload: None) + assert 'excluded by the message type filter' not in caplog.text + decoder.add_callback(MessageType.POSE, lambda header, payload: None) + assert 'Callback registered for POSE (10000), which is excluded by the message type filter.' in caplog.text + + +# Check that disabling payloads omits them from the results and the callback arguments. +def test_no_payload(): + test_bytes = P1_POSE_MESSAGE1 + P1_POSE_AUX_MESSAGE3 + decoder = FusionEngineDecoder(return_payload=False, return_bytes=True) + callback_args = [] + decoder.add_callback(None, lambda *args: callback_args.append(args)) + ret = decoder.on_data(test_bytes) + assert [len(r) for r in ret] == [2, 2] + assert [r[0].message_type for r in ret] == [MessageType.POSE, MessageType.POSE_AUX] + assert [bytes(r[1]) for r in ret] == [P1_POSE_MESSAGE1, P1_POSE_AUX_MESSAGE3] + assert [len(args) for args in callback_args] == [2, 2] + + def test_seq_backwards_warning(caplog): test_bytes = P1_POSE_MESSAGE1 + P1_POSE_MESSAGE2 + P1_POSE_MESSAGE1 decoder = FusionEngineDecoder(warn_on_error=True) diff --git a/python/tests/test_message_batch.py b/python/tests/test_message_batch.py new file mode 100644 index 00000000..1efe1d00 --- /dev/null +++ b/python/tests/test_message_batch.py @@ -0,0 +1,150 @@ +import pytest + +from fusion_engine_client.messages import * +from fusion_engine_client.parsers import FusionEngineEncoder, MessageBatch, MixedLogReader +from fusion_engine_client.utils.time_range import TimeRange + + +def _make_wrapper(system_time_sec, data_type, data): + message = InputDataWrapperMessage() + message.system_time_ns = int(round(system_time_sec * 1e9)) + message.data_type = data_type + message.data = data + return message + + +def _make_pose(p1_time_sec): + message = PoseMessage() + message.p1_time = Timestamp(p1_time_sec) + return message + + +@pytest.fixture +def data_path(tmpdir): + """Write a log mixing pose output, wrapped input of two types and two source IDs, and non-FusionEngine bytes.""" + encoder = FusionEngineEncoder() + entries = [ + (_make_pose(1.0), 0), + (_make_wrapper(100.00, InputDataType.M_TYPE_EXTERNAL_UNFRAMED_GNSS, b'first'), 0), + b'not a message', + (_make_wrapper(100.01, InputDataType.M_TYPE_RTCM3_POLARIS, b'other type'), 1), + (_make_pose(2.0), 0), + (_make_wrapper(100.50, InputDataType.M_TYPE_EXTERNAL_UNFRAMED_GNSS, b'second record'), 0), + (_make_pose(3.0), 1), + (_make_wrapper(103.00, InputDataType.M_TYPE_EXTERNAL_UNFRAMED_GNSS, b'third'), 1), + ] + path = tmpdir.join('test_file.p1log') + with open(path, 'wb') as f: + for entry in entries: + if isinstance(entry, bytes): + f.write(entry) + else: + message, source_id = entry + f.write(encoder.encode_message(message, source_identifier=source_id)) + return str(path) + + +def _read_iterator(path, **kwargs): + reader = MixedLogReader(path, return_header=True, return_payload=False, return_bytes=True, return_offset=True, + **kwargs) + return list(reader) + + +# Check that a batch contains the same messages, headers, and bytes as the iterator returns. +def test_matches_iterator(data_path): + expected = _read_iterator(data_path) + batch = MixedLogReader(data_path).read_batch() + + assert len(batch) == len(expected) == 7 + assert list(batch.offset) == [offset for _, _, offset in expected] + assert list(batch.message_type) == [header.message_type for header, _, _ in expected] + assert list(batch.sequence_number) == [header.sequence_number for header, _, _ in expected] + assert list(batch.source_identifier) == [header.source_identifier for header, _, _ in expected] + assert list(batch.size_bytes) == [len(data) for _, data, _ in expected] + for i, (header, data, _) in enumerate(expected): + assert bytes(batch.get_bytes(i)) == bytes(data) + assert bytes(batch.get_payload(i)) == bytes(data[MessageHeader.calcsize():]) + + +# Check that the batch honors the message type, source ID, time range, and max bytes filters like the iterator. +@pytest.mark.parametrize('kwargs', [ + {'message_types': [InputDataWrapperMessage]}, + {'message_types': [PoseMessage], 'invert_message_types': True}, + {'source_ids': [1]}, + {'time_range': TimeRange(start=1.5, end=2.5, absolute=True)}, +]) +def test_filters(data_path, kwargs): + expected = _read_iterator(data_path, **kwargs) + batch = MixedLogReader(data_path, **kwargs).read_batch() + assert list(batch.offset) == [offset for _, _, offset in expected] + + +def test_max_bytes(data_path): + # Let the reader index the whole file, then limit the read to the first 3 messages. + expected = _read_iterator(data_path) + _, data, offset = expected[2] + max_bytes = offset + len(data) + + reader = MixedLogReader(data_path) + reader.set_max_bytes(max_bytes) + batch = reader.read_batch() + assert list(batch.offset) == [offset for _, _, offset in expected[:3]] + + +# Check that read_batch() does not move the iterator. +def test_iterator_position_unchanged(data_path): + reader = MixedLogReader(data_path) + first = reader.read_next() + reader.read_batch() + second = reader.read_next() + assert first[0].sequence_number + 1 == second[0].sequence_number + + +# Check that a vectorized wrapper unpack matches the per-message unpack. +def test_unpack_array_input_data_wrapper(data_path): + batch = MixedLogReader(data_path, message_types=[InputDataWrapperMessage]).read_batch() + fields = batch.unpack_array(InputDataWrapperMessage) + + for i in range(len(batch)): + message = InputDataWrapperMessage() + message.unpack(bytes(batch.get_payload(i))) + assert fields['system_time_ns'][i] == message.system_time_ns + assert fields['data_type'][i] == message.data_type + assert bytes(InputDataWrapperMessage.get_data(batch.get_payload(i))) == message.data + + +# Check that selecting a subset with numpy indexing keeps the fields aligned. +def test_select(data_path): + batch = MixedLogReader(data_path, message_types=[InputDataWrapperMessage]).read_batch() + fields = batch.unpack_array(InputDataWrapperMessage) + selected = batch[fields['data_type'] == InputDataType.M_TYPE_EXTERNAL_UNFRAMED_GNSS] + + assert len(selected) == 3 + data = [bytes(InputDataWrapperMessage.get_data(selected.get_payload(i))) for i in range(len(selected))] + assert data == [b'first', b'second record', b'third'] + + +# Check that a scalar index is rejected, since it would not select an array of messages. +def test_select_scalar(data_path): + batch = MixedLogReader(data_path).read_batch() + with pytest.raises(TypeError): + _ = batch[0] + assert len(batch[0:1]) == 1 + + +# Check that unpacking a batch with other message types, or a class without unpack_array(), is an error. +def test_unpack_array_errors(data_path): + batch = MixedLogReader(data_path).read_batch() + with pytest.raises(ValueError): + batch.unpack_array(InputDataWrapperMessage) + + poses = batch[batch.message_type == MessageType.POSE] + with pytest.raises(NotImplementedError): + poses.unpack_array(PoseMessage) + + +# Check that a filter matching nothing produces an empty batch. +def test_empty(data_path): + batch = MixedLogReader(data_path, message_types=[EventNotificationMessage]).read_batch() + assert len(batch) == 0 + assert isinstance(batch, MessageBatch) diff --git a/python/tests/test_mixed_log_reader.py b/python/tests/test_mixed_log_reader.py index 2bd69bf5..84d387e7 100644 --- a/python/tests/test_mixed_log_reader.py +++ b/python/tests/test_mixed_log_reader.py @@ -139,6 +139,15 @@ def test_read_all(self, data_path): reader.rewind() self._check_results(reader, messages) + def test_read_headers_only(self, data_path): + messages = self._generate_mixed_data(data_path) + + # Read the headers without deserializing the payloads. + reader = MixedLogReader(str(data_path), return_header=True, return_payload=False, return_bytes=True) + results = list(reader) + assert [header.message_type for header, _ in results] == [m.get_type() for m in messages] + assert all(len(data) == header.get_message_size() for header, data in results) + def test_read_pose(self, data_path): messages = self._generate_mixed_data(data_path) expected_messages = [m for m in messages if isinstance(m, PoseMessage)]