PySpark の Testing 機能の１つである assertDataFrameEqual に対する基本的な検証結果

<!--
# PySpark の Testing 機能の assertDataFrameEqual に対する検証結果
-->

## 概要

この記事では、Spark 3.5 にて実装された PySpark の Testing 機能の１つである `assertDataFrameEqual` についての検証結果を共有します。assertDataFrameEqualはPySparkのテストユーティリティの一部で、2つのデータフレームが等しいかどうかを検証するためのメソッドです。検証したケースのコードを記事内で後述しています。

検証では、正常系（データに差異がない場合）の動作確認やデータに差異がある場合の動作確認などのケースを実施しました。しかし、大量データでの検証ではタイムアウトの問題が発生し、テーブル間でのデータ差異の確認には利用できなさそうでした。ただ、まだリリースされて間もないため、今後改善される可能性があります。

次の記事を参考にしました。

- PySpak における Testing 機能の概要
    - [Testing PySpark — PySpark 3.5.0 documentation (apache.org)](https://spark.apache.org/docs/latest/api/python/getting_started/testing_pyspark.html)
- assertDataFrameEqual のドキュメント
    - [pyspark.testing.assertDataFrameEqual — PySpark 3.5.0 documentation (apache.org)](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.testing.assertDataFrameEqual.html#pyspark.testing.assertDataFrameEqual)


## 事前準備

### ライブラリのimport

```python
import datetime

from pyspark.sql import Row
from decimal import Decimal

from pyspark.testing.utils import assertDataFrameEqual, assertSchemaEqual
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/e48a52f2-5eb3-00ae-79ae-6bb9f80fe119.png)

## メソッドの基本的な動作確認

### 正常系の動作確認

```python
test_schema_01 = """
id int,
name string,
weight float,
birth_date date
"""
test_data_01 = [
    Row(id=1, name="user_aaa", weight=53.5, birth_date=datetime.date(2020, 1, 1)),
    Row(id=2, name="user_bbb", weight=53.5, birth_date=datetime.date(2020, 1, 1)),
    Row(id=3, name="user_ccc", weight=53.5, birth_date=datetime.date(2020, 1, 1)),
]
test_df_01 = spark.createDataFrame(test_data_01, test_schema_01)
test_df_01.display()
```

```python
assertDataFrameEqual(
    test_df_01,
    test_df_01,
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/af9c871d-b87a-e851-a92d-ceb12e34db35.png)

### データに差異がある場合の動作確認

```python
exp_schema_01 = """
id int,
name string,
weight float,
birth_date date
"""
exp_data_01 = [
    Row(id=4, name="user_aaa", weight=53.5, birth_date=datetime.date(2020, 1, 1)),
    Row(id=2, name="user_bbb", weight=53.5, birth_date=datetime.date(2020, 1, 1)),
    Row(id=3, name="user_ccc", weight=53.5, birth_date=datetime.date(2020, 1, 1)),
]
exp_df_01 = spark.createDataFrame(exp_data_01, exp_schema_01)
exp_df_01.display()
```

```python
assertDataFrameEqual(
    test_df_01,
    exp_df_01,
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/b7c1a07e-a362-eb33-fdb8-94a239f686ee.png)

次のようなエラーが発生しました。

```text

assertDataFrameEqual(
    test_df_01,
    exp_df_01,
)
PySparkAssertionError: [DIFFERENT_ROWS] Results do not match: ( 100.00000 % )
*** actual ***
[31m- Row(id=1, name='user_aaa', weight=53.5, birth_date=datetime.date(2020, 1, 1))[0m
  Row(id=2, name='user_bbb', weight=53.5, birth_date=datetime.date(2020, 1, 1))
  Row(id=3, name='user_ccc', weight=53.5, birth_date=datetime.date(2020, 1, 1))


*** expected ***
  Row(id=2, name='user_bbb', weight=53.5, birth_date=datetime.date(2020, 1, 1))
  Row(id=3, name='user_ccc', weight=53.5, birth_date=datetime.date(2020, 1, 1))
[31m+ Row(id=4, name='user_aaa', weight=53.5, birth_date=datetime.date(2020, 1, 1))[0m
```

### スキーマに差異がある場合の動作確認

```python
exp_schema_02 = """
id int
"""
exp_data_02 = [
    Row(id=1),
]
exp_df_02 = spark.createDataFrame(exp_data_02, exp_schema_02)
exp_df_02.display()
```

```python
assertDataFrameEqual(
    test_df_01,
    exp_df_02,
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/20746af7-73b3-06f9-0041-cc4fa6f9aece.png)

次のようなエラーが発生しました。

```test
assertDataFrameEqual(
```

PySpark の Testing 機能には`assertSchemaEqual`というメソッドもあり、同様のエラーが発生しました。

```python
# assertSchemaEqual との比較
assertSchemaEqual(
    test_df_01.schema,
    exp_df_02.schema,
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/9f2ea91b-c616-b0be-fcd8-8720cef3d20a.png)

```text
--- actual
+++ expected
- StructType([StructField('id', IntegerType(), True), StructField('name', StringType(), True), StructField('weight', FloatType(), True), StructField('birth_date', DateType(), True)])
+ StructType([StructField('id', IntegerType(), True)])
```

### エラーハンドリング

try/exceptブロックによりエラーを取得したい場合には、次のようなコードにより可能です。         

```python
from pyspark.testing.utils import PySparkAssertionError

try:
    assertDataFrameEqual(test_df_01, exp_df_01)
except PySparkAssertionError:
    print("エラーです")
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/edefba11-b6ed-7d5a-4e57-f1477ec136b5.png)

## データ型の対応可否確認

### 正常系

```python
test_data_02 = [
    Row(
        string_column="AAA",
        byte_column=1,
        integer_column=1,
        bigint_column=1,
        float_column=12.300000190734863,
        double_column=12.3,
        numeric_column=Decimal("12"),
        boolean_column=True,
        date_column=datetime.date(2020, 1, 1),
        timestamp_column=datetime.datetime(2021, 1, 1, 0, 0),
        timestamp_ntz_column=datetime.datetime(2021, 1, 1, 0, 0),
        binary_column=bytearray(b"A"),
        struct_column=Row(struct_string_column="AAA", struct_int_column=1),
        array_column=["AAA", "BBB", "CCC"],
        map_column={"AAA": 1},
    )
]

test_schema_02 = """
--文字型
string_column string,

--整数型
byte_column byte,
integer_column integer,
bigint_column bigint,

--浮動小数点型
float_column float,
double_column double,
numeric_column numeric,

--真偽型
boolean_column boolean,

--日付時刻
date_column date,
timestamp_column timestamp,
timestamp_ntz timestamp_ntz,

--バイナリー型
binary_column binary,

--複合型
struct_column struct<
    struct_string_column :string,
    struct_int_column    :int
>,
array_column array<string>,
map_column map<string, int>
"""

test_df_02 = spark.createDataFrame(test_data_02, test_schema_02)
```

```python
assertDataFrameEqual(
    test_df_02,
    test_df_02,
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/9bfd8e07-3e07-3e5e-f371-30e7a8030c6e.png)

### 異常系

```python
exp_data_03 = [
    Row(
        string_column="ABC",
        byte_column=1,
        integer_column=1,
        bigint_column=1,
        float_column=12.300000190734863,
        double_column=12.3,
        numeric_column=Decimal("12"),
        boolean_column=True,
        date_column=datetime.date(2020, 1, 1),
        timestamp_column=datetime.datetime(2021, 1, 1, 0, 0),
        timestamp_ntz_column=datetime.datetime(2021, 1, 1, 0, 0),
        binary_column=bytearray(b"A"),
        struct_column=Row(struct_string_column="AAA", struct_int_column=1),
        array_column=["AAA", "BBB", "CCC"],
        map_column={"AAA": 1},
    )
]

exp_schema_03 = """
--文字型
string_column string,

--整数型
byte_column byte,
integer_column integer,
bigint_column bigint,

--浮動小数点型
float_column float,
double_column double,
numeric_column numeric,

--真偽型
boolean_column boolean,

--日付時刻
date_column date,
timestamp_column timestamp,
timestamp_ntz timestamp_ntz,

--バイナリー型
binary_column binary,

--複合型
struct_column struct<
    struct_string_column :string,
    struct_int_column    :int
>,
array_column array<string>,
map_column map<string, int>
"""

exp_df_03 = spark.createDataFrame(exp_data_03, exp_schema_03)
exp_df_03.display()
```

```python
assertDataFrameEqual(
    test_df_02,
    exp_df_03,
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/d53f6393-6317-08b3-9a7f-5b2c1b0db18d.png)

次のようなエラーが発生しました。

```text
PySparkAssertionError: [DIFFERENT_ROWS] Results do not match: ( 100.00000 % )
*** actual ***
[31m! Row(string_column='AAA', byte_column=1, integer_column=1, bigint_column=1, float_column=12.300000190734863, double_column=12.3, numeric_column=Decimal('12'), boolean_column=True, date_column=datetime.date(2020, 1, 1), timestamp_column=datetime.datetime(2021, 1, 1, 0, 0), timestamp_ntz=datetime.datetime(2021, 1, 1, 0, 0), binary_column=bytearray(b'A'), struct_column=Row(struct_string_column='AAA', struct_int_column=1), array_column=['AAA', 'BBB', 'CCC'], map_column={'AAA': 1})[0m


*** expected ***
[31m! Row(string_column='ABC', byte_column=1, integer_column=1, bigint_column=1, float_column=12.300000190734863, double_column=12.3, numeric_column=Decimal('12'), boolean_column=True, date_column=datetime.date(2020, 1, 1), timestamp_column=datetime.datetime(2021, 1, 1, 0, 0), timestamp_ntz=datetime.datetime(2021, 1, 1, 0, 0), binary_column=bytearray(b'A'), struct_column=Row(struct_string_column='AAA', struct_int_column=1), array_column=['AAA', 'BBB', 'CCC'], map_column={'AAA': 1})[0m
```

## 大量データの場合の動作確認

### 正常系

```python
filepath = "dbfs:/databricks-datasets/tpch/data-001/part/part.tbl"

schema = """
  p_partkey long,
  p_name string,
  p_mfgr string,
  p_brand string,
  p_type string,
  p_size int,
  p_container string,
  p_retailprice decimal(12, 2),
  p_comment string
"""

test_df_04 = spark.read.format("csv").schema(schema).option("sep", "|").load(filepath)
print(test_df_04.count())
```

```python
assertDataFrameEqual(
    test_df_04,
    test_df_04,
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/6d7cbba2-8f12-e8f9-a759-2d78003a175e.png)

### 異常系

```python
filepath = "dbfs:/databricks-datasets/tpch/data-001/part/part.tbl"

schema = """
  p_partkey long,
  p_name string,
  p_mfgr string,
  p_brand string,
  p_type string,
  p_size int,
  p_container string,
  p_retailprice decimal(12, 2),
  p_comment string
"""

exp_df_04 = spark.read.format("csv").schema(schema).option("sep", "|").load(filepath)
exp_df_04 = exp_df_04.sample(0.9999)
print(exp_df_04.count())
```

```python
assertDataFrameEqual(
    test_df_04,
    exp_df_04,
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/930880e4-e90a-55fb-77dd-7e81c86f1c2e.png)


次のようなメッセージが表示され、タイムアウトしてしまい、データの差異があることを確認できませんでした。

```text
TimeoutException: Futures timed out after [30 seconds]
```

## 結論

本記事では、PySparkの`assertDataFrameEqual`関数の使用方法とその検証結果について説明しました。この関数は、データフレームの等価性を確認するための強力なツールであり、データ型の違いなどに対しても正確に動作します。しかし、パフォーマンスについても考慮する必要があります。この関数を理解し、適切に使用することで、データフレームのテストがより簡単かつ効率的になります。