PySparkでspark-testing-baseを使ってみる

概要

spark用のテストライブラリのspark-testing-baseをPySparkでどのように使うかを軽く調べてみた。

調べたこと

• unittest2のTestCaseを継承しており、unittest2.main()でテスト実行できる
• クラス内部でspark contextを立ち上げており、テストでわざわざ自力で立ち上げる必要がない
  • 自分でspark contextを立ち上げていると、テスト実行時にエラーになる
• DataFrameの場合、スキーマと件数と各レコードの並びが一致するかを確認しているみたい

DataFrameのテストサンプル

from sparktestingbase.sqltestcase import SQLTestCase
import unittest2


class DataFrameTest(SQLTestCase):
    def test_expected_equal(self):
        data = [(1, 2, 3), (4, 5, 6), (7, 8, 9)]
        df = self.sqlCtx.createDataFrame(data, ['a', 'b', 'c'])
        self.assertDataFrameEqual(df, df)

    @unittest2.expectedFailure    
    def test_expected_unequal(self):
        data = [(1, 2, 3), (4, 5, 6), (7, 8, 9)]
        df1 = self.sqlCtx.createDataFrame(data, ['a', 'b', 'c'])
        df2 = self.sqlCtx.createDataFrame(data, ['a', 'b', 'd'])
        self.assertDataFrameEqual(df1, df2)


if __name__ == "__main__":
    unittest2.main()

関連

• Spark でUnit Testを書く