Polars中能否在方法链内执行验证与断言操作？

Polars目前没有内置的.assert()方法，但可以通过自定义函数或扩展实现方法链中的断言逻辑，满足长查询中提前终止执行的需求。以下是针对Eager和Lazy两种模式的实现方案：

Eager DataFrame 断言实现 #

利用.pipe()方法嵌入自定义断言函数，验证条件不满足时抛出AssertionError，满足则返回原DataFrame继续后续操作：

import polars as pl

df = pl.DataFrame({
    'A': ['X', 'X', 'Y', 'Z'],
})

def assert_condition(df, condition, message):
    if not df.select(condition).item():
        raise AssertionError(message)
    return df

# 验证A列至少有一个X
(
    df
    .pipe(assert_condition, (pl.col('A') == 'X').any(), "At least one element in A should be X")
    .filter(pl.col('A') == 'X')
)

# 验证A列至少有2个X
(
    df
    .pipe(assert_condition, (pl.col('A') == 'X').sum() > 1, "At least 2 elements in A should be X")
    .filter(pl.col('A') == 'X')
)

LazyFrame 断言实现 #

Lazy模式下可以用.map()方法嵌入验证逻辑，在执行计划触发时（如.collect()）进行断言检查：

def lazy_assert(lf, conditions):
    def validate(df):
        for cond in conditions:
            # 从条件的meta中获取自定义错误信息
            msg = cond.meta.get("message", "Assertion failed")
            if not df.select(cond).item():
                raise AssertionError(msg)
        return df
    return lf.map(validate)

(
    df
    .lazy()
    .pipe(lazy_assert, [
        ((pl.col('A') == 'X').sum() > 1).meta.set("message", "At least 2 elements in A should be X"),
        ((pl.col('A') == 'Y').sum() > 1).meta.set("message", "At least 2 elements in A should be Y")
    ])
    .filter(pl.col('A') == 'X')
    .collect()
)

进阶：自定义.assert()方法（谨慎使用） #

如果想要更贴近示例中的语法，可以通过Monkey Patch为Polars对象添加.assert()方法，但需注意可能与未来内置方法冲突：

# 为Eager DataFrame添加assert方法
def df_assert(self, condition, message=None):
    if not self.select(condition).item():
        raise AssertionError(message or "Assertion failed")
    return self

pl.DataFrame.assert = df_assert

# 为LazyFrame添加assert方法
def lf_assert(self, conditions):
    def validate(df):
        for cond in conditions:
            msg = cond.meta.get("message", "Assertion failed")
            if not df.select(cond).item():
                raise AssertionError(msg)
        return df
    return self.map(validate)

pl.LazyFrame.assert = lf_assert

# 使用自定义方法
(
    df
    .assert((pl.col('A') == 'X').any(), message="At least one element in A should be X")
    .filter(pl.col('A') == 'X')
)

(
    df
    .lazy()
    .assert([
        ((pl.col('A') == 'X').sum() > 1).meta.set("message", "At least 2 elements in A should be X"),
        ((pl.col('A') == 'Y').sum() > 1).meta.set("message", "At least 2 elements in A should be Y")
    ])
    .filter(pl.col('A') == 'X')
    .collect()
)

内容的提问来源于stack exchange，提问作者T.H Rice