Spark - добавить новый столбец во фрейм данных, столбцы которого соответствуют заданному шаблону scala

Question

Я хочу добавить к фрейму данных новый столбец карты с существующими столбцами, которые начинаются с заданного общего префикса.

Например, у меня есть ввод

{"Prefix_A": "v_A", "Prefix_B": "v_B", "Field": "v"}, {"Prefix_A": "v_A", "Prefix_B": "v_B", "Prefix_C": "v_C", "Field": "v"}

Я хочу объединить все поля с префиксом «Префикс_» и получить вывод

{"NewColumn": {"Prefix_A": "v_A", "Prefix_B": "v_B"}, "Field": "v"}, {"NewColumn": {"Prefix_A": "v_A", "Prefix_B": "v_B", "Prefix_C": "v_C"}, "Field": "v"}

Я хочу сделать это на лету, т.е. я не знаюстолбцы как данные без схемы JSON дампа.

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

Answer 1

Предполагая, что у вас есть входной файл data.json, содержащий ваши записи json, вы можете получить ожидаемый результат с помощью следующего кода:

import org.apache.spark.sql.functions.{col,struct}
import spark.implicits._

val df = spark.read.json("data.json")
val (prefixedColumns, otherColumns) = df.columns.partition(_.startsWith("Prefix"))
val transformedDf = df.select(
    struct(prefixedColumns.map(col):_*).as("NewColumn")
    +: otherColumns.map(col):_*)
transformedDf.write.json("output.json")

По сути, вы воссоздаете новый список столбцов, используя необходимыефункции, а затем применить их в select, используя нотацию scala :_* для преобразования последовательности в varargs