Preserve complete Arrow episodes and their timestep order

minari/dataset/_storages/arrow_storage.py

@@ -81,16 +81,7 @@ def get_episode_metadata(self, episode_indices: Iterable[int]) -> Iterable[Dict]
                 yield json.load(file)
 
     def get_episodes(self, episode_indices: Iterable[int]) -> Iterable[dict]:
-        dataset = pa.dataset.dataset(
-            [
-                pa.dataset.dataset(
-                    f"{self.data_path}/{ep_id}",
-                    format=self.FORMAT,
-                    ignore_prefixes=["_", ".", "metadata.json"],
-                )
-                for ep_id in episode_indices
-            ]
-        )
+        """Read complete episodes in the order provided by any index iterable."""
 
         def _to_dict(id, episode):
             return {
@@ -115,7 +106,16 @@ def _to_dict(id, episode):
                 ),
             }
 
-        return map(_to_dict, episode_indices, dataset.to_batches())
+        def _read_episode(episode_id):
+            dataset = pa.dataset.dataset(
+                f"{self.data_path}/{episode_id}",
+                format=self.FORMAT,
+                ignore_prefixes=["_", ".", "metadata.json"],
+            )
+            batch = dataset.to_table().combine_chunks().to_batches()[0]
+            return _to_dict(episode_id, batch)
+
+        return map(_read_episode, episode_indices)
 
     def update_episodes(self, episodes: Iterable[EpisodeBuffer]):
         total_steps = self.total_steps
@@ -160,6 +160,7 @@ def update_episodes(self, episodes: Iterable[EpisodeBuffer]):
                 format=self.FORMAT,
                 partitioning=["episode_id"],
                 existing_data_behavior="overwrite_or_ignore",
+                use_threads=False,
             )
 
             episode_metadata: dict = {"id": episode_id, "total_steps": len(rewards)}

tests/dataset/test_arrow_episode_reads.py

@@ -0,0 +1,48 @@
+import gymnasium as gym
+import numpy as np
+import pytest
+
+from minari.data_collector.episode_buffer import EpisodeBuffer
+from minari.dataset.minari_storage import MinariStorage
+
+
+def _storage(path, format, lengths):
+    storage = MinariStorage.new(
+        path,
+        observation_space=gym.spaces.Box(-1e6, 1e6, (1,)),
+        action_space=gym.spaces.Discrete(2),
+        data_format=format,
+    )
+    for index, length in enumerate(lengths):
+        storage.update_episodes(
+            [
+                EpisodeBuffer(
+                    observations=np.arange(length + 1, dtype=np.float32).reshape(-1, 1),
+                    actions=np.zeros(length, dtype=np.int64),
+                    rewards=[float(index + 1)] * length,
+                    terminations=[False] * length,
+                    truncations=[False] * length,
+                )
+            ]
+        )
+    return storage
+
+
+@pytest.mark.parametrize("format", ["arrow", "parquet"])
+def test_get_episodes_accepts_one_shot_indices(tmp_path, format):
+    storage = _storage(tmp_path, format, [2, 3])
+    episodes = list(storage.get_episodes(iter([1, 0])))
+    assert [episode["id"] for episode in episodes] == [1, 0]
+    assert [len(episode["rewards"]) for episode in episodes] == [3, 2]
+
+
+@pytest.mark.parametrize("format", ["arrow", "parquet"])
+def test_get_episodes_does_not_split_long_episodes(tmp_path, format):
+    length = 140000
+    storage = _storage(tmp_path, format, [length, 2])
+    episodes = list(storage.get_episodes([0, 1]))
+    assert [len(episode["rewards"]) for episode in episodes] == [length, 2]
+    np.testing.assert_array_equal(
+        episodes[0]["observations"].ravel(), np.arange(length + 1)
+    )
+    np.testing.assert_array_equal(episodes[1]["rewards"], [2, 2])