Spark - разделить строковый столбец, выходящий за разделитель, на одну часть

Question

У меня есть CSV-файл из двух строковых столбцов (термин, код).Столбец кода имеет специальный формат [num]-[two_letters]-[text], где text также может содержать тире -.Я хочу прочитать этот файл, используя Spark, в массив данных ровно из четырех столбцов (term, num, two_letters, text).

Input
+---------------------------------+
|  term  |          code          |
+---------------------------------+
| term01 |    12-AB-some text     |
| term02 | 130-CD-some-other-text |
+---------------------------------+


Output
+------------------------------------------+
|  term  | num | letters |       text      |
+------------------------------------------+
| term01 | 12  |   AB    |   some text     |
| term02 | 130 |   CD    | some-other-text |
+------------------------------------------+

Я могу разбить столбец code на три столбца, когда нет тирев его text части, но как мне найти решение, которое бы охватывало все случаи (например, получить весь текст после двух штрихов в одном столбце)?

Код для разделения столбца на три хорошо разъясненв ответе здесь

Answer 1

Вот как бы я это сделал, используя UDF:

case class MyData(num: Int, letters: String, text: String)
def udfSplit = udf(
  (input: String) => { 
    val res = input.split("-", 3) // limit=3 => pattern applied at most n - 1 times
    MyData(res(0).toInt, res(1), res(2))
  }
)

val df = spark.createDataFrame(
    Seq(
      ("term01", "12-AB-some text"), 
      ("term02", "130-CD-some-other-text")
    )
).toDF("term", "code")
df.show(false)
+------+----------------------+
|term  |code                  |
+------+----------------------+
|term01|12-AB-some text       |
|term02|130-CD-some-other-text|
+------+----------------------+

val res = df.withColumn("code", udfSplit($"code"))
res.show(false)
+------+------------------------+
|term  |code                    |
+------+------------------------+
|term01|[12,AB,some text]       |
|term02|[130,CD,some-other-text]|
+------+------------------------+

res.printSchema
root
 |-- term: string (nullable = true)
 |-- code: struct (nullable = true)
 |    |-- num: integer (nullable = false)
 |    |-- letters: string (nullable = true)
 |    |-- text: string (nullable = true)

res.select("term", "code.*").show(false)
+------+---+-------+---------------+
|term  |num|letters|text           |
+------+---+-------+---------------+
|term01|12 |AB     |some text      |
|term02|130|CD     |some-other-text|
+------+---+-------+---------------+

Answer 2

Вот один вариант с regexp_extract:

val df = Seq(("term01", "12-AB-some text"), ("term02", "130-CD-some-other-text")).toDF("term", "code")

// define the pattern that matches the string column
val p = "([0-9]+)-([a-zA-Z]{2})-(.*)"
// p: String = ([0-9]+)-([a-zA-Z]{2})-(.*)

// define the map from new column names to the group index in the pattern
val cols = Map("num" -> 1, "letters" -> 2, "text" -> 3)
// cols: scala.collection.immutable.Map[String,Int] = Map(num -> 1, letters -> 2, text -> 3)

// create the new columns on data frame
cols.foldLeft(df){ 
    case (df, (colName, groupIdx)) => df.withColumn(colName, regexp_extract($"code", p, groupIdx)) 
}.drop("code").show

+------+---+-------+---------------+
|  term|num|letters|           text|
+------+---+-------+---------------+
|term01| 12|     AB|      some text|
|term02|130|     CD|some-other-text|
+------+---+-------+---------------+