Delta LakeとSparkのバリアントデータ型をDatabricksで試す

# 導入

以下のDatabricks公式Blogにて、バリアントデータ型の導入が紹介されていました。

https://www.databricks.com/jp/blog/introducing-open-variant-data-type-delta-lake-and-apache-spark

Delta 4.0/Spark 4.0からバリアントサポートは含まれるようであり、DatabricksではDBR 15.3より利用可能になるとのこと。

既にDatabricks社員の方が試された記事も掲載されていました。

https://zenn.dev/akuwano/articles/c2c91c88b56536

DBRは15.3がベータとして公開されていることもあり、私も興味津々でしたので、少し試してみました。

検証環境はDatabricks on AWS、DBRは15.3ML(Beta)です。

:::note warn
上記の通りDBR15.3はベータリリースのため、将来的に下記コードは動かなくなる可能性があります。
:::


# とりあえず動かす

Pythonのノートブックを作って、JSON形式の**文字列**を含んだデータフレームを作成します。

意図的に異なる構造のデータにしたかったので、以下のように2種のJSON形式文字列を定義。

```python
json_str1 = '{"name":"Bob", "age":30, "cars": ["Ford", "BMW"], "pets": null}'
json_str2 = '{"id":"1234567890", "name":"Jhon Doe", "first_name":"Jhon", "last_name":"Doe", "address": {"street":"Main St", "city":"New York", "state":"NY"}}'
```

作成したJSON型文字列を使ってデータフレームを作成。

```python
df = spark.createDataFrame(
    [["001", json_str1], ["002", json_str2]],
    ["id", "json_str"],
)

# SQLで利用するため、Temporary Viewとしても定義
df.createOrReplaceTempView("json_table")
display(df)
```

以下のようなデータフレームができました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/d9711cd3-2c60-0c7c-76c2-ff1fdc467b2d.png)

ちなみに、アイコンを見て分かるように、全ての列は文字列型です。

念のため、スキーマ情報を出力。

```python
df.printSchema()
```

```text:出力
root
 |-- id: string (nullable = true)
 |-- json_str: string (nullable = true)
```

文字列ですね。

では、上記のブログにあるように、SQLで```parse_json```を使ってバリアント型へ変換してみます。

```sql
%sql
SELECT
  id,
  parse_json(json_str)
FROM
  json_table;
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/7e1de2cd-c25d-e00a-c7d8-8ce00b95818e.png)

JSON文字列がバリアント型へ変換されました。
また、異なる構造のJSONを同一列に含めましたが、問題なく変換されました。

従来は、構造の異なるJSON文字列を非文字列型データとして保持する際は、スキーマを合わせてStruct型で保管するか、フラットな構造の場合はDict型で保管するようなやり方だったかなと思います。
データによってはこういった変換や管理は手間でしたが、バリアント型を使えるようになることでより柔軟に半構造化データ列を管理できるようになりますね。

# Pysparkで使う

SQLではなくPython(Pyspark)でもやってみましょう。

Pysparkでも```parse_json```関数が提供されていますので、それを使って変換します。

```python
import pyspark.sql.functions as F

variant_df = df.select("id", F.parse_json("json_str").alias("json_obj"))

display(variant_df)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/9bd50e78-761d-9d64-89ea-209a3ca4c158.png)

SQLと同等の結果が得られました。

では、Variant列から要素を取り出すクエリを実行してみます。
SQLライクにパスナビゲーションを書く場合は以下のようになります。

```python
variant_df.selectExpr("id", "json_obj:name::string").display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/6a3460ba-1897-94fa-9598-5f75feeb25af.png)

パスナビゲーションは```列名```:```ドット区切りの要素名```::```変換後データタイプ```という書式の模様。

したがって、ネストした要素へのアクセスは以下のようにドット区切りで指定します。

```python
variant_df.selectExpr("id", "json_obj:address.city::string").display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/2be631dd-a968-c8be-2b6b-3695239d4b82.png)

なお、変換後データタイプ指定を省略すると、バリアント型のデータとして結果が得られます。

また、変換不能なデータタイプを指定すると、エラーになります。

```python
# INT型への変換はできない
variant_df.selectExpr("id", "json_obj:address.city::int").display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/e625bcb3-0bc2-d6be-61a2-a8df34b90aef.png)

よりPythonライクな記述をしたい場合は、```variant_get```関数を使うことが出来ます。

```python
variant_df.select("id", F.variant_get("json_obj", "$.name", "string")).display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/a760949d-9602-15a1-f17a-395af190c6f4.png)

```variant_get```は第一引数にバリアント型の列名、第2引数に要素のパスナビゲーション、第3引数に変換後データタイプを指定する関数です。詳細は以下を確認ください。

https://spark.apache.org/docs/4.0.0-preview1/api/python/reference/pyspark.sql/api/pyspark.sql.functions.variant_get.html

```variant_get```もデータタイプに変換不可能なものを指定すると、エラーになります。

このようなエラーを避けたい場合、```variant_get```の代わりに```try_variant_get```を利用することができます。

```python
variant_df.select("id", F.try_variant_get("json_obj", "$.name", "int")).display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/fb9b9f33-bd1c-978c-e442-e1121703201b.png)

```try_variant_get```はエラーが出ない代わりに、変換不能の場合は結果がnullとなります。

<br>

最後に、バリアント型列を含むデータフレームをDelta Tableとして保管します。

```python
# Delta Tableへ保管
variant_df.write.mode("overwrite").saveAsTable("training.default.variant_sample")
```

無事にできました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/a90a091f-24e5-c913-04fa-484fe0a7e731.png)

カタログエクスプローラ上でのプレビューも問題なく出来ます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/8022a01c-e56a-1d60-1377-e219658a2544.png)

詳細を見ると、Propertiesとして```delta.feature.variantType-preview=supported```が設定されていました。(今だけかな)

# まとめ

JSON形式で半構造データを取り扱う場合、かなりやり易くなったという印象です。
(スキーマを厳密に管理すべきユースケースにおいては注意が必要ですが)

最近だとLLMでJSON形式の構造化出力を行うことが増えてきそうだなと感じているのですが、そいういった出力結果の管理や利用の利便性向上につながるのではないかと期待しています。

かなり汎用的に使えそうな機能ですので、広範なユースケースでの適用が今後も増えて行きそうですね。