dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/dataset/models.py

@@ -0,0 +1,51 @@
+from enum import StrEnum
+from typing import Annotated, Literal, Union
+
+from pydantic import Field
+
+from dsgrid.data_models import DSGBaseModel
+from dsgrid.dimension.base_models import DimensionType
+
+
+class ValueFormat(StrEnum):
+    """Defines the format of value columns in a dataset."""
+
+    PIVOTED = "pivoted"
+    STACKED = "stacked"
+
+
+class TableFormat(StrEnum):
+    """Defines the table structure of a dataset."""
+
+    ONE_TABLE = "one_table"
+    TWO_TABLE = "two_table"
+
+
+# Keep old name as alias for backward compatibility during migration
+TableFormatType = ValueFormat
+
+
+class PivotedTableFormatModel(DSGBaseModel):
+    """Defines a pivoted table format where one dimension's records are columns."""
+
+    format_type: Literal[ValueFormat.PIVOTED] = ValueFormat.PIVOTED
+    pivoted_dimension_type: DimensionType = Field(
+        title="pivoted_dimension_type",
+        description="The dimension type whose records are columns that contain data values.",
+    )
+
+
+class StackedTableFormatModel(DSGBaseModel):
+    """Defines a stacked (unpivoted) table format with a single value column."""
+
+    format_type: Literal[ValueFormat.STACKED] = ValueFormat.STACKED
+
+
+# Alias for backward compatibility
+UnpivotedTableFormatModel = StackedTableFormatModel
+
+
+TableFormatModel = Annotated[
+    Union[PivotedTableFormatModel, StackedTableFormatModel],
+    Field(discriminator="format_type"),
+]

dsgrid/dataset/table_format_handler_base.py

@@ -0,0 +1,257 @@
+import abc
+import logging
+from typing import Iterable
+
+from dsgrid.config.project_config import ProjectConfig
+from dsgrid.dimension.base_models import DimensionCategory, DimensionType
+from dsgrid.query.query_context import QueryContext
+from dsgrid.query.models import (
+    AggregationModel,
+    ColumnModel,
+    ColumnType,
+    DatasetDimensionsMetadataModel,
+    DimensionMetadataModel,
+)
+from dsgrid.spark.types import DataFrame
+from dsgrid.utils.dataset import map_stacked_dimension, remove_invalid_null_timestamps
+from dsgrid.utils.spark import persist_intermediate_query
+from dsgrid.utils.timing import track_timing, timer_stats_collector
+
+
+logger = logging.getLogger(__name__)
+
+
+class TableFormatHandlerBase(abc.ABC):
+    """Base class for table format handers"""
+
+    def __init__(self, project_config: ProjectConfig, dataset_id: str | None = None):
+        self._project_config = project_config
+        self._dataset_id = dataset_id
+
+    def add_columns(
+        self,
+        df: DataFrame,
+        column_models: list[ColumnModel],
+        context: QueryContext,
+        value_columns: Iterable[str],
+    ) -> DataFrame:
+        """Add columns to the dataframe. For example, suppose the geography dimension is at
+        county resolution and the user wants to add a column for state.
+
+        Parameters
+        ----------
+        df : pyspark.sql.DataFrame
+        column_models : list
+        context : QueryContext
+        value_columns: Iterable[str]
+            Columns in the dataframe that contain load values.
+        """
+        columns = set(df.columns)
+        all_base_names = self.project_config.list_dimension_names(category=DimensionCategory.BASE)
+        for column in column_models:
+            name = column.dimension_name
+            if name in all_base_names or name in columns:
+                continue
+            supp_dim = self._project_config.get_dimension_with_records(name)
+            existing_metadata = context.get_dimension_metadata(
+                supp_dim.model.dimension_type, dataset_id=self._dataset_id
+            )
+            existing_base_metadata = [
+                x for x in existing_metadata if x.dimension_name in all_base_names
+            ]
+            if len(existing_base_metadata) != 1:
+                msg = (
+                    f"Bug: expected one base metadata object for {supp_dim.model.dimension_type}: "
+                    "{existing_base_metadata}"
+                )
+                raise Exception(msg)
+            base_dim_name = existing_base_metadata[0].dimension_name
+            if base_dim_name not in all_base_names:
+                msg = f"Bug: Expected {base_dim_name} to be a base dimension."
+                raise Exception(msg)
+            base_dim = self._project_config.get_dimension_with_records(base_dim_name)
+            records = self._project_config.get_base_to_supplemental_mapping_records(
+                base_dim, supp_dim
+            )
+
+            if column.function is not None:
+                # TODO #200: Do we want to allow this?
+                msg = f"Applying a SQL function to added column={name} is not supported yet"
+                raise NotImplementedError(msg)
+            expected_base_dim_cols = context.get_dimension_column_names_by_name(
+                supp_dim.model.dimension_type,
+                base_dim.model.name,
+                dataset_id=self._dataset_id,
+            )
+            if len(expected_base_dim_cols) > 1:
+                msg = "Bug: Non-time dimensions cannot have more than one base dimension column"
+                raise Exception(msg)
+            expected_base_dim_col = expected_base_dim_cols[0]
+            df = map_stacked_dimension(
+                df,
+                records,
+                expected_base_dim_col,
+                drop_column=False,
+                to_column=name,
+            )
+            if context.model.result.column_type == ColumnType.DIMENSION_NAMES:
+                assert supp_dim.model.dimension_type != DimensionType.TIME
+                column_names = [name]
+            else:
+                column_names = [expected_base_dim_col]
+            context.add_dimension_metadata(
+                supp_dim.model.dimension_type,
+                DimensionMetadataModel(dimension_name=name, column_names=column_names),
+                dataset_id=self.dataset_id,
+            )
+
+        if "fraction" in df.columns:
+            for col in value_columns:
+                df = df.withColumn(col, df[col] * df["fraction"])
+            df = df.drop("fraction")
+
+        return df
+
+    @abc.abstractmethod
+    def process_aggregations(
+        self, df: DataFrame, aggregations: list[AggregationModel], context: QueryContext
+    ) -> DataFrame:
+        """Aggregate the dimensional data as specified by aggregations.
+
+        Parameters
+        ----------
+        df : pyspark.sql.DataFrame
+        aggregations : AggregationModel
+        context : QueryContext
+
+        Returns
+        -------
+        pyspark.sql.DataFrame
+
+        """
+
+    @property
+    def project_config(self) -> ProjectConfig:
+        """Return the project config of the dataset being processed."""
+        return self._project_config
+
+    @property
+    def dataset_id(self) -> str | None:
+        """Return the ID of the dataset being processed."""
+        return self._dataset_id
+
+    def convert_columns_to_query_names(
+        self, df: DataFrame, dataset_id: str, context: QueryContext
+    ) -> DataFrame:
+        """Convert columns from dimension types to dimension query names."""
+        columns = set(df.columns)
+        for dim_type in DimensionType:
+            if dim_type == DimensionType.TIME:
+                time_dim = self._project_config.get_base_time_dimension()
+                df = time_dim.map_timestamp_load_data_columns_for_query_name(df)
+            elif dim_type.value in columns:
+                existing_col = dim_type.value
+                new_cols = context.get_dimension_column_names(dim_type, dataset_id=dataset_id)
+                assert len(new_cols) == 1, f"{dim_type=} {new_cols=}"
+                new_col = next(iter(new_cols))
+                if existing_col != new_col:
+                    df = df.withColumnRenamed(existing_col, new_col)
+                    logger.debug("Converted column from %s to %s", existing_col, new_col)
+
+        return df
+
+    def replace_ids_with_names(self, df: DataFrame) -> DataFrame:
+        """Replace dimension record IDs with names."""
+        assert not {"id", "name"}.intersection(df.columns), df.columns
+        orig = df
+        all_query_names = self._project_config.get_dimension_names_mapped_to_type()
+        for name in set(df.columns).intersection(all_query_names.keys()):
+            if all_query_names[name] != DimensionType.TIME:
+                # Time doesn't have records.
+                dim_config = self._project_config.get_dimension_with_records(name)
+                records = dim_config.get_records_dataframe().select("id", "name")
+                df = (
+                    df.join(records, on=df[name] == records["id"])
+                    .drop("id", name)
+                    .withColumnRenamed("name", name)
+                )
+        assert df.count() == orig.count(), f"counts changed {df.count()} {orig.count()}"
+        return df
+
+    @staticmethod
+    def _add_column_to_dim_type(
+        column: ColumnModel, dim_type: DimensionType, column_to_dim_type: dict[str, DimensionType]
+    ) -> None:
+        name = column.get_column_name()
+        if name in column_to_dim_type:
+            assert dim_type == column_to_dim_type[name], f"{name=} {column_to_dim_type}"
+        column_to_dim_type[name] = dim_type
+
+    def _build_group_by_columns(
+        self,
+        columns: list[ColumnModel],
+        context: QueryContext,
+        final_metadata: DatasetDimensionsMetadataModel,
+    ):
+        group_by_cols: list[str] = []
+        for column in columns:
+            dim = self._project_config.get_dimension(column.dimension_name)
+            dim_type = dim.model.dimension_type
+            match context.model.result.column_type:
+                case ColumnType.DIMENSION_TYPES:
+                    column_names = context.get_dimension_column_names_by_name(
+                        dim_type, column.dimension_name, dataset_id=self._dataset_id
+                    )
+                    if dim_type == DimensionType.TIME:
+                        group_by_cols += column_names
+                    else:
+                        group_by_cols.append(dim_type.value)
+                case ColumnType.DIMENSION_NAMES:
+                    column_names = [column.get_column_name()]
+                    expr = self._make_group_by_column_expr(column)
+                    group_by_cols.append(expr)
+                    if not isinstance(expr, str) or expr != column.dimension_name:
+                        # In this case we are replacing any existing query name with an expression
+                        # or alias, and so the old name must be removed.
+                        final_metadata.remove_metadata(dim_type, column.dimension_name)
+                case _:
+                    msg = f"Bug: unhandled: {context.model.result.column_type}"
+                    raise NotImplementedError(msg)
+            final_metadata.add_metadata(
+                dim_type,
+                DimensionMetadataModel(
+                    dimension_name=column.dimension_name, column_names=column_names
+                ),
+            )
+        return group_by_cols
+
+    @staticmethod
+    def _make_group_by_column_expr(column):
+        if column.function is None:
+            expr = column.dimension_name
+        else:
+            expr = column.function(column.dimension_name)
+            if column.alias is not None:
+                expr = expr.alias(column.alias)
+        return expr
+
+    @track_timing(timer_stats_collector)
+    def _remove_invalid_null_timestamps(self, df: DataFrame, orig_id, context: QueryContext):
+        if id(df) != orig_id:
+            # The table could have NULL timestamps that designate expected-missing data.
+            # Those rows could be obsolete after aggregating stacked dimensions.
+            # This is an expensive operation, so only do it if the dataframe changed.
+            value_columns = context.get_value_columns()
+            if not value_columns:
+                msg = "Bug: value_columns cannot be empty"
+                raise Exception(msg)
+            time_columns = context.get_dimension_column_names(
+                DimensionType.TIME, dataset_id=self._dataset_id
+            )
+            if time_columns:
+                # Persist the query up to this point to avoid multiple evaluations.
+                df = persist_intermediate_query(df, context.scratch_dir_context)
+                stacked_columns = set(df.columns) - value_columns.union(time_columns)
+                df = remove_invalid_null_timestamps(df, time_columns, stacked_columns)
+                logger.debug("Removed any rows with invalid null timestamps")
+        return df

dsgrid/dataset/table_format_handler_factory.py

@@ -0,0 +1,17 @@
+from dsgrid.dataset.models import ValueFormat
+from .table_format_handler_base import TableFormatHandlerBase
+
+from .unpivoted_table import UnpivotedTableHandler
+
+
+def make_table_format_handler(
+    value_format: ValueFormat, project_config, dataset_id=None
+) -> TableFormatHandlerBase:
+    """Return a format handler for the passed value format."""
+    match value_format:
+        case ValueFormat.STACKED:
+            handler = UnpivotedTableHandler(project_config, dataset_id=dataset_id)
+        case _:
+            raise NotImplementedError(str(value_format))
+
+    return handler

dsgrid/dataset/unpivoted_table.py

@@ -0,0 +1,121 @@
+import logging
+
+from dsgrid.common import VALUE_COLUMN
+from dsgrid.dimension.base_models import DimensionType
+from dsgrid.query.models import (
+    AggregationModel,
+    ColumnModel,
+    ColumnType,
+    DatasetDimensionsMetadataModel,
+)
+from dsgrid.query.query_context import QueryContext
+from dsgrid.spark.types import DataFrame
+from dsgrid.units.convert import convert_units_unpivoted
+from dsgrid.dataset.table_format_handler_base import TableFormatHandlerBase
+
+
+logger = logging.getLogger(__name__)
+
+
+class UnpivotedTableHandler(TableFormatHandlerBase):
+    """Implements behavior for tables stored in unpivoted format."""
+
+    def process_aggregations(
+        self, df: DataFrame, aggregations: list[AggregationModel], context: QueryContext
+    ):
+        orig_id = id(df)
+        df = self.process_stacked_aggregations(df, aggregations, context)
+        df = self._remove_invalid_null_timestamps(df, orig_id, context)
+        return df
+
+    def process_stacked_aggregations(
+        self, df, aggregations: list[AggregationModel], context: QueryContext
+    ):
+        """Aggregate the stacked dimensional data as specified by aggregations.
+
+        Parameters
+        ----------
+        df : pyspark.sql.DataFrame
+        aggregations : AggregationModel
+        context : QueryContext
+
+        Returns
+        -------
+        pyspark.sql.DataFrame
+
+        """
+        if not aggregations:
+            return df
+
+        final_metadata = DatasetDimensionsMetadataModel()
+        dim_type_to_base_query_name = self.project_config.get_dimension_type_to_base_name_mapping()
+        column_to_dim_type: dict[str, DimensionType] = {}
+        dropped_dimensions = set()
+        for agg in aggregations:
+            metric_query_name = None
+            columns: list[ColumnModel] = []
+            for dim_type, column in agg.iter_dimensions_to_keep():
+                assert dim_type not in dropped_dimensions, dim_type
+                columns.append(column)
+                self._add_column_to_dim_type(column, dim_type, column_to_dim_type)
+                if dim_type == DimensionType.METRIC:
+                    metric_query_name = column.dimension_name
+
+            if metric_query_name is None:
+                msg = f"Bug: A metric dimension is not included in {agg}"
+                raise Exception(msg)
+
+            dropped_dimensions.update(set(agg.list_dropped_dimensions()))
+            if not columns:
+                continue
+
+            df = self.add_columns(df, columns, context, [VALUE_COLUMN])
+            group_by_cols = self._build_group_by_columns(columns, context, final_metadata)
+            op = agg.aggregation_function
+            df = df.groupBy(*group_by_cols).agg(op(VALUE_COLUMN).alias(VALUE_COLUMN))
+
+            if metric_query_name not in dim_type_to_base_query_name[DimensionType.METRIC]:
+                to_dim = self.project_config.get_dimension_with_records(metric_query_name)
+                assert context.base_dimension_names.metric is not None
+                mapping = self.project_config.get_base_to_supplemental_config(
+                    self.project_config.get_dimension_with_records(
+                        context.base_dimension_names.metric
+                    ),
+                    to_dim,
+                )
+                from_dim_id = mapping.model.from_dimension.dimension_id
+                from_records = self.project_config.get_base_dimension_records_by_id(from_dim_id)
+                mapping_records = mapping.get_records_dataframe()
+                to_unit_records = to_dim.get_records_dataframe()
+                df = convert_units_unpivoted(
+                    df,
+                    _get_metric_column_name(context, metric_query_name),
+                    from_records,
+                    mapping_records,
+                    to_unit_records,
+                )
+
+            logger.debug(
+                "Aggregated dimensions with groupBy %s and operation %s",
+                group_by_cols,
+                op.__name__,
+            )
+
+        for dim_type in DimensionType:
+            metadata = final_metadata.get_metadata(dim_type)
+            if dim_type in dropped_dimensions and metadata:
+                metadata.clear()
+            context.replace_dimension_metadata(dim_type, metadata, dataset_id=self.dataset_id)
+        return df
+
+
+def _get_metric_column_name(context: QueryContext, metric_query_name):
+    match context.model.result.column_type:
+        case ColumnType.DIMENSION_TYPES:
+            metric_column = DimensionType.METRIC.value
+        case ColumnType.DIMENSION_NAMES:
+            metric_column = metric_query_name
+        case _:
+            msg = f"Bug: unhandled: {context.model.result.column_type}"
+            raise NotImplementedError(msg)
+    return metric_column

dsgrid/dimension/base_models.py

@@ -0,0 +1,230 @@
+"""Dimension types for dsgrid"""
+
+from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
+
+from pydantic import Field
+
+from dsgrid.exceptions import DSGInvalidDimension
+from dsgrid.data_models import DSGBaseModel, DSGEnum
+from dsgrid.utils.utilities import check_uniqueness
+
+
+class DimensionType(DSGEnum):
+    """Dimension types"""
+
+    METRIC = "metric"
+    GEOGRAPHY = "geography"
+    SECTOR = "sector"
+    SUBSECTOR = "subsector"
+    TIME = "time"
+    WEATHER_YEAR = "weather_year"
+    MODEL_YEAR = "model_year"
+    SCENARIO = "scenario"
+
+    def __lt__(self, other):
+        return self.value < other.value
+
+    @classmethod
+    def from_column(cls, column: str) -> "DimensionType":
+        try:
+            return cls(column)
+        except ValueError:
+            msg = f"column={column} is not expected or of a known dimension type."
+            raise DSGInvalidDimension(msg)
+
+    @staticmethod
+    def get_dimension_types_allowed_as_columns() -> set["DimensionType"]:
+        """Return the dimension types that may exist in the data table as columns."""
+        return {x for x in DimensionType if x != DimensionType.TIME}
+
+    @staticmethod
+    def get_allowed_dimension_column_names() -> set[str]:
+        """Return the dimension types that may exist in the data table as columns."""
+        return {x.value for x in DimensionType.get_dimension_types_allowed_as_columns()}
+
+
+class DimensionCategory(DSGEnum):
+    """Types of dimension categories in a project"""
+
+    BASE = "base"
+    SUBSET = "subset"
+    SUPPLEMENTAL = "supplemental"
+
+
+class DimensionRecordBaseModel(DSGBaseModel):
+    """Base class for all dsgrid dimension models"""
+
+    # TODO: add support/links for docs
+    id: str = Field(
+        title="ID",
+        description="Unique identifier within a dimension",
+    )
+    name: str = Field(
+        title="name",
+        description="User-defined name",
+    )
+
+
+class MetricDimensionBaseModel(DimensionRecordBaseModel):
+    """Base class for all metric dimensions (e.g. EnergyEndUse)"""
+
+
+class GeographyDimensionBaseModel(DimensionRecordBaseModel):
+    """Base class for all geography dimensions"""
+
+    time_zone: str | None = Field(
+        default=None,
+        title="Local Prevailing Time Zone",
+        description="""
+        These time zone information are used in reference to project time_zone
+        to convert between project time and local times as necessary.
+        All Prevailing time_zones account for daylight savings time.
+        If a location does not observe daylight savings, use Standard time_zones.
+        """,
+    )
+
+
+class ModelYearDimensionBaseModel(DimensionRecordBaseModel):
+    """Base class for all model year dimensions"""
+
+
+class ScenarioDimensionBaseModel(DimensionRecordBaseModel):
+    """Base class for all scenario dimensions"""
+
+
+class SectorDimensionBaseModel(DimensionRecordBaseModel):
+    """Base class for all sector dimensions"""
+
+
+class SubsectorDimensionBaseModel(DimensionRecordBaseModel):
+    """Base class for all subsector dimensions"""
+
+
+class WeatherYearDimensionBaseModel(DimensionRecordBaseModel):
+    """Base class for weather year dimensions"""
+
+
+_DIMENSION_TO_MODEL = {
+    DimensionType.METRIC: MetricDimensionBaseModel,
+    DimensionType.GEOGRAPHY: GeographyDimensionBaseModel,
+    DimensionType.SECTOR: SectorDimensionBaseModel,
+    DimensionType.SUBSECTOR: SubsectorDimensionBaseModel,
+    DimensionType.WEATHER_YEAR: WeatherYearDimensionBaseModel,
+    DimensionType.MODEL_YEAR: ModelYearDimensionBaseModel,
+    DimensionType.SCENARIO: ScenarioDimensionBaseModel,
+}
+
+
+class DatasetDimensionRequirements(DSGBaseModel):
+    """Defines the requirements for checking a dataset prior to registration."""
+
+    check_time_consistency: bool = True
+    check_dimension_associations: bool = True
+    require_all_dimension_types: bool = True
+
+
+def get_record_base_model(type_enum):
+    """Return the dimension model class for a DimensionType."""
+    dim_model = _DIMENSION_TO_MODEL.get(type_enum)
+    if dim_model is None:
+        msg = f"no mapping for {type_enum}"
+        raise DSGInvalidDimension(msg)
+    return dim_model
+
+
+def check_required_dimensions(dimensions, tag):
+    """Check that a project or dataset config contains all required dimensions.
+
+    Parameters
+    ----------
+    dimensions : list
+        list of DimensionReferenceModel
+    tag : str
+        User-defined string to include in exception messages
+
+    Raises
+    ------
+    ValueError
+        Raised if a required dimension is not provided.
+
+    """
+    dimension_types = {x.dimension_type for x in dimensions}
+    required_dim_types = set(DimensionType)
+    missing = required_dim_types.difference(dimension_types)
+    if missing:
+        msg = f"Required dimension(s) {missing} are not in {tag}."
+        raise ValueError(msg)
+
+    check_uniqueness((x.dimension_type for x in dimensions), tag)
+
+
+def check_required_dataset_dimensions(dimensions, requirements: DatasetDimensionRequirements, tag):
+    """Check that a dataset config contains all required dimensions.
+
+    Parameters
+    ----------
+    dimensions : list
+        list of DimensionReferenceModel
+    requirements : DatasetDimensionRequirements
+        Defines the dimension requirements in the dataset that must be checked.
+    tag : str
+        User-defined string to include in exception messages
+
+    Raises
+    ------
+    ValueError
+        Raised if a required dimension is not provided.
+
+    """
+    dimension_types = {x.dimension_type for x in dimensions}
+    if requirements.require_all_dimension_types:
+        required_dim_types = set(DimensionType)
+        missing = required_dim_types.difference(dimension_types)
+        if missing:
+            msg = f"Required dimension(s) {missing} are not in {tag}."
+            raise ValueError(msg)
+
+    check_uniqueness((x.dimension_type for x in dimensions), tag)
+
+
+def check_timezone_in_geography(dimension, err_msg=None):
+    """Check that a geography dimension contains valid time zones in records.
+
+    Parameters
+    ----------
+    dimension : DimensionModel
+    err_msg : str | None
+        Optional error message
+
+    Raises
+    ------
+    DSGInvalidDimension
+        Raised if a required dimension is not provided.
+
+    """
+    if dimension.dimension_type != DimensionType.GEOGRAPHY:
+        msg = (
+            f"Dimension has type {dimension.dimension_type}, "
+            "Can only check time_zone for Geography."
+        )
+        raise DSGInvalidDimension(msg)
+
+    if not hasattr(dimension.records[0], "time_zone"):
+        if err_msg is None:
+            err_msg = "These geography dimension records must include a time_zone column."
+        raise ValueError(err_msg)
+
+    record_tzs = {rec.time_zone for rec in dimension.records if rec.time_zone is not None}
+    invalid_tzs = []
+    for tz in record_tzs:
+        try:
+            ZoneInfo(tz)
+        except ZoneInfoNotFoundError:
+            invalid_tzs.append(tz)
+
+    if invalid_tzs:
+        msg = (
+            f"Geography dimension {dimension.dimension_id} has invalid time zone(s) in records "
+            f"{dimension.filename}: {invalid_tzs}. Use IANA time zone names only."
+        )
+        raise DSGInvalidDimension(msg)