typedspark 0.0.0__py3-none-any.whl

typedspark/__init__.py

@@ -0,0 +1,54 @@
+"""Typedspark: column-wise type annotations for pyspark DataFrames."""
+
+from typedspark._core.column import Column
+from typedspark._core.column_meta import ColumnMeta
+from typedspark._core.dataset import DataSet, DataSetImplements
+from typedspark._core.datatypes import (
+    ArrayType,
+    DayTimeIntervalType,
+    DecimalType,
+    MapType,
+    StructType,
+)
+from typedspark._core.literaltype import IntervalType
+from typedspark._schema.schema import MetaSchema, Schema
+from typedspark._transforms.structtype_column import structtype_column
+from typedspark._transforms.transform_to_schema import transform_to_schema
+from typedspark._utils.create_dataset import (
+    create_empty_dataset,
+    create_partially_filled_dataset,
+    create_structtype_row,
+)
+from typedspark._utils.databases import Catalogs, Database, Databases
+from typedspark._utils.load_table import create_schema, load_table
+from typedspark._utils.register_schema_to_dataset import (
+    register_schema_to_dataset,
+    register_schema_to_dataset_with_alias,
+)
+
+__all__ = [
+    "ArrayType",
+    "Catalogs",
+    "Column",
+    "ColumnMeta",
+    "Database",
+    "Databases",
+    "DataSet",
+    "DayTimeIntervalType",
+    "DecimalType",
+    "IntervalType",
+    "MapType",
+    "MetaSchema",
+    "DataSetImplements",
+    "Schema",
+    "StructType",
+    "create_empty_dataset",
+    "create_partially_filled_dataset",
+    "create_structtype_row",
+    "create_schema",
+    "load_table",
+    "register_schema_to_dataset",
+    "register_schema_to_dataset_with_alias",
+    "structtype_column",
+    "transform_to_schema",
+]

typedspark/_core/column.py

@@ -0,0 +1,97 @@
+"""Module containing classes and functions related to TypedSpark Columns."""
+
+from logging import warn
+from typing import Generic, Optional, TypeVar, Union, get_args, get_origin
+
+from pyspark.sql import Column as SparkColumn
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.functions import col
+from pyspark.sql.types import DataType
+
+from typedspark._core.datatypes import StructType
+
+T = TypeVar("T", bound=DataType)
+
+
+class EmptyColumn(SparkColumn):
+    """Column object to be instantiated when there is no active Spark session."""
+
+    def __init__(self, *args, **kwargs) -> None:  # pragma: no cover
+        pass
+
+
+class Column(SparkColumn, Generic[T]):
+    """Represents a ``Column`` in a ``Schema``. Can be used as:
+
+    .. code-block:: python
+
+        class A(Schema):
+            a: Column[IntegerType]
+            b: Column[StringType]
+    """
+
+    def __new__(
+        cls,
+        name: str,
+        dataframe: Optional[DataFrame] = None,
+        curid: Optional[int] = None,
+        dtype: Optional[T] = None,
+        parent: Union[DataFrame, "Column", None] = None,
+        alias: Optional[str] = None,
+    ):
+        """``__new__()`` instantiates the object (prior to ``__init__()``).
+
+        Here, we simply take the provided ``name``, create a pyspark
+        ``Column`` object and cast it to a typedspark ``Column`` object.
+        This allows us to bypass the pypsark ``Column`` constuctor in
+        ``__init__()``, which requires parameters that may be difficult
+        to access.
+        """
+        # pylint: disable=unused-argument
+
+        if dataframe is not None and parent is None:
+            parent = dataframe
+            warn("The use of Column(dataframe=...) is deprecated, use Column(parent=...) instead.")
+
+        column: SparkColumn
+        if SparkSession.getActiveSession() is None:
+            column = EmptyColumn()  # pragma: no cover
+        elif alias is not None:
+            column = col(f"{alias}.{name}")
+        elif parent is not None:
+            column = parent[name]
+        else:
+            column = col(name)
+
+        column.__class__ = Column
+        return column
+
+    def __init__(
+        self,
+        name: str,
+        dataframe: Optional[DataFrame] = None,
+        curid: Optional[int] = None,
+        dtype: Optional[T] = None,
+        parent: Union[DataFrame, "Column", None] = None,
+        alias: Optional[str] = None,
+    ):
+        # pylint: disable=unused-argument
+        self.str = name
+        self._dtype = dtype if dtype is not None else DataType
+        self._curid = curid
+
+    def __hash__(self) -> int:
+        return hash((self.str, self._curid))
+
+    @property
+    def dtype(self) -> T:
+        """Get the datatype of the column, e.g. Column[IntegerType] -> IntegerType."""
+        dtype = self._dtype
+
+        if get_origin(dtype) == StructType:
+            return StructType(
+                schema=get_args(dtype)[0],
+                parent=self,
+            )  # type: ignore
+
+        return dtype()  # type: ignore

typedspark/_core/column_meta.py

@@ -0,0 +1,26 @@
+"""Metadata for ``Column`` objects that can be accessed during runtime."""
+
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+
+@dataclass
+class ColumnMeta:
+    """Contains the metadata for a ``Column``. Used as:
+
+    .. code-block:: python
+
+        class A(Schema):
+            a: Annotated[
+                Column[IntegerType],
+                ColumnMeta(
+                    comment="This is a comment",
+                )
+            ]
+    """
+
+    comment: Optional[str] = None
+
+    def get_metadata(self) -> Optional[Dict[str, str]]:
+        """Returns the metadata of this column."""
+        return {"comment": self.comment} if self.comment else None

typedspark/_core/dataset.py

@@ -0,0 +1,302 @@
+"""Module containing classes and functions related to TypedSpark DataSets."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Callable, Generic, List, Literal, Optional, Type, TypeVar, Union, cast, overload
+
+from pyspark.sql import Column as SparkColumn
+from pyspark.sql import DataFrame
+from typing_extensions import Concatenate, ParamSpec
+
+from typedspark._core.validate_schema import validate_schema
+from typedspark._schema.schema import Schema
+
+_Schema = TypeVar("_Schema", bound=Schema)
+_Protocol = TypeVar("_Protocol", bound=Schema, covariant=True)
+_Implementation = TypeVar("_Implementation", bound=Schema, covariant=True)
+
+P = ParamSpec("P")
+_ReturnType = TypeVar("_ReturnType", bound=DataFrame)  # pylint: disable=C0103
+
+
+class DataSetImplements(DataFrame, Generic[_Protocol, _Implementation]):
+    """DataSetImplements allows us to define functions such as:
+
+    .. code-block:: python
+        class Age(Schema, Protocol):
+            age: Column[LongType]
+
+        def birthday(df: DataSetImplements[Age, T]) -> DataSet[T]:
+            return transform_to_schema(
+                df,
+                df.typedspark_schema,
+                {Age.age: Age.age + 1},
+            )
+
+    Such a function:
+    1. Takes as an input ``DataSetImplements[Age, T]``: a ``DataSet`` that implements the protocol
+       ``Age`` as ``T``.
+    2. Returns a ``DataSet[T]``: a ``DataSet`` of the same type as the one that was provided.
+
+    ``DataSetImplements`` should solely be used as a type annotation, it is never initialized."""
+
+    _schema_annotations: Type[_Implementation]
+
+    def __init__(self):
+        raise NotImplementedError(
+            "DataSetImplements should solely be used as a type annotation, it is never initialized."
+        )
+
+    @property
+    def typedspark_schema(self) -> Type[_Implementation]:
+        """Returns the ``Schema`` of the ``DataSet``."""
+        return self._schema_annotations
+
+    """The following functions are equivalent to their parents in ``DataFrame``, but since they
+    don't affect the ``Schema``, we can add type annotations here. We're omitting docstrings,
+    such that the docstring from the parent will appear."""
+
+    def alias(self, alias: str) -> DataSet[_Implementation]:
+        return DataSet[self._schema_annotations](super().alias(alias))  # type: ignore
+
+    def distinct(self) -> DataSet[_Implementation]:  # pylint: disable=C0116
+        return DataSet[self._schema_annotations](super().distinct())  # type: ignore
+
+    def filter(self, condition) -> DataSet[_Implementation]:  # pylint: disable=C0116
+        return DataSet[self._schema_annotations](super().filter(condition))  # type: ignore
+
+    @overload
+    def join(  # type: ignore
+        self,
+        other: DataFrame,
+        on: Optional[  # pylint: disable=C0103
+            Union[str, List[str], SparkColumn, List[SparkColumn]]
+        ] = ...,
+        how: None = ...,
+    ) -> DataFrame: ...  # pragma: no cover
+
+    @overload
+    def join(
+        self,
+        other: DataFrame,
+        on: Optional[  # pylint: disable=C0103
+            Union[str, List[str], SparkColumn, List[SparkColumn]]
+        ] = ...,
+        how: Literal["semi"] = ...,
+    ) -> DataSet[_Implementation]: ...  # pragma: no cover
+
+    @overload
+    def join(
+        self,
+        other: DataFrame,
+        on: Optional[  # pylint: disable=C0103
+            Union[str, List[str], SparkColumn, List[SparkColumn]]
+        ] = ...,
+        how: Optional[str] = ...,
+    ) -> DataFrame: ...  # pragma: no cover
+
+    def join(  # pylint: disable=C0116
+        self,
+        other: DataFrame,
+        on: Optional[  # pylint: disable=C0103
+            Union[str, List[str], SparkColumn, List[SparkColumn]]
+        ] = None,
+        how: Optional[str] = None,
+    ) -> DataFrame:
+        return super().join(other, on, how)  # type: ignore
+
+    def orderBy(self, *args, **kwargs) -> DataSet[_Implementation]:  # type: ignore  # noqa: N802, E501  # pylint: disable=C0116, C0103
+        return DataSet[self._schema_annotations](super().orderBy(*args, **kwargs))  # type: ignore
+
+    def transform(
+        self,
+        func: Callable[Concatenate[DataSet[_Implementation], P], _ReturnType],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> _ReturnType:
+        return super().transform(func, *args, **kwargs)  # type: ignore
+
+    @overload
+    def unionByName(  # noqa: N802  # pylint: disable=C0116, C0103
+        self,
+        other: DataSet[_Implementation],
+        allowMissingColumns: Literal[False] = ...,  # noqa: N803
+    ) -> DataSet[_Implementation]: ...  # pragma: no cover
+
+    @overload
+    def unionByName(  # noqa: N802  # pylint: disable=C0116, C0103
+        self,
+        other: DataFrame,
+        allowMissingColumns: bool = ...,  # noqa: N803
+    ) -> DataFrame: ...  # pragma: no cover
+
+    def unionByName(  # noqa: N802  # pylint: disable=C0116, C0103
+        self,
+        other: DataFrame,
+        allowMissingColumns: bool = False,  # noqa: N803
+    ) -> DataFrame:
+        res = super().unionByName(other, allowMissingColumns)
+        if isinstance(other, DataSet) and other._schema_annotations == self._schema_annotations:
+            return DataSet[self._schema_annotations](res)  # type: ignore
+        return res  # pragma: no cover
+
+
+class DataSet(DataSetImplements[_Schema, _Schema]):
+    """``DataSet`` subclasses pyspark ``DataFrame`` and hence has all the same
+    functionality, with in addition the possibility to define a schema.
+
+    .. code-block:: python
+
+        class Person(Schema):
+            name: Column[StringType]
+            age: Column[LongType]
+
+        def foo(df: DataSet[Person]) -> DataSet[Person]:
+            # do stuff
+            return df
+    """
+
+    def __new__(cls, dataframe: DataFrame) -> DataSet[_Schema]:
+        """``__new__()`` instantiates the object (prior to ``__init__()``).
+
+        Here, we simply take the provided ``df`` and cast it to a
+        ``DataSet``. This allows us to bypass the ``DataFrame``
+        constuctor in ``__init__()``, which requires parameters that may
+        be difficult to access. Subsequently, we perform schema validation, if
+        the schema annotations are provided.
+        """
+        dataframe = cast(DataSet, dataframe)
+        dataframe.__class__ = DataSet
+
+        # first we reset the schema annotations to None, in case they are inherrited through the
+        # passed DataFrame
+        dataframe._schema_annotations = None  # type: ignore
+
+        # then we use the class' schema annotations to validate the schema and add metadata
+        if hasattr(cls, "_schema_annotations"):
+            dataframe._schema_annotations = cls._schema_annotations  # type: ignore
+            dataframe._validate_schema()
+            dataframe._add_schema_metadata()
+
+        return dataframe  # type: ignore
+
+    def __init__(self, dataframe: DataFrame):
+        pass
+
+    def __class_getitem__(cls, item):
+        """Allows us to define a schema for the ``DataSet``.
+
+        To make sure that the DataSet._schema_annotations variable isn't reused globally, we
+        generate a subclass of the ``DataSet`` with the schema annotations as a class variable.
+        """
+        subclass_name = f"{cls.__name__}[{item.__name__}]"
+        subclass = type(subclass_name, (cls,), {"_schema_annotations": item})
+        return subclass
+
+    def _validate_schema(self) -> None:
+        """Validates the schema of the ``DataSet`` against the schema annotations."""
+        validate_schema(
+            self._schema_annotations.get_structtype(),
+            deepcopy(self.schema),
+            self._schema_annotations.get_schema_name(),
+        )
+
+    def _add_schema_metadata(self) -> None:
+        """Adds the ``ColumnMeta`` comments as metadata to the ``DataSet``.
+
+        Previously set metadata is deleted. Hence, if ``foo(dataframe: DataSet[A]) -> DataSet[B]``,
+        then ``DataSet[B]`` will not inherrit any metadata from ``DataSet[A]``.
+
+        Assumes validate_schema() in __setattr__() has been run.
+        """
+        for field in self._schema_annotations.get_structtype().fields:
+            self.schema[field.name].metadata = field.metadata
+
+    """The following functions are equivalent to their parents in ``DataSetImplements``. However,
+    to support functions like ``functools.reduce(DataSet.unionByName, datasets)``, we also add them
+    here. Unfortunately, this leads to some code redundancy, but we'll take that for granted."""
+
+    def alias(self, alias: str) -> DataSet[_Schema]:
+        return DataSet[self._schema_annotations](super().alias(alias))  # type: ignore
+
+    def distinct(self) -> DataSet[_Schema]:  # pylint: disable=C0116
+        return DataSet[self._schema_annotations](super().distinct())  # type: ignore
+
+    def filter(self, condition) -> DataSet[_Schema]:  # pylint: disable=C0116
+        return DataSet[self._schema_annotations](super().filter(condition))  # type: ignore
+
+    @overload
+    def join(  # type: ignore
+        self,
+        other: DataFrame,
+        on: Optional[  # pylint: disable=C0103
+            Union[str, List[str], SparkColumn, List[SparkColumn]]
+        ] = ...,
+        how: None = ...,
+    ) -> DataFrame: ...  # pragma: no cover
+
+    @overload
+    def join(
+        self,
+        other: DataFrame,
+        on: Optional[  # pylint: disable=C0103
+            Union[str, List[str], SparkColumn, List[SparkColumn]]
+        ] = ...,
+        how: Literal["semi"] = ...,
+    ) -> DataSet[_Schema]: ...  # pragma: no cover
+
+    @overload
+    def join(
+        self,
+        other: DataFrame,
+        on: Optional[  # pylint: disable=C0103
+            Union[str, List[str], SparkColumn, List[SparkColumn]]
+        ] = ...,
+        how: Optional[str] = ...,
+    ) -> DataFrame: ...  # pragma: no cover
+
+    def join(  # pylint: disable=C0116
+        self,
+        other: DataFrame,
+        on: Optional[  # pylint: disable=C0103
+            Union[str, List[str], SparkColumn, List[SparkColumn]]
+        ] = None,
+        how: Optional[str] = None,
+    ) -> DataFrame:
+        return super().join(other, on, how)  # type: ignore
+
+    def orderBy(self, *args, **kwargs) -> DataSet[_Schema]:  # type: ignore  # noqa: N802, E501  # pylint: disable=C0116, C0103
+        return DataSet[self._schema_annotations](super().orderBy(*args, **kwargs))  # type: ignore
+
+    def transform(
+        self,
+        func: Callable[Concatenate[DataSet[_Schema], P], _ReturnType],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> _ReturnType:
+        return super().transform(func, *args, **kwargs)  # type: ignore
+
+    @overload
+    def unionByName(  # noqa: N802  # pylint: disable=C0116, C0103
+        self,
+        other: DataSet[_Schema],
+        allowMissingColumns: Literal[False] = ...,  # noqa: N803
+    ) -> DataSet[_Schema]: ...  # pragma: no cover
+
+    @overload
+    def unionByName(  # noqa: N802  # pylint: disable=C0116, C0103
+        self,
+        other: DataFrame,
+        allowMissingColumns: bool = ...,  # noqa: N803
+    ) -> DataFrame: ...  # pragma: no cover
+
+    def unionByName(  # noqa: N802  # pylint: disable=C0116, C0103
+        self,
+        other: DataFrame,
+        allowMissingColumns: bool = False,  # noqa: N803
+    ) -> DataFrame:
+        res = super().unionByName(other, allowMissingColumns)
+        if isinstance(other, DataSet) and other._schema_annotations == self._schema_annotations:
+            return DataSet[self._schema_annotations](res)  # type: ignore
+        return res  # pragma: no cover

typedspark/_core/datatypes.py

@@ -0,0 +1,106 @@
+"""Here, we make our own definitions of ``MapType``, ``ArrayType`` and ``StructType`` in
+order to allow e.g. for ``ArrayType[StringType]``."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Dict, Generic, Type, TypeVar
+
+from pyspark.sql.types import DataType
+
+if TYPE_CHECKING:  # pragma: no cover
+    from typedspark._core.column import Column
+    from typedspark._schema.schema import Schema
+
+    _Schema = TypeVar("_Schema", bound=Schema)
+else:
+    _Schema = TypeVar("_Schema")
+
+_KeyType = TypeVar("_KeyType", bound=DataType)  # pylint: disable=invalid-name
+_ValueType = TypeVar("_ValueType", bound=DataType)  # pylint: disable=invalid-name
+_Precision = TypeVar("_Precision", bound=int)  # pylint: disable=invalid-name
+_Scale = TypeVar("_Scale", bound=int)  # pylint: disable=invalid-name
+_StartField = TypeVar("_StartField", bound=int)  # pylint: disable=invalid-name
+_EndField = TypeVar("_EndField", bound=int)  # pylint: disable=invalid-name
+
+
+class TypedSparkDataType(DataType):
+    """Base class for typedspark specific ``DataTypes``."""
+
+    @classmethod
+    def get_name(cls) -> str:
+        """Return the name of the type."""
+        return cls.__name__
+
+
+class StructTypeMeta(type):
+    """Initializes the schema attribute as None.
+
+    This allows for auto-complete in Databricks notebooks (uninitialized variables don't
+    show up in auto-complete there).
+    """
+
+    def __new__(cls, name: str, bases: Any, dct: Dict[str, Any]):
+        dct["schema"] = None
+        return super().__new__(cls, name, bases, dct)
+
+
+class StructType(Generic[_Schema], TypedSparkDataType, metaclass=StructTypeMeta):
+    """Allows for type annotations such as:
+
+    .. code-block:: python
+
+        class Job(Schema):
+            position: Column[StringType]
+            salary: Column[LongType]
+
+        class Person(Schema):
+            job: Column[StructType[Job]]
+    """
+
+    def __init__(
+        self,
+        schema: Type[_Schema],
+        parent: Column,
+    ) -> None:
+        self.schema = schema
+        self.schema._parent = parent
+
+
+class MapType(Generic[_KeyType, _ValueType], TypedSparkDataType):
+    """Allows for type annotations such as.
+
+    .. code-block:: python
+
+        class Basket(Schema):
+            items: Column[MapType[StringType, StringType]]
+    """
+
+
+class ArrayType(Generic[_ValueType], TypedSparkDataType):
+    """Allows for type annotations such as.
+
+    .. code-block:: python
+
+        class Basket(Schema):
+            items: Column[ArrayType[StringType]]
+    """
+
+
+class DecimalType(Generic[_Precision, _Scale], TypedSparkDataType):
+    """Allows for type annotations such as.
+
+    .. code-block:: python
+
+        class Numbers(Schema):
+            number: Column[DecimalType[Literal[10], Literal[0]]]
+    """
+
+
+class DayTimeIntervalType(Generic[_StartField, _EndField], TypedSparkDataType):
+    """Allows for type annotations such as.
+
+    .. code-block:: python
+
+        class TimeInterval(Schema):
+            interval: Column[DayTimeIntervalType[IntervalType.HOUR, IntervalType.SECOND]
+    """

typedspark/_core/literaltype.py

@@ -0,0 +1,41 @@
+"""Defines ``LiteralTypes``, e.g. ``IntervalType.DAY``, that map their class attribute
+to a ``Literal`` integer.
+
+Can be used for example in ``DayTimeIntervalType``.
+"""
+
+from typing import Dict, Literal
+
+
+class LiteralType:
+    """Base class for literal types, that map their class attribute to a Literal
+    integer."""
+
+    @classmethod
+    def get_dict(cls) -> Dict[str, str]:
+        """Returns a dictionary mapping e.g. "IntervalType.DAY" -> "Literal[0]"."""
+        dictionary = {}
+        for key, value in cls.__dict__.items():
+            if key.startswith("_"):
+                continue
+
+            key = f"{cls.__name__}.{key}"
+            value = str(value).replace("typing.", "")
+
+            dictionary[key] = value
+
+        return dictionary
+
+    @classmethod
+    def get_inverse_dict(cls) -> Dict[str, str]:
+        """Returns a dictionary mapping e.g. "Literal[0]" -> "IntervalType.DAY"."""
+        return {v: k for k, v in cls.get_dict().items()}
+
+
+class IntervalType(LiteralType):
+    """Interval types for ``DayTimeIntervalType``."""
+
+    DAY = Literal[0]
+    HOUR = Literal[1]
+    MINUTE = Literal[2]
+    SECOND = Literal[3]