dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/config/dimensions_config.py

@@ -0,0 +1,71 @@
+import logging
+from pathlib import Path
+
+from pydantic import field_validator, Field
+
+from dsgrid.data_models import DSGBaseModel
+from dsgrid.utils.utilities import check_uniqueness
+from .config_base import ConfigBase
+from .dimensions import DimensionModel, DimensionsListModel
+
+logger = logging.getLogger(__name__)
+
+
+class DimensionsConfigModel(DSGBaseModel):
+    """Represents multiple dimension models.
+
+    Used when registering multiple dimensions in one command.
+    """
+
+    dimensions: DimensionsListModel = Field(
+        title="dimensions",
+        description="Dimensions for submission to the dimension registry",
+    )
+
+    @field_validator("dimensions")
+    @classmethod
+    def check_files(cls, values: dict) -> dict:
+        """Validate dimension files are unique across all dimensions"""
+        check_uniqueness(
+            (x.filename for x in values if isinstance(x, DimensionModel) and x.filename),
+            "dimension record filename",
+        )
+        return values
+
+    @field_validator("dimensions")
+    @classmethod
+    def check_names(cls, values: dict) -> dict:
+        """Validate dimension names are unique across all dimensions."""
+        check_uniqueness(
+            [dim.name for dim in values],
+            "dimension record name",
+        )
+        return values
+
+
+class DimensionsConfig(ConfigBase):
+    """Provides an interface to a DimensionsConfigModel."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._src_dir = None
+
+    @staticmethod
+    def config_filename():
+        return "dimensions.json5"
+
+    @property
+    def config_id(self):
+        assert False, "not correct for this class"
+
+    @staticmethod
+    def model_class():
+        return DimensionsConfigModel
+
+    @classmethod
+    def load(cls, config_filename: Path, *args, **kwargs):
+        return super().load(config_filename, *args, **kwargs)
+
+    @classmethod
+    def load_from_model(cls, model):
+        return cls(model)

dsgrid/config/file_schema.py

@@ -0,0 +1,190 @@
+import logging
+from pathlib import Path
+from typing import Self
+
+from pydantic import Field, field_validator, model_validator
+
+from dsgrid.data_models import DSGBaseModel
+from dsgrid.dimension.base_models import DimensionType
+from dsgrid.exceptions import DSGInvalidDataset, DSGInvalidField
+from dsgrid.spark.functions import read_csv_duckdb, read_json, read_parquet
+from dsgrid.spark.types import DataFrame, DUCKDB_COLUMN_TYPES, SUPPORTED_TYPES
+from dsgrid.utils.scratch_dir_context import ScratchDirContext
+from dsgrid.utils.spark import write_dataframe
+from dsgrid.utils.utilities import check_uniqueness
+
+
+logger = logging.getLogger(__name__)
+
+
+class Column(DSGBaseModel):
+    name: str = Field(description="Name of the column")
+    dimension_type: DimensionType | None = Field(
+        default=None,
+        description="Dimension represented by the data in the column. Optional if this is a "
+        "time column or pivoted column. Required if the column represents a stacked dimension "
+        "but an alternate name is being used, such as 'county' instead of 'geography'. "
+        "dsgrid will rename any column that is set at runtime, writing out the result to the "
+        "registry's data directory. The original dataset is not modified.",
+    )
+    data_type: str | None = Field(
+        default=None, description="Type of the data in the column. If None, infer the type."
+    )
+
+    @field_validator("data_type")
+    @classmethod
+    def check_data_type(cls, data_type: str | None) -> str | None:
+        if data_type is None:
+            return None
+
+        type_upper = data_type.upper()
+        if type_upper not in SUPPORTED_TYPES:
+            supported_data_types = sorted(SUPPORTED_TYPES)
+            msg = f"{data_type=} is not one of {supported_data_types=}"
+            raise ValueError(msg)
+        return type_upper
+
+
+class FileSchema(DSGBaseModel):
+    """Defines the format of a data file (CSV, JSON, Parquet)."""
+
+    path: str | None = Field(description="Path to the file. Must be assigned during registration.")
+    columns: list[Column] = Field(
+        default=[], description="Custom schema for the columns in the file."
+    )
+    ignore_columns: list[str] = Field(
+        default=[],
+        description="List of column names to ignore (drop) when reading the file.",
+    )
+
+    @model_validator(mode="after")
+    def check_consistency(self) -> Self:
+        if len(self.columns) > 1:
+            check_uniqueness((x.name for x in self.columns), "column names")
+
+        # Check that ignore_columns don't overlap with columns
+        column_names = {x.name for x in self.columns}
+        ignore_set = set(self.ignore_columns)
+        overlap = column_names & ignore_set
+        if overlap:
+            msg = f"Columns cannot be in both 'columns' and 'ignore_columns': {overlap}"
+            raise ValueError(msg)
+
+        return self
+
+    def get_data_type_mapping(self) -> dict[str, str]:
+        """Return the mapping of column to data type."""
+        return {x.name: x.data_type for x in self.columns if x.data_type is not None}
+
+
+def read_data_file(
+    schema: FileSchema, scratch_dir_context: ScratchDirContext | None = None
+) -> DataFrame:
+    """Read a data file from a schema.
+
+    Parameters
+    ----------
+    schema : FileSchema
+        Schema defining the file path and column types.
+    scratch_dir_context : ScratchDirContext
+        Optional location to write temporary files.
+
+    Returns
+    -------
+    DataFrame
+        A Spark DataFrame containing the file data.
+    """
+    if schema.path is None:
+        msg = "File path is not assigned"
+        raise DSGInvalidDataset(msg)
+
+    path = Path(schema.path)
+    if not path.exists():
+        msg = f"{path} does not exist"
+        raise FileNotFoundError(msg)
+
+    expected_columns = {x.name for x in schema.columns}
+
+    match path.suffix:
+        case ".parquet":
+            df = read_parquet(path)
+        case ".csv":
+            column_schema = _get_column_schema(schema, DUCKDB_COLUMN_TYPES)
+            df = read_csv_duckdb(path, schema=column_schema)
+        case ".json":
+            df = read_json(path)
+        case _:
+            msg = f"Unsupported file type: {path.suffix}"
+            raise DSGInvalidDataset(msg)
+
+    actual_columns = set(df.columns)
+    diff = expected_columns.difference(actual_columns)
+    if diff:
+        msg = f"Expected columns {diff} are not in {actual_columns=}"
+        raise DSGInvalidDataset(msg)
+
+    df = _drop_ignored_columns(df, schema.ignore_columns)
+    renames = _get_column_renames(schema)
+    if renames:
+        df = _rename_columns(df, renames)
+        if scratch_dir_context is None:
+            renamed_path = path.with_stem(path.stem + "_renamed")
+            logger.warning(
+                "Creating temporary file at %s. Pass scratch_dir_context to avoid this.",
+                renamed_path,
+            )
+        else:
+            renamed_path = scratch_dir_context.get_temp_filename(suffix=path.suffix)
+        write_dataframe(df, renamed_path, overwrite=True)
+        schema.path = str(renamed_path)
+        for column in schema.columns:
+            if column.name in renames:
+                column.name = renames[column.name]
+                column.dimension_type = None
+    return df
+
+
+def _get_column_renames(schema: FileSchema) -> dict[str, str]:
+    """Return a mapping of columns to rename."""
+    mapping: dict[str, str] = {}
+    for column in schema.columns:
+        if column.dimension_type is not None and column.name != column.dimension_type.value:
+            mapping[column.name] = column.dimension_type.value
+    return mapping
+
+
+def _rename_columns(df: DataFrame, mapping: dict[str, str]) -> DataFrame:
+    for old_name, new_name in mapping.items():
+        df = df.withColumnRenamed(old_name, new_name)
+        logger.info("Renamed column %s to %s", old_name, new_name)
+    return df
+
+
+def _drop_ignored_columns(df: DataFrame, ignore_columns: list[str]) -> DataFrame:
+    if not ignore_columns:
+        return df
+
+    existing_columns = set(df.columns)
+    for col in ignore_columns:
+        if col in existing_columns:
+            df = df.drop(col)
+            logger.info("Dropped ignored column: %s", col)
+        else:
+            logger.warning("Ignored column '%s' not found in file", col)
+    return df
+
+
+def _get_column_schema(schema: FileSchema, backend_mapping: dict) -> dict[str, str] | None:
+    column_types = schema.get_data_type_mapping()
+    if not column_types:
+        return None
+
+    mapped_schema: dict[str, str] = {}
+    for key, val in column_types.items():
+        col_type = val.upper()
+        if col_type not in backend_mapping:
+            options = " ".join(sorted(backend_mapping.keys()))
+            msg = f"column type = {val} is not supported. {options=}"
+            raise DSGInvalidField(msg)
+        mapped_schema[key] = backend_mapping[col_type]
+    return mapped_schema

dsgrid/config/index_time_dimension_config.py

@@ -0,0 +1,80 @@
+import logging
+from datetime import datetime, timedelta
+from typing import Union
+
+import chronify
+import pandas as pd
+
+from dsgrid.time.types import IndexTimestampType
+from .dimensions import IndexTimeDimensionModel
+from .time_dimension_base_config import TimeDimensionBaseConfig
+from dsgrid.dimension.time import TimeIntervalType
+
+
+logger = logging.getLogger(__name__)
+
+
+class IndexTimeDimensionConfig(TimeDimensionBaseConfig):
+    """Provides an interface to a IndexTimeDimensionModel."""
+
+    @staticmethod
+    def model_class() -> IndexTimeDimensionModel:
+        return IndexTimeDimensionModel
+
+    def supports_chronify(self) -> bool:
+        return True
+
+    def to_chronify(
+        self,
+    ) -> Union[
+        chronify.IndexTimeRangeTZ, chronify.IndexTimeRangeNTZ, chronify.IndexTimeRangeWithTZColumn
+    ]:
+        time_cols = self.get_load_data_time_columns()
+        assert len(self._model.ranges) == 1
+        assert len(time_cols) == 1
+
+        # IndexTimeDimensionModel does not map to IndexTimeRangeNTZ and TZ at the moment
+        assert self.get_time_zone() is None
+        config = chronify.IndexTimeRangeWithTZColumn(
+            time_column=time_cols[0],
+            start=self._model.ranges[0].start,
+            length=self.get_lengths()[0],
+            start_timestamp=pd.Timestamp(self.get_start_times()[0]),
+            resolution=self.get_frequency(),
+            time_zone_column="time_zone",
+            measurement_type=self._model.measurement_type,
+            interval_type=self._model.time_interval_type,
+        )
+        return config
+
+    def get_frequency(self) -> timedelta:
+        freqs = [trange.frequency for trange in self.model.ranges]
+        if len(set(freqs)) > 1:
+            msg = f"IndexTimeDimensionConfig.get_frequency found multiple frequencies: {freqs}"
+            raise ValueError(msg)
+        return freqs[0]
+
+    def get_start_times(self) -> list[pd.Timestamp]:
+        """get represented start times"""
+        tz = self.get_tzinfo()
+        start_times = []
+        for trange in self.model.ranges:
+            start = datetime.strptime(trange.starting_timestamp, trange.str_format)
+            assert start.tzinfo is None
+            start_times.append(start.replace(tzinfo=tz))
+        return start_times
+
+    def get_lengths(self) -> list[int]:
+        return [trange.end - trange.start + 1 for trange in self.model.ranges]
+
+    def get_load_data_time_columns(self) -> list[str]:
+        return list(IndexTimestampType._fields)
+
+    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

dsgrid/config/input_dataset_requirements.py

@@ -0,0 +1,31 @@
+"""Defines dataset dimension requirements for a project."""
+
+from pydantic import conlist, Field
+
+from dsgrid.config.project_config import RequiredDimensionsModel, InputDatasetModel
+from dsgrid.data_models import DSGBaseModel
+
+
+class InputDatasetDimensionRequirementsModel(DSGBaseModel):
+    """Defines dataset dimension requirements."""
+
+    dataset_id: str
+    required_dimensions: RequiredDimensionsModel = Field(
+        title="required_dimensions",
+        description="Defines required record IDs that must exist for each dimension.",
+    )
+
+
+class InputDatasetDimensionRequirementsListModel(DSGBaseModel):
+    """Defines a list of dataset dimension requirements."""
+
+    dataset_dimension_requirements: conlist(
+        InputDatasetDimensionRequirementsModel, min_length=1
+    ) = Field(description="List of dataset dimension requirements")
+
+
+class InputDatasetListModel(DSGBaseModel):
+    datasets: conlist(InputDatasetModel, min_length=1) = Field(
+        title="datasets",
+        description="List of input datasets for the project.",
+    )

dsgrid/config/mapping_tables.py

@@ -0,0 +1,209 @@
+import csv
+import logging
+import os
+
+
+from pydantic import field_validator, Field, ValidationInfo, field_serializer
+
+from dsgrid.config.dimension_mapping_base import (
+    DimensionMappingBaseModel,
+    DimensionMappingDatasetToProjectBaseModel,
+    DimensionMappingPreRegisteredBaseModel,
+)
+from dsgrid.config.dimensions import DimensionReferenceModel
+from dsgrid.data_models import DSGBaseModel
+from dsgrid.utils.files import compute_file_hash
+from dsgrid.utils.utilities import convert_record_dicts_to_classes
+from .config_base import ConfigWithRecordFileBase
+
+
+logger = logging.getLogger(__name__)
+
+
+class MappingTableRecordModel(DSGBaseModel):
+    """Represents one record in dimension mapping record files. Maps one dimension to another."""
+
+    from_id: str = Field(
+        title="from_id",
+        description="Source mapping",
+    )
+    to_id: str | None = Field(
+        default=None,
+        title="to_id",
+        description="Destination mapping",
+    )
+    from_fraction: float = Field(
+        title="from_fraction",
+        description="Fraction of from_id to map to to_id",
+        default=1.0,
+    )
+
+    @field_validator("from_id", "to_id")
+    @classmethod
+    def check_to_id(cls, val):
+        if val == "":
+            return None
+        return val
+
+
+class MappingTableByNameModel(DimensionMappingPreRegisteredBaseModel):
+    """Attributes for a dimension mapping table for soon-to-be registered dimensions by name.
+    This will be converted to a MappingTableModel as soon as the dimensions are registered.
+    """
+
+    filename: str = Field(
+        title="filename",
+        alias="file",
+        description="Filename containing association table records.",
+    )
+
+
+class DatasetBaseToProjectMappingTableModel(DimensionMappingDatasetToProjectBaseModel):
+    """Attributes for a dimension mapping table to map soon-to-be-registered dataset base
+    dimensions to a project's dimensions. This will be converted to a MappingTableModel as soon as
+    the dimensions are registered.
+    """
+
+    filename: str = Field(
+        title="filename",
+        alias="file",
+        description="Filename containing association table records.",
+    )
+
+
+class DatasetBaseToProjectMappingTableListModel(DSGBaseModel):
+    """Represents the config file passed to register-and-submit-dataset command."""
+
+    mappings: list[DatasetBaseToProjectMappingTableModel]
+
+
+class MappingTableModel(DimensionMappingBaseModel):
+    """Attributes for a dimension mapping table"""
+
+    filename: str | None = Field(
+        title="filename",
+        alias="file",
+        default=None,
+        description="Filename containing association table records. Only assigned for user input "
+        "and output purposes. The registry database stores records in the mapping JSON document.",
+    )
+    file_hash: str | None = Field(
+        title="file_hash",
+        description="Hash of the contents of the file, computed by dsgrid.",
+        json_schema_extra={
+            "dsgrid_internal": True,
+        },
+        default=None,
+    )
+    records: list = Field(
+        title="records",
+        description="dimension mapping records in filename that get loaded at runtime",
+        json_schema_extra={
+            "dsgrid_internal": True,
+        },
+        default=[],
+    )
+
+    @field_validator("filename")
+    @classmethod
+    def check_filename(cls, filename):
+        """Validate record file"""
+        if filename is not None:
+            if filename:
+                if not os.path.isfile(filename):
+                    msg = f"{filename} does not exist"
+                    raise ValueError(msg)
+                if not filename.endswith(".csv"):
+                    msg = f"only CSV is supported: {filename}"
+                    raise ValueError(msg)
+        return filename
+
+    @field_validator("file_hash")
+    @classmethod
+    def compute_file_hash(cls, file_hash, info: ValidationInfo):
+        """Compute file hash."""
+        if "filename" not in info.data:
+            return file_hash
+
+        if not file_hash:
+            file_hash = compute_file_hash(info.data["filename"])
+        return file_hash
+
+    @field_validator("records")
+    @classmethod
+    def add_records(cls, records, info: ValidationInfo):
+        """Add records from the file."""
+        if "filename" not in info.data:
+            return records
+
+        if records:
+            if isinstance(records[0], dict):
+                records = convert_record_dicts_to_classes(records, MappingTableRecordModel)
+            return records
+
+        with open(info.data["filename"], encoding="utf-8-sig") as f_in:
+            return convert_record_dicts_to_classes(csv.DictReader(f_in), MappingTableRecordModel)
+
+    @field_serializer("filename")
+    def serialize_cls(self, val, _):
+        return None
+
+    @classmethod
+    def from_pre_registered_model(
+        cls,
+        model: MappingTableByNameModel | DatasetBaseToProjectMappingTableModel,
+        from_dimension: DimensionReferenceModel,
+        to_dimension: DimensionReferenceModel,
+    ):
+        return MappingTableModel(
+            mapping_type=model.mapping_type,
+            archetype=model.archetype,
+            from_dimension=from_dimension,
+            to_dimension=to_dimension,
+            description=model.description,
+            file=model.filename,
+            from_fraction_tolerance=model.from_fraction_tolerance,
+            to_fraction_tolerance=model.to_fraction_tolerance,
+        )
+
+
+class MappingTableConfig(ConfigWithRecordFileBase):
+    """Provides an interface to an MappingTableModel"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._dataframe = None
+
+    @staticmethod
+    def config_filename():
+        return "dimension_mapping.json5"
+
+    @property
+    def config_id(self):
+        return self.model.mapping_id
+
+    @staticmethod
+    def model_class():
+        return MappingTableModel
+
+    def get_unique_from_ids(self):
+        """Return the unique from IDs in an association table's records.
+
+        Returns
+        -------
+        set
+            set of str
+
+        """
+        return {x.from_id for x in self.model.records}
+
+    def get_unique_to_ids(self):
+        """Return the unique to IDs in an association table's records.
+
+        Returns
+        -------
+        set
+            set of str
+
+        """
+        return {x.to_id for x in self.model.records}

dsgrid/config/noop_time_dimension_config.py

@@ -0,0 +1,42 @@
+from datetime import timedelta
+
+from .dimensions import NoOpTimeDimensionModel
+from .time_dimension_base_config import TimeDimensionBaseConfig
+
+
+class NoOpTimeDimensionConfig(TimeDimensionBaseConfig):
+    """Provides an interface to an NoOpTimeDimensionModel."""
+
+    @staticmethod
+    def model_class() -> NoOpTimeDimensionModel:
+        return NoOpTimeDimensionModel
+
+    def check_dataset_time_consistency(self, load_data_df, time_columns) -> None:
+        pass
+
+    def get_frequency(self) -> timedelta:
+        return timedelta(days=0)
+
+    def get_time_ranges(self) -> list:
+        return []
+
+    def get_start_times(self) -> list:
+        return []
+
+    def get_lengths(self) -> list:
+        return []
+
+    def get_load_data_time_columns(self) -> list:
+        return []
+
+    def get_time_zone(self) -> None:
+        return None
+
+    def get_tzinfo(self) -> None:
+        return None
+
+    def get_time_interval_type(self) -> None:
+        return None
+
+    def list_expected_dataset_timestamps(self) -> list:
+        return []