Databricks にて Amazon が開発している大規模なデータ品質をテスト可能な PyDeequ を動かしてみた

## 概要

Databricks にて Amazon が開発している PyDeequ の動作確認を実施しましたので共有します。PyDeequ (python-deequ)の README に記載されているクイックスタートを Databricks の仕様に修正したコードで実行しました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/62c4164f-ec2e-3f10-9623-4742f2d7717b.png)

引用元：[awslabs/python-deequ: Python API for Deequ (github.com)](https://github.com/awslabs/python-deequ)

## クイックスタート on Databricks

### PyDeequ に対応した Spark バージョンのクラスター（`12.2 LTS (includes Apache Spark 3.3.2, Scala 2.12)`）にて deequ (`com.amazon.deequ:deequ:2.0.4-spark-3.3`)をインストール

 ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/f649faa7-792d-8641-d530-c0feb0a8f045.png)

### ライブラリのインストール、および、`SPARK_VERSION`変数の定義

```python
%pip install pydeequ==1.2.0 -q
dbutils.library.restartPython()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/f7784a9a-acd9-74d0-4908-0ba5ed030bc0.png)

```python
import os
os.environ["SPARK_VERSION"] = spark.version
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/eb622672-1ced-c262-613f-6a16f7442703.png)

### Set up a PySpark session

```python
from pyspark.sql import SparkSession, Row
import pydeequ

# Databricks では SparkSession の定義が不要
# spark = (SparkSession
#     .builder
#     .config("spark.jars.packages", pydeequ.deequ_maven_coord)
#     .config("spark.jars.excludes", pydeequ.f2j_maven_coord)
#     .getOrCreate())

df = spark.sparkContext.parallelize(
    [
        Row(a="foo", b=1, c=5),
        Row(a="bar", b=2, c=6),
        Row(a="baz", b=3, c=None),
    ]
).toDF()
```

```python
df.display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/e6fb8090-eef6-4cbf-48b1-960a963a6f80.png)

### Analyzers

```python
from pydeequ.analyzers import *

analysisResult = (
    AnalysisRunner(spark)
    .onData(df)
    .addAnalyzer(Size())
    .addAnalyzer(Completeness("b"))
    .run()
)

analysisResult_df = AnalyzerContext.successMetricsAsDataFrame(spark, analysisResult)
# analysisResult_df.show()
```

```python
analysisResult_df.display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/4b4445e3-d4c3-96e9-8ce3-a47182c97871.png)

### Profile

```python
from pyspark.sql.types import StringType

# String 型のカラムを除外
df2 = df.select([column for column in df.columns if df.schema[column].dataType != StringType()])
```

```python
from pydeequ.profiles import *

result = ColumnProfilerRunner(spark).onData(df2).run()

for col, profile in result.profiles.items():
    print(profile)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/0a7cdd75-6ec8-36ea-f8dd-ad91106b6d11.png)

### Constraint Suggestions

```python
from pydeequ.suggestions import *

suggestionResult = (
    ConstraintSuggestionRunner(spark).onData(df).addConstraintRule(DEFAULT()).run()
)

# Constraint Suggestions in JSON format
print(suggestionResult)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/4007e7ee-530a-63d6-5373-f34442c258f9.png)

### Constraint Verification

```python
from pydeequ.checks import *
from pydeequ.verification import *

check = Check(spark, CheckLevel.Warning, "Review Check")

checkResult = (
    VerificationSuite(spark)
    .onData(df)
    .addCheck(
        check.hasSize(lambda x: x >= 3)
        .hasMin("b", lambda x: x == 0)
        .isComplete("c")
        .isUnique("a")
        .isContainedIn("a", ["foo", "bar", "baz"])
        .isNonNegative("b")
    )
    .run()
)

checkResult_df = VerificationResult.checkResultsAsDataFrame(spark, checkResult)
# checkResult_df.show()
```

```python
checkResult_df.display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/767c3e97-a98f-7141-d191-abaa5355f303.png)

### Repository

```python
from pydeequ.repository import *
from pydeequ.analyzers import *

metrics_file = FileSystemMetricsRepository.helper_metrics_file(spark, "metrics.json")
repository = FileSystemMetricsRepository(spark, metrics_file)
key_tags = {"tag": "pydeequ hello world"}
resultKey = ResultKey(spark, ResultKey.current_milli_time(), key_tags)

analysisResult = (
    AnalysisRunner(spark)
    .onData(df)
    .addAnalyzer(ApproxCountDistinct("b"))
    .useRepository(repository)
    .saveOrAppendResult(resultKey)
    .run()
)
```

```python
result_metrep_df = (
    repository.load()
    .before(ResultKey.current_milli_time())
    .forAnalyzers([ApproxCountDistinct("b")])
    .getSuccessMetricsAsDataFrame()
)
```

```python
result_metrep_df.display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/5c5f47ca-bd4a-8092-032c-f592ea4d6906.png)


## 注意事項

### 1. pydeequ のインポート時のエラーへの対応方法

pydeequ のインポート時に次のようなエラーが発生する場合があります。

```text
RuntimeError: SPARK_VERSION environment variable is required. Supported values are: dict_keys(['3.3', '3.2', '3.1', '3.0', '2.4'])
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/c80a0606-e39d-7f59-6f9c-b8fa7fb5f97a.png)

`SPARK_VERSION`変数に spark のバージョンをセットすることで対応できます。

```python
import os
os.environ["SPARK_VERSION"] = spark.version
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/63419723-eac6-9ddd-83f6-1b1636bdbe02.png)

### Profile 時の StringColumnProfile エラーへの対応方法

```text
KeyError: 'StringColumnProfile'
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/fe07e21a-cd93-dd57-6183-d6a4a287b05c.png)

String 型のカラムがあるとエラーとなるため、そのカラムを除外することで対応できます。

```python
from pyspark.sql.types import StringType

# String 型のカラムを除外
df2 = df.select([column for column in df.columns if df.schema[column].dataType != StringType()])
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/a3ee7ed6-8b14-3d3c-54a3-b224aaa68774.png)