如何在 Pyspark 中的條件下動態鏈接？

Question

語境

數據框應具有category列，該列基於一組固定規則。 規則集變得相當大。

題

有沒有辦法使用元組列表（參見下面的示例）來動態鏈接when條件以實現與底部硬編碼解決方案相同的結果。

# Potential list of rule definitions
category_rules = [
    ('A', 8, 'small'),
    ('A', 30, 'large'),
    ('B', 5, 'small'),
    # Group, size smaller value --> Category
    # and so on ... e.g.,
]

例子

這是一個重現性的玩具示例。 由組和 id 組成的數據框應該添加列category ，這取決於group列的內容。 規則列表如上一節所示。

輸入數據

df = df.withColumn(
    'category',
    F.when(
        (F.col('group') == 'A')
        & (F.col('size') < 8),
        F.lit('small')
    ).when(
        (F.col('group') == 'A')
        & (F.col('size') < 30),
        F.lit('large')
    ).when(
        (F.col('group') == 'B')
        & (F.col('size') < 5),
        F.lit('small')
    ).otherwise(
        F.lit('unkown')
    )
)

+-----+-----+----+--------+
|group|   id|size|category|
+-----+-----+----+--------+
|    A|45345|   5|   small|
|    C|55345|   5|  unkown|
|    A|35345|  10|   large|
|    B|65345|   4|   small|
+-----+-----+----+--------+

硬編碼解決方案

df = df.withColumn( 'category', F.when( (F.col('group') == 'A') & (F.col('size') < 8), F.lit('small') ).when( (F.col('group') == 'A') & (F.col('size') < 30), F.lit('large') ).when( (F.col('group') == 'B') & (F.col('size') < 5), F.lit('small') ).otherwise( F.lit('unkown') ) )

 +-----+-----+----+--------+ |group| id|size|category| +-----+-----+----+--------+ | A|45345| 5| small| | C|55345| 5| unkown| | A|35345| 10| large| | B|65345| 4| small| +-----+-----+----+--------+

[編輯 1] 添加更復雜的條件來解釋為什么需要鏈接。

Answer 1

基於dataframe api的解決方案：

cond = F.when(F.col('group') == category_rules[0][0], F.lit(category_rules[0][1]))
for c in category_rules[1:]:
    cond = cond.when(F.col('group') == c[0], F.lit(c[1]))
cond = cond.otherwise('unknown')

df.withColumn("category", cond).show()

Answer 2

您可以使用字符串插值來創建表達式，例如：

CASE 
   WHEN (group = 'A') THEN 'small' 
   WHEN (group = 'B') THEN 'large'
   ELSE 'unkown'
END

然后在 Spark 表達式中使用它：

from pyspark.sql.functions import expr

data = [('A', '45345'), ('C', '55345'), ('A', '35345'), ('B', '65345')]
df = spark.createDataFrame(data, ['group', 'id'])

category_rules = [('A', 'small'), ('B', 'large')]

when_cases = [f"WHEN (group = '{r[0]}') THEN '{r[1]}'" for r in category_rules]

rules_expr = "CASE " + " ".join(when_cases) + " ELSE 'unkown' END"
# CASE WHEN (group = 'A') THEN 'small' WHEN (group = 'B') THEN 'large' ELSE 'unkown' END

df.withColumn('category', expr(rules_expr)).show()

# +-----+-----+--------+
# |group|   id|category|
# +-----+-----+--------+
# |    A|45345|   small|
# |    C|55345|  unkown|
# |    A|35345|   small|
# |    B|65345|   large|
# +-----+-----+--------+

Answer 3

我希望這個解決方案適合你：

使用您使用“group”和“category”列定義的元組列表創建一個新數據框： category_rules = [('A', 'small'),('B', 'large'), etc] 這將是你的“lookup_dataframe”

lookup_df = spark.createDataFrame(category_rules , ['group', 'category'])

然后，您可以在“組”列上加入兩個數據框，因此對於具有組值的每一行，都將獲得您從 lookup_df 加入的列中的類別值。

df = df.join(lookup_dataframe, ['group'], 'left')

通過進行左連接，如果您的 df（右側）中有一個未包含在 lookup_df 中的組值，例如“C”，它將具有空值。