Cache dataset in-memory

Author: samedii

datastream/dataset.py

@@ -487,6 +487,27 @@ def zip(datasets: List[Dataset]) -> Dataset[Tuple]:
             ))
         )
 
+    def cache(self, key_column):
+        '''Cache dataset in-memory based on key column.'''
+        from functools import lru_cache
+
+        key_mapping = dict(zip(
+            self.dataframe[key_column],
+            range(len(self)),
+        ))
+
+        @lru_cache(maxsize=None)
+        def only_key(key):
+            return self.get_item(self.dataframe, key_mapping[key])
+
+        return Dataset(
+            dataframe=self.dataframe,
+            length=self.length,
+            get_item=lambda dataframe, index: only_key(
+                dataframe.iloc[index][key_column]
+            ),
+        )
+
 
 def test_equal():
     dataset1 = Dataset.from_subscriptable([4, 7, 12])

datastream/datastream.py

@@ -253,6 +253,13 @@ def multi_sample(self: Datastream[T], n: int) -> Datastream[T]:
             MultiSampler.from_number(n, self.dataset),
         )
 
+    def cache(self, key_column):
+        '''Cache dataset in-memory. See :func:`Dataset.cache` for details.'''
+        return Datastream(
+            self.dataset.cache(key_column),
+            self.sampler,
+        )
+
 
 def test_datastream_merge():