PySparkでlistからOR条件を生成する

概要

PySparkでlistの中に入っている値を使ってOR条件を生成したい。

例えばlistに[('a', 1), ('b', 2)]のようなtupleが入っていて、 (x='a' and y=1) or (x='b' and y=2) のような条件を生成する。

バージョン情報

• Spark 2.4.3

サンプルデータ

カラムx(char)とカラムy(int)の2つのカラムを持つDataFrameを生成する。

from pyspark.sql import Row

df = spark.createDataFrame([
    Row(x='A', y=1),
    Row(x='B', y=2),
    Row(x='C', y=3),
])
df.collect()
    #=> [Row(x='A', y=1), Row(x='B', y=2), Row(x='C', y=3)]

単一カラムでOR条件

単一のカラムであれば難しいことはなくisinでリストを渡せば済む。

from pyspark.sql import functions as F

df.where(F.col('x').isin(['A', 'C'])) \
    .collect()
    #=> [Row(x='A', y=1), Row(x='C', y=3)]

複数カラムでOR条件

冒頭で例に出した (x='a' and y=1) or (x='b' and y=2) のような条件について。

条件はこんな風に宣言して連結できる。

cond1 = (F.col('x') == 'A') & (F.col('y') == 1)
cond2 = (F.col('x') == 'B') & (F.col('y') == 2)

df.where(cond1 | cond2) \
    .collect()
    #=> [Row(x='A', y=1), Row(x='B', y=2)]

[('a', 1), ('b', 2)] のような変数に入っている場合は、 |= で連結していけば条件が作れる。

arr = [('a', 1), ('b', 2)]
cond = None
for x, y in arr:
    cond |= (F.col('x') == x) & (F.col('y') == y)

df.where(cond) \
    .collect()
    #=>  [Row(x='A', y=1), Row(x='B', y=2)]

改定履歴

Author: Masato Watanabe, Date: 2019-09-03, 記事投稿