如何对 pyspark `withColumn` 动作进行单元测试 - Python 3？答案

【问题标题】：How to unittest pyspark `withColumn` action - Python 3?如何对 pyspark `withColumn` 动作进行单元测试 - Python 3？
【发布时间】：2021-09-16 14:15:18
【问题描述】：

我正在学习pyspark，我有一个函数：

import re

def function_1(string):
    new_string = re.sub(r"!", " ", string)
    return new_string


udf_function_1 = udf(lambda s: function_1(s), StringType())


def function_2(data):
    new_data = data \
        .withColumn("column_1", udf_function_1("column_1"))
    return new_data

我的问题是如何在 Python 中为function_2() 编写单元测试。

【问题讨论】：

标签： python unit-testing pyspark apache-spark-sql python-unittest

【解决方案1】：

您究竟想在function_2 中测试什么？

以下是保存在名为sample_test.py 的文件中的简单测试。我使用了pytest，但您可以在 unittest 中修改非常相似的代码。

# sample_test.py

from pyspark import sql


spark = sql.SparkSession.builder \
        .appName("local-spark-session") \
        .getOrCreate()
        
def test_create_session():
    assert isinstance(spark, sql.SparkSession) == True
    assert spark.sparkContext.appName == 'local-spark-session'

def test_spark_version():
    assert spark.version == '3.1.2'

运行测试...

C:\Users\user\Desktop>pytest -v sample_test.py
============================================= test session starts =============================================
platform win32 -- Python 3.6.7, pytest-6.2.5, py-1.10.0, pluggy-1.0.0 -- c:\users\user\appdata\local\programs\python\python36\python.exe
cachedir: .pytest_cache
rootdir: C:\Users\user\Desktop
collected 2 items

sample_test.py::test_create_session PASSED                                                               [ 50%]
sample_test.py::test_spark_version PASSED                                                                [100%]

============================================== 2 passed in 4.81s ==============================================

【讨论】：