Merge pull request #166 from alliander-opensource/feature/lazy-excel-loading

Lazy excel loading

Author: bramstoeller

src/power_grid_model_io/data_stores/excel_file_store.py

@@ -13,6 +13,7 @@
 
 from power_grid_model_io.data_stores.base_data_store import BaseDataStore
 from power_grid_model_io.data_types import TabularData
+from power_grid_model_io.data_types.tabular_data import LazyDataFrame
 
 
 class ExcelFileStore(BaseDataStore[TabularData]):
@@ -24,7 +25,7 @@ class ExcelFileStore(BaseDataStore[TabularData]):
     same values) or renamed.
     """
 
-    __slots__ = ("_file_paths", "_header_rows")
+    __slots__ = ("_file_paths", "_excel_files", "_header_rows")
 
     _unnamed_pattern: re.Pattern = re.compile(r"Unnamed: \d+_level_\d+")
 
@@ -34,6 +35,7 @@ def __init__(self, file_path: Optional[Path] = None, **extra_paths: Path):
         # Create a dictionary of all supplied file paths:
         # {"": file_path, extra_name[0]: extra_path[0], extra_name[1]: extra_path[1], ...}
         self._file_paths: Dict[str, Path] = {}
+        self._excel_files: Dict[str, pd.ExcelFile] = {}
         if file_path is not None:
             self._file_paths[""] = file_path
         for name, path in extra_paths.items():
@@ -62,21 +64,37 @@ def load(self) -> TabularData:
         have no prefix, while the tables of all the extra files will be prefixed with the name of the key word argument
         as supplied in the constructor.
         """
-        data: Dict[str, pd.DataFrame] = {}
+        data: Dict[str, LazyDataFrame] = {}
         for name, path in self._file_paths.items():
-            with path.open(mode="rb") as file_pointer:
-                spreadsheet = pd.read_excel(io=file_pointer, sheet_name=None, header=self._header_rows)
-            for sheet_name, sheet_data in spreadsheet.items():
-                sheet_data = self._remove_unnamed_column_placeholders(data=sheet_data)
-                sheet_data = self._handle_duplicate_columns(data=sheet_data, sheet_name=sheet_name)
+            self._excel_files[name] = pd.ExcelFile(path)
+            for sheet_name in self._excel_files[name].sheet_names:
+                loader = self._load_sheet_wrapper(name, sheet_name)
                 if name:
                     sheet_name = f"{name}.{sheet_name}"
                 if sheet_name in data:
                     raise ValueError(f"Duplicate sheet name '{sheet_name}'")
-                data[sheet_name] = sheet_data
-
+                data[sheet_name] = loader
         return TabularData(**data)
 
+    def _load_sheet_wrapper(self, name: str, sheet_name: str):
+        """
+        Load a single Excel sheet as a Pandas DataFrame.
+
+        Args:
+            name: the name of the file (empty string for the main sheet)
+            sheet_name: the name of the sheet
+
+        Returns: The contents the specified Excel sheet.
+        """
+
+        def wrapper():
+            sheet_data = self._excel_files[name].parse(sheet_name, header=self._header_rows)
+            sheet_data = self._remove_unnamed_column_placeholders(data=sheet_data)
+            sheet_data = self._handle_duplicate_columns(data=sheet_data, sheet_name=sheet_name)
+            return sheet_data
+
+        return wrapper
+
     def save(self, data: TabularData) -> None:
         """
         Store tabular data as one or more Excel file.

tests/unit/data_stores/test_excel_file_store.py

@@ -15,7 +15,7 @@
 from power_grid_model_io.data_stores.excel_file_store import ExcelFileStore
 from power_grid_model_io.data_types.tabular_data import TabularData
 
-from ...utils import assert_log_exists
+from ...utils import MockExcelFile, assert_log_exists
 
 PandasExcelData = Dict[str, pd.DataFrame]
 
@@ -103,55 +103,57 @@ def test_files__read_only():
 
 @patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._handle_duplicate_columns")
 @patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._remove_unnamed_column_placeholders")
-@patch("power_grid_model_io.data_stores.excel_file_store.Path.open", mock_open())
-@patch("power_grid_model_io.data_stores.excel_file_store.pd.read_excel")
+@patch("power_grid_model_io.data_stores.excel_file_store.pd.ExcelFile")
 def test_load(
-    mock_read_excel: MagicMock,
+    mock_excel_file: MagicMock,
     mock_remove_unnamed_column_placeholders: MagicMock,
     mock_handle_duplicate_columns: MagicMock,
     objects_excel: PandasExcelData,
 ):
-    # Arrange
     fs = ExcelFileStore(file_path=Path("input_data.xlsx"))
-    mock_read_excel.return_value = objects_excel
+    mock_excel_file.return_value = MockExcelFile(objects_excel)
     mock_remove_unnamed_column_placeholders.side_effect = noop
     mock_handle_duplicate_columns.side_effect = noop
 
     # Act
     data = fs.load()
 
     # Assert
-    mock_read_excel.assert_called_once()
+    mock_excel_file.assert_called_once()
+    pd.testing.assert_frame_equal(data["Nodes"], objects_excel["Nodes"])
+    pd.testing.assert_frame_equal(data["Lines"], objects_excel["Lines"])
     assert mock_remove_unnamed_column_placeholders.call_args_list[0] == call(data=objects_excel["Nodes"])
     assert mock_remove_unnamed_column_placeholders.call_args_list[1] == call(data=objects_excel["Lines"])
     assert mock_handle_duplicate_columns.call_args_list[0] == call(data=objects_excel["Nodes"], sheet_name="Nodes")
     assert mock_handle_duplicate_columns.call_args_list[1] == call(data=objects_excel["Lines"], sheet_name="Lines")
-    pd.testing.assert_frame_equal(data["Nodes"], objects_excel["Nodes"])
-    pd.testing.assert_frame_equal(data["Lines"], objects_excel["Lines"])
 
 
 @patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._handle_duplicate_columns")
 @patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._remove_unnamed_column_placeholders")
-@patch("power_grid_model_io.data_stores.excel_file_store.Path.open", mock_open())
-@patch("power_grid_model_io.data_stores.excel_file_store.pd.read_excel")
+@patch("power_grid_model_io.data_stores.excel_file_store.pd.ExcelFile")
 def test_load__extra(
-    mock_read_excel: MagicMock,
+    mock_excel_file: MagicMock,
     mock_remove_unnamed_column_placeholders: MagicMock,
     mock_handle_duplicate_columns: MagicMock,
     objects_excel: PandasExcelData,
     specs_excel: PandasExcelData,
 ):
+
     # Arrange
     fs = ExcelFileStore(Path("input_data.xlsx"), foo=Path("foo_types.xlsx"))
-    mock_read_excel.side_effect = (objects_excel, specs_excel)
+    mock_excel_file.side_effect = (MockExcelFile(objects_excel), MockExcelFile(specs_excel))
     mock_remove_unnamed_column_placeholders.side_effect = noop
     mock_handle_duplicate_columns.side_effect = noop
 
     # Act
     data = fs.load()
 
     # Assert
-    assert mock_read_excel.call_count == 2
+    assert mock_excel_file.call_count == 2
+    pd.testing.assert_frame_equal(data["Nodes"], objects_excel["Nodes"])
+    pd.testing.assert_frame_equal(data["Lines"], objects_excel["Lines"])
+    pd.testing.assert_frame_equal(data["foo.Colors"], specs_excel["Colors"])
+    pd.testing.assert_frame_equal(data["foo.Lines"], specs_excel["Lines"])
     assert mock_remove_unnamed_column_placeholders.call_args_list[0] == call(data=objects_excel["Nodes"])
     assert mock_remove_unnamed_column_placeholders.call_args_list[1] == call(data=objects_excel["Lines"])
     assert mock_remove_unnamed_column_placeholders.call_args_list[2] == call(data=specs_excel["Colors"])
@@ -160,26 +162,21 @@ def test_load__extra(
     assert mock_handle_duplicate_columns.call_args_list[1] == call(data=objects_excel["Lines"], sheet_name="Lines")
     assert mock_handle_duplicate_columns.call_args_list[2] == call(data=specs_excel["Colors"], sheet_name="Colors")
     assert mock_handle_duplicate_columns.call_args_list[3] == call(data=specs_excel["Lines"], sheet_name="Lines")
-    pd.testing.assert_frame_equal(data["Nodes"], objects_excel["Nodes"])
-    pd.testing.assert_frame_equal(data["Lines"], objects_excel["Lines"])
-    pd.testing.assert_frame_equal(data["foo.Colors"], specs_excel["Colors"])
-    pd.testing.assert_frame_equal(data["foo.Lines"], specs_excel["Lines"])
 
 
 @patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._handle_duplicate_columns")
 @patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._remove_unnamed_column_placeholders")
-@patch("power_grid_model_io.data_stores.excel_file_store.Path.open", mock_open())
-@patch("power_grid_model_io.data_stores.excel_file_store.pd.read_excel")
+@patch("power_grid_model_io.data_stores.excel_file_store.pd.ExcelFile")
 def test_load__extra__duplicate_sheet_name(
-    mock_read_excel: MagicMock,
+    mock_excel_file: MagicMock,
     mock_remove_unnamed_column_placeholders: MagicMock,
     mock_handle_duplicate_columns: MagicMock,
 ):
     # Arrange
     foo_data = {"bar.Nodes": pd.DataFrame()}
     bar_data = {"Nodes": pd.DataFrame()}
     fs = ExcelFileStore(Path("foo.xlsx"), bar=Path("bar.xlsx"))
-    mock_read_excel.side_effect = (foo_data, bar_data)
+    mock_excel_file.side_effect = (MockExcelFile(foo_data), MockExcelFile(bar_data))
     mock_remove_unnamed_column_placeholders.side_effect = noop
     mock_handle_duplicate_columns.side_effect = noop
 

tests/unit/data_stores/test_vision_excel_file_store.py

@@ -7,16 +7,16 @@
 from power_grid_model_io.data_stores.vision_excel_file_store import VisionExcelFileStore
 
 
-@patch("power_grid_model_io.data_stores.excel_file_store.pd.read_excel")
+@patch("power_grid_model_io.data_stores.excel_file_store.pd.ExcelFile")
 @patch("power_grid_model_io.data_stores.excel_file_store.Path.open", mock_open())
-def test_header_rows(read_excel_mock: MagicMock):
+def test_header_rows(mock_excel_file: MagicMock):
     # Arrange
     store = VisionExcelFileStore(file_path=Path("dummy.xlsx"))
-    read_excel_mock.return_value = {}
+    mock_excel_file.return_value.sheet_names = ["foo"]
 
     # Act
-    store.load()
+    data = store.load()
+    data["foo"]
 
     # Assert
-    read_excel_mock.assert_called_once()
-    assert read_excel_mock.call_args_list[0].kwargs["header"] == [0, 1]
+    mock_excel_file.return_value.parse.assert_called_once_with("foo", header=[0, 1])

tests/utils.py

@@ -256,3 +256,15 @@ def __len__(self):
 
     def __getitem__(self, item: str):
         return MockVal(pd.Series(name=item, dtype=np.float64))
+
+
+class MockExcelFile:
+    def __init__(self, data: Dict[str, pd.DataFrame]):
+        self.data = data
+
+    @property
+    def sheet_names(self) -> List[str]:
+        return list(self.data.keys())
+
+    def parse(self, sheet_name: str, **_kwargs) -> pd.DataFrame:
+        return self.data[sheet_name]