web-dev-qa-db-ja.com

PySpark：2つの条件と3つの結果を持つwithColumn（）

私はSparkとPySparkを使用しています。次の擬似コードと同等の結果を達成しようとしています。

df = df.withColumn('new_column', 
    IF fruit1 == fruit2 THEN 1, ELSE 0. IF fruit1 IS NULL OR fruit2 IS NULL 3.)

私はPySparkでこれをしようとしていますが、構文についてはわかりません。ポインタはありますか？ expr()を調べましたが、機能させることができませんでした。

dfはpyspark.sql.dataframe.DataFrameであることに注意してください。

apache-sparkhivepysparkapache-spark-sqlhiveql
27
user2205916

これを実装するには、いくつかの効率的な方法があります。必要なインポートから始めましょう：

from pyspark.sql.functions import col, expr, when

Expr内でHive IF関数を使用できます。

new_column_1 = expr(
    """IF(fruit1 IS NULL OR fruit2 IS NULL, 3, IF(fruit1 = fruit2, 1, 0))"""
)

またはwhen + otherwise：

new_column_2 = when(
    col("fruit1").isNull() | col("fruit2").isNull(), 3
).when(col("fruit1") == col("fruit2"), 1).otherwise(0)

最後に、次のトリックを使用できます。

from pyspark.sql.functions import coalesce, lit

new_column_3 = coalesce((col("fruit1") == col("fruit2")).cast("int"), lit(3))

サンプルデータの場合：

df = sc.parallelize([
    ("orange", "Apple"), ("kiwi", None), (None, "banana"), 
    ("mango", "mango"), (None, None)
]).toDF(["fruit1", "fruit2"])

これは次のように使用できます。

(df
    .withColumn("new_column_1", new_column_1)
    .withColumn("new_column_2", new_column_2)
    .withColumn("new_column_3", new_column_3))

結果は次のとおりです。

+------+------+------------+------------+------------+
|fruit1|fruit2|new_column_1|new_column_2|new_column_3|
+------+------+------------+------------+------------+
|orange| Apple|           0|           0|           0|
|  kiwi|  null|           3|           3|           3|
|  null|banana|           3|           3|           3|
| mango| mango|           1|           1|           1|
|  null|  null|           3|           3|           3|
+------+------+------------+------------+------------+
46
user6910411

次のようにUDFを使用する必要があります

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf

def func(fruit1, fruit2):
    if fruit1 == None or fruit2 == None:
        return 3
    if fruit1 == fruit2:
        return 1
    return 0

func_udf = udf(func, IntegerType())
df = df.withColumn('new_column',func_udf(df['fruit1'], df['fruit2']))
15
David