"Pandera检查多列"

Question

我正在编写一个Python脚本，使用Pandera库验证Pandas DataFrame中的列。我遇到的问题是，我可以逐列进行验证，但未能找到如何在验证某一列的值时，考虑数据框中另一列的值的方法。

具体来说，我想检查一列中是否存在某些值，如果存在，我还想检查该行另一列的值是否包含特定文本。

import numpy as np
import pandas as pd
import pandera as pa

dataframe = pd.DataFrame({'column_A': ['ABC company', 'BBB company', 'ABC company', 'CCC company'],
                          'column_B': ['1000', np.nan, '2000', np.nan]
                          })

schema = pa.DataFrameSchema(
    columns={
        'column_A': pa.Column(pa.String),
        'column_B': pa.Column(pa.String, nullable=True,
                              checks=pa.Check(
                                  lambda df: (df['column_A'].str.contains('ABC')) & (~df['column_B'].isna()))
                              )
    }
)

schema.validate(dataframe)

这是我尝试解决此问题的方法，但Pandera并不知道在验证'column_B'列时，我需要同时查看多个序列（series）。

ysth · Answer

我明白，您需要调整布尔掩码并隔离校验条件以避免出现KeyError异常。请参考以下代码：

schema = pa.DataFrameSchema(
    columns={
        "column_A": pa.Column(pa.String),
        "column_B": pa.Column(pa.String, nullable=True),
    },
    checks=[
        pa.Check(
            lambda df: (
                df["column_A"].str.contains("ABC")
                & (b_notna := df["column_B"].notna())
            ).where(b_notna, True),
            ignore_na=False
        )
    ]
)

# 输出结果
print(schema.validate(dataframe))
# 此处显示已验证的数据框内容
#      column_A column_B
# 0  ABC company     1000
# 1  BBB company      NaN
# 2  ABC company     2000
# 3  CCC company      NaN

# 如果我们将'column_A'的第一项设置为"A_BC"，则校验将失败：

# 失败案例：
#      列名   索引   失败情况
# 0  column_A      0  A_BC company
# 1  column_B      0          1000

Pshemo · Answer

如果您想在数据框级别执行校验（参见文档中的“# 定义一个宽泛的数据框级校验”部分），您需要在列之外定义您的测试：

import numpy as np
import pandas as pd
import pandera as pa

dataframe = pd.DataFrame({'column_A': ['ABC公司', 'BBB公司', 'ABC公司', 'CCC公司'],
                          'column_B': ['1000', np.nan, '2000', np.nan]
                          })

# 定义针对整个数据框的测试
check_AB = pa.Check(
    lambda df: (df['column_A'].str.contains('ABC')) & (~df['column_B'].isna()),
    name='check_AB'  # 校验名称
)

schema = pa.DataFrameSchema(
    columns={
        'column_A': pa.Column(pa.String),
        'column_B': pa.Column(pa.String, nullable=True)
    },
    checks=check_AB  # <- 在数据框层面应用的校验
)

# 输出：
# 当运行 schema.validate(dataframe) 时，
# 显示如下：
# ...
# <Check check_AB>
# 失败案例：
#      列名   索引   失败情况
# 0  column_A      1  BBB公司
# 1  column_A      3  CCC公司