Fix validation of saved multi-year datasets

changelog.d/2416.fixed.md

@@ -0,0 +1 @@
+- Fixed loading local multi-year H5 datasets saved by `UKMultiYearDataset.save` by distinguishing yearly groups from single-year pandas tables.

docs/book/usage/simulations.md

@@ -297,6 +297,10 @@ print(sim.calculate("household_net_income", 2026))
 `Simulation` and `Microsimulation` can also load H5 files from local paths,
 Hugging Face URLs, or Google Cloud Storage URLs:
 
+Files written by either `UKSingleYearDataset.save(path)` or
+`UKMultiYearDataset.save(path)` can be loaded with `Simulation(dataset=path)`.
+The loader detects single-year tables and multi-year groups from the H5 layout.
+
 ```python
 sim = Microsimulation(
     dataset="gs://policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.55.10"

policyengine_uk/data/dataset_schema.py

@@ -22,7 +22,8 @@ def validate_file_path(file_path: str, raise_exception: bool = True):
                 raise FileNotFoundError(f"File not found: {file_path}")
             return False
 
-        # Check if the file contains time_period, person, benunit, and household datasets
+        # Pandas stores each table as an annotated HDF5 group. A multi-year
+        # file has unannotated parent groups here, with tables below the years.
         with h5py.File(file_path, "r") as f:
             required_datasets = [
                 "time_period",
@@ -31,7 +32,11 @@ def validate_file_path(file_path: str, raise_exception: bool = True):
                 "household",
             ]
             for dataset in required_datasets:
-                if dataset not in f:
+                if (
+                    dataset not in f
+                    or not isinstance(f[dataset], h5py.Group)
+                    or "pandas_type" not in f[dataset].attrs
+                ):
                     if raise_exception:
                         raise ValueError(
                             f"Dataset '{dataset}' not found in the file: {file_path}"
@@ -298,26 +303,30 @@ def validate_file_path(file_path: str, raise_exception: bool = False):
             else:
                 return False
 
-        # Check if the file contains datasets for multiple years
+        # h5py lists only immediate children: the entity groups contain the
+        # year-named pandas tables written by save, even for a single year.
         with h5py.File(file_path, "r") as f:
             for required_dataset in ["person", "benunit", "household"]:
-                if not any(f"{required_dataset}" in key for key in f.keys()):
+                if required_dataset not in f:
                     if raise_exception:
                         raise ValueError(
                             f"Dataset '{required_dataset}' not found in the file: {file_path}"
                         )
                     else:
                         return False
 
-                # Check that there is at least one dataset year in the folder (keys include e.g. person/2025)
-
-                # Check that there is at least one dataset year in the folder
-                years_found = False
-                for key in f.keys():
-                    parts = key.split("/")
-                    if len(parts) >= 2 and required_dataset == parts[0]:
-                        years_found = True
-                        break
+                group = f[required_dataset]
+                years_found = (
+                    isinstance(group, h5py.Group)
+                    and "pandas_type" not in group.attrs
+                    and any(
+                        len(year) == 4
+                        and year.isdigit()
+                        and isinstance(group[year], h5py.Group)
+                        and "pandas_type" in group[year].attrs
+                        for year in group
+                    )
+                )
 
                 if not years_found:
                     if raise_exception:

policyengine_uk/tests/test_simulation_dataset_sources.py

@@ -1,7 +1,11 @@
 from pathlib import Path
+from tempfile import TemporaryDirectory
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
+from hypothesis import example, given, settings
+from hypothesis import strategies as st
 
 import policyengine_uk.simulation as simulation_module
 from policyengine_uk import Simulation
@@ -144,48 +148,80 @@ def test_dataset_source_rejects_unsupported_remote_urls():
         )
 
 
-def test_build_from_file_loads_local_multi_year_dataset(monkeypatch, tmp_path):
-    cache_key = "hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5"
-    dataset_path = tmp_path / "dataset.h5"
-    built = {}
-
-    class FakeMultiYearDataset:
-        @staticmethod
-        def validate_file_path(file_path, raise_exception=False):
-            built["validated_file_path"] = file_path
-            built["raise_exception"] = raise_exception
-            return True
-
-        def __init__(self, file_path):
-            self.file_path = file_path
-
-    def fake_build_from_multi_year_dataset(dataset):
-        built["dataset"] = dataset
+@pytest.mark.usefixtures("cloned_uk_tax_benefit_system")
+def test_build_from_file_loads_local_multi_year_dataset(tmp_path):
+    """Files written by either UK schema dispatch without validator stubs."""
+    from policyengine_uk.data import UKMultiYearDataset, UKSingleYearDataset
+    from policyengine_uk.data.economic_assumptions import extend_single_year_dataset
 
-    simulation_module._url_dataset_cache.pop(cache_key, None)
-    monkeypatch.setattr(simulation_module, "UKMultiYearDataset", FakeMultiYearDataset)
-    monkeypatch.setattr(
-        simulation_module, "_pre_encode_enum_columns", lambda dataset, tbs: None
+    single = _single_person_dataset()
+    multi = extend_single_year_dataset(
+        single, _current_law_system().parameters, end_year=2028
     )
+    multi_path = str(tmp_path / "multi.h5")
+    multi.save(multi_path)
+
+    assert UKMultiYearDataset.validate_file_path(multi_path, True)
+    assert not UKSingleYearDataset.validate_file_path(multi_path, False)
+    with pytest.raises(ValueError, match="Dataset 'time_period' not found in the file"):
+        UKSingleYearDataset.validate_file_path(multi_path)
+    simulation = Simulation(dataset=Path(multi_path))
+    assert simulation.dataset.years == multi.years
+    for year in multi.years:
+        np.testing.assert_allclose(
+            simulation.calculate("employment_income_before_lsr", year),
+            multi[year].person.employment_income,
+            rtol=1e-6,
+        )
 
-    simulation = Simulation.__new__(Simulation)
-    simulation.tax_benefit_system = _current_law_system()
-    monkeypatch.setattr(
-        simulation,
-        "build_from_multi_year_dataset",
-        fake_build_from_multi_year_dataset,
+    # Keep the existing single-year save/load route covered alongside the fix.
+    single_path = str(tmp_path / "single.h5")
+    single.save(single_path)
+    assert UKSingleYearDataset.validate_file_path(single_path)
+    assert not UKMultiYearDataset.validate_file_path(single_path)
+    with pytest.raises(ValueError, match="No yearly data found for 'person'"):
+        UKMultiYearDataset.validate_file_path(single_path, True)
+    single_simulation = Simulation(dataset=single_path)
+    np.testing.assert_array_equal(
+        single_simulation.calculate("employment_income_before_lsr", 2025),
+        single.person.employment_income,
     )
 
-    try:
-        Simulation.build_from_file(simulation, Path(dataset_path), cache_key=cache_key)
 
-        assert built["validated_file_path"] == str(dataset_path)
-        assert built["raise_exception"] is False
-        assert built["dataset"].file_path == str(dataset_path)
-        assert simulation.dataset is built["dataset"]
-        assert simulation_module._url_dataset_cache[cache_key] is built["dataset"]
-    finally:
-        simulation_module._url_dataset_cache.pop(cache_key, None)
+@pytest.mark.usefixtures("cloned_uk_tax_benefit_system")
+@settings(max_examples=8, deadline=None)
+@example(years={2025})
+@given(years=st.sets(st.integers(2020, 2035), min_size=1, max_size=4))
+def test_saved_datasets_have_exactly_one_schema_and_load(years):
+    """Both layouts, including a multi-year layout with just one year."""
+    from policyengine_uk.data import UKMultiYearDataset, UKSingleYearDataset
+
+    singles = []
+    for year in sorted(years):
+        single = _single_person_dataset()
+        single.time_period = str(year)
+        single.person["employment_income"] = float(year * 10)
+        singles.append(single)
+
+    with TemporaryDirectory() as directory:
+        for index, dataset in enumerate(
+            [*singles, UKMultiYearDataset(datasets=singles)]
+        ):
+            path = str(Path(directory) / f"dataset-{index}.h5")
+            dataset.save(path)
+            single_valid = UKSingleYearDataset.validate_file_path(path, False)
+            multi_valid = UKMultiYearDataset.validate_file_path(path, False)
+            assert single_valid != multi_valid
+            assert multi_valid == isinstance(dataset, UKMultiYearDataset)
+            simulation = Simulation(dataset=path)
+            expected = (
+                dataset.datasets if multi_valid else {int(dataset.time_period): dataset}
+            )
+            for year, original in expected.items():
+                np.testing.assert_array_equal(
+                    simulation.calculate("employment_income_before_lsr", year),
+                    original.person.employment_income,
+                )
 
 
 def _single_person_dataset():
@@ -204,7 +240,16 @@ def _single_person_dataset():
         }
     )
     benunit = pd.DataFrame({"benunit_id": [1]})
-    household = pd.DataFrame({"household_id": [1]})
+    household = pd.DataFrame(
+        {
+            "household_id": [1],
+            "household_weight": [1_000.0],
+            "region": ["LONDON"],
+            "tenure_type": ["RENT_PRIVATELY"],
+            "rent": [12_000.0],
+            "council_tax": [1_500.0],
+        }
+    )
     return UKSingleYearDataset(
         person=person, benunit=benunit, household=household, fiscal_year=2025
     )