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, gem2, gem3 и item содержит количество элементов. Теперь, в зависимости от дефектного элемента, мне нужно спроектировать счетчик данного дефектного элемента из item как новый столбец с именем count.

Например, если столбец 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)

Answer 1

Вы также можете решить эту проблему с помощью более классического подхода, используя встроенные в SQL функции when-case:

import sparkSession.implicits._

val defectiveItems = Seq(
(1, "gem1", Map("gem1" -> 10, "gem2" -> 0, "gem3" -> 0)),
(2, "gem1", Map("gem1" -> 15, "gem2" -> 0, "gem3" -> 0)),
(3, "gem1", Map("gem1" -> 33, "gem2" -> 0, "gem3" -> 0)),
(4, "gem3", Map("gem1" -> 0, "gem2" -> 0, "gem3" -> 2))
).toDF("Id", "defectiveItem", "item")
import org.apache.spark.sql.functions._
val datasetWithCount = defectiveItems.withColumn("count", when($"defectiveItem" === "gem1", $"item.gem1").otherwise(when($"defectiveItem" === "gem2", $"item.gem2").otherwise($"item.gem3")))

println("All items="+datasetWithCount.collectAsList())

Будет напечатано:

All items=[[1,gem1,Map(gem1 -> 10, gem2 -> 0, gem3 -> 0),10], [2,gem1,Map(gem1 -> 15, gem2 -> 0, gem3 -> 0),15], [3,gem1,Map(gem1 -> 33, gem2 -> 0, gem3 -> 0),33], [4,gem3,Map(gem1 -> 0, gem2 -> 0, gem3 -> 2),2]]

Используя собственные решения, вы можете воспользоваться преимуществами внутренней оптимизации Spark для планов выполнения.

Answer 2

Вы можете получить желаемый выходной кадр данных, используя функцию 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)

Надеюсь, ответ полезен