如何在 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”，它将具有空值。