使用PySpark动态生成CASE WHEN语句实现复杂数据映射

本文介绍如何使用PySpark基于DataFrame中的数据动态生成`CASE WHEN`语句，以实现复杂的数据映射逻辑。通过将映射规则存储在DataFrame中，并根据这些规则构建SQL表达式，可以灵活地处理包含通配符的映射关系，从而避免复杂的JOIN操作，提升数据处理效率。

在PySpark中，有时需要根据DataFrame中的多列值组合来生成结果，并且这些组合与结果的映射关系存储在另一个DataFrame中。当映射关系中包含通配符时，传统的JOIN操作可能难以实现。本文将介绍一种使用动态生成的CASE WHEN语句来解决此问题的方法。

问题描述

假设我们有两个DataFrame：

• df：包含需要进行映射的数据，例如col1, col2, col3等列。
• mapping_table：包含映射规则，其中某些列的值可能是通配符*，表示该列的值不影响结果。

我们的目标是根据mapping_table中的规则，将df中的每一行映射到一个结果值。

解决方案：动态生成CASE WHEN语句

该解决方案的核心思想是将mapping_table转换为一个CASE WHEN语句，然后使用expr函数将其应用到df上。

以下是实现步骤：

1. 构建CASE WHEN语句

   首先，我们需要遍历mapping_table中的每一行，并根据每一行的数据构建一个WHEN子句。如果某个列的值是*，则忽略该列。

   

   爱改写

   AI写作和改写润色工具

   87

   查看详情

   from pyspark.sql import SparkSession
   from pyspark.sql.functions import expr
   
   # 创建 SparkSession
   spark = SparkSession.builder.appName("dynamic_case_when").getOrCreate()
   
   # 示例数据
   map_data = [('a', 'b', 'c', 'good'), ('a', 'a', '*', 'very good'), 
             ('b', 'd', 'c', 'bad'), ('a', 'b', 'a', 'very good'),
             ('c', 'c', '*', 'very bad'), ('a', 'b', 'b', 'bad')]
   
   columns = ["col1", "col2", 'col3', 'result']
   
   mapping_table = spark.createDataFrame(map_data, columns)
   
   data =[[('a', 'b', 'c')], [('a', 'a', 'b')], 
           [('c', 'c', 'a')], [('c', 'c', 'b')],
           [('a', 'b', 'b')], [('a', 'a', 'd')]
         ]
   
   columns = ["col1", "col2", 'col3']
   df = spark.createDataFrame(data, columns)
   df = df.selectExpr(
       "_1.col1 as col1",
       "_1.col2 as col2",
       "_1.col3 as col3"
   )
   
   ressql = 'case '
   for m in map_data:
       p = [f"{p[0]} = '{p[1]}'" for p in zip(columns, m[:3]) if p[1] != "*"]
       ressql = ressql + ' when ' + ' and '.join(p) + f" then '{m[3]}'"
   ressql = ressql + ' end'
   
   print(ressql)

   登录后复制

   上述代码中，我们首先创建了一个CASE语句的开头case。然后，我们遍历map_data中的每一行m。对于每一行，我们使用列表推导式[f"{p[0]} = '{p[1]}'" for p in zip(columns, m[:3]) if p[1] != "*"]来构建一个条件列表p。这个列表包含所有非通配符列的条件。最后，我们将这些条件用and连接起来，并添加到CASE语句中，同时添加对应的结果m[3]。

2. 应用CASE WHEN语句

   使用expr函数将生成的CASE WHEN语句应用到df上，创建一个新的result列。

   from pyspark.sql import functions as F
   
   df = df.withColumn('result', F.expr(ressql))
   df.show()

   登录后复制

   F.expr(ressql)会将字符串ressql解析为一个SQL表达式，并将其应用到DataFrame df上。withColumn函数会在DataFrame中添加一个新的列result，其值是根据CASE WHEN语句计算出来的。

示例代码

完整的示例代码如下：

from pyspark.sql import SparkSession
from pyspark.sql.functions import expr

# 创建 SparkSession
spark = SparkSession.builder.appName("dynamic_case_when").getOrCreate()

# 示例数据
map_data = [('a', 'b', 'c', 'good'), ('a', 'a', '*', 'very good'), 
          ('b', 'd', 'c', 'bad'), ('a', 'b', 'a', 'very good'),
          ('c', 'c', '*', 'very bad'), ('a', 'b', 'b', 'bad')]

columns = ["col1", "col2", 'col3', 'result']

mapping_table = spark.createDataFrame(map_data, columns)

data =[[('a', 'b', 'c')], [('a', 'a', 'b')], 
        [('c', 'c', 'a')], [('c', 'c', 'b')],
        [('a', 'b', 'b')], [('a', 'a', 'd')]
      ]

columns = ["col1", "col2", 'col3']
df = spark.createDataFrame(data, columns)
df = df.selectExpr(
    "_1.col1 as col1",
    "_1.col2 as col2",
    "_1.col3 as col3"
)

ressql = 'case '
for m in map_data:
    p = [f"{p[0]} = '{p[1]}'" for p in zip(columns, m[:3]) if p[1] != "*"]
    ressql = ressql + ' when ' + ' and '.join(p) + f" then '{m[3]}'"
ressql = ressql + ' end'

from pyspark.sql import functions as F

df = df.withColumn('result', F.expr(ressql))
df.show()

# 关闭 SparkSession
spark.stop()

注意事项

• 性能：动态生成CASE WHEN语句的方法在mapping_table非常大时可能会影响性能。在这种情况下，可以考虑使用其他方法，例如广播变量和UDF。
• SQL注入：如果mapping_table中的数据来自外部源，需要注意SQL注入的风险。应该对数据进行适当的转义或验证。
• 数据类型：确保df和mapping_table中列的数据类型一致，否则可能会导致错误。

总结

本文介绍了一种使用PySpark动态生成CASE WHEN语句来解决复杂数据映射问题的方法。该方法可以灵活地处理包含通配符的映射关系，避免复杂的JOIN操作。但是，需要注意性能和安全问题。在实际应用中，需要根据具体情况选择合适的解决方案。

以上就是使用PySpark动态生成CASE WHEN语句实现复杂数据映射的详细内容，更多请关注php中文网其它相关文章！