pyspark - 根据同一行中的其他列动态选择列内容答案

【问题标题】：pyspark - Dynamically select column content based on other column from the same rowpyspark - 根据同一行中的其他列动态选择列内容
【发布时间】：2022-01-12 11:46:25
【问题描述】：

我的数据框看起来像：

categoryName	catA	catB
catA	0.25	0.75
catB	0.5	0.5

其中 categoryName 为 String 类型，而 cat* 为 Double。我想添加将包含列中的值的列，该列的名称位于 categoryName 列中：

categoryName	catA	catB	score
catA	0.25	0.75	0.25
catB	0.5	0.7	0.7

在第一行'score'的值来自列名'catA' 在列名“catB”的第二行“分数”值中谢谢

【问题讨论】：

标签： dataframe pyspark

【解决方案1】：

一种方法是根据每行的列名和值创建一个映射，然后使用在所需列中定义的值访问该映射。

这很酷的是它可以为任意多的列工作。

例子：

from pyspark.sql import SparkSession
import pyspark.sql.functions as F

data = [
    {"categoryName": "catA", "catA": 0.25, "catB": 0.75},
    {"categoryName": "catB", "catA": 0.5, "catB": 0.7},
]

spark = SparkSession.builder.getOrCreate()
df = spark.createDataFrame(data)
df = (
    df.withColumn(
        "map", F.expr("map(" + ",".join([f"'{c}', {c}" for c in df.columns]) + ")")
    )
    .withColumn("score", F.expr("map[categoryName]"))
    .drop("map")
)

结果：

+----+----+------------+-----+                                                  
|catA|catB|categoryName|score|
+----+----+------------+-----+
|0.25|0.75|catA        |0.25 |
|0.5 |0.7 |catB        |0.7  |
+----+----+------------+-----+

【讨论】：