Patito

Patito通过结合pydantic和polars，提供了一种编写现代、类型注释的数据框逻辑的简单方法。

Patito提供了一个简单的方法来声明pydantic数据模型，这些模型同时充当polars数据框的模式。这些模式可以用于

👮 简单且高效的数据框验证。
🧪 简单生成有效的模拟数据框进行测试。
🐍 以面向对象的方式检索和表示单个行。
🧠 为代码库中的核心数据模型提供一个单一的事实来源。

Patito对polars提供了第一类支持，这是一个"用Rust编写的闪电般快速的数据帧库"。

安装

pip install patito

文档

Patio的完整文档可以在这里找到。

👮 数据验证

Patito允许您通过创建patito.Model的类型注释子类来指定数据框中每列的类型

# models.py
from typing import Literal

import patito as pt


class Product(pt.Model):
    product_id: int = pt.Field(unique=True)
    temperature_zone: Literal["dry", "cold", "frozen"]
    is_for_sale: bool

类 Product 表示数据框的 模式，而 Product 的实例表示数据框的单行。Patito 可以高效地验证任意数据框的内容，并提供易于理解的错误信息。

import polars as pl

df = pl.DataFrame(
    {
        "product_id": [1, 1, 3],
        "temperature_zone": ["dry", "dry", "oven"],
    }
)
try:
    Product.validate(df)
except pt.exceptions.DataFrameValidationError as exc:
    print(exc)
# 3 validation errors for Product
# is_for_sale
#   Missing column (type=type_error.missingcolumns)
# product_id
#   2 rows with duplicated values. (type=value_error.rowvalue)
# temperature_zone
#   Rows with invalid values: {'oven'}. (type=value_error.rowvalue)

🧪 合成有效的测试数据

Patito 鼓励您严格验证数据框输入，从而确保在运行时正确无误。但是，强制正确性会带来摩擦，尤其是在测试期间。以下函数为例

import polars as pl

def num_products_for_sale(products: pl.DataFrame) -> int:
    Product.validate(products)
    return products.filter(pl.col("is_for_sale")).height

以下测试会因为 patito.exceptions.DataFrameValidationError 而失败

def test_num_products_for_sale():
    products = pl.DataFrame({"is_for_sale": [True, True, False]})
    assert num_products_for_sale(products) == 2

为了使测试通过，我们需要为 temperature_zone 和 product_id 列添加有效的虚拟数据。这会迅速引入大量样板代码到所有涉及数据框的测试中，从而模糊了每个测试实际测试的内容。因此，Patito 提供了 examples 构造函数来生成完全符合给定模型模式的测试数据。

Product.examples({"is_for_sale": [True, True, False]})
# shape: (3, 3)
# ┌─────────────┬──────────────────┬────────────┐
# │ is_for_sale ┆ temperature_zone ┆ product_id │
# │ ---         ┆ ---              ┆ ---        │
# │ bool        ┆ str              ┆ i64        │
# ╞═════════════╪══════════════════╪════════════╡
# │ true        ┆ dry              ┆ 0          │
# ├╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
# │ true        ┆ dry              ┆ 1          │
# ├╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
# │ false       ┆ dry              ┆ 2          │
# └─────────────┴──────────────────┴────────────┘

examples() 方法接受与常规数据框构造函数相同的参数，主要区别在于它会为任何未指定的列填充有效的虚拟数据。因此，测试可以重写为

def test_num_products_for_sale():
    products = Product.examples({"is_for_sale": [True, True, False]})
    assert num_products_for_sale(products) == 2

🖼️ 模型感知的数据框类

Patito 提供了 patito.DataFrame 类，该类扩展了 polars.DataFrame，以提供与 patito.Model 相关的实用方法。数据框的模式可以在运行时通过调用 patito.DataFrame.set_model(model) 指定，之后将提供一系列上下文相关的函数。

• DataFrame.validate() - 验证给定的数据框并返回自身。
• DataFrame.drop() - 删除所有未在模型中指定为字段的冗余列。
• DataFrame.cast() - 将不与给定类型注解兼容的列转换为指定类型。当指定 Field(dtype=...) 时，即使在不兼容的情况下，也会强制使用指定的数据类型。
• DataFrame.get(predicate) - 从数据帧中检索一行，作为模型的一个实例。如果过滤器谓词不正好返回一行，则会引发异常。
• DataFrame.fill_null(strategy="defaults") - 根据模型模式上设置的默认值填充缺失值。
• DataFrame.derive() - 带有 Field(derived_from=...) 注解的模型字段表示应该通过某些任意的 polars 表达式定义一个列。如果 derived_from 被指定为字符串，则给定值将被解释为带有 polars.col() 的列名。在调用 DataFrame.derive() 时，根据 derived_from 表达式创建并填充这些列。

以下示例将最好地说明这些方法

from typing import Literal

import patito as pt
import polars as pl


class Product(pt.Model):
    product_id: int = pt.Field(unique=True)
    # Specify a specific dtype to be used
    popularity_rank: int = pt.Field(dtype=pl.UInt16)
    # Field with default value "for-sale"
    status: Literal["draft", "for-sale", "discontinued"] = "for-sale"
    # The eurocent cost is extracted from the Euro cost string "€X.Y EUR"
    eurocent_cost: int = pt.Field(
        derived_from=100 * pl.col("cost").str.extract(r"€(\d+\.+\d+)").cast(float).round(2)
    )


products = pt.DataFrame(
    {
        "product_id": [1, 2],
        "popularity_rank": [2, 1],
        "status": [None, "discontinued"],
        "cost": ["€2.30 EUR", "€1.19 EUR"],
    }
)
product = (
    products
    # Specify the schema of the given data frame
    .set_model(Product)
    # Derive the `eurocent_cost` int column from the `cost` string column using regex
    .derive()
    # Drop the `cost` column as it is not part of the model
    .drop()
    # Cast the popularity rank column to an unsigned 16-bit integer and cents to an integer
    .cast()
    # Fill missing values with the default values specified in the schema
    .fill_null(strategy="defaults")
    # Assert that the data frame now complies with the schema
    .validate()
    # Retrieve a single row and cast it to the model class
    .get(pl.col("product_id") == 1)
)
print(repr(product))
# Product(product_id=1, popularity_rank=2, status='for-sale', eurocent_cost=230)

每个 Patito 模型自动获得一个 .DataFrame 属性，这是一个自定义的数据帧子类，其中在实例化时调用 .set_model()。换句话说，pt.DataFrame(...).set_model(Product) 等同于 Product.DataFrame(...)。

🐍 将行表示为类

数据帧非常适合在 一组 对象上执行矢量化操作。但是，当需要检索单行并对其操作时，数据帧结构自然就不足了。Patito 允许你在模型上定义的方法中嵌入行级逻辑。

# models.py
import patito as pt

class Product(pt.Model):
    product_id: int = pt.Field(unique=True)
    name: str

    @property
    def url(self) -> str:
        return (
            "https://example.com/no/products/"
            f"{self.product_id}-"
            f"{self.name.lower().replace(' ', '-')}"
        )

可以使用 from_row() 方法从数据帧的单行实例化类

products = pl.DataFrame(
    {
        "product_id": [1, 2],
        "name": ["Skimmed milk", "Eggs"],
    }
)
milk_row = products.filter(pl.col("product_id" == 1))
milk = Product.from_row(milk_row)
print(milk.url)
# https://example.com/no/products/1-skimmed-milk

如果你通过使用 patito.DataFrame.set_model() 或直接使用 Product.DataFrame 将 Product 模型与 DataFrame “连接”，则可以使用 .get() 方法将数据帧过滤到单个行，并将其转换为相应的模型类

products = Product.DataFrame(
    {
        "product_id": [1, 2],
        "name": ["Skimmed milk", "Eggs"],
    }
)
milk = products.get(pl.col("product_id") == 1)
print(milk.url)
# https://example.com/no/products/1-skimmed-milk