Databricksにおけるpandas function API

[pandas function APIs \| Databricks on AWS](https://docs.databricks.com/spark/latest/spark-sql/pandas-function-apis.html) [2022/6/27時点]の翻訳です。

:::note warn
本書は抄訳であり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
:::

pandas function APIを用いることで、pandasインスタンスを入力として受け取り、pandasインスタンスを出力するネイティブのPython関数をPySparkデータフレームに直接適用できるようになります。function APIは[pandasユーザー定義関数](https://qiita.com/taka_yayoi/items/b65197128ee698d87910)と同様に、データを転送し、pandasがデータを操作する際に[Apache Arrow](https://arrow.apache.org/)を使用します。しかし、pandas function APIではPython型ヒントはオプションです。

pandas function APIには3つのタイプがあります。

- Grouped map
- Map
- Cogrouped map

pandas function APIは、pandas UDFの実行で使用されるのと同じ内部ロジックを活用します。このため、PyArrow、サポートされるSQLの対応、設定などはpandas UDFと同じ特性を共有します。

詳細に関しては、[Spark3\.0における新機能: Pandas UDFとPython型ヒント](https://qiita.com/ktmrmshk/items/24aa8467b906c56c332f)をご覧ください。

# Grouped map

「split-apply-combine」パターンを実装するために`groupBy().applyInPandas()`を通じて、グループ分けされたデータの変換を行います。split-apply-combineは3つのステップから構成されます。

- `DataFrame.groupBy`を用いてデータをグループに分割します。
- それぞれのグループに関数を適用します。関数の入力、出力は両方とも`pandas.DataFrame`です。入力データにはそれぞれのグループの全ての行と列が含まれます。
- 新たな`DataFrame`に結果を統合します。

`groupBy().applyInPandas()`を使用するには、以下を定義する必要があります。
- それぞれのグループに対する処理を定義するPython関数
- 出力`DataFrame`のスキーマを定義する`StructType`オブジェクトあるいは文字列

返却される`pandas.DataFrame`のカラムのラベルは、出力スキーマが文字列として定義された場合はフィールド名が一致する必要があり、文字列でない場合には、integerのインデックスなどフィールのデータ型がマッチする必要があります。`pandas.DataFrame`を構成する際にどのようにカラムにラベルをつけるのかに関しては、[pandas\.DataFrame](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html)を参照ください。

関数を適用する前に、グループの全データがメモリーにロードされます。特にグループのサイズに偏りがある場合、アウトオブメモリー例外を引き起こすことがあります。設定[maxRecordsPerBatch](https://spark.apache.org/docs/latest/sql-pyspark-pandas-with-arrow.html#setting-arrow-batch-size)はグループには適用されないので、ご自身の手で、グループ分けされたデータがメモリーに収まるようにする必要があります。

以下の例では、グループのそれぞれの値から平均値を引き算するためにどのように`groupby().applyInPandas()`を使用するのかを示しています。

```py:Python
df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
    ("id", "v"))

def subtract_mean(pdf):
    # pdf is a pandas.DataFrame
    v = pdf.v
    return pdf.assign(v=v - v.mean())

df.groupby("id").applyInPandas(subtract_mean, schema="id long, v double").show()
# +---+----+
# | id|   v|
# +---+----+
# |  1|-0.5|
# |  1| 0.5|
# |  2|-3.0|
# |  2|-1.0|
# |  2| 4.0|
# +---+----+
```

詳細な使い方に関しては、[pyspark\.sql\.GroupedData\.applyInPandas](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.GroupedData.applyInPandas.html#pyspark-sql-groupeddata-applyinpandas)をご覧ください。

# Map

`DataFrame.mapInPandas()`を用いることでpandasのインスタンスを用いたmapオペレーションを実行し、`pandas.DataFrame`のイテレータを現在のPySparkデータフレームを表現する`pandas.DataFrame`の別のイテレータに変換し、結果をPySparkデータフレームとして返却します。

内部の関数は、`pandas.DataFrame`のイテレータを受け取り、`pandas.DataFrame`のイテレータを返却します。Series to Series pandas UDFのような幾つかのpandas UDFと異なり、任意の長さの出力を返却することができます。

以下の例では`mapInPandas()`の使い方を説明しています。

```py:Python
df = spark.createDataFrame([(1, 21), (2, 30)], ("id", "age"))

def filter_func(iterator):
    for pdf in iterator:
        yield pdf[pdf.id == 1]

df.mapInPandas(filter_func, schema=df.schema).show()
# +---+---+
# | id|age|
# +---+---+
# |  1| 21|
# +---+---+
```

詳細については、[pyspark\.sql\.DataFrame\.mapInPandas](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.mapInPandas.html#pyspark-sql-dataframe-mapinpandas)をご覧ください。

# Cogrouped map

pandasインスタンスを用いたCogrouped mapオペレーションには、共通のキーでグループ結合される2つのPySpark`Dataframe`に対して`DataFrame.groupby().cogroup().applyInPandas()`を使用し、それぞれの共通グループにPython関数が適用されます。以下のステップから構成されます。

- それぞれのデータフレームでキーを同じくするグループのデータがシャッフルされ、同じ共通グループにまとめられます。
- それぞれの共通グループに関数を適用します。関数の入力は`pandas.DataFrame`です(オプションとしてキーを表現するタプル)。関数の出力は`pandas.DataFrame`です。
- 全てのグループの`pandas.DataFrame`を新たなPySpark `DataFrame`で結合します。

`groupBy().cogroup().applyInPandas()`を使用するには、以下を定義する必要があります。

- それぞれの共通グループに対する計算処理を定義するPython関数
- 出力`DataFrame`のスキーマを定義する`StructType`オブジェクトあるいは文字列

返却される`pandas.DataFrame`のカラムのラベルは、出力スキーマが文字列として定義された場合はフィールド名が一致する必要があり、文字列でない場合には、integerのインデックスなどフィールのデータ型がマッチする必要があります。`pandas.DataFrame`を構成する際にどのようにカラムにラベルをつけるのかに関しては、[pandas\.DataFrame](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html)を参照ください。

関数を適用する前に、グループの全データがメモリーにロードされます。特にグループのサイズに偏りがある場合、アウトオブメモリー例外を引き起こすことがあります。設定[maxRecordsPerBatch](https://spark.apache.org/docs/latest/sql-pyspark-pandas-with-arrow.html#setting-arrow-batch-size)はグループには適用されないので、ご自身の手で、グループ分けされたデータがメモリーに収まるようにする必要があります。

以下の例では、2つのデータセット間の`asof join`を実行するために、どのように`groupby().cogroup().applyInPandas()`を使用するのかを示しています。

```py:Python
import pandas as pd

df1 = spark.createDataFrame(
    [(20000101, 1, 1.0), (20000101, 2, 2.0), (20000102, 1, 3.0), (20000102, 2, 4.0)],
    ("time", "id", "v1"))

df2 = spark.createDataFrame(
    [(20000101, 1, "x"), (20000101, 2, "y")],
    ("time", "id", "v2"))

def asof_join(l, r):
    return pd.merge_asof(l, r, on="time", by="id")

df1.groupby("id").cogroup(df2.groupby("id")).applyInPandas(
    asof_join, schema="time int, id int, v1 double, v2 string").show()
# +--------+---+---+---+
# |    time| id| v1| v2|
# +--------+---+---+---+
# |20000101|  1|1.0|  x|
# |20000102|  1|3.0|  x|
# |20000101|  2|2.0|  y|
# |20000102|  2|4.0|  y|
# +--------+---+---+---+
```

詳細な使い方については、[pyspark\.sql\.PandasCogroupedOps\.applyInPandas](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.PandasCogroupedOps.applyInPandas.html#pyspark-sql-pandascogroupedops-applyinpandas)をご覧ください。

### Databricks 無料トライアル

[Databricks 無料トライアル](https://databricks.com/jp/try-databricks)