如何在 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列的内容。规则列表如上一节所示。

输入数据

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

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

硬编码解决方案

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

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')
    )
)

[编辑 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()