diff options
| -rw-r--r-- | data/hdf5.py | 5 |
1 files changed, 5 insertions, 0 deletions
diff --git a/data/hdf5.py b/data/hdf5.py index d848023..6e2f9a4 100644 --- a/data/hdf5.py +++ b/data/hdf5.py @@ -19,6 +19,11 @@ class TaxiDataset(H5PYDataset): def data_path(self): return os.path.join(data.path, self.filename) + def extract(self, request): + if not self.load_in_memory: + raise ValueError('extract called on a dataset not loaded in memory') + return dict(zip(self.sources, self.get_data(None, request))) + class TaxiStream(DataStream): def __init__(self, which_set, filename='data.hdf5', iteration_scheme=None, **kwargs): dataset = TaxiDataset(which_set, filename, **kwargs) |