Индекс элемента массива в scala

Question

У меня есть таблица со столбцом массива, например:

+-Name-+
 array
  0: {"given_name":"B. A.", "surname":"Name1"}
  1: {"given_name":"A.", "surname":"Name2"}
  2: {"given_name":"C." "surname":"Name3"}

Я хотел бы добавить еще один элемент "index", начиная с 1, в массив, чтобы найти последовательность автора, например

+-Name-+
array
  0: {"given_name":"B. A.", "surname":"Name1", "index":"1"}
  1: {"given_name":"A.", "surname":"Name2", "index":"2"}
  2: {"given_name":"C." "surname":"Name3", "index":"3"}

Как это сделать в Scala, ваша помощь очень ценится.

Answer 1

Вот один из подходов, использующий UDF, который отображает каждый элемент столбца с типом массива для включения также индекса элемента:

import org.apache.spark.sql.Row
import org.apache.spark.sql.functions._
import spark.implicits._

case class Name(given_name: String, surname: String)
case class NameIdx(given_name: String, surname: String, index: Int)

val df = Seq(
  Seq(Name("John", "Doe"), Name("Jane", "Smith"), Name("Mike", "Davis")),
  Seq(Name("Rachel", "Smith"), Name("Steve", "Thompson"))
).toDF("name")

val addIndex = udf((names: Seq[Row]) => names.map{
  case name @ Row(gn: String, sn: String) => NameIdx(gn, sn, names.indexOf(name) + 1)
})

df.select(addIndex($"name").as("name")).show(false)
// +----------------------------------------------+
// |name                                          |
// +----------------------------------------------+
// |[[John,Doe,1], [Jane,Smith,2], [Mike,Davis,3]]|
// |[[Rachel,Smith,1], [Steve,Thompson,2]]        |
// +----------------------------------------------+

Чтобы получить значения JSON, примените to_json следующим образом:

df.select(to_json(addIndex($"name")).as("name")).show(false)
// +-----------------------------------------------------------------------------------------------------------------------------------------------------+
// |name                                                                                                                                                 |
// +-----------------------------------------------------------------------------------------------------------------------------------------------------+
// |[{"given_name":"John","surname":"Doe","index":1},{"given_name":"Jane","surname":"Smith","index":2},{"given_name":"Mike","surname":"Davis","index":3}]|
// |[{"given_name":"Rachel","surname":"Smith","index":1},{"given_name":"Steve","surname":"Thompson","index":2}]                                          |
// +-----------------------------------------------------------------------------------------------------------------------------------------------------+