如何使用断言函数比较PySpark DataFrame中的两列

针对你用Spark DataFrame替代Pandas assert_array_almost_equal的需求，这里提供几种实用实现方式：

1. 精确匹配检查（适用于字符串/完全相等场景） #

如果你的expected和Curated是字符串列（比如处理后的邮箱），可以直接过滤两列不相等的行，通过断言确保不匹配行数为0，同时输出不匹配详情：

from pyspark.sql.functions import col

# 找出两列不相等的行
mismatched_rows = df1.filter(col("expected") != col("Curated"))
# 获取不匹配行数
mismatch_count = mismatched_rows.count()

# 断言：存在不匹配行则抛出错误并打印详情
assert mismatch_count == 0, f"发现{mismatch_count}条不匹配记录：\n{mismatched_rows.show(truncate=False)}"

2. 近似相等检查（适用于数值型列） #

如果是数值列需要近似匹配（对应assert_array_almost_equal的场景），可以设置容忍度，判断两列差值的绝对值是否在允许范围内：

from pyspark.sql.functions import abs, lit

tolerance = 1e-6  # 自定义容忍度阈值
mismatched_rows = df1.filter(abs(col("expected") - col("Curated")) > lit(tolerance))
mismatch_count = mismatched_rows.count()

assert mismatch_count == 0, f"发现{mismatch_count}条超出容忍度的记录：\n{mismatched_rows.show(truncate=False)}"

3. 封装成可复用的断言函数 #

如果需要多次使用，可以把检查逻辑封装成函数：

from pyspark.sql.functions import col, abs, lit

def assert_columns_equal(df, col1, col2, tolerance=None):
    if tolerance is None:
        # 精确匹配逻辑
        mismatched = df.filter(col(col1) != col(col2))
    else:
        # 近似匹配逻辑
        mismatched = df.filter(abs(col(col1) - col(col2)) > lit(tolerance))
    
    count = mismatched.count()
    assert count == 0, f"列{col1}和{col2}存在{count}条不匹配记录：\n{mismatched.show(truncate=False)}"

# 调用示例（精确匹配邮箱列）
assert_columns_equal(df1, "expected", "Curated")

# 调用示例（数值列近似匹配）
# assert_columns_equal(df1, "expected", "Curated", tolerance=1e-6)

结合你的原有代码，修改后的完整示例：

from pyspark.sql.functions import col

df1 = context.spark.read.option("header", True).csv("./test/input/test/Book1.csv") 
df1 = df1.withColumn("Curated", dataclean.clean_email(col("email")))
df1.show()

# 执行断言检查
mismatched_rows = df1.filter(col("expected") != col("Curated"))
mismatch_count = mismatched_rows.count()
assert mismatch_count == 0, f"发现{mismatch_count}条不匹配记录：\n{mismatched_rows.show(truncate=False)}"

内容的提问来源于stack exchange，提问作者Xi12