python - PySpark 如何使用 col 函数引用另一列的值

Question

我有一个包含以下列的数据框

field1 , field2 , field3 , field_name

示例数据

"a1", "b1", "c1", "field1"
"a2", "b2", "c2", "field2"
"a3", "b3", "c3", "field3"

我想将新列“fieldvalue”添加到数据框中，以便它包含与“fieldname”列内容对应的列中的值

因此上面的第一行将有 fieldvalue = "a1"，因为 fieldname 包含“field1”

输出数据框应如下所示

field1, field2, field3 , fieldname, fieldvalue

数据

"a1", "b1", "c1", "field1", "a1"
"a2", "b2", "c2", "field2", "b2"
"a3", "b3", "c3", "field3", "c3"

我尝试使用以下语法

df1 = df1.withColumn("fieldValue", func.col(func.col("fieldName")))

但它失败并出现以下错误，因为 func.col 需要一个列，而不是一个字符串

方法 col([class org.apache.spark.sql.Column]) 不存在

Answer 1

使用col您只能使用不依赖于实际行的值的常量值作为参数。原因是用于执行该函数的 DAG 是在 Spark 查看数据之前创建的。 col 只需从数据框中选择一列，并且该列对于所有行都必须相同。

相反，每一行都可以转换为 array ，然后使用 find_in_set 计算所需列的数组内的索引最后从数组中取出所需的值:

from pyspark.sql import functions as F

df.withColumn("array", F.array(df.columns)) \
  .withColumn("index", F.expr(f"find_in_set(field_name, '{','.join(df.columns)}') - 1")) \
  .withColumn("fieldvalue", F.expr("array[index]")) \
  .show()

输出:

+------+------+------+----------+--------------------+-----+----------+
|field1|field2|field3|field_name|               array|index|fieldvalue|
+------+------+------+----------+--------------------+-----+----------+
|    a1|    b1|    c1|    field1|[a1, b1, c1, field1]|    0|        a1|
|    a2|    b2|    c2|    field2|[a2, b2, c2, field2]|    1|        b2|
|    a3|    b3|    c3|    field3|[a3, b3, c3, field3]|    2|        c3|
+------+------+------+----------+--------------------+-----+----------+

使用 udf 可以实现相同的结果，但 Spark SQL 函数通常比 udf 更快，尤其是在使用 Python 时。

编辑:如果array不起作用，同样的想法可以用于map :

m = "map(" + ",".join([f"'{c}', {c}" for c in df.columns]) + ")"

df.withColumn("map", F.expr(m)) \
  .withColumn("fieldvalue", F.expr("map[field_name]")) \
  .show()