patito 0.4.3__py3-none-any.whl → 0.5.0__py3-none-any.whl

patito/polars.py

@@ -4,10 +4,10 @@ from __future__ import annotations
 from typing import (
     TYPE_CHECKING,
     Any,
+    Collection,
     Generic,
-    List,
+    Iterable,
     Optional,
-    Sequence,
     Type,
     TypeVar,
     Union,
@@ -15,6 +15,7 @@ from typing import (
 )
 
 import polars as pl
+from polars.type_aliases import IntoExpr
 from pydantic import create_model
 from typing_extensions import Literal
 
@@ -71,9 +72,10 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
         predicate_pushdown: bool = True,
         projection_pushdown: bool = True,
         simplify_expression: bool = True,
-        string_cache: bool = False,
         no_optimization: bool = False,
         slice_pushdown: bool = True,
+        common_subplan_elimination: bool = True,
+        streaming: bool = False,
     ) -> "DataFrame[ModelType]":  # noqa: DAR101, DAR201
         """
         Collect into a DataFrame.
@@ -86,9 +88,10 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
             predicate_pushdown=predicate_pushdown,
             projection_pushdown=projection_pushdown,
             simplify_expression=simplify_expression,
-            string_cache=string_cache,
             no_optimization=no_optimization,
             slice_pushdown=slice_pushdown,
+            common_subplan_elimination=common_subplan_elimination,
+            streaming=streaming,
         )
         if getattr(self, "model", False):
             cls = DataFrame._construct_dataframe_model_class(model=self.model)
@@ -207,11 +210,8 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             │ i64  ┆ str    │
             ╞══════╪════════╡
             │ 1    ┆ A      │
-            ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
             │ 1    ┆ B      │
-            ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
             │ 2    ┆ A      │
-            ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
             │ 2    ┆ B      │
             └──────┴────────┘
             >>> casted_classes = classes.cast()
@@ -223,11 +223,8 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             │ u16  ┆ cat    │
             ╞══════╪════════╡
             │ 1    ┆ A      │
-            ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
             │ 1    ┆ B      │
-            ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
             │ 2    ┆ A      │
-            ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
             │ 2    ┆ B      │
             └──────┴────────┘
             >>> casted_classes.validate()
@@ -290,7 +287,11 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
                 columns.append(pl.col(column).cast(default_dtypes[column]))
         return self.with_columns(columns)
 
-    def drop(self: DF, name: Optional[Union[str, List[str]]] = None) -> DF:
+    def drop(
+        self: DF,
+        columns: Optional[Union[str, Collection[str]]] = None,
+        *more_columns: str,
+    ) -> DF:
         """
         Drop one or more columns from the dataframe.
 
@@ -299,9 +300,10 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         :ref:`DataFrame.set_model <DataFrame.set_model>`, are dropped.
 
         Args:
-            name: A single column string name, or list of strings, indicating
+            columns: A single column string name, or list of strings, indicating
                 which columns to drop. If not specified, all columns *not*
                 specified by the associated dataframe model will be dropped.
+            more_columns: Additional named columns to drop.
 
         Returns:
             DataFrame[Model]: New dataframe without the specified columns.
@@ -319,13 +321,12 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             │ i64      │
             ╞══════════╡
             │ 1        │
-            ├╌╌╌╌╌╌╌╌╌╌┤
             │ 2        │
             └──────────┘
 
         """
-        if name is not None:
-            return super().drop(name)
+        if columns is not None:
+            return self._from_pydf(super().drop(columns)._df)
         else:
             return self.drop(list(set(self.columns) - set(self.model.columns)))
 
@@ -416,7 +417,6 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             │ i64 ┆ i64 ┆ i64        │
             ╞═════╪═════╪════════════╡
             │ 1   ┆ 1   ┆ 2          │
-            ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
             │ 2   ┆ 2   ┆ 4          │
             └─────┴─────┴────────────┘
         """
@@ -426,11 +426,11 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
                 derived_from = props["derived_from"]
                 dtype = self.model.dtypes[column_name]
                 if isinstance(derived_from, str):
-                    df = df.with_column(
+                    df = df.with_columns(
                         pl.col(derived_from).cast(dtype).alias(column_name)
                     )
                 elif isinstance(derived_from, pl.Expr):
-                    df = df.with_column(derived_from.cast(dtype).alias(column_name))
+                    df = df.with_columns(derived_from.cast(dtype).alias(column_name))
                 else:
                     raise TypeError(
                         "Can not derive dataframe column from type "
@@ -447,6 +447,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             ]
         ] = None,
         limit: Optional[int] = None,
+        matches_supertype: bool = True,
     ) -> DF:
         """
         Fill null values using a filling strategy, literal, or ``Expr``.
@@ -461,6 +462,8 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
                 provided.
             limit: The number of consecutive null values to forward/backward fill.
                 Only valid if ``strategy`` is ``"forward"`` or ``"backward"``.
+            matches_supertype: Fill all matching supertype of the fill ``value``.
+
 
         Returns:
             DataFrame[Model]: A new dataframe with nulls filled in according to the
@@ -483,13 +486,18 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             │ str    ┆ i64   │
             ╞════════╪═══════╡
             │ apple  ┆ 10    │
-            ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
             │ banana ┆ 19    │
             └────────┴───────┘
         """
         if strategy != "defaults":  # pragma: no cover
-            return cast(  # type: ignore[redundant-cast]
-                DF, super().fill_null(value=value, strategy=strategy, limit=limit)
+            return cast(  # pyright: ignore[redundant-cast]
+                DF,
+                super().fill_null(
+                    value=value,
+                    strategy=strategy,
+                    limit=limit,
+                    matches_supertype=matches_supertype,
+                ),
             )
         return self.with_columns(
             [
@@ -596,7 +604,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             create_model(  # type: ignore
                 "UntypedRow",
                 __base__=Model,
-                **pydantic_annotations,
+                **pydantic_annotations,  # pyright: ignore
             ),
         )
 
@@ -651,15 +659,17 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             ...     b: str = pt.Field(derived_from="source_of_b")
             ...
             >>> csv_file = io.StringIO("a,source_of_b\n1,1")
-            >>> CSVModel.DataFrame.read_csv(csv_file).drop()
-            shape: (1, 2)
-            ┌─────┬─────┐
-            │ a   ┆ b   │
-            │ --- ┆ --- │
-            │ f64 ┆ str │
-            ╞═════╪═════╡
-            │ 1.0 ┆ 1   │
-            └─────┴─────┘
+
+
+            # >>> CSVModel.DataFrame.read_csv(csv_file).drop()
+            # shape: (1, 2)
+            # ┌─────┬─────┐
+            # │ a   ┆ b   │
+            # │ --- ┆ --- │
+            # │ f64 ┆ str │
+            # ╞═════╪═════╡
+            # │ 1.0 ┆ 1   │
+            # └─────┴─────┘
         """
         kwargs.setdefault("dtypes", cls.model.dtypes)
         if not kwargs.get("has_header", True) and "columns" not in kwargs:
@@ -670,27 +680,24 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
     # --- Type annotation overrides ---
     def filter(  # noqa: D102
         self: DF,
-        predicate: Union[pl.Expr, str, pl.Series, list[bool], np.ndarray[Any, Any]],
+        predicate: Union[
+            pl.Expr, str, pl.Series, list[bool], np.ndarray[Any, Any], bool
+        ],
     ) -> DF:
         return cast(DF, super().filter(predicate=predicate))
 
     def select(  # noqa: D102
         self: DF,
-        exprs: Union[str, pl.Expr, pl.Series, Sequence[Union[str, pl.Expr, pl.Series]]],
+        *exprs: Union[IntoExpr, Iterable[IntoExpr]],
+        **named_exprs: IntoExpr,
     ) -> DF:
-        return cast(DF, super().select(exprs=exprs))  # type: ignore[redundant-cast]
-
-    def with_column(self: DF, column: Union[pl.Series, pl.Expr]) -> DF:  # noqa: D102
-        return cast(DF, super().with_column(column=column))
+        return cast(  # pyright: ignore[redundant-cast]
+            DF, super().select(*exprs, **named_exprs)
+        )
 
     def with_columns(  # noqa: D102
         self: DF,
-        exprs: Union[
-            pl.Expr,
-            pl.Series,
-            Sequence[Union[pl.Expr, pl.Series]],
-            None,
-        ] = None,
-        **named_exprs: Union[pl.Expr, pl.Series],
+        *exprs: Union[IntoExpr, Iterable[IntoExpr]],
+        **named_exprs: IntoExpr,
     ) -> DF:
-        return cast(DF, super().with_columns(exprs=exprs, **named_exprs))
+        return cast(DF, super().with_columns(*exprs, **named_exprs))

patito/pydantic.py

@@ -19,6 +19,7 @@ from typing import (
 )
 
 import polars as pl
+from polars.datatypes import PolarsDataType
 from pydantic import BaseConfig, BaseModel, Field, create_model  # noqa: F401
 from pydantic.main import ModelMetaclass as PydanticModelMetaclass
 from typing_extensions import Literal, get_args
@@ -110,7 +111,7 @@ class ModelMetaclass(PydanticModelMetaclass):
 
     @property
     def dtypes(  # type: ignore
-        cls: Type[ModelType],
+        cls: Type[ModelType],  # pyright: ignore
     ) -> dict[str, Type[pl.DataType]]:
         """
         Return the polars dtypes of the dataframe.
@@ -129,18 +130,16 @@ class ModelMetaclass(PydanticModelMetaclass):
             ...     price: float
             ...
             >>> Product.dtypes
-            {'name': <class 'polars.datatypes.Utf8'>, \
-'ideal_temperature': <class 'polars.datatypes.Int64'>, \
-'price': <class 'polars.datatypes.Float64'>}
+            {'name': Utf8, 'ideal_temperature': Int64, 'price': Float64}
         """
         return {
             column: valid_dtypes[0] for column, valid_dtypes in cls.valid_dtypes.items()
         }
 
     @property
-    def valid_dtypes(  # type: ignore  # noqa: C901
-        cls: Type[ModelType],
-    ) -> dict[str, List[Type[pl.DataType]]]:
+    def valid_dtypes(  # type: ignore
+        cls: Type[ModelType],  # pyright: ignore
+    ) -> dict[str, List[Union[pl.PolarsDataType, pl.List]]]:
         """
         Return a list of polars dtypes which Patito considers valid for each field.
 
@@ -164,82 +163,91 @@ class ModelMetaclass(PydanticModelMetaclass):
             ...     float_column: float
             ...
             >>> pprint(MyModel.valid_dtypes)
-            {'bool_column': [<class 'polars.datatypes.Boolean'>],
-             'float_column': [<class 'polars.datatypes.Float64'>,
-                              <class 'polars.datatypes.Float32'>],
-             'int_column': [<class 'polars.datatypes.Int64'>,
-                            <class 'polars.datatypes.Int32'>,
-                            <class 'polars.datatypes.Int16'>,
-                            <class 'polars.datatypes.Int8'>,
-                            <class 'polars.datatypes.UInt64'>,
-                            <class 'polars.datatypes.UInt32'>,
-                            <class 'polars.datatypes.UInt16'>,
-                            <class 'polars.datatypes.UInt8'>],
-             'str_column': [<class 'polars.datatypes.Utf8'>]}
+            {'bool_column': [Boolean],
+             'float_column': [Float64, Float32],
+             'int_column': [Int64, Int32, Int16, Int8, UInt64, UInt32, UInt16, UInt8],
+             'str_column': [Utf8]}
         """
         valid_dtypes = {}
         for column, props in cls._schema_properties().items():
-            if "dtype" in props:
-                valid_dtypes[column] = [
-                    props["dtype"],
-                ]
-            elif "enum" in props and props["type"] == "string":
-                valid_dtypes[column] = [pl.Categorical, pl.Utf8]
-            elif "type" not in props:
+            column_dtypes: List[Union[PolarsDataType, pl.List]]
+            if props.get("type") == "array":
+                array_props = props["items"]
+                item_dtypes = cls._valid_dtypes(props=array_props)
+                if item_dtypes is None:
+                    raise NotImplementedError(
+                        f"No valid dtype mapping found for column '{column}'."
+                    )
+                column_dtypes = [pl.List(dtype) for dtype in item_dtypes]
+            else:
+                column_dtypes = cls._valid_dtypes(props=props)  # pyright: ignore
+
+            if column_dtypes is None:
                 raise NotImplementedError(
                     f"No valid dtype mapping found for column '{column}'."
                 )
-            elif props["type"] == "integer":
-                valid_dtypes[column] = [
-                    pl.Int64,
-                    pl.Int32,
-                    pl.Int16,
-                    pl.Int8,
-                    pl.UInt64,
-                    pl.UInt32,
-                    pl.UInt16,
-                    pl.UInt8,
-                ]
-            elif props["type"] == "number":
-                if props.get("format") == "time-delta":
-                    valid_dtypes[column] = [
-                        pl.Duration,
-                    ]  # pyright: reportPrivateImportUsage=false
-                else:
-                    valid_dtypes[column] = [pl.Float64, pl.Float32]
-            elif props["type"] == "boolean":
-                valid_dtypes[column] = [
-                    pl.Boolean,
-                ]
-            elif props["type"] == "string":
-                string_format = props.get("format")
-                if string_format is None:
-                    valid_dtypes[column] = [
-                        pl.Utf8,
-                    ]
-                elif string_format == "date":
-                    valid_dtypes[column] = [
-                        pl.Date,
-                    ]
-                # TODO: Find out why this branch is not being hit
-                elif string_format == "date-time":  # pragma: no cover
-                    valid_dtypes[column] = [
-                        pl.Datetime,
-                    ]
-            elif props["type"] == "null":
-                valid_dtypes[column] = [
-                    pl.Null,
-                ]
-            else:  # pragma: no cover
-                raise NotImplementedError(
-                    f"No valid dtype mapping found for column '{column}'"
-                )
+            valid_dtypes[column] = column_dtypes
 
         return valid_dtypes
 
+    @staticmethod
+    def _valid_dtypes(  # noqa: C901
+        props: Dict,
+    ) -> Optional[List[pl.PolarsDataType]]:
+        """
+        Map schema property to list of valid polars data types.
+
+        Args:
+            props: Dictionary value retrieved from BaseModel._schema_properties().
+
+        Returns:
+            List of valid dtypes. None if no mapping exists.
+        """
+        if "dtype" in props:
+            return [
+                props["dtype"],
+            ]
+        elif "enum" in props and props["type"] == "string":
+            return [pl.Categorical, pl.Utf8]
+        elif "type" not in props:
+            return None
+        elif props["type"] == "integer":
+            return [
+                pl.Int64,
+                pl.Int32,
+                pl.Int16,
+                pl.Int8,
+                pl.UInt64,
+                pl.UInt32,
+                pl.UInt16,
+                pl.UInt8,
+            ]
+        elif props["type"] == "number":
+            if props.get("format") == "time-delta":
+                return [pl.Duration]
+            else:
+                return [pl.Float64, pl.Float32]
+        elif props["type"] == "boolean":
+            return [pl.Boolean]
+        elif props["type"] == "string":
+            string_format = props.get("format")
+            if string_format is None:
+                return [pl.Utf8]
+            elif string_format == "date":
+                return [pl.Date]
+            # TODO: Find out why this branch is not being hit
+            elif string_format == "date-time":  # pragma: no cover
+                return [pl.Datetime]
+            else:
+                return None  # pragma: no cover
+        elif props["type"] == "null":
+            return [pl.Null]
+        else:  # pragma: no cover
+            return None
+
     @property
     def valid_sql_types(  # type: ignore  # noqa: C901
-        cls: Type[ModelType],
+        cls: Type[ModelType],  # pyright: ignore
     ) -> dict[str, List["DuckDBSQLType"]]:
         """
         Return a list of DuckDB SQL types which Patito considers valid for each field.
@@ -302,7 +310,7 @@ class ModelMetaclass(PydanticModelMetaclass):
                 from patito.duckdb import _enum_type_name
 
                 # fmt: off
-                valid_dtypes[column] = [
+                valid_dtypes[column] = [  # pyright: ignore
                     _enum_type_name(field_properties=props),  # type: ignore
                     "VARCHAR", "CHAR", "BPCHAR", "TEXT", "STRING",
                 ]
@@ -374,7 +382,7 @@ class ModelMetaclass(PydanticModelMetaclass):
 
     @property
     def sql_types(  # type: ignore
-        cls: Type[ModelType],
+        cls: Type[ModelType],  # pyright: ignore
     ) -> dict[str, str]:
         """
         Return compatible DuckDB SQL types for all model fields.
@@ -405,7 +413,7 @@ class ModelMetaclass(PydanticModelMetaclass):
 
     @property
     def defaults(  # type: ignore
-        cls: Type[ModelType],
+        cls: Type[ModelType],  # pyright: ignore
     ) -> dict[str, Any]:
         """
         Return default field values specified on the model.
@@ -432,7 +440,7 @@ class ModelMetaclass(PydanticModelMetaclass):
 
     @property
     def non_nullable_columns(  # type: ignore
-        cls: Type[ModelType],  # pyright: reportGeneralTypeIssues=false
+        cls: Type[ModelType],  # pyright: ignore
     ) -> set[str]:
         """
         Return names of those columns that are non-nullable in the schema.
@@ -456,7 +464,7 @@ class ModelMetaclass(PydanticModelMetaclass):
 
     @property
     def nullable_columns(  # type: ignore
-        cls: Type[ModelType],  # pyright: reportGeneralTypeIssues=false
+        cls: Type[ModelType],  # pyright: ignore
     ) -> set[str]:
         """
         Return names of those columns that are nullable in the schema.
@@ -480,7 +488,7 @@ class ModelMetaclass(PydanticModelMetaclass):
 
     @property
     def unique_columns(  # type: ignore
-        cls: Type[ModelType],
+        cls: Type[ModelType],  # pyright: ignore
     ) -> set[str]:
         """
         Return columns with uniqueness constraint.
@@ -529,14 +537,18 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
     defaults: ClassVar[Dict[str, Any]]
 
-    @classmethod
+    @classmethod  # type: ignore[misc]
     @property
-    def DataFrame(cls: Type[ModelType]) -> Type[DataFrame[ModelType]]:
+    def DataFrame(
+        cls: Type[ModelType],
+    ) -> Type[DataFrame[ModelType]]:  # pyright: ignore  # noqa
         """Return DataFrame class where DataFrame.set_model() is set to self."""
 
-    @classmethod
+    @classmethod  # type: ignore[misc]
     @property
-    def LazyFrame(cls: Type[ModelType]) -> Type[LazyFrame[ModelType]]:
+    def LazyFrame(
+        cls: Type[ModelType],
+    ) -> Type[LazyFrame[ModelType]]:  # pyright: ignore
         """Return DataFrame class where DataFrame.set_model() is set to self."""
 
     @classmethod
@@ -570,7 +582,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
             >>> df = pl.DataFrame(
             ...     [["1", "product name", "1.22"]],
-            ...     columns=["product_id", "name", "price"],
+            ...     schema=["product_id", "name", "price"],
             ... )
             >>> Product.from_row(df)
             Product(product_id=1, name='product name', price=1.22)
@@ -582,7 +594,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
         elif _PANDAS_AVAILABLE and isinstance(row, pd.DataFrame):
             dataframe = pl.DataFrame._from_pandas(row)
         elif _PANDAS_AVAILABLE and isinstance(row, pd.Series):  # type: ignore[unreachable]
-            return cls(**dict(row.iteritems()))  # type: ignore[unreachable]
+            return cls(**dict(row.items()))  # type: ignore[unreachable]
         else:
             raise TypeError(f"{cls.__name__}.from_row not implemented for {type(row)}.")
         return cls._from_polars(dataframe=dataframe, validate=validate)
@@ -622,7 +634,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
             >>> df = pl.DataFrame(
             ...     [["1", "product name", "1.22"]],
-            ...     columns=["product_id", "name", "price"],
+            ...     schema=["product_id", "name", "price"],
             ... )
             >>> Product._from_polars(df)
             Product(product_id=1, name='product name', price=1.22)
@@ -977,7 +989,6 @@ class Model(BaseModel, metaclass=ModelMetaclass):
             │ str       ┆ cat              ┆ i64        │
             ╞═══════════╪══════════════════╪════════════╡
             │ product A ┆ dry              ┆ 0          │
-            ├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
             │ product B ┆ dry              ┆ 1          │
             └───────────┴──────────────────┴────────────┘
         """
@@ -1461,8 +1472,8 @@ class FieldDoc:
         product_id
           2 rows with duplicated values. (type=value_error.rowvalue)
         price
-          Polars dtype <class 'polars.datatypes.Int64'> \
-          does not match model field type. (type=type_error.columndtype)
+          Polars dtype Int64 does not match model field type. \
+          (type=type_error.columndtype)
         brand_color
           2 rows with out of bound values. (type=value_error.rowvalue)
     """

patito/sql.py

@@ -45,7 +45,7 @@ class Case:
 
         Examples:
             >>> import patito as pt
-            >>> db = pt.Database()
+            >>> db = pt.duckdb.Database()
             >>> relation = db.to_relation("select 1 as a union select 2 as a")
             >>> case_statement = pt.sql.Case(
             ...     on_column="a",
@@ -53,7 +53,7 @@ class Case:
             ...     default="three",
             ...     as_column="b",
             ... )
-            >>> relation.select(f"*, {case_statement}").to_df()
+            >>> relation.select(f"*, {case_statement}").order(by="a").to_df()
             shape: (2, 2)
             ┌─────┬─────┐
             │ a   ┆ b   │
@@ -61,7 +61,6 @@ class Case:
             │ i64 ┆ str │
             ╞═════╪═════╡
             │ 1   ┆ one │
-            ├╌╌╌╌╌┼╌╌╌╌╌┤
             │ 2   ┆ two │
             └─────┴─────┘
         """