scala - SparkSQL：如何根据列名选择列值

Question

我正在使用具有以下架构的数据框：

root
 |-- Id: integer (nullable = true)
 |-- defectiveItem: string (nullable = true)
 |-- item: struct (nullable = true)
 |    |-- gem1: integer (nullable = true)
 |    |-- gem2: integer (nullable = true)
 |    |-- gem3: integer (nullable = true)

defectiveItem列包含 , 中的值，gem1并包含项目的计数。现在根据缺陷项，我需要将给定缺陷项的计数从一个名为的新列中投影出来。gem2gem3itemitemcount

例如，如果defectiveItem列包含gem1并且item包含{"gem1":3,"gem2":4,"gem3":5}结果count列应该包含 3。

生成的架构应如下所示：

root
     |-- Id: integer (nullable = true)
     |-- defectiveItem: string (nullable = true)
     |-- item: struct (nullable = true)
     |    |-- gem1: integer (nullable = true)
     |    |-- gem2: integer (nullable = true)
     |    |-- gem3: integer (nullable = true)
     |-- count: integer (nullable = true)

标签： scalaapache-sparkapache-spark-sql

Answer 1

udf您可以使用以下函数获取所需的输出数据框

import org.apache.spark.sql.functions._
def getItemUdf = udf((defectItem: String, item: Row)=> item.getAs[Int](defectItem))

df.withColumn("count", getItemUdf(col("defectiveItem"), col("item"))).show(false)

希望回答有用