tech.quantco.comA QuantCo Engineering Blog BlogAboutJobsquantco.com Home Blog About Jobs quantco.com 发布于 2025年4月16日 星期三

dataframely — 一个声明式的、🐻‍❄️-native 的数据框验证库

作者

在 QuantCo，我们不断尝试提高代码库的质量，以确保它们易于维护。最近，这通常涉及到将数据管道从 pandas 迁移到 polars，以实现显著的性能提升。

2023年底，我们开始着手改造一个长期运行项目中庞大的遗留代码库。在改造过程中，我们意识到现有的数据框处理代码存在一个根本缺陷：列名、数据类型、值范围以及其他不变性——从代码中无法直接看出。

因此，理解函数行为的典型方法是在客户端基础设施上执行它——因为只有在那里才能获得实际数据。然后，我们会手动单步执行每个 pandas 转换，以检查每次更改前后数据。显然，这既繁琐、容易出错，效率也很低。

一旦我们在 polars 中重写了一系列转换，由于缺乏静态类型检查或运行时数据框内容验证，bug 很难被发现。为了确保正确性，我们经常需要在大型数据集上端到端地运行整个管道——这需要大量的时间和计算资源。

最终，我们意识到需要一种更好的方法来描述、验证和推断数据管道中数据框的内容。我们希望在阅读代码时就能清楚地了解不变性，并在运行时强制执行这些不变性，以确保正确性。

数据框验证的救赎之道

解决这个问题的一个自然方法是使用数据框验证库。早在 2023 年，就已经存在一些 Python 库，允许定义数据框模式并验证数据框是否符合这些模式，即是否满足预定义的期望。

在某些项目中，我们已经在使用 pandera（https://github.com/unionai-oss/pandera），这是一个广为人知的开源库，用于验证 pandas 数据框。不幸的是，在 2023 年，pandera 没有任何 polars 支持，而一个值得关注的 polars 原生替代方案，即 patito（https://github.com/JakobGM/patito），还处于起步阶段，无法被认为是生产就绪的。

然而，即使在今天，我们在使用 pandera 和 patito 时仍然遇到一些限制。我们认为这些限制是其范围和设计固有的，无法通过向这些项目贡献代码来轻松解决——尽管如此，我们仍然积极地这样做（例如，我们维护 pandera 的 conda-forge feedstock）。

具体来说，pandera 和 patito 缺少对以下内容的支持：

隆重推出 dataframely：一个 Polars 原生数据框验证库

为了弥补这些库的不足，我们开发了 dataframely。dataframely 是一个声明式数据框验证库，对 polars 数据框提供一流的支持。其目的是使使用 polars 编写的数据管道 (1) 更加健壮，通过确保数据符合预期；(2) 更加可读，通过向数据框类型提示添加模式信息。

光说不练假把式，让我们看一些代码示例。

定义模式

要开始使用 dataframely，首先要定义一个模式。在 QuantCo，我们经常处理保险索赔——例如，我们可以为包含医院发票的数据框创建一个模式：

import dataframely as dy
class InvoiceSchema(dy.Schema):
  invoice_id = dy.String(primary_key=True)
  admission_date = dy.Date(nullable=False)
  discharge_date = dy.Date(nullable=False)
  amount = dy.Decimal(nullable=False, min_exclusive=Decimal(0))
  @dy.rule()
  def discharge_after_admission() -> pl.Expr:
    return pl.col("discharge_date") >= pl.col("admission_date")

虽然我们可以根据数据框的列及其数据类型来描述数据框，但我们也可以对列级别以及跨列级别对期望进行编码。例如，我们可以将一列（或多列）指定为主键，或者定义一个自定义验证规则，该规则在_跨越_列的范围内生效。

验证数据框

一旦我们定义了一个模式，我们可以将一个 pl.DataFrame 或 pl.LazyFrame 传递给它的 validate 类方法，以验证其内容是否与模式定义匹配。如果我们要自动将列类型强制转换为模式中指定的类型，我们可以传递 cast=True。

invoices = pl.DataFrame({
  "invoice_id": ["001", "002", "003"],
  "admission_date": [date(2025, 1, 1), date(2025, 1, 5), date(2025, 1, 1)],
  "discharge_date": [date(2025, 1, 4), date(2025, 1, 7), date(2025, 1, 1)],
  "amount": [1000.0, 200.0, 400.0]
})
validated: dy.DataFrame[InvoiceSchema] = InvoiceSchema.validate(invoices, cast=True)

如果 invoices 中的任何行无效，即在单个列或整个模式上定义的任何规则的计算结果为 False，则会引发验证异常。否则，如果 invoices 中的所有行都有效，validate 将返回一个经过验证的数据框，其类型为 dy.DataFrame[InvoiceSchema]。

重要的是，dy.DataFrame[InvoiceSchema] 只是一个类型构造，而在运行时，您仍然处理的是 pl.DataFrame。这样做的好处是，dataframely 可以逐步采用，任何 dy.DataFrame[...] 都可以轻松地传递给接受 pl.DataFrame 的方法（反之亦然，通过使用 type: ignore 注释）。

然而，最大的好处是：泛型数据框类型立即告诉代码的读者他们可以期望在数据框中找到什么数据。这显著提高了 mypy 对于基于数据框的代码的有用性：类型检查器现在可以确保传递给方法的数据框满足与其内容相关的某些前提条件 - 而不会在运行时产生隐藏的性能损失。

验证数据框组

通常，数据框（或者更确切地说是“表”）是相互依赖的，而正确的数据验证需要考虑共享一个公共主键的多个表。dataframely 使您可以定义 "集合"，用于包含一组数据框，并在集合级别定义验证规则。为了创建集合，我们首先为诊断数据框引入第二个模式：

class DiagnosisSchema(dy.Schema):
  invoice_id = dy.String(primary_key=True)
  diagnosis_code = dy.String(primary_key=True, regex=r"[A-Z][0-9]{2,4}")
  is_main = dy.Bool(nullable=False)
  @dy.rule(group_by=["invoice_id"])
  def exactly_one_main_diagnosis() -> pl.Expr:
    return pl.col("is_main").sum() == 1

然后，我们可以创建一个集合，将发票和属于这些发票的诊断信息捆绑在一起：

# Introduce a collection for groups of schema-validated data frames
class HospitalClaims(dy.Collection):
  invoices: dy.LazyFrame[InvoiceSchema]
  diagnoses: dy.LazyFrame[DiagnosisSchema]
  @dy.filter()
  def at_least_one_diagnosis_per_invoice(self) -> pl.LazyFrame:
    return self.invoices.join(
      self.diagnoses.select(pl.col("invoice_id").unique()),
      on="invoice_id",
      how="inner",
    )

请注意，我们如何通过使用 @dy.filter 装饰器跨集合成员添加验证，从而进一步定义我们对集合内容的期望。

如果我们调用集合上的 validate，如果任何输入数据框不满足其模式定义，或者 如果集合上的过滤器导致删除任何输入数据框中的至少一行，则将引发验证异常。

软验证和验证失败内省

虽然调用 validate 对于确保正确性很有用，但在生产管道中，我们通常不希望在运行时引发异常。为此，dataframely 提供了 filter 方法来对模式和集合执行 "软验证"。filter 返回通过验证的行和一个额外的 FailureInfo 对象，用于检查无效行：

good, failure = InvoiceSchema.filter(invoices, cast=True)
# Inspect the reasons for the failed rows
failure.counts()
# Inspect the co-occurrences of validation failures
failure.cooccurrence_counts()

由于 filter 不会引发异常，我们可以安全地在生产代码中使用它，并将无效行记录下来以便以后检查。

其他功能

在使用 dataframely 的过程中，我们意识到，定义模式并因此对数据框内容进行编码，除了运行验证之外，还具有各种好处。例如，如果我们要将数据框写入数据库，我们可以自动导出表的 SQL 模式。另一种可能性是自动生成符合模式的单元测试样本数据，从而使测试作者可以专注于测试内容，而不是冗长地创建数据框。要了解所有可能性，请查看 API 文档。

实践经验

理解数据框的结构和内容在使用表格数据时至关重要——这是我们在 QuantCo 构建高度健壮的数据管道的核心要求。dataframely 已经使我们更接近这个目标：今天，我们已成功地在多个团队的日常工作中使用 dataframely，这些团队跨越多个客户，用于分析和生产管道。

我们的数据科学家和工程师喜欢 dataframely，因为它

我们很高兴开源 dataframely 并与数据工程社区分享它。如果您正在使用复杂的数据管道并希望提高可靠性、生产力和安心，我们相信您也会喜欢它。

请在 GitHub 上查看 dataframely，并告诉我们您的想法！

标签

pythonpolarsdata-validation

上一篇文章

调试自动化的 Conda-Forge Feedstock 更新 ← 返回博客 githubGitHublinkedinLinkedinblueskyBlueskymastodonMastodon tech.quantco.com • 法律信息 • © 2025