pytest断言用于pyspark数据帧比较

Question

我有2个pyspark数据帧，如附件中所示。 Expected_df和actual_df

在单元测试中，我试图检查两者是否相等。

我的代码是

expected = map(lambda row: row.asDict(), expected_df.collect()) 
actual = map(lambda row: row.asDict(), actaual_df.collect()) 
assert expected = actual 

由于两个df相同，但行顺序不同，因此断言在这里失败。 比较此类df的最佳方法是什么。

Answer 1

确保排序顺序的另一种方法是：

from pandas.testing import assert_frame_equal

def assert_frame_with_sort(results, expected, key_columns):
    results_sorted = results.sort_values(by=key_columns).reset_index(drop=True)
    expected_sorted = expected.sort_values(by=key_columns).reset_index(drop=True)
    assert_frame_equal(results_sorted, expected_sorted)

Answer 2

不幸的是，如果不对任何列（特别是在键列上）应用排序，就无法完成此操作，原因是没有任何guarantee for ordering of records in a DataFrame。您无法预测记录将在数据框中显示的顺序。下面的方法对我来说很好：

expected = expected_df.orderBy('period_start_time').collect()
actual = actaual_df.orderBy('period_start_time').collect() 
assert expected = actual

Answer 3

尝试使用“ ==”代替“ =“。 断言期望==实际

Answer 4

这是在一些pyspark文档中完成的：

assert sorted(expected_df.collect()) == sorted(actaual_df.collect())

Answer 5

您可以尝试pyspark-test

https://pypi.org/project/pyspark-test/0.1.0/

这是受pyspark的panadas测试模块启发的。

用法很简单

from pyspark_test import assert_pyspark_df_equal

assert_pyspark_df_equal(df_1, df_2)

除了只比较数据框外，就像熊猫测试模块一样，它还接受许多可选参数，您可以在文档中进行检查。