时间戳

首页 > 解决方案 > 如何模拟对pyspark sql函数的内部调用

python - 如何模拟对pyspark sql函数的内部调用

问题描述

得到以下 pyspark 代码：

import pyspark.sql.functions as F

null_or_unknown_count = df.sample(0.01).filter(
    F.col('env').isNull() | (F.col('env') == 'Unknown')
).count()

在测试代​​码中，数据框被模拟，所以我试图为这个调用设置 return_value，如下所示：

from unittest import mock
from unittest.mock import ANY

@mock.patch('pyspark.sql.DataFrame', spec=pyspark.sql.DataFrame)
def test_null_or_unknown_validation(self, mock_df):
    mock_df.sample(0.01).filter(ANY).count.return_value = 250

但这失败了，原因如下：

File "/usr/local/lib/python3.7/site-packages/pyspark/sql/functions.py", line 44, in _
  jc = getattr(sc._jvm.functions, name)(col._jc if isinstance(col, Column) else col)
AttributeError: 'NoneType' object has no attribute '_jvm'

也试过mock_df.sample().filter().count.return_value = 250了，这给出了同样的错误。

如何F.col('env').isNull() | (F.col('env') == 'Unknown')正确模拟过滤器？

标签： pythonapache-sparkpysparkmockingpython-unittest

解决方案

感谢我工作中的聪明同事，这就是答案。我们必须模拟pyspark.sql.functions.col然后设置一个return_value。

@mock.patch('pyspark.sql.functions.col')
@mock.patch('pyspark.sql.DataFrame', spec=pyspark.sql.DataFrame)
def test_null_or_unknown_validation(self, mock_df, mock_functions):
    mock_functions.isNull.return_value = True # (or False also works)
    mock_df.sample(0.01).filter(ANY).count.return_value = 250

使用mock_df.sample().filter().count.return_value = 250也可以正常工作。

推荐阅读