主页 / coding / 问题 / 79341299
Accepted
smurphy
Asked: 2025-01-09 11:25:08 +0800 CST

根据另一列中提供的列名数组创建具有值的列

  • 772

我想创建一个新列，其中包含列中列出的列名的值数组lookup。

示例输入

input_df = spark.createDataFrame([
    Row(id=123, alert=1, operation=1, lookup=[]),
    Row(id=234, alert=0, operation=0, lookup=['alert']),
    Row(id=345, alert=1, operation=0, lookup=['operation']),
    Row(id=456, alert=0, operation=1, lookup=['alert', 'operation']),
])

预期输出

ID 警报 手术 抬头 查找值
123 1 1 [] []
234 0 0 [alert] [0]
345 1 0 [operation] [0]
456 0 1 [alert, operation] [0, 1]

我尝试过

input_df.withColumn("lookup_values", F.transform(F.col("lookup"), lambda x: input_df[f'{x}'])).show()

失败并出现以下错误：

AnalysisException：[UNRESOLVED_COLUMN.WITH_SUGGESTION] 无法解析名称为 的列或函数参数Column<'x_1'>。您是指下列之一吗？[ id, alert, operation, lookup]。

这个错误令人惊讶，因为下面的代码虽然没有产生预期的结果，但却没有产生错误：

input_df.withColumn("lookup_values", F.transform(F.col("lookup"), lambda x: input_df['alert'])).show()
ID 警报 手术 抬头 查找值
123 1 1 [] []
234 0 0 [alert] [0]
345 1 0 [operation] [1]
456 0 1 [alert, operation] [0, 0]
apache-spark

2 个回答

  1. 一种方法是将整行传递到 UDF，然后根据列将查找值放入列表中lookup：

    @func.udf(returnType=ArrayType(IntegerType()))
def lookup_values_udf(row):
    return [row[field] for field in row["lookup"]]

input_df.withColumn(
    "lookup_values", 
    lookup_values_udf(func.struct([func.col(col) for col in input_df.columns]))
).show(
    10, False
)

+---+-----+---------+------------------+-------------+
|id |alert|operation|lookup            |lookup_values|
+---+-----+---------+------------------+-------------+
|123|1    |1        |[]                |[]           |
|234|0    |0        |[alert]           |[0]          |
|345|1    |0        |[operation]       |[0]          |
|456|0    |1        |[alert, operation]|[0, 1]       |
+---+-----+---------+------------------+-------------+
    • 1
  2. Best Answer

    这是不使用 UDF 而是使用内置函数的答案。处理大量数据时速度应该会更快：

    from pyspark.sql import functions as F

input_df.withColumn(
    "lookup_values",
    F.create_map(
        [F.lit("alert"), F.col("alert"), F.lit("operation"), F.col("operation")]
    ),
).withColumn(
    "lookup_values", 
    F.transform(F.col("lookup"), lambda x: F.col("lookup_values")[x])
).display()

    +---+-----+---------+------------------+-------------+
| id|alert|operation|            lookup|lookup_values|
+---+-----+---------+------------------+-------------+
|123|    1|        1|                []|           []|
|234|    0|        0|           [alert]|          [0]|
|345|    1|        0|       [operation]|          [0]|
|456|    0|        1|[alert, operation]|       [0, 1]|
+---+-----+---------+------------------+-------------+
    • 1

相关问题