Как смоделировать внутренний вызов функции pyspark sql - PullRequest
Новая
       79

Как смоделировать внутренний вызов функции pyspark sql

6 голосов
/ 02 ноября 2019

Получил следующий фрагмент кода 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')??

1 Ответ

2 голосов
/ 08 ноября 2019

Спасибо моему умному коллеге по работе, вот ответ. Мы должны смоделировать pyspark.sql.functions.col и затем установить возвращаемое значение. 

@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 также отлично работает.

...