isinで特定の番号の範囲を取得しようとしたらSparkRuntimeExceptionが発生した

はじめに

PySparkで作成したDataFrameのデータのうち、
idが20番台のデータを取得したいコードを書いたところ、SparkRuntimeExceptionが発生してしまったので備忘メモとして残します。

やりたいこと

• 作成したDataFrameのレコードの中で、idカラムが20番台(20~29)のレコードを取得したい！
• 20, 21, 22,...という形で指定するのではなく範囲を指定して取得したい！

環境

使用した環境

・Google Colaboratory
・Python: 3.10.12
・PySpark: 3.5.0

今回使用したコード

以下のようにfilterとisinを使用して20番台のレコードを取得しようとするとSparkRuntimeExceptionが発生してしまいます。

# 20番台の範囲を指定したいのでrangeを使用したがSparkRuntimeExceptionが発生してしまう
sdf.filter(
F.col("id").isin(range(20, 30))
)

エラー内容

---------------------------------------------------------------------------
SparkRuntimeException                     Traceback (most recent call last)
<ipython-input-10-3db942e00f8f> in <cell line: 1>()
      1 sdf.filter(
----> 2 F.col("id").isin(range(20, 30))
      3 )

4 frames
/usr/local/lib/python3.10/dist-packages/pyspark/errors/exceptions/captured.py in deco(*a, **kw)
    183                 # Hide where the exception came from that shows a non-Pythonic
    184                 # JVM exception message.
--> 185                 raise converted from None
    186             else:
    187                 raise

SparkRuntimeException: [UNSUPPORTED_FEATURE.LITERAL_TYPE] The feature is not supported: Literal for '[20, 21, 22, 23, 24, 25, 26, 27, 28, 29]' of class java.util.ArrayList.

原因

公式ドキュメントよりrange型は使用できる型ではないようです。

解決方法

rangeの外側をlistで囲って型を変換することで実行できるようになりました。

df.filter(
F.col("id").isin(list(range(20, 30)))
).show()

出力結果

+---+---+---+
| id|  a|  b|
+---+---+---+
| 20|  n|123|
| 21|  o|234|
| 22|  p|234|
| 23|  q|123|
| 24|  r|123|
| 25|  s|123|
| 26|  t|123|
| 27|  u|123|
| 28|  v|123|
| 29|  w|123|
+---+---+---+