PySpark書いてて混乱した、ArrayType、MapType、StructTypeを整理する

概要

表題の通りです。

データ構造としてどの型をどう使うのが適切なのか整理したかったので今回の記事にしてみました。

具体例だけ見て納得したい方はこちら。
あとは追加で、それぞれのデータ型の展開の仕方も書いてあるので気になる方は直接こちらへ。

整理

特徴を列挙するとこんな感じ。

ArrayType(配列)

• 配列の各要素はデータ型が同じ必要がある
• 配列は順序を保有する
• 同じ要素を複数回含むことができる

MapType(マップ型)

• キーバリューでデータを保存する
• 値は任意のデータ型を持てる
• 検索が高速

StructType(構造体型)

• 各フィールドに名前がつけられる
• フィールドにアクセスするときはドット演算子やフィールド名でアクセス可能

ここだけ見ると、MapとStructの用途の違いがいまいちピンとこないので、それぞれ似ている特徴のデータ型2つを比較して、さらに具体的な用途まで落とし込んでみたいと思います。

それぞれの違いに注目

MapとStructの違い

• MapTypeはキーと値のペアを動的に保持するのでキー自体がデータの一部としてみなされる
• StructTypeはフィールドが固定されており、フィールド名があらかじめ決まっている。定義されたフィールドに対してデータの構造が明確に固定される

ArrayとStructの違い

• Arrayは同じ方の複数の値を保持するが、個々の要素に名前やラベルがない
• Structは異なる型を持つ複数のフィールドを保持でき、フィールドに名前をつけられる

その違い、具体的にどういうことやねん

MapTypeの具体例

まず、MapTypeの用途を考えてみます。例えばユーザーのレビュー情報を保存するデータ構造を考えるとき、各商品ごとに評価やコメントなどの属性がキーとして動的に関連付けられます。

動的にと言っているくらいなので、データごとに異なるキーがあっても対応できます。

例えば下のように、user_1, user_2まではrating, commentしかキーはありませんでしたが、新たにrecommendというキーを生やしたり、逆にcommentというキーがなかったりがあり得ます。

schema = MapType(StringType(), StringType())

data = [
    ("user_1", {"rating": "5", "comment": "Great product"}),
    ("user_2", {"rating": "3", "comment": "Average quality"}),
    ("user_3", {"rating": "4", "comment": "Good value", "recommend": "yes"}),
    ("user_4", {"rating": "2", "recommend": "no"})
]

df = spark.createDataFrame(data, ["user_id", "review_details"])
df.show(truncate=False)

データ構造として柔軟に持ちたいときは嬉しそう。

とはいえデータの一貫性を保つのが大変なので、例えば常にrecommendがあるかどうかを確認する処理を入れておくみたいな気遣いが必要そうですね。

df = df.withColumn(
    "recommend", 
    when(df["review_details"].getItem("recommend").isNull(), "unknown")
    .otherwise(df["review_details"].getItem("recommend"))
)

StructTypeの具体例

Mapとは違い、あらかじめ定義されたフィールド名とデータ型を持ちます。ここで、商品情報を保持するスキーマを考えたとき、このようなデータの利用例が考えられます。

from pyspark.sql.types import StructType, StructField, StringType, FloatType

schema = StructType([
    StructField("item_id", StringType(), True),
    StructField("name", StringType(), True),
    StructField("price", FloatType(), True),
    StructField("category", StringType(), True)
])

data = [
    ("001", "Apple", 0.99, "Fruit"),
    ("002", "Banana", 0.79, "Fruit")
]

df = spark.createDataFrame(data, schema=schema)
df.show(truncate=False)

Mapと違うのは、各フィールドのデータの型が事前に決まっていたり、動的なキーが存在しないことです。明らかに構造を定義できるときは適してそう。

あとは、よくネストされた形でStructをみることも多い様な気がします。

schema = StructType([
    StructField("item_id", StringType(), True),
    StructField("details", StructType([
        StructField("name", StringType(), True),
        StructField("price", FloatType(), True),
        StructField("category", StringType(), True)
    ]))
])

data = [
    ("001", ("Apple", 0.99, "Fruit")),
    ("002", ("Banana", 0.79, "Fruit"))
]

df = spark.createDataFrame(data, schema=schema)
df.show(truncate=False)

商品の詳細をネストされた構造体で定義したものです。

データにアクセスするときはドット表記やブラケット(df["item_id"])みたいな形でアクセスできます。

データの展開の仕方の違い

arrayとmapだと、explodeで展開できます。structはexplodeで展開できないのですが、selectで愚直に展開することができます。

structの場合

例えば以下のようなスキーマがある場合

root
 |-- user: struct (nullable = true)
 |    |-- name: string (nullable = true)
 |    |-- address: struct (nullable = true)
 |    |    |-- city: string (nullable = true)
 |    |    |-- state: string (nullable = true)

df.select("user.name", "user.address.city", "user.address.state").show()

これで展開できます。各フィールドを直接指定する感じですね。なので withColumnを利用しても展開することができます。

arrayの場合

いきなりですが、arrayがexplodeで展開できるのはいいとして、structのデータ構造をarrayで持っている場合のデータの展開方法を考えます。

下記のようなデータを定義してみます。

from pyspark.sql import Row
from pyspark.sql.functions import explode

data = [
    Row(order_id="001", items=[
        Row(item_id="A001", content_type="Type1", date_modified="2024-09-01", no=1),
        Row(item_id="A002", content_type="Type2", date_modified="2024-09-02", no=2)
    ]),
    Row(order_id="002", items=[
        Row(item_id="A003", content_type="Type3", date_modified="2024-09-03", no=3)
    ])
]

df = spark.createDataFrame(data)
df.show(truncate=False)

この形のデータ構造に対して、まずはarrayをならす必要があるので explodeします。

df_exploded = df.select("order_id", explode("items").alias("item"))
df_exploded.show(truncate=False)

展開するとこんな感じですが、各orderの詳細を見るためには結局order_idでgroupbyする必要があります。今回はどうデータを使うかは考えずに展開していきます。

ここまでくるとitemはstruct型なので、先ほどのように各フィールドを指定すれば展開できます。

df_final = df_exploded.select(
    "order_id",
    "item.item_id",
    "item.content_type",
    "item.date_modified",
    "item.no"
)
df_final.show(truncate=False)

まとめ

データ型	主な用途	特徴	例
ArrayType	同じ型の複数の値をリストで保持する場合	同じデータ型の要素が順序付きで格納される	["apple", "banana"]
MapType	キーと値のペアを保持し、キーで高速に検索したい場合	キーが一意で、キーと値がペアになって格納される。そしてキーを動的に生成可能	{"age": 30, "name": "A"}
StructType	レコード形式で複数の属性を持つデータを保持する場合	固定された名前を持つフィールドを複数持ち、フィールドごとに型が異なる	{"name": "apple", "price": 1.99}

ちなみに

上記で紹介したのはかなりシンプルなデータ構造であり、データは複雑に持たないに限るのですが、、、

とはいえ現実には array<struct> や structが入り混じってるデータ構造もザラにあります。データは無闇にexplodeするのではなく、使える手札をきちんと理解している上でデータ要件に適切なデータ型を定義していきたいですね。