Merge pull request #1747 from PolicyEngine/codex/support-local-h5-datasets

Support local dataset file paths

Author: anth-volk

changelog.d/local-dataset-filepaths.md

@@ -0,0 +1 @@
+- Allow UK simulations to load local dataset file paths, including materialized H5 files passed by upstream dataset resolvers.

policyengine_uk/simulation.py

@@ -105,7 +105,13 @@ def __init__(
         scenario: Optional[Scenario] = None,
         situation: Optional[Dict] = None,
         dataset: Optional[
-            Union[pd.DataFrame, str, UKSingleYearDataset, UKMultiYearDataset]
+            Union[
+                pd.DataFrame,
+                str,
+                os.PathLike,
+                UKSingleYearDataset,
+                UKMultiYearDataset,
+            ]
         ] = None,
         trace: bool = False,
         reform: Union[Dict, Type[Reform]] = None,
@@ -148,8 +154,8 @@ def __init__(
         # Build simulation from appropriate source
         if situation is not None:
             self.build_from_situation(situation)
-        elif isinstance(dataset, str):
-            self.build_from_url(dataset)
+        elif isinstance(dataset, (str, os.PathLike)):
+            self.build_from_dataset_source(dataset)
         elif isinstance(dataset, pd.DataFrame):
             self.build_from_dataframe(dataset)
         elif isinstance(dataset, Dataset):
@@ -159,7 +165,7 @@ def __init__(
         elif isinstance(dataset, UKMultiYearDataset):
             self.build_from_multi_year_dataset(dataset)
         elif dataset is None:
-            self.build_from_url(get_default_dataset_url())
+            self.build_from_dataset_source(get_default_dataset_url())
         else:
             raise ValueError(f"Unsupported dataset type: {dataset.__class__}")
 
@@ -260,6 +266,52 @@ def build_from_situation(self, situation: Dict) -> None:
             self, fiscal_year=self.default_input_period
         )
 
+    def build_from_dataset_source(
+        self, dataset_source: Union[str, os.PathLike]
+    ) -> None:
+        """Build simulation from a supported dataset source string."""
+        dataset_source = os.fspath(dataset_source)
+        if dataset_source.startswith("hf://"):
+            self.build_from_url(dataset_source)
+            return
+        if "://" in dataset_source:
+            raise ValueError(
+                "Only HuggingFace dataset URLs are supported directly by "
+                "policyengine-uk. Download or materialize other dataset "
+                "sources to a local file path before passing them to Simulation."
+            )
+        self.build_from_file(dataset_source)
+
+    def build_from_file(
+        self,
+        dataset_file: Union[str, os.PathLike],
+        *,
+        cache_key: Optional[str] = None,
+    ) -> None:
+        """Build simulation from a local dataset file path."""
+        dataset_file = os.fspath(dataset_file)
+
+        if UKMultiYearDataset.validate_file_path(dataset_file, False):
+            multi_year_dataset = UKMultiYearDataset(dataset_file)
+        elif UKSingleYearDataset.validate_file_path(dataset_file, False):
+            multi_year_dataset = extend_single_year_dataset(
+                UKSingleYearDataset(dataset_file),
+                self.tax_benefit_system.parameters,
+            )
+        else:
+            dataset = Dataset.from_file(dataset_file, self.default_input_period)
+            self.build_from_dataset(dataset)
+            return
+
+        # Pre-encode string enum columns to int16 once before caching so
+        # subsequent loads skip the expensive astype(str) + searchsorted path.
+        _pre_encode_enum_columns(multi_year_dataset, self.tax_benefit_system)
+        if cache_key is not None:
+            _url_dataset_cache[cache_key] = multi_year_dataset
+
+        self.build_from_multi_year_dataset(multi_year_dataset)
+        self.dataset = multi_year_dataset
+
     def build_from_url(self, url: str) -> None:
         """Build simulation from a HuggingFace dataset URL.
 
@@ -269,7 +321,7 @@ def build_from_url(self, url: str) -> None:
         Raises:
             ValueError: If URL is not a HuggingFace URL
         """
-        if "hf://" not in url:
+        if not url.startswith("hf://"):
             raise ValueError(f"Non-HuggingFace URLs are currently not supported.")
 
         # Return early from in-memory cache if available: skips HDF5 reading,
@@ -295,26 +347,7 @@ def build_from_url(self, url: str) -> None:
             version=version,
         )
 
-        # Determine dataset type and build accordingly
-        if UKMultiYearDataset.validate_file_path(dataset_file, False):
-            multi_year_dataset = UKMultiYearDataset(dataset_file)
-        elif UKSingleYearDataset.validate_file_path(dataset_file, False):
-            multi_year_dataset = extend_single_year_dataset(
-                UKSingleYearDataset(dataset_file),
-                self.tax_benefit_system.parameters,
-            )
-        else:
-            dataset = Dataset.from_file(dataset_file, self.default_input_period)
-            self.build_from_dataset(dataset)
-            return
-
-        # Pre-encode string enum columns to int16 once before caching so
-        # subsequent loads skip the expensive astype(str) + searchsorted path.
-        _pre_encode_enum_columns(multi_year_dataset, self.tax_benefit_system)
-        _url_dataset_cache[url] = multi_year_dataset
-
-        self.build_from_multi_year_dataset(multi_year_dataset)
-        self.dataset = multi_year_dataset
+        self.build_from_file(dataset_file, cache_key=url)
 
     def build_from_dataframe(self, df: pd.DataFrame) -> None:
         """Build simulation from a pandas DataFrame.

policyengine_uk/tests/test_simulation_dataset_sources.py

@@ -0,0 +1,110 @@
+from pathlib import Path
+
+import pytest
+
+import policyengine_uk.simulation as simulation_module
+from policyengine_uk import Simulation
+
+
+class _StopDatasetLoad(Exception):
+    pass
+
+
+@pytest.mark.parametrize(
+    "dataset_source_factory",
+    [str, lambda dataset_path: dataset_path],
+    ids=["str", "pathlike"],
+)
+def test_simulation_routes_local_dataset_path_to_file_loader(
+    monkeypatch, tmp_path, dataset_source_factory
+):
+    captured = {}
+    dataset_path = tmp_path / "dataset.h5"
+
+    def fake_build_from_file(self, dataset_file, *, cache_key=None):
+        captured["dataset_file"] = dataset_file
+        captured["cache_key"] = cache_key
+        raise _StopDatasetLoad
+
+    monkeypatch.setattr(Simulation, "build_from_file", fake_build_from_file)
+
+    with pytest.raises(_StopDatasetLoad):
+        Simulation(dataset=dataset_source_factory(dataset_path))
+
+    assert captured == {
+        "dataset_file": str(dataset_path),
+        "cache_key": None,
+    }
+
+
+def test_dataset_source_routes_huggingface_urls_to_url_loader(monkeypatch):
+    captured = {}
+    simulation = Simulation.__new__(Simulation)
+    url = "hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5"
+
+    def fake_build_from_url(dataset_url):
+        captured["url"] = dataset_url
+
+    def fake_build_from_file(dataset_file, *, cache_key=None):
+        raise AssertionError("HuggingFace URLs should not be treated as local files.")
+
+    monkeypatch.setattr(simulation, "build_from_url", fake_build_from_url)
+    monkeypatch.setattr(simulation, "build_from_file", fake_build_from_file)
+
+    Simulation.build_from_dataset_source(simulation, url)
+
+    assert captured == {"url": url}
+
+
+def test_dataset_source_rejects_unsupported_remote_urls():
+    simulation = Simulation.__new__(Simulation)
+
+    with pytest.raises(ValueError, match="Only HuggingFace dataset URLs"):
+        Simulation.build_from_dataset_source(
+            simulation,
+            "gs://policyengine-uk-data-private/enhanced_frs_2023_24.h5",
+        )
+
+
+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
+
+    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
+    )
+
+    simulation = Simulation.__new__(Simulation)
+    simulation.tax_benefit_system = object()
+    monkeypatch.setattr(
+        simulation,
+        "build_from_multi_year_dataset",
+        fake_build_from_multi_year_dataset,
+    )
+
+    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)