Pyspark中的条件如何动态链接？答案

【问题标题】：How to dynamically chain when conditions in Pyspark?Pyspark中的条件如何动态链接？
【发布时间】：2021-01-30 03:26:59
【问题描述】：

上下文

数据框应具有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|
+-----+-----+-----+

硬编码解决方案

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] 添加更复杂的条件来解释为什么需要链接。

【问题讨论】：

标签： python dataframe apache-spark pyspark

【解决方案1】：

我希望这个解决方案适合你：

使用您使用列“组”和“类别”定义的元组列表创建一个新数据框： category_rules = [('A', 'small'),('B', 'large'), 等等] 这将是您的“lookup_dataframe”

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

然后，您可以在“组”列上左连接两个数据框，因此对于具有组值的每一行，将从 lookup_df 获得您连接的列中的类别值。

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

通过进行左连接，如果您的 df（右侧）中有一个未包含在 lookup_df 中的组值，例如“C”，它将具有空值。

【讨论】：

谢谢。那将是一个基本的解决方案。但我想知道如果有一组更复杂的条件，你将如何做到这一点。例如，除了group 之外，一些数值需要匹配一组相关条件。将更新问题以使这一点更清楚。

【解决方案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|
# +-----+-----+--------+

【讨论】：

如更新后的问题所示，这也应该适用于多种复杂条件。由于这基本上是 SQL，我想知道您是否发现建议的解决方案存在任何安全问题，例如 SQL 注入？
大数据存储中的 SQL 注入？我认为这很困难，但绝对不是不可能的。尽管这取决于您如何生成 category_rules 列表，因为这是任何潜在注入的唯一来源。它还支持复杂的条件是的

【解决方案3】：

基于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()

【讨论】：