Fix streaming Parquet imports and video fallback

tests/datasets/test_episode_parquet_reader.py

@@ -20,7 +20,7 @@ pytest.importorskip("pyarrow", reason="pyarrow is required (install lerobot[data
 import pyarrow as pa
 import pyarrow.parquet as pq
 
-from lerobot.datasets.episode_parquet import EpisodeParquetReader
+from lerobot.streaming.episode_parquet import EpisodeParquetReader
 
 
 def _table(episodes: list[int]) -> pa.Table:

tests/datasets/test_episode_video_streaming.py

@@ -103,7 +103,14 @@ def test_parser_accepts_co64_chunk_offsets():
     np.testing.assert_array_equal(mp4.sample_offsets, np.array([10_000, 10_050, 10_025]))
 
 
-def _fake_cache(monkeypatch, tmp_path, *, byte_budget=8, max_open_decoders=1):
+def _fake_cache(
+    monkeypatch,
+    tmp_path,
+    *,
+    byte_budget=8,
+    max_open_decoders=1,
+    video_backend="torchcodec",
+):
     manifest = EpisodeVideoManifest(video_keys=["camera"], files=[], spans={})
     cache = EpisodeByteCache(
         manifest,
@@ -112,6 +119,7 @@ def _fake_cache(monkeypatch, tmp_path, *, byte_budget=8, max_open_decoders=1):
         workers=1,
         open_decoders=False,
         max_open_decoders=max_open_decoders,
+        video_backend=video_backend,
     )
     monkeypatch.setattr(
         cache,
@@ -161,6 +169,54 @@ def test_decoder_count_has_independent_limit(monkeypatch, tmp_path):
         assert cache.open_decoder_count == 1
 
 
+def test_decoder_eviction_and_cache_shutdown_close_backend_resources(monkeypatch, tmp_path):
+    opened = []
+
+    class FakeDecoder:
+        def __init__(self):
+            self.closed = False
+
+        def close(self):
+            self.closed = True
+
+    def open_decoder(_data):
+        decoder = FakeDecoder()
+        opened.append(decoder)
+        return decoder
+
+    monkeypatch.setattr("lerobot.streaming.episode_cache.open_video_decoder", open_decoder)
+    with _fake_cache(monkeypatch, tmp_path, byte_budget=20, max_open_decoders=1) as cache:
+        cache.get_decoder(0, "camera")
+        cache.get_decoder(1, "camera")
+
+        assert opened[0].closed
+        assert not opened[1].closed
+
+    assert opened[1].closed
+
+
+def test_decoder_falls_back_to_pyav_when_torchcodec_rejects_mini_mp4(monkeypatch, tmp_path):
+    opened_backends = []
+
+    class FakeDecoder:
+        pass
+
+    def open_decoder(_data, frame_mappings=None, *, backend="torchcodec"):
+        assert frame_mappings is None
+        opened_backends.append(backend)
+        if backend == "torchcodec":
+            raise ValueError("No valid stream found")
+        return FakeDecoder()
+
+    monkeypatch.setattr("lerobot.streaming.episode_cache.open_video_decoder", open_decoder)
+    with _fake_cache(monkeypatch, tmp_path, video_backend="torchcodec") as cache:
+        decoder = cache.get_decoder(0, "camera")
+
+        assert isinstance(decoder, FakeDecoder)
+        assert opened_backends == ["torchcodec", "pyav"]
+        assert cache.decoder_fallback_count == 1
+
+
 def test_releasing_episode_allows_immediate_eviction(monkeypatch, tmp_path):
     with _fake_cache(monkeypatch, tmp_path, byte_budget=5) as cache:
         cache.retain_episode(0)

tests/datasets/test_streaming_factory.py

@@ -33,6 +33,7 @@ def test_factory_wires_production_streaming_settings(monkeypatch):
     dataset_config = DatasetConfig(
         repo_id="owner/dataset",
         streaming=True,
+        video_backend="pyav",
         streaming_data_root="memory://payload",
         streaming_episode_pool_size=7,
         streaming_prefetch_episodes=3,
@@ -54,5 +55,6 @@ def test_factory_wires_production_streaming_settings(monkeypatch):
     assert captured["kwargs"]["prefetch_episodes"] == 3
     assert captured["kwargs"]["byte_budget_gb"] == 2.5
     assert captured["kwargs"]["max_num_shards"] == 1
+    assert captured["kwargs"]["video_backend"] == "pyav"
     assert captured["kwargs"]["return_uint8"] is True
     assert captured["kwargs"]["repeat"] is True

tests/datasets/test_streaming_production.py

@@ -61,19 +61,26 @@ def test_streaming_matches_map_style_with_exact_coverage(tmp_path: Path, lerobot
         _assert_item_equal(sample, map_dataset[int(sample["index"])])
 
 
-def test_streaming_rgb_video_matches_map_style(tmp_path: Path, lerobot_dataset_factory) -> None:
+@pytest.mark.parametrize("video_backend", ["torchcodec", "pyav"])
+def test_streaming_rgb_video_matches_map_style(
+    tmp_path: Path,
+    lerobot_dataset_factory,
+    video_backend: str,
+) -> None:
     root = tmp_path / "dataset"
     map_dataset = lerobot_dataset_factory(
         root=root,
         repo_id=DUMMY_REPO_ID,
         total_episodes=2,
         total_frames=20,
+        video_backend=video_backend,
     )
     streaming = StreamingLeRobotDataset(
         DUMMY_REPO_ID,
         root=root,
         shuffle=False,
         buffer_size=2,
+        video_backend=video_backend,
     )
 
     for sample in streaming: