"""Tests for rerun.experimental.LazyChunkStream and RrdReader.""" from __future__ import annotations import subprocess from typing import TYPE_CHECKING import pyarrow.compute as pc import pytest import rerun as rr from inline_snapshot import snapshot as inline_snapshot from rerun.experimental import Chunk, DeriveLens, LazyChunkStream, MutateLens, RrdReader, Selector from .conftest import TEST_APP_ID as APP_ID, TEST_RECORDING_ID as RECORDING_ID if TYPE_CHECKING: from pathlib import Path # --------------------------------------------------------------------------- # RrdReader basics # --------------------------------------------------------------------------- def test_rrd_reader_properties(test_rrd_path: Path) -> None: reader = RrdReader(test_rrd_path) recs = reader.recordings() assert len(recs) == 1 assert recs[0].application_id == APP_ID assert recs[0].recording_id == RECORDING_ID def test_rrd_reader_file_not_found(tmp_path: Path) -> None: with pytest.raises(ValueError, match="not found"): RrdReader(tmp_path / "nonexistent.rrd") # --------------------------------------------------------------------------- # to_chunks / iter # --------------------------------------------------------------------------- def test_to_chunks(test_rrd_path: Path) -> None: """to_chunks() returns Chunk objects with expected properties.""" chunks = RrdReader(test_rrd_path).stream().to_chunks() assert len(chunks) > 0 for chunk in chunks: assert isinstance(chunk, Chunk) assert chunk.num_rows > 0 chunk.format(redact=True) def test_iter(test_rrd_path: Path) -> None: stream = RrdReader(test_rrd_path).stream() collected = stream.to_chunks() stream2 = RrdReader(test_rrd_path).stream() iterated = list(stream2) assert len(iterated) == len(collected) # --------------------------------------------------------------------------- # collect # --------------------------------------------------------------------------- def test_collect_returns_chunk_store(test_rrd_path: Path) -> None: """collect() returns a ChunkStore with correct entity paths.""" store = RrdReader(test_rrd_path).stream().collect() paths = store.schema().entity_paths() assert "/robots/arm" in paths assert "/cameras/front" in paths assert "/config" in paths # --------------------------------------------------------------------------- # Identity roundtrip # --------------------------------------------------------------------------- def test_identity_roundtrip(test_rrd_path: Path, tmp_path: Path) -> None: reader = RrdReader(test_rrd_path) out = tmp_path / "roundtrip.rrd" reader.stream().write_rrd(out, application_id=APP_ID, recording_id=RECORDING_ID) original = reader.store() roundtripped = RrdReader(out).store() assert original.schema() == roundtripped.schema() # Strong check: `rerun rrd compare` verifies semantic equality of the data. process = subprocess.run( ["rerun", "rrd", "compare", "--unordered", str(test_rrd_path), str(out)], check=False, capture_output=True, ) if process.returncode != 0: print(process.stdout.decode("utf-8")) print(process.stderr.decode("utf-8")) assert process.returncode == 0, f"RRD compare failed: {process.stderr.decode('utf-8')}" # --------------------------------------------------------------------------- # filter # --------------------------------------------------------------------------- def test_filter_content(test_rrd_path: Path) -> None: # Single string store = RrdReader(test_rrd_path).stream().filter(content="/robots/**").collect() assert store.schema().entity_paths() == ["/robots/arm"] # List of strings store2 = RrdReader(test_rrd_path).stream().filter(content=["/robots/**", "/cameras/**"]).collect() assert store2.schema().entity_paths() == ["/cameras/front", "/robots/arm"] def test_filter_is_static(test_rrd_path: Path) -> None: static_store = RrdReader(test_rrd_path).stream().filter(is_static=True).collect() assert static_store.schema().entity_paths() == ["/config"] temporal_store = RrdReader(test_rrd_path).stream().filter(is_static=False).collect() assert "/config" not in temporal_store.schema().entity_paths() def test_filter_has_timeline(test_rrd_path: Path) -> None: store = RrdReader(test_rrd_path).stream().filter(has_timeline="other_timeline").collect() # Only /robots/arm has the other_timeline assert store.schema().entity_paths() == ["/robots/arm"] def test_filter_component(test_rrd_path: Path) -> None: store = RrdReader(test_rrd_path).stream().filter(components="Points3D:positions").collect() assert store.schema().entity_paths() == ["/robots/arm"] # colors column should be stripped — only positions remains arm_cols = store.schema().columns_for(entity_path="/robots/arm") assert len(arm_cols) == 1 assert arm_cols[0].component == "Points3D:positions" def test_component_slice_gets_new_chunk_id(test_rrd_path: Path) -> None: """Slicing by component must produce chunks with fresh IDs, not reuse the original.""" original_ids = {c.id for c in RrdReader(test_rrd_path).stream().to_chunks()} # filter keeps only the matching column -> sliced chunk filtered = RrdReader(test_rrd_path).stream().filter(components="Points3D:positions").to_chunks() for chunk in filtered: assert chunk.id not in original_ids, "filter(components=...) must assign a new ChunkId" # NOLINT # drop keeps the non-matching columns -> also a sliced chunk dropped = RrdReader(test_rrd_path).stream().drop(components="Points3D:positions").to_chunks() for chunk in dropped: if chunk.entity_path == "/robots/arm": assert chunk.id not in original_ids, "drop(components=...) must assign a new ChunkId" # NOLINT def test_filter_multiple_components(test_rrd_path: Path) -> None: """filter(components=[A, B]) keeps both columns when present (OR semantics).""" store = RrdReader(test_rrd_path).stream().filter(components=["Points3D:positions", "Points3D:colors"]).collect() assert store.schema().entity_paths() == ["/robots/arm"] arm_cols = store.schema().columns_for(entity_path="/robots/arm") assert len(arm_cols) == 2 def test_filter_multiple_components_partial(test_rrd_path: Path) -> None: """filter(components=[A, Z]) where Z doesn't exist: keep A only.""" store = RrdReader(test_rrd_path).stream().filter(components=["Points3D:positions", "Nonexistent:foo"]).collect() assert store.schema().entity_paths() == ["/robots/arm"] def test_filter_multiple_components_none_present(test_rrd_path: Path) -> None: """filter(components=[Z1, Z2]) where neither exist: empty store.""" store = RrdReader(test_rrd_path).stream().filter(components=["Nonexistent:a", "Nonexistent:b"]).collect() assert store.schema().entity_paths() == [] def test_drop_multiple_components(test_rrd_path: Path) -> None: """drop(components=[A, B]) removes both columns.""" store = RrdReader(test_rrd_path).stream().drop(components=["Points3D:positions", "Points3D:colors"]).collect() # /robots/arm had only those two components, so it should be gone assert "/robots/arm" not in store.schema().entity_paths() def test_split_multiple_components(test_rrd_path: Path) -> None: """split(components=[A, B]): matched gets A+B, complement gets rest.""" stream = RrdReader(test_rrd_path).stream() matched, complement = stream.split(components=["Points3D:positions", "Points3D:colors"]) matched_store = matched.collect() complement_store = complement.collect() assert matched_store.schema().entity_paths() == ["/robots/arm"] assert len(complement_store.schema().entity_paths()) > 0 # --------------------------------------------------------------------------- # drop # --------------------------------------------------------------------------- def test_drop(test_rrd_path: Path) -> None: store = RrdReader(test_rrd_path).stream().drop(content="/robots/**").collect() paths = store.schema().entity_paths() assert "/robots/arm" not in paths assert "/cameras/front" in paths assert "/config" in paths # --------------------------------------------------------------------------- # split / merge # --------------------------------------------------------------------------- def test_split_merge_roundtrip(test_rrd_path: Path) -> None: original = RrdReader(test_rrd_path).stream().collect() stream2 = RrdReader(test_rrd_path).stream() static_branch, temporal_branch = stream2.split(is_static=True) merged = LazyChunkStream.merge(static_branch, temporal_branch).collect() assert original.schema() == merged.schema() def test_split_drop_one_branch(test_rrd_path: Path) -> None: """Consuming only one branch of a split should not hang.""" stream = RrdReader(test_rrd_path).stream() matching, _non_matching = stream.split(content="/robots/**") store = matching.collect() assert store.schema().entity_paths() == ["/robots/arm"] # --------------------------------------------------------------------------- # from_iter # --------------------------------------------------------------------------- def test_from_iter(test_rrd_path: Path) -> None: original = RrdReader(test_rrd_path).stream().to_chunks() roundtripped = LazyChunkStream.from_iter(original).collect() assert roundtripped.schema() == RrdReader(test_rrd_path).store().schema() # --------------------------------------------------------------------------- # Composition # --------------------------------------------------------------------------- def test_chained_filters(test_rrd_path: Path) -> None: store = RrdReader(test_rrd_path).stream().filter(is_static=False).filter(content="/robots/**").collect() assert store.schema().entity_paths() == ["/robots/arm"] # --------------------------------------------------------------------------- # Dangling split branch # --------------------------------------------------------------------------- def test_dangling_split_matched_only(test_rrd_path: Path, capfd: pytest.CaptureFixture[str]) -> None: """Using only the matched branch of a split should work (degenerated to filter) and warn.""" stream = RrdReader(test_rrd_path).stream() matched, _unmatched = stream.split(content="/robots/**") store = matched.collect() assert store.schema().entity_paths() == ["/robots/arm"] captured = capfd.readouterr() assert "only one branch" in captured.err.lower(), ( f"Expected a warning about dangling split branch on stderr, got: {captured.err!r}" ) def test_dangling_split_unmatched_only(test_rrd_path: Path, capfd: pytest.CaptureFixture[str]) -> None: """Using only the unmatched branch of a split should work (degenerated to drop) and warn.""" stream = RrdReader(test_rrd_path).stream() _matched, unmatched = stream.split(content="/robots/**") store = unmatched.collect() assert "/robots/arm" not in store.schema().entity_paths() captured = capfd.readouterr() assert "only one branch" in captured.err.lower(), ( f"Expected a warning about dangling split branch on stderr, got: {captured.err!r}" ) # --------------------------------------------------------------------------- # Move semantics # --------------------------------------------------------------------------- def test_stream_consumed_after_filter(test_rrd_path: Path) -> None: """A stream consumed by filter() cannot be used again as a builder input.""" stream = RrdReader(test_rrd_path).stream() _filtered = stream.filter(is_static=True) with pytest.raises(ValueError, match="already been consumed"): stream.drop(is_static=False) def test_stream_consumed_after_split(test_rrd_path: Path) -> None: """A stream consumed by split() cannot be used again as a builder input.""" stream = RrdReader(test_rrd_path).stream() _a, _b = stream.split(is_static=True) with pytest.raises(ValueError, match="already been consumed"): stream.filter(content="/foo/**") def test_merge_same_stream_twice(test_rrd_path: Path) -> None: """Passing the same stream to merge twice is an error.""" stream = RrdReader(test_rrd_path).stream() a, _b = stream.split(is_static=True) with pytest.raises(ValueError, match="already been consumed"): LazyChunkStream.merge(a, a) def test_merge_indirect_reuse(test_rrd_path: Path) -> None: """A stream used as split upstream and also passed directly to merge is an error.""" stream = RrdReader(test_rrd_path).stream() a, b = stream.split(is_static=True) _b1, b2 = b.split(content="/robots/**") # b was consumed by the second split, so passing it to merge should fail. with pytest.raises(ValueError, match="already been consumed"): LazyChunkStream.merge(b, b2, a) def test_terminal_does_not_consume(test_rrd_path: Path) -> None: """Terminals (collect, write_rrd, iter) borrow without consuming.""" stream = RrdReader(test_rrd_path).stream() store1 = stream.collect() store2 = stream.collect() assert store1.schema() == store2.schema() chunks = list(stream) assert len(chunks) > 0 # --------------------------------------------------------------------------- # Lenses # --------------------------------------------------------------------------- def test_lenses_identity(test_rrd_path: Path) -> None: """A lens with Selector('.') passes through the struct component data unchanged.""" lens = MutateLens("Imu:accel", Selector(".")) store = RrdReader(test_rrd_path).stream().filter(content="/sensors/**").lenses(lens).collect() assert store.summary() == inline_snapshot( "/sensors/imu rows=2 static=False timelines=['my_index'] cols=['Imu:accel', 'my_index']" ) def test_lenses_field_selector(test_rrd_path: Path) -> None: """A lens with Selector('.x') extracts a struct field and reinterprets it as a Rerun Scalar.""" lens = DeriveLens("Imu:accel").to_component(rr.Scalars.descriptor_scalars(), Selector(".x")) store = RrdReader(test_rrd_path).stream().filter(content="/sensors/**").lenses(lens).collect() assert store.summary() == inline_snapshot( "/sensors/imu rows=2 static=False timelines=['my_index'] cols=['Scalars:scalars', 'my_index']" ) # Verify the extracted values are correct chunks = RrdReader(test_rrd_path).stream().filter(content="/sensors/**").lenses(lens).to_chunks() rb = chunks[0].to_record_batch() scalars = rb.column("Scalars:scalars") assert scalars.to_pylist() == [[0.1], [0.4]] def test_lenses_multiple_outputs(test_rrd_path: Path) -> None: """A single lens can produce multiple output groups at different entity paths.""" lenses = [ DeriveLens("Imu:accel", output_entity="/out/x").to_component(rr.Scalars.descriptor_scalars(), Selector(".x")), DeriveLens("Imu:accel", output_entity="/out/z").to_component(rr.Scalars.descriptor_scalars(), Selector(".z")), ] store = RrdReader(test_rrd_path).stream().filter(content="/sensors/**").lenses(lenses).collect() assert store.summary() == inline_snapshot("""\ /out/x rows=2 static=False timelines=['my_index'] cols=['Scalars:scalars', 'my_index'] /out/z rows=2 static=False timelines=['my_index'] cols=['Scalars:scalars', 'my_index']\ """) def test_lenses_drop_unmatched(test_rrd_path: Path) -> None: """With drop_unmatched (default), unmatched chunks are not forwarded.""" lens = DeriveLens("nonexistent:Component:foo").to_component("out:Component:bar", Selector(".")) store = RrdReader(test_rrd_path).stream().lenses(lens, output_mode="drop_unmatched").collect() assert store.summary() == inline_snapshot("") def test_lenses_forward_unmatched(test_rrd_path: Path) -> None: """With forward_unmatched, transformed chunks replace originals and unmatched chunks pass through.""" lens = DeriveLens("Imu:accel", output_entity="/transformed").to_component( rr.Scalars.descriptor_scalars(), Selector(".x") ) store = ( RrdReader(test_rrd_path) .stream() .lenses(lens, output_mode="forward_unmatched") .drop(content="/__properties/**") .collect() ) assert store.summary() == inline_snapshot("""\ /cameras/front rows=1 static=False timelines=['my_index'] cols=['TextLog:text', 'my_index'] /config rows=1 static=True timelines=[] cols=['TextLog:text'] /robots/arm rows=2 static=False timelines=['my_index', 'other_timeline'] cols=['Points3D:colors', 'Points3D:positions', 'my_index', 'other_timeline'] /transformed rows=2 static=False timelines=['my_index'] cols=['Scalars:scalars', 'my_index']\ """) def test_lenses_forward_all(test_rrd_path: Path) -> None: """With forward_all, both transformed and original data are forwarded.""" lens = DeriveLens("Imu:accel", output_entity="/transformed").to_component( rr.Scalars.descriptor_scalars(), Selector(".x") ) store = ( RrdReader(test_rrd_path) .stream() .lenses(lens, output_mode="forward_all") .drop(content="/__properties/**") .collect() ) assert store.summary() == inline_snapshot("""\ /cameras/front rows=1 static=False timelines=['my_index'] cols=['TextLog:text', 'my_index'] /config rows=1 static=True timelines=[] cols=['TextLog:text'] /robots/arm rows=2 static=False timelines=['my_index', 'other_timeline'] cols=['Points3D:colors', 'Points3D:positions', 'my_index', 'other_timeline'] /sensors/imu rows=2 static=False timelines=['my_index'] cols=['Imu:accel', 'my_index'] /transformed rows=2 static=False timelines=['my_index'] cols=['Scalars:scalars', 'my_index']\ """) def test_lenses_consumes_stream(test_rrd_path: Path) -> None: """Calling .lenses() consumes the stream (move semantics).""" lens = DeriveLens("Imu:accel").to_component(rr.Scalars.descriptor_scalars(), Selector(".x")) stream = RrdReader(test_rrd_path).stream() _transformed = stream.lenses(lens) with pytest.raises(ValueError, match="already been consumed"): stream.filter(is_static=True) def test_lenses_chained_with_filter(test_rrd_path: Path) -> None: """Lenses can be composed with filter in a pipeline.""" lens = DeriveLens("Imu:accel").to_component(rr.Scalars.descriptor_scalars(), Selector(".z")) store = RrdReader(test_rrd_path).stream().filter(content="/sensors/**").lenses(lens).collect() assert store.summary() == inline_snapshot( "/sensors/imu rows=2 static=False timelines=['my_index'] cols=['Scalars:scalars', 'my_index']" ) def test_lenses_content_filter_match(test_rrd_path: Path) -> None: """With `content` set to a matching path, lenses apply only to those chunks; others pass through.""" lens = DeriveLens("Imu:accel", output_entity="/transformed").to_component( rr.Scalars.descriptor_scalars(), Selector(".x") ) store = ( RrdReader(test_rrd_path) .stream() .lenses(lens, output_mode="drop_unmatched", content="/sensors/**") .drop(content="/__properties/**") .collect() ) # /sensors/imu was matched by content -> lens applied -> produced /transformed. # All other chunks pass through unchanged regardless of `drop_unmatched`. assert store.summary() == inline_snapshot("""\ /cameras/front rows=1 static=False timelines=['my_index'] cols=['TextLog:text', 'my_index'] /config rows=1 static=True timelines=[] cols=['TextLog:text'] /robots/arm rows=2 static=False timelines=['my_index', 'other_timeline'] cols=['Points3D:colors', 'Points3D:positions', 'my_index', 'other_timeline'] /transformed rows=2 static=False timelines=['my_index'] cols=['Scalars:scalars', 'my_index']\ """) def test_lenses_content_filter_excludes_lens_target(test_rrd_path: Path) -> None: """Chunks outside the `content` scope bypass the lens and pass through, regardless of output_mode.""" lens = DeriveLens("Imu:accel", output_entity="/transformed").to_component( rr.Scalars.descriptor_scalars(), Selector(".x") ) # Content scope only includes /robots/**, so /sensors/imu is bypassed entirely # (the lens never sees it). /robots/arm is in scope but the lens doesn't match, # so under drop_unmatched it's dropped. store = ( RrdReader(test_rrd_path) .stream() .lenses(lens, output_mode="drop_unmatched", content="/robots/**") .drop(content="/__properties/**") .collect() ) # /sensors/imu passes through (out of content scope). No /transformed (lens never matched). paths = store.schema().entity_paths() assert "/transformed" not in paths assert "/sensors/imu" in paths assert "/robots/arm" not in paths # in scope, but dropped by drop_unmatched def test_lenses_invalid_output_mode(test_rrd_path: Path) -> None: """Invalid output_mode string raises ValueError.""" lens = DeriveLens("Points3D:positions").to_component("Points3D:positions", Selector(".")) with pytest.raises(ValueError, match="Unknown output_mode"): RrdReader(test_rrd_path).stream().lenses(lens, output_mode="invalid") # type: ignore[arg-type] def test_lenses_time_extraction(test_rrd_path: Path) -> None: """A lens can extract a timestamp field from a struct component as a new timeline.""" lens = ( DeriveLens("Imu:accel") .to_component(rr.Scalars.descriptor_scalars(), Selector(".x")) .to_timeline("sensor_time", "timestamp_ns", Selector(".timestamp")) ) store = RrdReader(test_rrd_path).stream().filter(content="/sensors/**").lenses(lens).collect() assert store.summary() == inline_snapshot( "/sensors/imu rows=2 static=False timelines=['my_index', 'sensor_time'] cols=['Scalars:scalars', 'my_index', 'sensor_time']" ) chunks = RrdReader(test_rrd_path).stream().filter(content="/sensors/**").lenses(lens).to_chunks() rb = chunks[0].to_record_batch() scalars = rb.column("Scalars:scalars") assert scalars.to_pylist() == [[0.1], [0.4]] sensor_time = rb.column("sensor_time") assert [t.value for t in sensor_time.to_pylist()] == [1000000000, 2000000000] def test_lenses_dynamic_selector(test_rrd_path: Path) -> None: """A lens with a dynamic selector uses .pipe() to transform data with a Python callable.""" selector = Selector(".x").pipe(lambda arr: pc.multiply(arr, 2.0)) lens = DeriveLens("Imu:accel").to_component(rr.Scalars.descriptor_scalars(), selector) store = RrdReader(test_rrd_path).stream().filter(content="/sensors/**").lenses(lens).collect() assert store.summary() == inline_snapshot( "/sensors/imu rows=2 static=False timelines=['my_index'] cols=['Scalars:scalars', 'my_index']" ) chunks = RrdReader(test_rrd_path).stream().filter(content="/sensors/**").lenses(lens).to_chunks() rb = chunks[0].to_record_batch() scalars = rb.column("Scalars:scalars") assert scalars.to_pylist() == [[0.2], [0.8]] # --------------------------------------------------------------------------- # map / flat_map # --------------------------------------------------------------------------- def test_map_identity(test_rrd_path: Path) -> None: """map(identity) preserves all chunks.""" original = RrdReader(test_rrd_path).stream().to_chunks() mapped = RrdReader(test_rrd_path).stream().map(lambda c: c).to_chunks() assert len(mapped) == len(original) def test_flat_map_identity(test_rrd_path: Path) -> None: """flat_map(lambda c: [c]) preserves all chunks.""" original = RrdReader(test_rrd_path).stream().to_chunks() mapped = RrdReader(test_rrd_path).stream().flat_map(lambda c: [c]).to_chunks() assert len(mapped) == len(original) def test_flat_map_drop_all(test_rrd_path: Path) -> None: """flat_map(lambda c: []) produces an empty stream.""" chunks = RrdReader(test_rrd_path).stream().flat_map(lambda _c: []).to_chunks() assert len(chunks) == 0 def test_map_after_filter(test_rrd_path: Path) -> None: """Map composes after filter.""" store = RrdReader(test_rrd_path).stream().filter(content="/robots/**").map(lambda c: c).collect() assert store.schema().entity_paths() == ["/robots/arm"] def test_map_in_split_branch(test_rrd_path: Path) -> None: """Map works on a branch produced by split.""" stream = RrdReader(test_rrd_path).stream() matched, _non_matched = stream.split(content="/robots/**") mapped = matched.map(lambda c: c) store = mapped.collect() assert store.schema().entity_paths() == ["/robots/arm"] def test_map_error_propagation(test_rrd_path: Path) -> None: """A map callable that raises propagates the exception.""" def _raise(_c: Chunk) -> Chunk: raise ValueError("test error from map") with pytest.raises(ValueError, match="test error from map"): RrdReader(test_rrd_path).stream().map(_raise).to_chunks() def test_flat_map_error_propagation(test_rrd_path: Path) -> None: """A flat_map callable that raises propagates the exception.""" def _raise(_c: Chunk) -> list[Chunk]: raise ValueError("test error from flat_map") with pytest.raises(ValueError, match="test error from flat_map"): RrdReader(test_rrd_path).stream().flat_map(_raise).to_chunks() def test_map_consumed_stream(test_rrd_path: Path) -> None: """Calling map on a consumed stream raises ValueError.""" stream = RrdReader(test_rrd_path).stream() _filtered = stream.filter(is_static=True) with pytest.raises(ValueError, match="already been consumed"): stream.map(lambda c: c) def test_flat_map_consumed_stream(test_rrd_path: Path) -> None: """Calling flat_map on a consumed stream raises ValueError.""" stream = RrdReader(test_rrd_path).stream() _filtered = stream.filter(is_static=True) with pytest.raises(ValueError, match="already been consumed"): stream.flat_map(lambda c: [c]) def test_map_multiple_executions(test_rrd_path: Path) -> None: """Calling to_chunks() twice on a mapped stream produces the same results.""" stream = RrdReader(test_rrd_path).stream().map(lambda c: c) first = stream.to_chunks() second = stream.to_chunks() assert len(first) == len(second)