dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/config/registration_models.py

@@ -0,0 +1,188 @@
+"""Contains data models to control bulk registration of projects and datasets."""
+
+from pathlib import Path
+from typing import Any, Iterable
+
+from pydantic import Field, ValidationInfo, field_validator, model_validator
+
+from dsgrid.data_models import DSGBaseModel
+from dsgrid.dimension.base_models import DimensionType
+from dsgrid.utils.files import load_data
+
+
+class ProjectRegistrationModel(DSGBaseModel):
+    """Defines a project to be registered."""
+
+    project_id: str = Field(description="Project ID")
+    config_file: Path = Field(description="Path to project.json5")
+    log_message: str | None = Field(
+        default=None,
+        description="Log message to use when registering the project. Defaults to an auto-generated message.",
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def fix_paths(cls, data: dict[str, Any]) -> dict[str, Any]:
+        _fix_paths(data, ("config_file",))
+        return data
+
+    @field_validator("log_message")
+    def fix_log_message(cls, log_message: str | None, info: ValidationInfo) -> str | None:
+        if log_message is None and "project_id" in info.data:
+            log_message = f"Register project {info.data['project_id']}"
+        return log_message
+
+
+class DatasetRegistrationModel(DSGBaseModel):
+    """Defines a dataset to be registered."""
+
+    dataset_id: str = Field(description="Dataset ID")
+    config_file: Path = Field(description="Path to dataset.json5")
+    replace_dimension_names_with_ids: bool = Field(
+        description="Replace the dimension entries with IDs of dimensions in the database "
+        "with matching names. Typically only useful for tests.",
+        default=False,
+    )
+    log_message: str | None = Field(
+        default=None,
+        description="Log message to use when registering the dataset. Defaults to an auto-generated message.",
+    )
+
+    @field_validator("log_message")
+    def fix_log_message(cls, log_message: str | None, info: ValidationInfo) -> str | None:
+        if log_message is None and "dataset_id" in info.data:
+            log_message = f"Register dataset {info.data['dataset_id']}"
+        return log_message
+
+    @model_validator(mode="before")
+    @classmethod
+    def fix_paths(cls, data: dict[str, Any]) -> dict[str, Any]:
+        _fix_paths(data, ("config_file",))
+        return data
+
+
+class DatasetSubmissionModel(DSGBaseModel):
+    """Defines how a dataset should be submitted to a project."""
+
+    dataset_id: str
+    project_id: str
+    dimension_mapping_file: Path | None = Field(
+        description="Path to file containing mappings of dataset-to-project dimensions",
+        default=None,
+    )
+    dimension_mapping_references_file: Path | None = Field(
+        description="Path to file containing references to mappings of dataset-to-project dimensions",
+        default=None,
+    )
+    replace_dimension_mapping_names_with_ids: bool = Field(
+        description="Replace the dimension mapping entries with IDs of dimension mappings "
+        "in the database with matching names. Typically only useful for tests.",
+        default=False,
+    )
+    autogen_reverse_supplemental_mappings: set[DimensionType] = Field(
+        description="Dimensions on which to attempt create reverse mappings from supplemental dimensions.",
+        default=set(),
+    )
+    log_message: str | None = Field(
+        default=None,
+        description="Log message to use when submitting the dataset. Defaults to an auto-generated message.",
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def fix_autogen_reverse_supplemental_mappings(cls, data: dict[str, Any]) -> dict[str, Any]:
+        if "autogen_reverse_supplemental_mappings" in data:
+            data["autogen_reverse_supplemental_mappings"] = {
+                DimensionType(x) for x in data["autogen_reverse_supplemental_mappings"]
+            }
+        return data
+
+    @field_validator("log_message")
+    def fix_log_message(cls, log_message: str | None, info: ValidationInfo) -> str | None:
+        if log_message is None and "dataset_id" in info.data:
+            log_message = (
+                f"Submit dataset {info.data['dataset_id']} to project {info.data['project_id']}"
+            )
+        return log_message
+
+
+class SubmittedDatasetsJournal(DSGBaseModel):
+    """Defines a dataset that was successfully submitted to a project."""
+
+    dataset_id: str
+    project_id: str
+
+
+class RegistrationJournal(DSGBaseModel):
+    """Defines projects and datasets that were succesfully registered."""
+
+    registered_projects: list[str] = []
+    registered_datasets: list[str] = []
+    submitted_datasets: list[SubmittedDatasetsJournal] = []
+
+    def add_dataset(self, dataset_id: str) -> None:
+        assert dataset_id not in self.registered_datasets, dataset_id
+        self.registered_datasets.append(dataset_id)
+
+    def add_project(self, project_id: str) -> None:
+        assert project_id not in self.registered_projects, project_id
+        self.registered_projects.append(project_id)
+
+    def add_submitted_dataset(self, dataset_id: str, project_id: str) -> None:
+        entry = SubmittedDatasetsJournal(dataset_id=dataset_id, project_id=project_id)
+        assert entry not in self.submitted_datasets, entry
+        self.submitted_datasets.append(entry)
+
+    def has_entries(self) -> bool:
+        return (
+            bool(self.registered_projects)
+            or bool(self.registered_datasets)
+            or bool(self.submitted_datasets)
+        )
+
+
+class RegistrationModel(DSGBaseModel):
+    """Defines a list of projects and datasets to be registered."""
+
+    projects: list[ProjectRegistrationModel] = Field(description="List of projects to register.")
+    datasets: list[DatasetRegistrationModel] = Field(description="List of datasets to register.")
+    dataset_submissions: list[DatasetSubmissionModel] = Field(
+        description="List of datasets to be submitted to projects."
+    )
+
+    def filter_by_journal(self, journal: RegistrationJournal) -> "RegistrationModel":
+        """Return a new instance of RegistrationModel by filtering an existing instance with
+        a journal.
+        """
+        projects = list(
+            filter(lambda x: x.project_id not in journal.registered_projects, self.projects)
+        )
+        datasets = list(
+            filter(lambda x: x.dataset_id not in journal.registered_datasets, self.datasets)
+        )
+        dataset_submissions = list(
+            filter(
+                lambda x: SubmittedDatasetsJournal(
+                    dataset_id=x.dataset_id, project_id=x.project_id
+                )
+                not in journal.submitted_datasets,
+                self.dataset_submissions,
+            )
+        )
+        return RegistrationModel(
+            projects=projects,
+            datasets=datasets,
+            dataset_submissions=dataset_submissions,
+        )
+
+
+def _fix_paths(data: dict[str, Any], fields: Iterable[str]) -> None:
+    for field in fields:
+        val = data.get(field)
+        if isinstance(val, str):
+            data[field] = Path(val)
+
+
+def create_registration(input_file: Path):
+    """Create registration inputs."""
+    return RegistrationModel(**load_data(input_file))

dsgrid/config/representative_period_time_dimension_config.py

@@ -0,0 +1,194 @@
+import abc
+import logging
+from datetime import timedelta
+from typing import Type, Any, Union
+
+import chronify
+
+from dsgrid.dimension.time import RepresentativePeriodFormat, TimeIntervalType
+from dsgrid.time.types import (
+    OneWeekPerMonthByHourType,
+    OneWeekdayDayAndOneWeekendDayPerMonthByHourType,
+)
+from .dimensions import RepresentativePeriodTimeDimensionModel
+from .time_dimension_base_config import TimeDimensionBaseConfig
+
+
+logger = logging.getLogger(__name__)
+
+
+class RepresentativePeriodTimeDimensionConfig(TimeDimensionBaseConfig):
+    """Provides an interface to an RepresentativePeriodTimeDimensionModel."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # We expect the list of required formats to grow.
+        # It's possible that one function (or set of functions) can handle all permutations
+        # of parameters. We can make that determination once we have requirements for more
+        # formats.
+        match self.model.format:
+            case RepresentativePeriodFormat.ONE_WEEK_PER_MONTH_BY_HOUR:
+                self._format_handler = OneWeekPerMonthByHourHandler()
+            case RepresentativePeriodFormat.ONE_WEEKDAY_DAY_AND_ONE_WEEKEND_DAY_PER_MONTH_BY_HOUR:
+                self._format_handler = OneWeekdayDayAndWeekendDayPerMonthByHourHandler()
+            case _:
+                msg = self.model.format.value
+                raise NotImplementedError(msg)
+
+    def supports_chronify(self) -> bool:
+        return True
+
+    def to_chronify(
+        self,
+    ) -> Union[chronify.RepresentativePeriodTimeTZ, chronify.RepresentativePeriodTimeNTZ]:
+        if len(self._model.ranges) != 1:
+            msg = (
+                "Mapping RepresentativePeriodTime with chronify is only supported with one range: "
+                f"{self._model.ranges}"
+            )
+            raise NotImplementedError(msg)
+        range_ = self._model.ranges[0]
+        if range_.start != 1 or range_.end != 12:
+            msg = (
+                "Mapping RepresentativePeriodTime with chronify is only supported with a full year: "
+                f"{range_}"
+            )
+            raise NotImplementedError(msg)
+        # RepresentativePeriodTimeDimensionModel does not map to NTZ at the moment
+        if isinstance(self._format_handler, OneWeekPerMonthByHourHandler) or isinstance(
+            self._format_handler, OneWeekdayDayAndWeekendDayPerMonthByHourHandler
+        ):
+            return chronify.RepresentativePeriodTimeTZ(
+                measurement_type=self._model.measurement_type,
+                interval_type=self._model.time_interval_type,
+                time_format=chronify.RepresentativePeriodFormat(self._model.format.value),
+                time_zone_column="time_zone",
+            )
+
+        msg = f"Cannot chronify time_config for {self._format_handler}"
+        raise NotImplementedError(msg)
+
+    @staticmethod
+    def model_class() -> RepresentativePeriodTimeDimensionModel:
+        return RepresentativePeriodTimeDimensionModel
+
+    def get_frequency(self) -> timedelta:
+        return self._format_handler.get_frequency()
+
+    def get_start_times(self) -> list[Any]:
+        return self._format_handler.get_start_times(self.model.ranges)
+
+    def get_lengths(self) -> list[int]:
+        return self._format_handler.get_lengths(self.model.ranges)
+
+    def get_load_data_time_columns(self) -> list[str]:
+        return self._format_handler.get_load_data_time_columns()
+
+    def get_time_zone(self) -> None:
+        return None
+
+    def get_tzinfo(self) -> None:
+        return None
+
+    def get_time_interval_type(self) -> TimeIntervalType:
+        return self.model.time_interval_type
+
+
+class RepresentativeTimeFormatHandlerBase(abc.ABC):
+    """Provides implementations for different representative time formats."""
+
+    @staticmethod
+    @abc.abstractmethod
+    def get_representative_time_type() -> Type:
+        """Return the time type representing the data."""
+
+    @abc.abstractmethod
+    def get_frequency(self):
+        """Return the frequency.
+
+        Returns
+        -------
+        timedelta
+
+        """
+
+    @staticmethod
+    @abc.abstractmethod
+    def get_load_data_time_columns():
+        """Return the required timestamp columns in the load data table.
+
+        Returns
+        -------
+        list
+
+        """
+
+
+class OneWeekPerMonthByHourHandler(RepresentativeTimeFormatHandlerBase):
+    """Handler for format with hourly data that includes one week per month."""
+
+    @staticmethod
+    def get_representative_time_type() -> OneWeekPerMonthByHourType:
+        return OneWeekPerMonthByHourType
+
+    def get_frequency(self) -> timedelta:
+        return timedelta(hours=1)
+
+    @staticmethod
+    def get_start_times(ranges) -> list[OneWeekPerMonthByHourType]:
+        """Get the starting combination of (month, day_of_week, hour) based on sorted order"""
+        start_times = []
+        for model in ranges:
+            start_times.append(OneWeekPerMonthByHourType(month=model.start, day_of_week=0, hour=0))
+        return start_times
+
+    @staticmethod
+    def get_lengths(ranges) -> list[int]:
+        """Get the number of unique combinations of (month, day_of_week, hour)"""
+        lengths = []
+        for model in ranges:
+            n_months = model.end - model.start + 1
+            lengths.append(n_months * 7 * 24)
+        return lengths
+
+    @staticmethod
+    def get_load_data_time_columns() -> list[str]:
+        return list(OneWeekPerMonthByHourType._fields)
+
+
+class OneWeekdayDayAndWeekendDayPerMonthByHourHandler(RepresentativeTimeFormatHandlerBase):
+    """Handler for format with hourly data that includes one weekday day and one weekend day
+    per month.
+    """
+
+    @staticmethod
+    def get_representative_time_type() -> OneWeekdayDayAndOneWeekendDayPerMonthByHourType:
+        return OneWeekdayDayAndOneWeekendDayPerMonthByHourType
+
+    def get_frequency(self) -> timedelta:
+        return timedelta(hours=1)
+
+    @staticmethod
+    def get_start_times(ranges) -> list[OneWeekdayDayAndOneWeekendDayPerMonthByHourType]:
+        """Get the starting combination of (month, hour, is_weekday) based on sorted order"""
+        start_times = []
+        for model in ranges:
+            start_times.append(
+                OneWeekdayDayAndOneWeekendDayPerMonthByHourType(
+                    month=model.start, hour=0, is_weekday=False
+                )
+            )
+        return start_times
+
+    @staticmethod
+    def get_lengths(ranges) -> list[int]:
+        """Get the number of unique combinations of (month, hour, is_weekday)"""
+        lengths = []
+        for model in ranges:
+            n_months = model.end - model.start + 1
+            lengths.append(n_months * 24 * 2)
+        return lengths
+
+    @staticmethod
+    def get_load_data_time_columns() -> list[str]:
+        return list(OneWeekdayDayAndOneWeekendDayPerMonthByHourType._fields)

dsgrid/config/simple_models.py

@@ -0,0 +1,49 @@
+"""Defines simplified data models for testing and filtering."""
+
+from pydantic import field_validator, model_validator, Field
+
+from dsgrid.data_models import DSGBaseModel
+from dsgrid.dimension.base_models import DimensionType
+
+
+class DimensionSimpleModel(DSGBaseModel):
+    dimension_type: DimensionType
+    dimension_name: str | None = None
+    record_ids: list[str]
+
+
+class DimensionsSimpleModel(DSGBaseModel):
+    base_dimensions: list[DimensionSimpleModel]
+    supplemental_dimensions: list[DimensionSimpleModel] = Field(default=[])
+
+    @field_validator("base_dimensions")
+    @classmethod
+    def check_base_dimensions(cls, base_dimensions):
+        dimension_types = {x.dimension_type for x in base_dimensions}
+        if len(dimension_types) != len(base_dimensions):
+            msg = "base_dimensions cannot contain duplicate dimension types"
+            raise ValueError(msg)
+        return base_dimensions
+
+    @model_validator(mode="after")
+    def check_supplemental_dimensions(self) -> "DimensionsSimpleModel":
+        for dim in self.supplemental_dimensions:
+            if dim.dimension_name is None:
+                msg = f"supplemental dimensions must define dimension_name: {dim}"
+                raise ValueError(msg)
+        return self
+
+
+class DatasetSimpleModel(DSGBaseModel):
+    dataset_id: str
+    dimensions: list[DimensionSimpleModel]
+
+
+class ProjectSimpleModel(DSGBaseModel):
+    project_id: str
+    dimensions: DimensionsSimpleModel
+
+
+class RegistrySimpleModel(DSGBaseModel):
+    projects: list[ProjectSimpleModel]
+    datasets: list[DatasetSimpleModel]

dsgrid/config/supplemental_dimension.py

@@ -0,0 +1,29 @@
+"""Defines a supplemental dimension."""
+
+from typing import Annotated
+from pydantic import Field
+
+
+from dsgrid.data_models import DSGBaseModel
+from .dimensions import DimensionModel
+from .mapping_tables import MappingTableByNameModel
+
+
+class SupplementalDimensionModel(DimensionModel):
+    """Defines a supplemental dimension."""
+
+    mapping: MappingTableByNameModel = Field(
+        description="Defines how the supplemental dimension will be mapped to the project's base "
+        "dimension.",
+        title="mapping",
+    )
+
+
+class SupplementalDimensionsListModel(DSGBaseModel):
+    """Defines a list of supplemental dimensions."""
+
+    supplemental_dimensions: Annotated[
+        list[SupplementalDimensionModel], Field(min_length=1)
+    ] = Field(
+        description="List of supplemental dimensions and mappings to be registered",
+    )

dsgrid/config/time_dimension_base_config.py

@@ -0,0 +1,192 @@
+import abc
+import logging
+from datetime import tzinfo
+from typing import Any
+import pandas as pd
+
+import chronify
+
+from .dimension_config import DimensionBaseConfigWithoutFiles
+from dsgrid.dimension.time import (
+    TimeIntervalType,
+    TimeBasedDataAdjustmentModel,
+)
+from dsgrid.dimension.time_utils import (
+    build_time_ranges,
+)
+from dsgrid.config.dimensions import TimeRangeModel
+
+from dsgrid.spark.types import (
+    DataFrame,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class TimeDimensionBaseConfig(DimensionBaseConfigWithoutFiles, abc.ABC):
+    """Base class for all time dimension configs"""
+
+    def supports_chronify(self) -> bool:
+        """Return True if the config can be converted to chronify."""
+        return False
+
+    # @abc.abstractmethod
+    def to_chronify(self) -> chronify.TimeBaseModel:
+        """Return the chronify version of the time model."""
+        # This is likely temporary until we can use chronify models directly.
+        msg = f"{type(self)}.to_chronify"
+        raise NotImplementedError(msg)
+
+    def check_dataset_time_consistency(self, load_data_df, time_columns: list[str]) -> None:
+        """Check consistency of the load data with the time dimension.
+
+        Parameters
+        ----------
+        load_data_df : pyspark.sql.DataFrame
+        time_columns : list[str]
+
+        Raises
+        ------
+        DSGInvalidDataset
+            Raised if the dataset is inconsistent with the time dimension.
+        """
+        msg = f"{type(self)}.check_dataset_time_consistency is not implemented"
+        raise NotImplementedError(msg)
+
+    def build_time_dataframe(self) -> DataFrame:
+        """Build time dimension as specified in config in a spark dataframe.
+
+        Returns
+        -------
+        pyspark.sql.DataFrame
+        """
+        msg = f"{self.__class__.__name__}.build_time_dataframe is not implemented"
+        raise NotImplementedError(msg)
+
+    @abc.abstractmethod
+    def get_load_data_time_columns(self) -> list[str]:
+        """Return the required timestamp columns in the load data table.
+
+        Returns
+        -------
+        list
+        """
+
+    def list_load_data_columns_for_query_name(self) -> list[str]:
+        """Return the time columns expected in the load data table for this dimension's query name.
+
+        Returns
+        -------
+        list[str]
+        """
+        # This may need to be re-implemented by child classes.
+        return [self.model.name]
+
+    def map_timestamp_load_data_columns_for_query_name(self, df) -> DataFrame:
+        """Map the timestamp columns in the load data table to those specified by the query name.
+
+        Parameters
+        ----------
+        df : pyspark.sql.DataFrame
+
+        Returns
+        -------
+        pyspark.sql.DataFrame
+        """
+        time_cols = self.get_load_data_time_columns()
+        if len(time_cols) > 1:
+            msg = (
+                "Handling of multiple time columns needs to be implemented in the child class: "
+                f"{type(self)}: {time_cols=}"
+            )
+            raise NotImplementedError(msg)
+
+        time_col = time_cols[0]
+        if time_col not in df.columns:
+            return df
+        return df.withColumnRenamed(time_col, self.model.name)
+
+    def get_time_ranges(self) -> list[Any]:
+        """Return time ranges with time_zone applied.
+
+        Returns
+        -------
+        list
+            list of DatetimeRange
+        """
+        msg = f"{type(self)}.get_time_ranges is not implemented"
+        raise NotImplementedError(msg)
+
+    @abc.abstractmethod
+    def get_start_times(self) -> list[Any]:
+        """Return the list of starting timestamp (with tzinfo) for this dimension.
+        One per time range.
+
+        Returns
+        -------
+        list[Any]
+        """
+
+    @abc.abstractmethod
+    def get_lengths(self) -> list[int]:
+        """Return the list of time range length (number of time steps) for this dimension.
+        One per time range.
+
+        Returns
+        -------
+        list[Any]
+        """
+
+    @abc.abstractmethod
+    def get_time_zone(self) -> str | None:
+        """Return a time zone instance for this dimension."""
+
+    def get_time_zones(self) -> list[str]:
+        """Return a list of time zones for this dimension."""
+        if self.get_time_zone():
+            return [self.get_time_zone()]
+        return []
+
+    @abc.abstractmethod
+    def get_tzinfo(self) -> tzinfo | None:
+        """Return a tzinfo instance for this dimension.
+
+        Returns
+        -------
+        tzinfo | None
+        """
+
+    @abc.abstractmethod
+    def get_time_interval_type(self) -> TimeIntervalType:
+        """Return the time interval type for this dimension.
+
+        Returns
+        -------
+        TimeIntervalType
+        """
+
+    def list_expected_dataset_timestamps(
+        self,
+        time_based_data_adjustment: TimeBasedDataAdjustmentModel | None = None,
+    ) -> list[tuple]:
+        """Return a list of the timestamps expected in the load_data table.
+        Parameters
+        ----------
+        time_based_data_adjustmen : TimeBasedDataAdjustmentModel | None
+
+        Returns
+        -------
+        list
+            List of tuples of columns representing time in the load_data table.
+
+        """
+        msg = f"{type(self)}.list_expected_dataset_timestamps is not implemented"
+        raise NotImplementedError(msg)
+
+    def _build_time_ranges(
+        self,
+        time_ranges: list[TimeRangeModel],
+        tz: str | None = None,
+    ) -> list[tuple[pd.Timestamp, pd.Timestamp, pd.Timedelta]]:
+        return build_time_ranges(time_ranges, tz=tz)