Lazy loading CSV files (or tables in general)

Signed-off-by: Bram Stoeller <[email protected]>

src/power_grid_model_io/data_stores/csv_dir_store.py

@@ -6,7 +6,7 @@
 """
 
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Callable, Dict, List
 
 import pandas as pd
 
@@ -26,17 +26,24 @@ class CsvDirStore(BaseDataStore[TabularData]):
 
     def __init__(self, dir_path: Path, **csv_kwargs):
         super().__init__()
-        self._dir_path = dir_path
+        self._dir_path = Path(dir_path)
         self._csv_kwargs: Dict[str, Any] = csv_kwargs
         self._header_rows: List[int] = [0]
 
     def load(self) -> TabularData:
         """
-        Load all CSV files in a directory as tabular data.
+        Create a lazy loader for all CSV files in a directory and store them in a TabularData instance.
         """
-        data: Dict[str, pd.DataFrame] = {}
+
+        def lazy_csv_loader(csv_path: Path) -> Callable[[], pd.DataFrame]:
+            def csv_loader():
+                return pd.read_csv(filepath_or_buffer=csv_path, header=self._header_rows, **self._csv_kwargs)
+
+            return csv_loader
+
+        data: Dict[str, Callable[[], pd.DataFrame]] = {}
         for path in self._dir_path.glob("*.csv"):
-            data[path.stem] = pd.read_csv(filepath_or_buffer=path, header=self._header_rows, **self._csv_kwargs)
+            data[path.stem] = lazy_csv_loader(path)
 
         return TabularData(**data)
 

src/power_grid_model_io/data_types/tabular_data.py

@@ -6,7 +6,7 @@
 which supports unit conversions and value substitutions
 """
 
-from typing import Dict, Iterable, Optional, Tuple, Union
+from typing import Callable, Dict, Generator, Iterable, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -15,14 +15,16 @@
 from power_grid_model_io.mappings.unit_mapping import UnitMapping
 from power_grid_model_io.mappings.value_mapping import ValueMapping
 
+LazyDataFrame = Callable[[], pd.DataFrame]
+
 
 class TabularData:
     """
     The TabularData class is a wrapper around Dict[str, Union[pd.DataFrame, np.ndarray]],
     which supports unit conversions and value substitutions
     """
 
-    def __init__(self, **tables: Union[pd.DataFrame, np.ndarray]):
+    def __init__(self, **tables: Union[pd.DataFrame, np.ndarray, LazyDataFrame]):
         """
         Tabular data can either be a collection of pandas DataFrames and/or numpy structured arrays.
         The key word arguments will define the keys of the data.
@@ -34,12 +36,12 @@ def __init__(self, **tables: Union[pd.DataFrame, np.ndarray]):
             **tables: A collection of pandas DataFrames and/or numpy structured arrays
         """
         for table_name, table_data in tables.items():
-            if not isinstance(table_data, (pd.DataFrame, np.ndarray)):
+            if not isinstance(table_data, (pd.DataFrame, np.ndarray)) and not callable(table_data):
                 raise TypeError(
                     f"Invalid data type for table '{table_name}'; "
                     f"expected a pandas DataFrame or NumPy array, got {type(table_data).__name__}."
                 )
-        self._data: Dict[str, Union[pd.DataFrame, np.ndarray]] = tables
+        self._data: Dict[str, Union[pd.DataFrame, np.ndarray, LazyDataFrame]] = tables
         self._units: Optional[UnitMapping] = None
         self._substitution: Optional[ValueMapping] = None
         self._log = structlog.get_logger(type(self).__name__)
@@ -73,7 +75,7 @@ def get_column(self, table_name: str, column_name: str) -> pd.Series:
         Returns:
             The required column, with unit conversions and value substitutions applied
         """
-        table_data = self._data[table_name]
+        table_data = self[table_name]
 
         # If the index 'column' is requested, but no column called 'index' exist,
         # return the index of the dataframe as if it were an actual column.
@@ -176,6 +178,8 @@ def __getitem__(self, table_name: str) -> Union[pd.DataFrame, np.ndarray]:
 
         Returns: The 'raw' table data
         """
+        if callable(self._data[table_name]):
+            self._data[table_name] = self._data[table_name]()
         return self._data[table_name]
 
     def keys(self) -> Iterable[str]:
@@ -187,13 +191,14 @@ def keys(self) -> Iterable[str]:
 
         return self._data.keys()
 
-    def items(self) -> Iterable[Tuple[str, Union[pd.DataFrame, np.ndarray]]]:
+    def items(self) -> Generator[Tuple[str, Union[pd.DataFrame, np.ndarray]], None, None]:
         """
         Mimic the dictionary .items() function
 
-        Returns: An iterator over the table names and the raw table data
+        Returns: An generator of the table names and the raw table data
         """
 
         # Note: PyCharm complains about the type, but it is correct, as an ItemsView extends from
         # AbstractSet[Tuple[_KT_co, _VT_co]], which actually is compatible with Iterable[_KT_co, _VT_co]
-        return self._data.items()
+        for key in self._data:
+            yield key, self[key]