dataframely 1.2.1__tar.gz → 1.3.1__tar.gz

{dataframely-1.2.1 → dataframely-1.3.1}/.copier-answers.yml

@@ -1,7 +1,7 @@
 # This file is managed by Copier; DO NOT EDIT OR REMOVE.
 _commit: v0.3.0
 _src_path: https://github.com/quantco/copier-template-python-open-source
-add_autobump_workflow: true
+add_autobump_workflow: false
 author_email: oliver.borchert@quantco.com
 author_name: Oliver Borchert
 github_url: https://github.com/quantco/dataframely

{dataframely-1.2.1 → dataframely-1.3.1}/.github/workflows/build.yml

@@ -78,7 +78,7 @@ jobs:
       id-token: write
     environment: pypi
     steps:
-      - uses: actions/download-artifact@95815c38cf2ff2164869cbab79da8d1f422bc89e # v4.2.1
+      - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0
         with:
           path: dist
           merge-multiple: true

{dataframely-1.2.1 → dataframely-1.3.1}/Cargo.toml

@@ -12,3 +12,7 @@ pyo3 = { version = "0.24", features = ["abi3-py311", "extension-module"] }
 rand = { version = "0.9", features = ["std_rng"] }
 regex-syntax = "0.8"
 thiserror = "2.0"
+
+[profile.release]
+codegen-units = 1
+lto = true

{dataframely-1.2.1 → dataframely-1.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataframely
-Version: 1.2.1
+Version: 1.3.1
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12

{dataframely-1.2.1 → dataframely-1.3.1}/dataframely/__init__.py

@@ -18,6 +18,7 @@ from ._typing import DataFrame, LazyFrame
 from .collection import Collection
 from .columns import (
     Any,
+    Array,
     Bool,
     Column,
     Date,
@@ -90,5 +91,6 @@ __all__ = [
     "String",
     "Struct",
     "List",
+    "Array",
     "Object",
 ]

{dataframely-1.2.1 → dataframely-1.3.1}/dataframely/_base_schema.py

@@ -98,7 +98,7 @@ class SchemaMeta(ABCMeta):
             )
 
         # 2) Check that the columns referenced in the group rules exist.
-        for name, rule in result.rules.items():
+        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:
@@ -106,7 +106,7 @@ class SchemaMeta(ABCMeta):
                         sorted(f"'{col}'" for col in missing_columns)
                     )
                     raise ImplementationError(
-                        f"Group validation rule '{name}' has been implemented "
+                        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}."
                     )
@@ -116,16 +116,16 @@ class SchemaMeta(ABCMeta):
         # empty data frame and checking for the evaluated dtypes.
         if len(result.rules) > 0:
             lf_empty = pl.LazyFrame(
-                schema={name: col.dtype for name, col in result.columns.items()}
+                schema={col_name: col.dtype for col_name, col in result.columns.items()}
             )
             # NOTE: For some reason, `polars` does not yield correct dtypes when calling
             #  `collect_schema()`
             schema = with_evaluation_rules(lf_empty, result.rules).collect().schema
-            for name, rule in result.rules.items():
-                dtype = schema[name]
+            for rule_name, rule in result.rules.items():
+                dtype = schema[rule_name]
                 if not isinstance(dtype, pl.Boolean):
                     raise RuleImplementationError(
-                        name, dtype, isinstance(rule, GroupRule)
+                        rule_name, dtype, isinstance(rule, GroupRule)
                     )
 
         return super().__new__(mcs, name, bases, namespace, *args, **kwargs)

{dataframely-1.2.1 → dataframely-1.3.1}/dataframely/_typing.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, Generic, ParamSpec, TypeVar
+from typing import TYPE_CHECKING, Any, Concatenate, Generic, ParamSpec, TypeVar
 
 import polars as pl
 
@@ -48,6 +48,14 @@ class DataFrame(pl.DataFrame, Generic[S]):
     def lazy(self, *args: Any, **kwargs: Any) -> LazyFrame[S]:
         raise NotImplementedError  # pragma: no cover
 
+    def pipe(
+        self,
+        function: Callable[Concatenate[DataFrame[S], P], R],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> R:
+        raise NotImplementedError  # pragma: no cover
+
     @inherit_signature(pl.DataFrame.rechunk)
     def rechunk(self, *args: Any, **kwargs: Any) -> DataFrame[S]:
         raise NotImplementedError  # pragma: no cover
@@ -92,6 +100,14 @@ class LazyFrame(pl.LazyFrame, Generic[S]):
     def lazy(self, *args: Any, **kwargs: Any) -> LazyFrame[S]:
         raise NotImplementedError  # pragma: no cover
 
+    def pipe(
+        self,
+        function: Callable[Concatenate[LazyFrame[S], P], R],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> R:
+        raise NotImplementedError  # pragma: no cover
+
     @inherit_signature(pl.LazyFrame.set_sorted)
     def set_sorted(self, *args: Any, **kwargs: Any) -> LazyFrame[S]:
         raise NotImplementedError  # pragma: no cover

{dataframely-1.2.1 → dataframely-1.3.1}/dataframely/columns/__init__.py

@@ -3,6 +3,7 @@
 
 from ._base import Column
 from .any import Any
+from .array import Array
 from .bool import Bool
 from .datetime import Date, Datetime, Duration, Time
 from .decimal import Decimal
@@ -17,6 +18,7 @@ from .struct import Struct
 __all__ = [
     "Column",
     "Any",
+    "Array",
     "Bool",
     "Date",
     "Datetime",

dataframely-1.3.1/dataframely/columns/array.py

@@ -0,0 +1,105 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Sequence
+from typing import Any, Literal
+
+import polars as pl
+
+from dataframely._compat import pa, sa, sa_TypeEngine
+from dataframely.random import Generator
+
+from ._base import Column
+from .struct import Struct
+
+
+class Array(Column):
+    """A fixed-shape array column."""
+
+    def __init__(
+        self,
+        inner: Column,
+        shape: int | tuple[int, ...],
+        *,
+        nullable: bool = True,
+        # polars doesn't yet support grouping by arrays,
+        # see https://github.com/pola-rs/polars/issues/22574
+        primary_key: Literal[False] = False,
+        check: Callable[[pl.Expr], pl.Expr] | None = None,
+        alias: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ):
+        """
+        Args:
+            inner: The inner column type. No validation rules on the inner type are supported yet.
+            shape: The shape of the array.
+            nullable: Whether this column may contain null values.
+            primary_key: Whether this column is part of the primary key of the schema.
+                Not yet supported for the Array type.
+            check: A custom check to run for this column. Must return a non-aggregated
+                boolean expression.
+            alias: An overwrite for this column's name which allows for using a column
+                name that is not a valid Python identifier. Especially note that setting
+                this option does _not_ allow to refer to the column with two different
+                names, the specified alias is the only valid name.
+            metadata: A dictionary of metadata to attach to the column.
+        """
+        if inner.primary_key or (
+            isinstance(inner, Struct)
+            and any(col.primary_key for col in inner.inner.values())
+        ):
+            raise ValueError(
+                "`primary_key=True` is not yet supported for inner types of the Array type."
+            )
+
+        # We disallow validation rules on the inner type since Polars arrays currently don't support .eval(). Converting
+        # to a list and calling .list.eval() is possible, however, since the shape can have multiple axes, the recursive
+        # conversion could have significant performance impact. Hence, we simply disallow inner validation rules.
+        # Another option would be to allow validation rules only for sampling, but not enforce them.
+        if inner.validation_rules(pl.lit(None)):
+            raise ValueError(
+                "Validation rules on the inner type of Array are not yet supported."
+            )
+
+        super().__init__(
+            nullable=nullable,
+            primary_key=False,
+            check=check,
+            alias=alias,
+            metadata=metadata,
+        )
+        self.inner = inner
+        self.shape = shape if isinstance(shape, tuple) else (shape,)
+
+    @property
+    def dtype(self) -> pl.DataType:
+        return pl.Array(self.inner.dtype, self.shape)
+
+    def sqlalchemy_dtype(self, dialect: sa.Dialect) -> sa_TypeEngine:
+        # NOTE: We might want to add support for PostgreSQL's ARRAY type or use JSON in the future.
+        raise NotImplementedError("SQL column cannot have 'Array' type.")
+
+    def _pyarrow_dtype_of_shape(self, shape: Sequence[int]) -> pa.DataType:
+        if shape:
+            size, *rest = shape
+            return pa.list_(self._pyarrow_dtype_of_shape(rest), size)
+        else:
+            return self.inner.pyarrow_dtype
+
+    @property
+    def pyarrow_dtype(self) -> pa.DataType:
+        return self._pyarrow_dtype_of_shape(self.shape)
+
+    def _sample_unchecked(self, generator: Generator, n: int) -> pl.Series:
+        # Sample the inner elements in a flat series
+        n_elements = n * math.prod(self.shape)
+        all_elements = self.inner.sample(generator, n_elements)
+
+        # Finally, apply a null mask
+        return generator._apply_null_mask(
+            all_elements.reshape((n, *self.shape)),
+            null_probability=self._null_probability,
+        )

{dataframely-1.2.1 → dataframely-1.3.1}/dataframely/columns/list.py

@@ -147,4 +147,7 @@ class List(Column):
                 chain([0], element_lengths.cum_sum()), element_lengths
             )
         ]
-        return pl.Series(list_elements)
+        # Finally, apply a null mask
+        return generator._apply_null_mask(
+            pl.Series(list_elements), null_probability=self._null_probability
+        )

{dataframely-1.2.1 → dataframely-1.3.1}/dataframely/columns/struct.py

@@ -96,10 +96,14 @@ class Struct(Column):
         return pa.struct({name: col.pyarrow_dtype for name, col in self.inner.items()})
 
     def _sample_unchecked(self, generator: Generator, n: int) -> pl.Series:
-        return (
+        series = (
             pl.DataFrame(
                 {name: col.sample(generator, n) for name, col in self.inner.items()}
             )
             .select(pl.struct(pl.all()))
             .to_series()
         )
+        # Apply a null mask.
+        return generator._apply_null_mask(
+            series, null_probability=self._null_probability
+        )

{dataframely-1.2.1 → dataframely-1.3.1}/dataframely/mypy.py

@@ -16,6 +16,7 @@ from mypy.nodes import (
     MemberExpr,
     NameExpr,
     StrExpr,
+    TupleExpr,
     TypeInfo,
     Var,
 )
@@ -189,7 +190,28 @@ def _convert_dy_column_to_dtype(
                 ],
             )
         return api.named_type("builtins.list")
-    if column_type == "Any":
+    if column_type == "Array":
+        if isinstance(column_args[0], CallExpr) and isinstance(
+            column_args[0].callee, MemberExpr | NameExpr
+        ):
+            inner_type = _convert_dy_column_to_dtype(
+                api,
+                column_args[0].callee.name,
+                column_args[0].args,
+            )
+            # If the array has more than one dimension, return a list of lists of the inner type.
+            if len(column_args) > 1 and isinstance(column_args[1], TupleExpr):
+                for _ in range(len(column_args[1].items) - 1):
+                    inner_type = api.named_type(
+                        "builtins.list",
+                        [inner_type],
+                    )
+            return api.named_type(
+                "builtins.list",
+                [inner_type],
+            )
+        return api.named_type("builtins.list")
+    if column_type == "Any" or column_type == "Object":
         return AnyType(TypeOfAny.explicit)
     # If we can't infer the type, we default to `Any`.
     # This is, for example, the case for self-defined types, e.g., via `functools.partial`.

{dataframely-1.2.1 → dataframely-1.3.1}/dataframely/random.py

@@ -47,7 +47,7 @@ class Generator:
         Returns:
             A seed of type ``uint32``.
         """
-        return self.numpy_generator.integers(0, 2**32)
+        return self.numpy_generator.integers(0, 2**32, dtype=int)
 
     # ------------------------------------ NUMBERS ----------------------------------- #
 
@@ -392,8 +392,10 @@ class Generator:
     def _apply_null_mask(self, series: pl.Series, null_probability: float) -> pl.Series:
         if null_probability == 0:
             return series
-        null_mask = self.numpy_generator.random(series.len()) < null_probability
-        return series.scatter(np.where(null_mask)[0], None)
+        null_mask = (
+            pl.Series(self.numpy_generator.random(series.len())) > null_probability
+        )
+        return pl.select(pl.when(null_mask).then(series)).to_series()
 
 
 # --------------------------------------- UTILS -------------------------------------- #

{dataframely-1.2.1 → dataframely-1.3.1}/dataframely/testing/__init__.py

@@ -6,6 +6,7 @@ from .const import (
     COLUMN_TYPES,
     FLOAT_COLUMN_TYPES,
     INTEGER_COLUMN_TYPES,
+    NO_VALIDATION_COLUMN_TYPES,
     SUPERTYPE_COLUMN_TYPES,
 )
 from .factory import create_collection, create_collection_raw, create_schema
@@ -18,6 +19,7 @@ __all__ = [
     "FLOAT_COLUMN_TYPES",
     "INTEGER_COLUMN_TYPES",
     "SUPERTYPE_COLUMN_TYPES",
+    "NO_VALIDATION_COLUMN_TYPES",
     "create_collection",
     "create_collection_raw",
     "create_schema",

{dataframely-1.2.1 → dataframely-1.3.1}/dataframely/testing/const.py

@@ -47,3 +47,8 @@ SUPERTYPE_COLUMN_TYPES: list[type[dc.Column]] = [
 ALL_COLUMN_TYPES: list[type[dc.Column]] = (
     [dc.Any] + COLUMN_TYPES + SUPERTYPE_COLUMN_TYPES
 )
+
+# The following is a list of column types that, when created with default parameter values, add no validation rules.
+NO_VALIDATION_COLUMN_TYPES: list[type[dc.Column]] = [
+    t for t in ALL_COLUMN_TYPES if t not in FLOAT_COLUMN_TYPES
+]