dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/data_models.py

@@ -0,0 +1,155 @@
+"""Base functionality for all Pydantic data models used in dsgrid"""
+
+import logging
+from enum import Enum
+from pathlib import Path
+from typing import Any, Self
+
+from pydantic import ConfigDict, BaseModel, Field, ValidationError
+
+from dsgrid.exceptions import DSGInvalidParameter
+from dsgrid.utils.files import in_other_dir, load_data, dump_data
+
+
+logger = logging.getLogger(__name__)
+
+
+def make_model_config(**kwargs) -> ConfigDict:
+    """Return a Pydantic config"""
+    return ConfigDict(
+        str_strip_whitespace=True,
+        validate_assignment=True,
+        validate_default=True,
+        extra="forbid",
+        use_enum_values=False,
+        arbitrary_types_allowed=True,
+        populate_by_name=True,
+        **kwargs,
+    )
+
+
+class DSGBaseModel(BaseModel):
+    """Base data model for all dsgrid data models"""
+
+    model_config = make_model_config()
+
+    @classmethod
+    def load(cls, filename):
+        """Load a data model from a file.
+        Temporarily changes to the file's parent directory so that Pydantic
+        validators can load relative file paths within the file.
+
+        Parameters
+        ----------
+        filename : str
+
+        """
+        filename = Path(filename)
+        if not filename.is_file():
+            msg = f"{filename} is not a file"
+            raise DSGInvalidParameter(msg)
+
+        with in_other_dir(filename.parent):
+            try:
+                return cls(**load_data(filename.name))
+            except ValidationError:
+                logger.exception("Failed to validate %s", filename)
+                raise
+
+    @classmethod
+    def get_fields_with_extra_attribute(cls, attribute):
+        fields = set()
+        for f, attrs in cls.model_fields.items():
+            if attrs.json_schema_extra.get(attribute):
+                fields.add(f)
+        return fields
+
+    def model_dump(self, *args, by_alias=True, **kwargs):
+        return super().model_dump(*args, by_alias=by_alias, **self._handle_kwargs(**kwargs))
+
+    def model_dump_json(self, *args, by_alias=True, **kwargs):
+        return super().model_dump_json(*args, by_alias=by_alias, **self._handle_kwargs(**kwargs))
+
+    @staticmethod
+    def _handle_kwargs(**kwargs):
+        return {k: v for k, v in kwargs.items() if k not in ("by_alias",)}
+
+    def serialize(self, *args, **kwargs) -> dict[str, Any]:
+        return self.model_dump(*args, mode="json", **kwargs)
+
+    @classmethod
+    def from_file(cls, filename: Path) -> Self:
+        """Deserialize the model from a file. Unlike the load method,
+        this does not change directories.
+        """
+        return cls(**load_data(filename))
+
+    def to_file(self, filename: Path) -> None:
+        """Serialize the model to a file."""
+        data = self.serialize()
+        dump_data(data, filename, indent=2)
+
+
+class DSGBaseDatabaseModel(DSGBaseModel):
+    """Base model for all configs stored in the database."""
+
+    id: int | None = Field(
+        default=None,
+        description="Registry database ID",
+        json_schema_extra={
+            "dsgrid_internal": True,
+        },
+    )
+    version: str | None = Field(
+        default=None,
+        title="version",
+        description="Version, generated by dsgrid",
+        json_schema_extra={
+            "dsgrid_internal": True,
+            "updateable": False,
+        },
+    )
+
+
+class EnumValue:
+    """Class to define a DSGEnum value"""
+
+    def __init__(self, value, description, **kwargs):
+        self.value = value
+        self.description = description
+        for kwarg, val in kwargs.items():
+            self.__setattr__(kwarg, val)
+
+
+class DSGEnum(Enum):
+    """dsgrid Enum class"""
+
+    def __new__(cls, *args):
+        obj = object.__new__(cls)
+        assert len(args) in (1, 2)
+        if isinstance(args[0], EnumValue):
+            obj._value_ = args[0].value
+            obj.description = args[0].description
+            for attr, val in args[0].__dict__.items():
+                if attr not in ("value", "description"):
+                    setattr(obj, attr, val)
+        elif len(args) == 2:
+            obj._value_ = args[0]
+            obj.description = args[1]
+        else:
+            obj._value_ = args[0]
+            obj.description = None
+        return obj
+
+    @classmethod
+    def format_for_docs(cls):
+        """Returns set of formatted enum values for docs."""
+        return str([e.value for e in cls]).replace("'", "``")
+
+    @classmethod
+    def format_descriptions_for_docs(cls):
+        """Returns formatted dict of enum values and descriptions for docs."""
+        desc = {}
+        for e in cls:
+            desc[f"``{e.value}``"] = f"{e.description}"
+        return desc

dsgrid/dataset/dataset.py

@@ -0,0 +1,123 @@
+"""Provides access to a dataset."""
+
+import abc
+import logging
+
+from sqlalchemy import Connection
+
+from dsgrid.config.dataset_config import DatasetConfig
+from dsgrid.config.dataset_schema_handler_factory import make_dataset_schema_handler
+from dsgrid.config.dimension_mapping_base import DimensionMappingReferenceListModel
+from dsgrid.config.project_config import ProjectConfig
+from dsgrid.query.query_context import QueryContext
+from dsgrid.dataset.dataset_schema_handler_base import DatasetSchemaHandlerBase
+from dsgrid.registry.data_store_interface import DataStoreInterface
+from dsgrid.registry.dimension_mapping_registry_manager import DimensionMappingRegistryManager
+from dsgrid.registry.dimension_registry_manager import DimensionRegistryManager
+from dsgrid.spark.types import DataFrame
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetBase(abc.ABC):
+    """Base class for datasets"""
+
+    def __init__(self, schema_handler: DatasetSchemaHandlerBase):
+        self._config = schema_handler.config
+        self._handler = schema_handler
+        self._id = schema_handler.config.model.dataset_id
+        # Can't use dashes in view names. This will need to be handled when we implement
+        # queries based on dataset ID.
+
+    @property
+    def config(self) -> DatasetConfig:
+        return self._config
+
+    @property
+    def dataset_id(self) -> str:
+        return self._id
+
+    @property
+    def handler(self) -> DatasetSchemaHandlerBase:
+        return self._handler
+
+
+class Dataset(DatasetBase):
+    """Represents a dataset used within a project."""
+
+    @classmethod
+    def load(
+        cls,
+        config: DatasetConfig,
+        dimension_mgr: DimensionRegistryManager,
+        dimension_mapping_mgr: DimensionMappingRegistryManager,
+        store: DataStoreInterface,
+        mapping_references: list[DimensionMappingReferenceListModel],
+        conn: Connection | None = None,
+    ):
+        """Load a dataset from a store.
+
+        Parameters
+        ----------
+        config : DatasetConfig
+        dimension_mgr : DimensionRegistryManager
+        dimension_mapping_mgr : DimensionMappingRegistryManager
+        mapping_references: list[DimensionMappingReferenceListModel]
+
+        Returns
+        -------
+        Dataset
+
+        """
+        return cls(
+            make_dataset_schema_handler(
+                conn,
+                config,
+                dimension_mgr,
+                dimension_mapping_mgr,
+                store=store,
+                mapping_references=mapping_references,
+            )
+        )
+
+    def make_project_dataframe(
+        self, query: QueryContext, project_config: ProjectConfig
+    ) -> DataFrame:
+        return self._handler.make_project_dataframe(query, project_config)
+
+
+class StandaloneDataset(DatasetBase):
+    """Represents a dataset used outside of a project."""
+
+    @classmethod
+    def load(
+        cls,
+        config: DatasetConfig,
+        dimension_mgr: DimensionRegistryManager,
+        dimension_mapping_mgr: DimensionMappingRegistryManager,
+        store: DataStoreInterface,
+        mapping_references: list[DimensionMappingReferenceListModel] | None = None,
+        conn: Connection | None = None,
+    ):
+        """Load a dataset from a store.
+
+        Parameters
+        ----------
+        config : DatasetConfig
+        dimension_mgr : DimensionRegistryManager
+
+        Returns
+        -------
+        Dataset
+
+        """
+        return cls(
+            make_dataset_schema_handler(
+                conn,
+                config,
+                dimension_mgr,
+                dimension_mapping_mgr,
+                store=store,
+                mapping_references=mapping_references,
+            )
+        )

dsgrid/dataset/dataset_expression_handler.py

@@ -0,0 +1,86 @@
+import operator
+
+from dsgrid.exceptions import DSGInvalidOperation
+from dsgrid.spark.functions import join_multiple_columns
+from dsgrid.spark.types import DataFrame
+from dsgrid.utils.py_expression_eval import Parser
+
+
+class DatasetExpressionHandler:
+    """Abstracts SQL expressions for dataset combinations with mathematical expressions."""
+
+    def __init__(self, df: DataFrame, dimension_columns: list[str], value_columns: list[str]):
+        self.df = df
+        self.dimension_columns = dimension_columns
+        self.value_columns = value_columns
+
+    def _op(self, other, op):
+        orig_self_count = self.df.count()
+        orig_other_count = other.df.count()
+        if orig_self_count != orig_other_count:
+            msg = (
+                f"{op=} requires that the datasets have the same length "
+                f"{orig_self_count=} {orig_other_count=}"
+            )
+            raise DSGInvalidOperation(msg)
+
+        def renamed(col):
+            return col + "_other"
+
+        other_df = other.df
+        for column in self.value_columns:
+            other_df = other_df.withColumnRenamed(column, renamed(column))
+        df = join_multiple_columns(self.df, other_df, self.dimension_columns)
+
+        for column in self.value_columns:
+            other_column = renamed(column)
+            df = df.withColumn(column, op(getattr(df, column), getattr(df, other_column)))
+
+        df = df.select(*self.df.columns)
+        joined_count = df.count()
+        if joined_count != orig_self_count:
+            msg = (
+                f"join for operation {op=} has a different row count than the original. "
+                f"{orig_self_count=} {joined_count=}"
+            )
+            raise DSGInvalidOperation(msg)
+
+        return DatasetExpressionHandler(df, self.dimension_columns, self.value_columns)
+
+    def __add__(self, other):
+        return self._op(other, operator.add)
+
+    def __mul__(self, other):
+        return self._op(other, operator.mul)
+
+    def __sub__(self, other):
+        return self._op(other, operator.sub)
+
+    def __or__(self, other):
+        if self.df.columns != other.df.columns:
+            msg = (
+                "Union is only allowed when datasets have identical columns: "
+                f"{self.df.columns=} vs {other.df.columns=}"
+            )
+            raise DSGInvalidOperation(msg)
+        return DatasetExpressionHandler(
+            self.df.union(other.df), self.dimension_columns, self.value_columns
+        )
+
+
+def evaluate_expression(expr: str, dataset_mapping: dict[str, DatasetExpressionHandler]):
+    """Evaluates an expresion containing dataset IDs.
+
+    Parameters
+    ----------
+    expr : str
+        Dataset combination expression, such as "dataset1 | dataset2"
+    dataset_mapping : dict[str, DatasetExpressionHandler]
+        Maps dataset ID to dataset. Each dataset_id in expr must be present in the mapping.
+
+    Returns
+    -------
+    DatasetExpressionHandler
+
+    """
+    return Parser().parse(expr).evaluate(dataset_mapping)

dsgrid/dataset/dataset_mapping_manager.py

@@ -0,0 +1,121 @@
+import logging
+from pathlib import Path
+from typing import Self
+
+from dsgrid.query.dataset_mapping_plan import (
+    DatasetMappingPlan,
+    MapOperation,
+    MapOperationCheckpoint,
+)
+from dsgrid.spark.types import DataFrame
+from dsgrid.utils.files import delete_if_exists
+from dsgrid.utils.spark import read_dataframe, write_dataframe
+from dsgrid.utils.scratch_dir_context import ScratchDirContext
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetMappingManager:
+    """Manages the mapping operations for a dataset."""
+
+    def __init__(
+        self,
+        dataset_id: str,
+        plan: DatasetMappingPlan,
+        scratch_dir_context: ScratchDirContext,
+        checkpoint: MapOperationCheckpoint | None = None,
+    ):
+        self._dataset_id = dataset_id
+        self._plan = plan
+        self._scratch_dir_context = scratch_dir_context
+        self._checkpoint = checkpoint
+        self._checkpoint_file: Path | None = None
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, *args, **kwargs) -> None:
+        # Don't cleanup if an exception occurred.
+        self.cleanup()
+
+    @property
+    def plan(self) -> DatasetMappingPlan:
+        """Return the mapping plan for the dataset."""
+        return self._plan
+
+    @property
+    def scratch_dir_context(self) -> ScratchDirContext:
+        """Return the scratch_dir_context."""
+        return self._scratch_dir_context
+
+    def try_read_checkpointed_table(self) -> DataFrame | None:
+        """Read the checkpointed table for the dataset, if it exists."""
+        if self._checkpoint is None:
+            return None
+        return read_dataframe(self._checkpoint.persisted_table_filename)
+
+    def get_completed_mapping_operations(self) -> set[str]:
+        """Return the names of completed mapping operations."""
+        return (
+            set() if self._checkpoint is None else set(self._checkpoint.completed_operation_names)
+        )
+
+    def has_completed_operation(self, op: MapOperation) -> bool:
+        """Return True if the mapping operation has been completed."""
+        return op.name in self.get_completed_mapping_operations()
+
+    def persist_table(self, df: DataFrame, op: MapOperation) -> DataFrame:
+        """Persist the intermediate table to the filesystem and return the persisted DataFrame."""
+        persisted_file = self._scratch_dir_context.get_temp_filename(
+            suffix=".parquet", add_tracked_path=False
+        )
+        write_dataframe(df, persisted_file)
+        self.save_checkpoint(persisted_file, op)
+        logger.info("Persisted mapping operation name=%s to %s", op.name, persisted_file)
+        return read_dataframe(persisted_file)
+
+    def save_checkpoint(self, persisted_table: Path, op: MapOperation) -> None:
+        """Save a checkpoint after persisting an operation to the filesystem."""
+        completed_operation_names: list[str] = []
+        for mapping_op in self._plan.list_mapping_operations():
+            completed_operation_names.append(mapping_op.name)
+            if mapping_op.name == op.name:
+                break
+
+        checkpoint = MapOperationCheckpoint(
+            dataset_id=self._dataset_id,
+            completed_operation_names=completed_operation_names,
+            persisted_table_filename=persisted_table,
+            mapping_plan_hash=self._plan.compute_hash(),
+        )
+        checkpoint_filename = self._scratch_dir_context.get_temp_filename(
+            suffix=".json", add_tracked_path=False
+        )
+        checkpoint.to_file(checkpoint_filename)
+        if self._checkpoint_file is not None and not self._plan.keep_intermediate_files:
+            assert self._checkpoint is not None, self._checkpoint
+            logger.info("Remove previous checkpoint: %s", self._checkpoint_file)
+            delete_if_exists(self._checkpoint_file)
+            delete_if_exists(self._checkpoint.persisted_table_filename)
+
+        self._checkpoint = checkpoint
+        self._checkpoint_file = checkpoint_filename
+        logger.info("Saved checkpoint in %s", self._checkpoint_file)
+
+    def cleanup(self) -> None:
+        """Cleanup the intermediate files. Call if the operation completed succesfully."""
+        if self._plan.keep_intermediate_files:
+            logger.info("Keeping intermediate files for dataset %s", self._dataset_id)
+            return
+
+        if self._checkpoint_file is not None:
+            logger.info("Removing checkpoint filename %s", self._checkpoint_file)
+            delete_if_exists(self._checkpoint_file)
+            self._checkpoint_file = None
+        if self._checkpoint is not None:
+            logger.info(
+                "Removing persisted intermediate table filename %s",
+                self._checkpoint.persisted_table_filename,
+            )
+            delete_if_exists(self._checkpoint.persisted_table_filename)
+            self._checkpoint = None