dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/registry/filesystem_data_store.py

@@ -0,0 +1,150 @@
+import logging
+from pathlib import Path
+from typing import Self
+
+from dsgrid.registry.data_store_interface import DataStoreInterface
+from dsgrid.spark.functions import coalesce
+from dsgrid.spark.types import DataFrame
+from dsgrid.utils.files import delete_if_exists
+from dsgrid.utils.spark import read_dataframe, write_dataframe, write_dataframe_and_auto_partition
+
+
+TABLE_FILENAME = "table.parquet"
+LOOKUP_TABLE_FILENAME = "lookup_table.parquet"
+MISSING_ASSOCIATIONS_TABLE_FILENAME = "missing_associations_table.parquet"
+# We used to write these filenames. Keep support for old registries, for now.
+ALT_TABLE_FILENAME = "load_data.parquet"
+ALT_LOOKUP_TABLE_FILENAME = "load_data_lookup.parquet"
+
+logger = logging.getLogger(__name__)
+
+
+class FilesystemDataStore(DataStoreInterface):
+    """Data store that stores tables in Parquet files on the local or remote filesystem."""
+
+    @classmethod
+    def create(cls, base_path: Path) -> Self:
+        base_path.mkdir(exist_ok=True)
+        return cls(base_path)
+
+    @classmethod
+    def load(cls, base_path: Path) -> Self:
+        if not base_path.exists():
+            msg = f"Base path {base_path} does not exist. Cannot load FilesystemDataStore."
+            raise FileNotFoundError(msg)
+        return cls(base_path)
+
+    def read_table(self, dataset_id: str, version: str) -> DataFrame:
+        filename = self._table_filename(dataset_id, version)
+        if not filename.exists():
+            filename = self._alt_table_filename(dataset_id, version)
+        if not filename.exists():
+            msg = f"Table does not exist for dataset {dataset_id}, version {version} at {filename.parent}."
+            raise FileNotFoundError(msg)
+        return read_dataframe(filename)
+
+    def replace_table(self, df: DataFrame, dataset_id: str, version: str) -> None:
+        filename = self._get_existing_table_filename(dataset_id, version)
+        if filename is None:
+            self.write_table(df, dataset_id, version)
+            return
+        self._replace_table(df, filename)
+
+    def read_lookup_table(self, dataset_id: str, version: str) -> DataFrame:
+        filename = self._get_existing_lookup_table_filename(dataset_id, version)
+        if filename is None:
+            msg = f"Table does not exist for dataset {dataset_id}, version {version}."
+            raise FileNotFoundError(msg)
+        return read_dataframe(filename)
+
+    def replace_lookup_table(self, df: DataFrame, dataset_id: str, version: str) -> None:
+        filename = self._get_existing_lookup_table_filename(dataset_id, version)
+        if filename is None:
+            self.write_lookup_table(df, dataset_id, version)
+            return
+        self._replace_table(df, filename)
+
+    def read_missing_associations_tables(
+        self, dataset_id: str, version: str
+    ) -> dict[str, DataFrame]:
+        assoc_dir = self._missing_associations_dir(dataset_id, version)
+        if not assoc_dir.exists():
+            return {}
+        return {x.stem: read_dataframe(x) for x in assoc_dir.iterdir()}
+
+    def write_table(
+        self, df: DataFrame, dataset_id: str, version: str, overwrite: bool = False
+    ) -> None:
+        filename = self._table_filename(dataset_id, version)
+        filename.parent.mkdir(parents=True, exist_ok=True)
+        write_dataframe_and_auto_partition(df, filename)
+
+    def write_lookup_table(
+        self, df: DataFrame, dataset_id: str, version: str, overwrite: bool = False
+    ) -> None:
+        filename = self._lookup_table_filename(dataset_id, version)
+        filename.parent.mkdir(parents=True, exist_ok=True)
+        write_dataframe(coalesce(df, 1), filename, overwrite=overwrite)
+
+    def write_missing_associations_tables(
+        self, dfs: dict[str, DataFrame], dataset_id: str, version: str, overwrite: bool = False
+    ) -> None:
+        for name, df in dfs.items():
+            filename = self._missing_associations_table_filename(name, dataset_id, version)
+            filename.parent.mkdir(parents=True, exist_ok=True)
+            write_dataframe_and_auto_partition(df, filename)
+
+    def remove_tables(self, dataset_id: str, version: str) -> None:
+        delete_if_exists(self._base_dir(dataset_id, version))
+
+    @property
+    def _data_dir(self) -> Path:
+        return self.base_path / "data"
+
+    def _base_dir(self, dataset_id: str, version: str) -> Path:
+        return self._data_dir / dataset_id / version
+
+    def _lookup_table_filename(self, dataset_id: str, version: str) -> Path:
+        return self._data_dir / dataset_id / version / LOOKUP_TABLE_FILENAME
+
+    def _missing_associations_dir(self, dataset_id: str, version: str) -> Path:
+        return self._data_dir / dataset_id / version / "missing_associations"
+
+    def _missing_associations_table_filename(
+        self, name: str, dataset_id: str, version: str
+    ) -> Path:
+        return self._missing_associations_dir(dataset_id, version) / f"{name}.parquet"
+
+    def _table_filename(self, dataset_id: str, version: str) -> Path:
+        return self._data_dir / dataset_id / version / TABLE_FILENAME
+
+    def _alt_lookup_table_filename(self, dataset_id: str, version: str) -> Path:
+        return self._data_dir / dataset_id / version / ALT_LOOKUP_TABLE_FILENAME
+
+    def _alt_table_filename(self, dataset_id: str, version: str) -> Path:
+        return self._data_dir / dataset_id / version / ALT_TABLE_FILENAME
+
+    def _get_existing_lookup_table_filename(self, dataset_id: str, version: str) -> Path | None:
+        filename = self._lookup_table_filename(dataset_id, version)
+        if filename.exists():
+            return filename
+        alt_filename = self._alt_lookup_table_filename(dataset_id, version)
+        if alt_filename.exists():
+            return alt_filename
+        return None
+
+    def _get_existing_table_filename(self, dataset_id: str, version: str) -> Path | None:
+        filename = self._table_filename(dataset_id, version)
+        if filename.exists():
+            return filename
+        alt_filename = self._alt_table_filename(dataset_id, version)
+        if alt_filename.exists():
+            return alt_filename
+        return None
+
+    @staticmethod
+    def _replace_table(df: DataFrame, filename: Path) -> None:
+        tmp_name = filename.parent / f"{filename.stem}_tmp.parquet"
+        write_dataframe(df, tmp_name)
+        delete_if_exists(filename)
+        tmp_name.rename(filename)

dsgrid/registry/filter_registry_manager.py

@@ -0,0 +1,123 @@
+import logging
+
+from sqlalchemy import Connection
+
+from dsgrid.config.simple_models import RegistrySimpleModel
+from dsgrid.config.dataset_schema_handler_factory import make_dataset_schema_handler
+from dsgrid.spark.functions import is_dataframe_empty
+from dsgrid.utils.timing import track_timing, timer_stats_collector
+from .registry_manager import RegistryManager
+
+
+logger = logging.getLogger(__name__)
+
+
+class FilterRegistryManager(RegistryManager):
+    """Specialized RegistryManager that performs filtering operations."""
+
+    @track_timing(timer_stats_collector)
+    def filter(self, simple_model: RegistrySimpleModel, conn: Connection | None = None):
+        """Filter the registry as described by simple_model.
+
+        Parameters
+        ----------
+        simple_model : RegistrySimpleModel
+            Filter all configs and data according to this model.
+        """
+        if conn is None:
+            with self.project_manager.db.engine.begin() as conn:
+                self._filter(conn, simple_model)
+        else:
+            self._filter(conn, simple_model)
+
+    def _filter(self, conn: Connection, simple_model: RegistrySimpleModel):
+        project_ids_to_keep = {x.project_id for x in simple_model.projects}
+        to_remove = [
+            x for x in self._project_mgr.list_ids(conn=conn) if x not in project_ids_to_keep
+        ]
+        for project_id in to_remove:
+            self._project_mgr.remove(project_id, conn=conn)
+
+        dataset_ids_to_keep = {x.dataset_id for x in simple_model.datasets}
+        dataset_ids_to_remove = set(self._dataset_mgr.list_ids(conn=conn)) - dataset_ids_to_keep
+        for dataset_id in dataset_ids_to_remove:
+            self._dataset_mgr.remove(dataset_id, conn=conn)
+
+        modified_dims = set()
+        modified_dim_records = {}
+
+        def handle_dimension(simple_dim, dim):
+            records = dim.get_records_dataframe()
+            df = records.filter(records.id.isin(simple_dim.record_ids))
+            filtered_records = [x.asDict() for x in df.collect()]
+            modified_dims.add(dim.model.dimension_id)
+            modified_dim_records[dim.model.dimension_id] = {
+                x.id for x in df.select("id").distinct().collect()
+            }
+            return filtered_records
+
+        logger.info("Filter project dimensions")
+        for project in simple_model.projects:
+            changed_project = False
+            project_config = self._project_mgr.get_by_id(project.project_id, conn=conn)
+            indices_to_remove = []
+            for i, dataset in enumerate(project_config.model.datasets):
+                if dataset.dataset_id in dataset_ids_to_remove:
+                    indices_to_remove.append(i)
+            for index in reversed(indices_to_remove):
+                project_config.model.datasets.pop(index)
+                changed_project = True
+            for simple_dim in project.dimensions.base_dimensions:
+                for dim in project_config.list_base_dimensions(
+                    dimension_type=simple_dim.dimension_type
+                ):
+                    dim.model.records = handle_dimension(simple_dim, dim)
+                    self.dimension_manager.db.replace(conn, dim.model)
+
+            for simple_dim in project.dimensions.supplemental_dimensions:
+                for dim in project_config.list_supplemental_dimensions(simple_dim.dimension_type):
+                    if dim.model.name == simple_dim.dimension_name:
+                        dim.model.records = handle_dimension(simple_dim, dim)
+                        self.dimension_manager.db.replace(conn, dim.model)
+            if changed_project:
+                self.project_manager.db.replace(conn, project_config.model)
+
+        logger.info("Filter dataset dimensions")
+        for dataset in simple_model.datasets:
+            logger.info("Filter dataset %s", dataset.dataset_id)
+            dataset_config = self._dataset_mgr.get_by_id(dataset.dataset_id, conn=conn)
+            for simple_dim in dataset.dimensions:
+                dim = dataset_config.get_dimension(simple_dim.dimension_type)
+                dim.model.records = handle_dimension(simple_dim, dim)
+                self.dimension_manager.db.replace(conn, dim.model)
+            handler = make_dataset_schema_handler(
+                conn,
+                dataset_config,
+                self._dimension_mgr,
+                self._dimension_mapping_mgr,
+                store=self._data_store,
+            )
+            handler.filter_data(dataset.dimensions, self._data_store)
+
+        logger.info("Filter dimension mapping records")
+        for mapping in self._dimension_mapping_mgr.iter_configs():
+            records = None
+            changed = False
+            from_id = mapping.model.from_dimension.dimension_id
+            to_id = mapping.model.to_dimension.dimension_id
+            if from_id in modified_dims or to_id in modified_dims:
+                records = mapping.get_records_dataframe()
+                if from_id in modified_dims:
+                    records = records.filter(records.from_id.isin(modified_dim_records[from_id]))
+                    changed = True
+                if to_id in modified_dims:
+                    records = records.filter(records.to_id.isin(modified_dim_records[to_id]))
+                    changed = True
+
+            # TODO: probably need to remove a dimension mapping if it is empty
+            if records is not None and changed and not is_dataframe_empty(records):
+                mapping.model.records = [x.asDict() for x in records.collect()]
+                self.dimension_mapping_manager.db.replace(conn, mapping.model)
+                logger.info(
+                    "Filtered dimension mapping records from ID %s", mapping.model.mapping_id
+                )

dsgrid/registry/project_config_generator.py

@@ -0,0 +1,57 @@
+import logging
+from pathlib import Path
+from typing import Iterable
+
+from chronify.utils.path_utils import check_overwrite
+
+from dsgrid.dimension.time import TimeDimensionType
+from dsgrid.exceptions import DSGInvalidParameter
+from dsgrid.utils.files import dump_data
+from dsgrid.config.project_config import make_unvalidated_project_config
+
+
+logger = logging.getLogger(__name__)
+
+
+def generate_project_config(
+    project_id: str,
+    dataset_ids: Iterable[str],
+    metric_types: Iterable[str],
+    name: str | None = None,
+    description: str | None = None,
+    time_type: TimeDimensionType = TimeDimensionType.DATETIME,
+    output_directory: Path | None = None,
+    overwrite: bool = False,
+):
+    """Generate project config files and filesystem skeleton."""
+    if not metric_types:
+        msg = "At least one metric type must be passed"
+        raise DSGInvalidParameter(msg)
+    output_dir = (output_directory or Path()) / project_id
+    check_overwrite(output_dir, overwrite)
+    output_dir.mkdir()
+    project_dir = output_dir / "project"
+    project_dir.mkdir()
+    project_file = project_dir / "project.json5"
+    datasets_dir = output_dir / "datasets"
+    datasets_dir.mkdir()
+    (datasets_dir / "historical").mkdir()
+    (datasets_dir / "modeled").mkdir()
+    dimensions_dir = project_dir / "dimensions"
+    dimensions_dir.mkdir()
+    (dimensions_dir / "subset").mkdir()
+    (dimensions_dir / "supplemental").mkdir()
+    (project_dir / "dimension_mappings").mkdir()
+
+    config = make_unvalidated_project_config(
+        project_id,
+        dataset_ids,
+        metric_types,
+        name=name,
+        description=description,
+        time_type=time_type,
+    )
+    dump_data(config, project_file, indent=2)
+    logger.info(
+        "Created project directory structure at %s with config file %s", output_dir, project_file
+    )