Databricks (Spark) にて Spark dataframe と Pandas 変換時の PyArrow が利用されないことによるエラー事象への対応方法

## 概要

Databricks (Spark) にて Spark dataframe と Pandas 変換時の PyArrow が利用されないことによるエラーに遭遇したため、事象の再現方法と対応方法を共有します。

Spark dataframe と Pandas 変換時の PyArrow が利用されない場合（`spark.sql.execution.arrow.pyspark.enabled`が`false`の場合）には、ドライバーノードとワーカーノードでデータの通信が行われ、`spark.rpc.message.maxSize`の上限を超えることで次のエラーが発生します。メッセージでは`spark.rpc.message.maxSize`の上限を上げることが提案されていますが、その上限を上げる対応方法を実施しませんでした。

```text
org.apache.spark.SparkException: Job aborted due to stage failure: Serialized task 29:0 was 277567143 bytes, which exceeds max allowed: spark.rpc.message.maxSize (268435456 bytes). Consider increasing spark.rpc.message.maxSize or using broadcast variables for large values.
```

Databricks では、基本的には`spark.sql.execution.arrow.pyspark.enabled`の設定値がデフォルトで`true`となっておりますが、特定の状況下では次のような警告とともに`spark.sql.execution.arrow.pyspark.enabled`の設定値が`false`となります。

```text
/databricks/spark/python/pyspark/sql/pandas/conversion.py:124: UserWarning: toPandas attempted Arrow optimization because 'spark.sql.execution.arrow.pyspark.enabled' is set to true; however, failed by the reason below:
  A null type field may not be non-nullable
Attempting non-optimization as 'spark.sql.execution.arrow.pyspark.fallback.enabled' is set to true.
  warn(msg)
```

Array 型のカラムを追加する際に`array`メソッドをそのまま指定している場合に、上記の警告が表示されることがわかりました。その対応方法として、明示的に CAST する方法で警告が表示されなくなりました。

```python
from pyspark.sql.functions import col, array

df = spark.createDataFrame([(1,)], 'col_1 long')
df = df.withColumn('array_column', array())  ## lit([]) でも同様の結果となる

pdf = df.toPandas()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/5627060a-414f-4590-53fb-20b51d8e28a2.png)

```python
from pyspark.sql.functions import col, array

df = spark.createDataFrame([(1,)], 'col_1 long')
df = df.withColumn('array_column', array().cast('array<void>'))
```

```python
pdf = df.toPandas()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/3949decf-e8c7-bdf9-0fff-2b6a49700a77.png)


本記事では、調査手順、エラーの再現方法、及び、その対応方法を共有します。

## 調査手順

spark のレポジトリーをベースに Github Copilot Chat で本件について質問したところ、次のような回答が返ってきました。私が遭遇したケースでは、PySparkとArrowの間に互換性の問題があったのではないかと仮説をたてました。

```text
下記の警告がArrow最適化が何らかの理由により表示されるケースを箇条書きで教えて。日本語で回答して

/databricks/spark/python/pyspark/sql/pandas/conversion.py:124: UserWarning: toPandas attempted Arrow optimization because 'spark.sql.execution.arrow.pyspark.enabled' is set to true;
```

```text
spark.sql.execution.arrow.pyspark.enabledがtrueに設定されているため、Arrow最適化が試みられ、以下のようなケースで警告が表示される可能性があります：

- Arrowがインストールされていない場合。
- PySparkとArrowの間に互換性の問題がある場合。
- データ形式がArrowによる最適化をサポートしていない場合。
- Arrowの最適化がメモリ制限を超える可能性がある大規模なデータフレームに対して使用された場合。

具体的な警告メッセージは問題の詳細によりますが、上記のいずれかのシナリオが原因である可能性が高いです。
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/d0141450-c56f-71e8-4cb4-9817673ba259.png)

続けて、エラーに関するコードを Github Copilot Chat で提示してもらいました。リンク先のコードを確認したところ、`to_arrow_schema`メソッドがデータ型の変換を実施していることがわかりました。

```text
そのエラーに関係するコードを提示して
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/711d5754-7976-6032-8c63-77b34c68db81.png)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/9f5c7341-b467-dcf2-e5e0-984b768e4f84.png)

引用元：[spark/python/pyspark/sql/pandas/conversion.py at 76b1c122cb7d77e8f175b25b935b9296a669d5d8 · apache/spark (github.com)](https://github.com/apache/spark/blob/76b1c122cb7d77e8f175b25b935b9296a669d5d8/python/pyspark/sql/pandas/conversion.py#L33-L129)

`to_arrow_schema`メソッド内で利用されている`to_arrow_type`における ArrayType の変換に関するコードを確認しました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/b0dc2294-7e16-891b-e56f-2b0ced29ed14.png)

引用元：[spark/python/pyspark/sql/pandas/types.py at 76b1c122cb7d77e8f175b25b935b9296a669d5d8 · apache/spark (github.com)](https://github.com/apache/spark/blob/76b1c122cb7d77e8f175b25b935b9296a669d5d8/python/pyspark/sql/pandas/types.py#L94-L108)

PyArrow 変換時のエラーが発生するケースとして、次のようなコードとなる場合（`NullType`かつ`nullable`が`False`）であることを特定できました。`A null type field may not be non-nullable`という`spark.sql.execution.arrow.pyspark.enabled`が`false`となった際の警告文内に記載されている一部のメッセージと一致しました。`nullable`を`True`にするためには、カラム追加時に CAST すればよいことが、後述する対応方法を把握することができました。

```python
import pyarrow as pa

print(
    pa.field(
        "element",
        pa.null(),
        False,
    )
)
```

```text
ValueError: A null type field may not be non-nullable
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/fe405336-8bad-6678-6848-cf2b400b93c3.png)

## 事前準備

### テストスキーマとテーブルの準備

```python
%sql
CREATE SCHEMA if NOT EXISTS hive_metastore.oom_test;
CREATE OR REPLACE TABLE hive_metastore.oom_test.partsupp (
  ps_partkey long,
  ps_suppkey long,
  ps_availqty int,
  ps_supplycost float,
  ps_comment string
);
```

```python
filepath = "dbfs:/databricks-datasets/tpch/data-001/partsupp/partsupp.tbl"
tgt_tbl = "hive_metastore.oom_test.partsupp"


schema = """
ps_partkey long,
ps_suppkey long,
ps_availqty int,
ps_supplycost float,
ps_comment string
"""

df = (spark
       .read
       .format("csv")
       .schema(schema)
       .option("sep", "|")
       .load(filepath)
    )

df.write.mode("overwrite").saveAsTable(tgt_tbl)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/a2029282-5f13-6026-5f64-75fae6380615.png)

## `spark.sql.execution.arrow.pyspark.enabled`が`false`の場合にエラーになることを確認

### `spark.sql.execution.arrow.pyspark.enabled`が`True`の場合の動作確認

正常終了しました。

```python
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
print(spark.conf.get("spark.sql.execution.arrow.pyspark.enabled"))

tgt_tbl = "hive_metastore.oom_test.partsupp"

df = spark.table(tgt_tbl)
pdf = df.toPandas()

new_col = f"ps_comment_1"
pdf[new_col] = pdf['ps_comment'] + f"_1"

df_from_pandas = spark.createDataFrame(pdf)
print(df_from_pandas.count())
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/dc046c65-a8a6-a5c6-eaa3-fcca04dbe666.png)

### `spark.sql.execution.arrow.pyspark.enabled`が`False`の場合の動作確認

エラーが発生しました。

```python
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "false")
print(spark.conf.get("spark.sql.execution.arrow.pyspark.enabled"))

tgt_tbl = "hive_metastore.oom_test.partsupp"

df = spark.table(tgt_tbl)
pdf = df.toPandas()

new_col = f"ps_comment_1"
pdf[new_col] = pdf['ps_comment'] + f"_1"

df_from_pandas = spark.createDataFrame(pdf)
print(df_from_pandas.count())
```

```text
org.apache.spark.SparkException: Job aborted due to stage failure: Serialized task 29:0 was 277567143 bytes, which exceeds max allowed: spark.rpc.message.maxSize (268435456 bytes). Consider increasing spark.rpc.message.maxSize or using broadcast variables for large values.
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/945820b3-ddca-95bc-41c6-1d13e8e005b7.png)

### `spark.sql.execution.arrow.pyspark.enabled`が`False`、かつ、Pandas のデータサイズが大きくなる処理を実行する場合

Spark ドライバーが異常終了します。

```text
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "false")
print(spark.conf.get("spark.sql.execution.arrow.pyspark.enabled"))

tgt_tbl = "hive_metastore.oom_test.partsupp"

df = spark.table(tgt_tbl)
pdf = df.toPandas()

for i in range(1, 6):
    new_col = f"ps_comment_{i}"
    pdf[new_col] = pdf['ps_comment'] + f"_{i}"

df_from_pandas = spark.createDataFrame(pdf)
print(df_from_pandas.count())
```

```text
The spark driver has stopped unexpectedly and is restarting. Your notebook will be automatically reattached.
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/defa5302-f578-d58a-c40f-cdf5a584b3a9.png)

## `spark.sql.execution.arrow.pyspark.enabled`が`false`となるケースと対応方法

### `spark.sql.execution.arrow.pyspark.enabled`が`false`となるケース

前述の警告が表示されます。

```python
from pyspark.sql.functions import col, array

df = spark.createDataFrame([(1,)], 'col_1 long')
df = df.withColumn('array_column', array())  ## lit([]) でも同様の結果となる

pdf = df.toPandas()
```

```text
/databricks/spark/python/pyspark/sql/pandas/conversion.py:124: UserWarning: toPandas attempted Arrow optimization because 'spark.sql.execution.arrow.pyspark.enabled' is set to true; however, failed by the reason below:
  A null type field may not be non-nullable
Attempting non-optimization as 'spark.sql.execution.arrow.pyspark.fallback.enabled' is set to true.
  warn(msg)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/84dcb77f-2645-e762-f3b8-34d3fd30a9b0.png)

上記コードの結果に基づき PyArrow への変換を実施すると、エラーとなりました。

```python
print(f"fields:       {df.schema.fields}")
print(f"dataType:     {df.schema.fields[1].dataType.elementType}")
print(f"containsNull: {df.schema.fields[1].dataType.containsNull}")


import pyarrow as pa

print(
    pa.field(
        "element",
        pa.null(),  # NullType
        df.schema.fields[1].dataType.containsNull,
    )
)
```

```text
ValueError: A null type field may not be non-nullable
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/8ccaf2e9-2969-8112-23c5-241ec8ac28cf.png)

### データ型を明示することによる対応方法

CAST によりデータ変換を実施することで、警告が表示されなくなりました。

```python
from pyspark.sql.functions import col, array

df = spark.createDataFrame([(1,)], 'col_1 long')
df = df.withColumn('array_column', array().cast('array<void>'))
```

```python
pdf = df.toPandas()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/3949decf-e8c7-bdf9-0fff-2b6a49700a77.png)

上記コードの結果に基づき PyArrow への変換を実施すると、想定通りに変換できました。

```python
print(f"fields:       {df.schema.fields}")
print(f"dataType:     {df.schema.fields[1].dataType.elementType}")
print(f"containsNull: {df.schema.fields[1].dataType.containsNull}")
```

```pyton
import pyarrow as pa

print(
    pa.field(
        "element",
        pa.null(),  # NullType
        df.schema.fields[1].dataType.containsNull,
    )
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/16b222d1-176a-f166-5938-0b2fb86dd1bd.png)

## まとめ

この記事は、Databricks (Spark) で Spark dataframe と Pandas の変換時に PyArrow が利用されないことによるエラー事象について説明しました。このエラーは、spark.sql.execution.arrow.pyspark.enabledがfalseの場合に発生します。この設定がfalseになると、ドライバーノードとワーカーノード間でデータの通信が行われ、その結果、spark.rpc.message.maxSizeの上限を超える可能性があります。

記事では、この問題の再現方法と対応策を詳しく説明しています。具体的には、Array 型のカラムを追加する際にarrayメソッドをそのまま指定すると、上記の警告が表示されることが確認されました。これを解決するための対応策として、カラムを追加する際にデータ型を明示的に CAST する方法が提案しています。