dataframely 1.7.3__cp310-abi3-win_amd64.whl

dataframely/__init__.py

@@ -0,0 +1,100 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+import importlib.metadata
+import warnings
+
+try:
+    __version__ = importlib.metadata.version(__name__)
+except importlib.metadata.PackageNotFoundError as e:  # pragma: no cover
+    warnings.warn(f"Could not determine version of {__name__}\n{e!s}", stacklevel=2)
+    __version__ = "unknown"
+
+from . import random
+from ._base_collection import CollectionMember
+from ._filter import filter
+from ._rule import rule
+from ._typing import DataFrame, LazyFrame, Validation
+from .collection import Collection, deserialize_collection
+from .columns import (
+    Any,
+    Array,
+    Bool,
+    Column,
+    Date,
+    Datetime,
+    Decimal,
+    Duration,
+    Enum,
+    Float,
+    Float32,
+    Float64,
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    Integer,
+    List,
+    Object,
+    String,
+    Struct,
+    Time,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+)
+from .config import Config
+from .failure import FailureInfo
+from .functional import (
+    concat_collection_members,
+    filter_relationship_one_to_at_least_one,
+    filter_relationship_one_to_one,
+)
+from .schema import Schema, deserialize_schema, read_parquet_metadata_schema
+
+__all__ = [
+    "random",
+    "filter",
+    "rule",
+    "DataFrame",
+    "LazyFrame",
+    "Collection",
+    "CollectionMember",
+    "deserialize_collection",
+    "Config",
+    "FailureInfo",
+    "concat_collection_members",
+    "filter_relationship_one_to_at_least_one",
+    "filter_relationship_one_to_one",
+    "Schema",
+    "deserialize_schema",
+    "read_parquet_metadata_schema",
+    "Any",
+    "Bool",
+    "Column",
+    "Date",
+    "Datetime",
+    "Decimal",
+    "Duration",
+    "Time",
+    "Enum",
+    "Float",
+    "Float32",
+    "Float64",
+    "Int8",
+    "Int16",
+    "Int32",
+    "Int64",
+    "Integer",
+    "UInt8",
+    "UInt16",
+    "UInt32",
+    "UInt64",
+    "String",
+    "Struct",
+    "List",
+    "Array",
+    "Object",
+    "Validation",
+]

dataframely/_base_collection.py

@@ -0,0 +1,343 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import textwrap
+import typing
+from abc import ABCMeta
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from typing import Annotated, Any, cast, get_args, get_origin
+
+import polars as pl
+from typing_extensions import Self
+
+from ._filter import Filter
+from ._typing import LazyFrame as TypedLazyFrame
+from .exc import AnnotationImplementationError, ImplementationError
+from .schema import Schema
+
+_MEMBER_ATTR = "__dataframely_members__"
+_FILTER_ATTR = "__dataframely_filters__"
+
+
+@dataclass(kw_only=True)
+class CollectionMember:
+    """An annotation class that configures different behavior for a collection member.
+
+    Members:
+        ignored_in_filters: Indicates that a member should be ignored in the
+            ``@dy.filter`` methods of a collection. This also affects the computation
+            of the shared primary key in the collection.
+
+    Example:
+        .. code:: python
+
+            class MyCollection(dy.Collection):
+                a: dy.LazyFrame[MySchema1]
+                b: dy.LazyFrame[MySchema2]
+
+                ignored_member: Annotated[
+                    dy.LazyFrame[MySchema3],
+                    dy.CollectionMember(ignored_in_filters=True)
+                ]
+
+                @dy.filter
+                def my_filter(self) -> pl.DataFrame:
+                    return self.a.join(self.b, on="shared_key")
+    """
+
+    #: Whether the member should be ignored in the filter method.
+    ignored_in_filters: bool = False
+    #: Whether the member's non-primary key columns should be inlined for sampling.
+    #: This means that value overrides are supplied on the top-level rather than in
+    #: a subkey with the member's name. Only valid if the member's primary key matches
+    #: the collection's common primary key. Two members that share common column names
+    #: may not both be inlined for sampling.
+    inline_for_sampling: bool = False
+
+
+# --------------------------------------- UTILS -------------------------------------- #
+
+
+def _common_primary_keys(columns: Iterable[type[Schema]]) -> set[str]:
+    return set.intersection(*[set(schema.primary_keys()) for schema in columns])
+
+
+# ------------------------------------------------------------------------------------ #
+#                                    COLLECTION META                                   #
+# ------------------------------------------------------------------------------------ #
+
+
+@dataclass
+class MemberInfo(CollectionMember):
+    """Information about a member of a collection."""
+
+    #: The schema of the member.
+    schema: type[Schema]
+    #: Whether the member is optional.
+    is_optional: bool
+
+
+@dataclass
+class Metadata:
+    """Utility class to gather members and filters associated with a collection."""
+
+    members: dict[str, MemberInfo] = field(default_factory=dict)
+    filters: dict[str, Filter] = field(default_factory=dict)
+
+    def update(self, other: Self) -> None:
+        self.members.update(other.members)
+        self.filters.update(other.filters)
+
+
+class CollectionMeta(ABCMeta):
+    def __new__(
+        mcs,  # noqa: N804
+        name: str,
+        bases: tuple[type[object], ...],
+        namespace: dict[str, Any],
+        *args: Any,
+        **kwargs: Any,
+    ) -> CollectionMeta:
+        result = Metadata()
+        for base in bases:
+            result.update(mcs._get_metadata_recursively(base))
+        result.update(mcs._get_metadata(namespace))
+        namespace[_MEMBER_ATTR] = result.members
+        namespace[_FILTER_ATTR] = result.filters
+
+        # We now have all necessary information about filters and members. We want to
+        # check some preconditions to not run into issues later...
+
+        non_ignored_member_schemas = [
+            m.schema for m in result.members.values() if not m.ignored_in_filters
+        ]
+
+        # 1) Check that there are overlapping primary keys that allow the application
+        # of filters.
+        if len(non_ignored_member_schemas) > 0 and len(result.filters) > 0:
+            if len(_common_primary_keys(non_ignored_member_schemas)) == 0:
+                raise ImplementationError(
+                    "Members of a collection must have an overlapping primary key "
+                    "but did not find any."
+                )
+
+        # 2) Check that filter names do not overlap with any column or rule names
+        if len(result.members) > 0:
+            taken = set.union(
+                *(
+                    set(member.schema.column_names())
+                    for member in result.members.values()
+                ),
+                *(
+                    set(member.schema._validation_rules())
+                    for member in result.members.values()
+                ),
+            )
+            intersection = taken & set(result.filters)
+            if len(intersection) > 0:
+                raise ImplementationError(
+                    "Filters defined on the collection must not be named the same as any "
+                    "column or rule in any of the member frames but found "
+                    f"{len(intersection)} such filters: {sorted(intersection)}."
+                )
+
+        # 3) Check that inlining for sampling is configured correctly.
+        if len(non_ignored_member_schemas) > 0:
+            common_primary_keys = _common_primary_keys(non_ignored_member_schemas)
+            inlined_columns: set[str] = set()
+            for member, info in result.members.items():
+                if info.inline_for_sampling:
+                    if set(info.schema.primary_keys()) != common_primary_keys:
+                        raise ImplementationError(
+                            f"Member '{member}' is inlined for sampling but its primary "
+                            "key is a superset of the common primary key. Such a member "
+                            "must not be inlined to be able to provide multiple values "
+                            "for a single combination of the common primary key."
+                        )
+                    non_primary_key_columns = (
+                        set(info.schema.column_names()) - common_primary_keys
+                    )
+                    if len(inlined_columns & non_primary_key_columns):
+                        raise ImplementationError(
+                            f"At least one column name of member '{member}' clashes "
+                            "with a column name of another member that is inlined for "
+                            "sampling."
+                        )
+                    inlined_columns.update(non_primary_key_columns)
+
+        return super().__new__(mcs, name, bases, namespace, *args, **kwargs)
+
+    @staticmethod
+    def _get_metadata_recursively(kls: type[object]) -> Metadata:
+        result = Metadata()
+        for base in kls.__bases__:
+            result.update(CollectionMeta._get_metadata_recursively(base))
+        result.update(CollectionMeta._get_metadata(kls.__dict__))  # type: ignore
+        return result
+
+    @staticmethod
+    def _get_metadata(source: dict[str, Any]) -> Metadata:
+        result = Metadata()
+
+        # Get all members via the annotations
+        if "__annotations__" in source:
+            for attr, kls in source["__annotations__"].items():
+                result.members[attr] = CollectionMeta._derive_member_info(
+                    attr, kls, CollectionMember()
+                )
+
+        # Get all filters by traversing the source
+        for attr, value in {
+            k: v for k, v in source.items() if not k.startswith("__")
+        }.items():
+            if isinstance(value, Filter):
+                result.filters[attr] = value
+
+        return result
+
+    @staticmethod
+    def _derive_member_info(
+        attr: str, type_annotation: Any, collection_member: CollectionMember
+    ) -> MemberInfo:
+        origin = get_origin(type_annotation)
+
+        if origin is None:
+            # `None` annotation is not allowed
+            raise AnnotationImplementationError(attr, type_annotation)
+        elif origin == Annotated:
+            # Maybe happy path: annotated member, dispatch recursively
+            annotation_args = cast(list[Any], get_args(type_annotation))
+            if len(annotation_args) > 2:
+                raise AnnotationImplementationError(attr, type_annotation)
+            if not isinstance(annotation_args[1], CollectionMember):
+                raise AnnotationImplementationError(attr, type_annotation)
+            return CollectionMeta._derive_member_info(
+                attr, annotation_args[0], annotation_args[1]
+            )
+        elif origin == typing.Union:
+            # Happy path: optional member
+            union_args = get_args(type_annotation)
+            if len(union_args) != 2:
+                raise AnnotationImplementationError(attr, type_annotation)
+            if not any(get_origin(arg) is None for arg in union_args):
+                raise AnnotationImplementationError(attr, type_annotation)
+
+            [not_none_arg] = [arg for arg in union_args if get_origin(arg) is not None]
+            if not issubclass(get_origin(not_none_arg), TypedLazyFrame):
+                raise AnnotationImplementationError(attr, type_annotation)
+
+            return MemberInfo(
+                schema=get_args(not_none_arg)[0],
+                is_optional=True,
+                ignored_in_filters=collection_member.ignored_in_filters,
+                inline_for_sampling=collection_member.inline_for_sampling,
+            )
+        elif issubclass(origin, TypedLazyFrame):
+            # Happy path: required member
+            return MemberInfo(
+                schema=get_args(type_annotation)[0],
+                is_optional=False,
+                ignored_in_filters=collection_member.ignored_in_filters,
+                inline_for_sampling=collection_member.inline_for_sampling,
+            )
+        else:
+            # Some other unknown annotation
+            raise AnnotationImplementationError(attr, type_annotation)
+
+    def __repr__(cls) -> str:
+        parts = [f'[Collection "{cls.__class__.__name__}"]']
+        parts.append(textwrap.indent("Members:", prefix=" " * 2))
+        for name, member in cls.members().items():  # type: ignore
+            parts.append(
+                textwrap.indent(
+                    f'- "{name}": {member.schema.__name__}'
+                    f"(optional={member.is_optional}, "
+                    f"ignored_in_filters={member.ignored_in_filters}, "
+                    f"inline_for_sampling={member.inline_for_sampling})",
+                    prefix=" " * 4,
+                )
+            )
+        if filters := cls._filters():  # type: ignore
+            parts.append(textwrap.indent("Filters:", prefix=" " * 2))
+            for name, member in filters.items():
+                parts.append(textwrap.indent(f'- "{name}":', prefix=" " * 4))
+                parts.append(
+                    textwrap.indent(
+                        f"{member.logic(cls.create_empty()).explain()}",  # type: ignore
+                        prefix=" " * 8,
+                    )
+                )
+        parts.append("")  # Add line break at the end
+        return "\n".join(parts)
+
+
+class BaseCollection(metaclass=CollectionMeta):
+    """Internal utility abstraction to reference collections without introducing
+    cyclical dependencies."""
+
+    @classmethod
+    def members(cls) -> dict[str, MemberInfo]:
+        """Information about the members of the collection."""
+        return getattr(cls, _MEMBER_ATTR)
+
+    @classmethod
+    def member_schemas(cls) -> dict[str, type[Schema]]:
+        """The schemas of all members of the collection."""
+        return {name: member.schema for name, member in cls.members().items()}
+
+    @classmethod
+    def required_members(cls) -> set[str]:
+        """The names of all required members of the collection."""
+        return {
+            name for name, member in cls.members().items() if not member.is_optional
+        }
+
+    @classmethod
+    def optional_members(cls) -> set[str]:
+        """The names of all optional members of the collection."""
+        return {name for name, member in cls.members().items() if member.is_optional}
+
+    @classmethod
+    def ignored_members(cls) -> set[str]:
+        """The names of all members of the collection that are ignored in filters."""
+        return {
+            name for name, member in cls.members().items() if member.ignored_in_filters
+        }
+
+    @classmethod
+    def non_ignored_members(cls) -> set[str]:
+        """The names of all members of the collection that are not ignored in filters
+        (default)."""
+        return {
+            name
+            for name, member in cls.members().items()
+            if not member.ignored_in_filters
+        }
+
+    @classmethod
+    def common_primary_keys(cls) -> list[str]:
+        """The primary keys shared by non ignored members of the collection."""
+        return sorted(
+            _common_primary_keys(
+                [
+                    member.schema
+                    for member in cls.members().values()
+                    if not member.ignored_in_filters
+                ]
+            )
+        )
+
+    @classmethod
+    def _filters(cls) -> dict[str, Filter[Self]]:
+        return getattr(cls, _FILTER_ATTR)
+
+    def to_dict(self) -> dict[str, pl.LazyFrame]:
+        """Return a dictionary representation of this collection."""
+        return {
+            member: getattr(self, member)
+            for member in self.member_schemas()
+            if getattr(self, member) is not None
+        }

dataframely/_base_schema.py

@@ -0,0 +1,192 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import textwrap
+from abc import ABCMeta
+from copy import copy
+from dataclasses import dataclass, field
+from typing import Any
+
+import polars as pl
+from typing_extensions import Self
+
+from ._rule import GroupRule, Rule
+from .columns import Column
+from .exc import ImplementationError
+
+_COLUMN_ATTR = "__dataframely_columns__"
+_RULE_ATTR = "__dataframely_rules__"
+
+# --------------------------------------- UTILS -------------------------------------- #
+
+
+def _build_rules(
+    custom: dict[str, Rule], columns: dict[str, Column]
+) -> dict[str, Rule]:
+    # NOTE: Copy here to prevent in-place modification of the custom rules
+    rules: dict[str, Rule] = copy(custom)
+
+    # Add primary key validation to the list of rules if applicable
+    primary_keys = _primary_keys(columns)
+    if len(primary_keys) > 0:
+        rules["primary_key"] = Rule(~pl.struct(primary_keys).is_duplicated())
+
+    # Add column-specific rules
+    column_rules = {
+        f"{col_name}|{rule_name}": Rule(expr)
+        for col_name, column in columns.items()
+        for rule_name, expr in column.validation_rules(pl.col(col_name)).items()
+    }
+    rules.update(column_rules)
+
+    return rules
+
+
+def _primary_keys(columns: dict[str, Column]) -> list[str]:
+    return list(k for k, col in columns.items() if col.primary_key)
+
+
+# ------------------------------------------------------------------------------------ #
+#                                      SCHEMA META                                     #
+# ------------------------------------------------------------------------------------ #
+
+
+@dataclass
+class Metadata:
+    """Utility class to gather columns and rules associated with a schema."""
+
+    columns: dict[str, Column] = field(default_factory=dict)
+    rules: dict[str, Rule] = field(default_factory=dict)
+
+    def update(self, other: Self) -> None:
+        self.columns.update(other.columns)
+        self.rules.update(other.rules)
+
+
+class SchemaMeta(ABCMeta):
+    def __new__(
+        mcs,  # noqa: N804
+        name: str,
+        bases: tuple[type[object], ...],
+        namespace: dict[str, Any],
+        *args: Any,
+        **kwargs: Any,
+    ) -> SchemaMeta:
+        result = Metadata()
+        for base in bases:
+            result.update(mcs._get_metadata_recursively(base))
+        result.update(mcs._get_metadata(namespace))
+        namespace[_COLUMN_ATTR] = result.columns
+        namespace[_RULE_ATTR] = result.rules
+
+        # At this point, we already know all columns and custom rules. We want to run
+        # some checks...
+
+        # 1) Check that the column names clash with none of the rule names. To this end,
+        # we assume that users cast dtypes, i.e. additional rules for dtype casting
+        # are also checked.
+        all_column_names = set(result.columns)
+        all_rule_names = set(_build_rules(result.rules, result.columns).keys()) | set(
+            f"{col}|dtype" for col in result.columns
+        )
+        common_names = all_column_names & all_rule_names
+        if len(common_names) > 0:
+            common_list = ", ".join(sorted(f"'{col}'" for col in common_names))
+            raise ImplementationError(
+                "Rules and columns must not be named equally but found "
+                f"{len(common_names)} overlaps: {common_list}."
+            )
+
+        # 2) Check that the columns referenced in the group rules exist.
+        for rule_name, rule in result.rules.items():
+            if isinstance(rule, GroupRule):
+                missing_columns = set(rule.group_columns) - set(result.columns)
+                if len(missing_columns) > 0:
+                    missing_list = ", ".join(
+                        sorted(f"'{col}'" for col in missing_columns)
+                    )
+                    raise ImplementationError(
+                        f"Group validation rule '{rule_name}' has been implemented "
+                        f"incorrectly. It references {len(missing_columns)} columns "
+                        f"which are not in the schema: {missing_list}."
+                    )
+
+        return super().__new__(mcs, name, bases, namespace, *args, **kwargs)
+
+    def __getattribute__(cls, name: str) -> Any:
+        val = super().__getattribute__(name)
+        # Dynamically set the name of the column if it is a `Column` instance.
+        if isinstance(val, Column):
+            val._name = val.alias or name
+        return val
+
+    @staticmethod
+    def _get_metadata_recursively(kls: type[object]) -> Metadata:
+        result = Metadata()
+        for base in kls.__bases__:
+            result.update(SchemaMeta._get_metadata_recursively(base))
+        result.update(SchemaMeta._get_metadata(kls.__dict__))  # type: ignore
+        return result
+
+    @staticmethod
+    def _get_metadata(source: dict[str, Any]) -> Metadata:
+        result = Metadata()
+        for attr, value in {
+            k: v for k, v in source.items() if not k.startswith("__")
+        }.items():
+            if isinstance(value, Column):
+                result.columns[value.alias or attr] = value
+            if isinstance(value, Rule):
+                # We must ensure that custom rules do not clash with internal rules.
+                if attr == "primary_key":
+                    raise ImplementationError(
+                        "Custom validation rule must not be named `primary_key`."
+                    )
+                result.rules[attr] = value
+        return result
+
+    def __repr__(cls) -> str:
+        parts = [f'[Schema "{cls.__name__}"]']
+        parts.append(textwrap.indent("Columns:", prefix=" " * 2))
+        for name, col in cls.columns().items():
+            parts.append(textwrap.indent(f'- "{name}": {col!r}', prefix=" " * 4))
+        if validation_rules := cls._schema_validation_rules():
+            parts.append(textwrap.indent("Rules:", prefix=" " * 2))
+            for name, rule in validation_rules.items():
+                parts.append(textwrap.indent(f'- "{name}": {rule!r}', prefix=" " * 4))
+        parts.append("")  # Add line break at the end
+        return "\n".join(parts)
+
+
+class BaseSchema(metaclass=SchemaMeta):
+    """Internal utility abstraction to reference schemas without introducing cyclical
+    dependencies."""
+
+    @classmethod
+    def column_names(cls) -> list[str]:
+        """The column names of this schema."""
+        return list(getattr(cls, _COLUMN_ATTR).keys())
+
+    @classmethod
+    def columns(cls) -> dict[str, Column]:
+        """The column definitions of this schema."""
+        columns: dict[str, Column] = getattr(cls, _COLUMN_ATTR)
+        for name in columns.keys():
+            # Dynamically set the name of the columns.
+            columns[name]._name = name
+        return columns
+
+    @classmethod
+    def primary_keys(cls) -> list[str]:
+        """The primary key columns in this schema (possibly empty)."""
+        return _primary_keys(cls.columns())
+
+    @classmethod
+    def _validation_rules(cls) -> dict[str, Rule]:
+        return _build_rules(cls._schema_validation_rules(), cls.columns())
+
+    @classmethod
+    def _schema_validation_rules(cls) -> dict[str, Rule]:
+        return getattr(cls, _RULE_ATTR)

dataframely/_compat.py

@@ -0,0 +1,39 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+
+from typing import Any
+
+
+class _DummyModule:  # pragma: no cover
+    def __init__(self, module: str) -> None:
+        self.module = module
+
+    def __getattr__(self, name: str) -> Any:
+        raise ValueError(f"Module '{self.module}' is not installed.")
+
+
+# ------------------------------------ SQLALCHEMY ------------------------------------ #
+
+try:
+    import sqlalchemy as sa
+    import sqlalchemy.dialects.mssql as sa_mssql
+    from sqlalchemy.sql.type_api import TypeEngine as sa_TypeEngine
+except ImportError:  # pragma: no cover
+    sa = _DummyModule("sqlalchemy")  # type: ignore
+    sa_mssql = _DummyModule("sqlalchemy")  # type: ignore
+
+    class sa_TypeEngine:  # type: ignore # noqa: N801
+        pass
+
+
+# -------------------------------------- PYARROW ------------------------------------- #
+
+try:
+    import pyarrow as pa
+except ImportError:  # pragma: no cover
+    pa = _DummyModule("pyarrow")
+
+# ------------------------------------------------------------------------------------ #
+
+__all__ = ["sa", "sa_mssql", "sa_TypeEngine", "pa"]