Объедините произвольное количество столбцов в новый столбец типа Array в Pyspark

Question

У меня есть фрейм данных pyspark, содержащий N столбцов, содержащих целые числа. Некоторые поля также могут быть пустыми. Например:

+---+-----+-----+
| id| f_1 | f_2 |
+---+-----+-----+
|  1| null| null|
|  2|123  | null|
|  3|124  |127  |
+---+-----+-----+

Я хочу объединить все столбцы с префиксом f в массив pyspark в новом столбце. Например:

+---+---------+
| id| combined|
+---+---------+
|  1| []      |
|  2|[123]    |
|  3|[124,127]|
+---+---------+

Ближе, что мне удалось получить, это:

features_filtered = features.select(F.concat(* features.columns[1:]).alias('combined')) 

, который возвращает ноль (я предполагаю, из-за нулей в начальном фрейме данных). Из того, что я искал, я хотел бы использовать .coalesce() или, может быть, .fillna() для обработки / удаления нулей, но мне не удалось заставить его работать.

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

Answer 1

В pyspark можно сделать как

df = df.withColumn("combined_array", f.array(*[i for i in df.columns if i.startswith('f')]))
      .withColumn("combined", expr('''FILTER(combined_array, x -> x is not null)'''))

Answer 2

Попробуйте это - (В scala, но может быть реализовано в python с минимальным изменением)

Загрузить данные

 val data =
      """
        |id| f_1 | f_2
        | 1| null| null
        | 2|123  | null
        | 3|124  |127
      """.stripMargin
    val stringDS = data.split(System.lineSeparator())
      .map(_.split("\\|").map(_.replaceAll("""^[ \t]+|[ \t]+$""", "")).mkString(","))
      .toSeq.toDS()
    val df = spark.read
      .option("sep", ",")
      .option("inferSchema", "true")
      .option("header", "true")
      .option("nullValue", "null")
      .csv(stringDS)
    df.printSchema()
    df.show(false)

    /**
      * root
      * |-- id: integer (nullable = true)
      * |-- f_1: integer (nullable = true)
      * |-- f_2: integer (nullable = true)
      *
      * +---+----+----+
      * |id |f_1 |f_2 |
      * +---+----+----+
      * |1  |null|null|
      * |2  |123 |null|
      * |3  |124 |127 |
      * +---+----+----+
      */

Преобразовать в массив

    df.withColumn("array", array(df.columns.filter(_.startsWith("f")).map(col): _*))
      .withColumn("combined", expr("FILTER(array, x -> x is not null)"))
      .show(false)

    /**
      * +---+----+----+----------+----------+
      * |id |f_1 |f_2 |array     |combined  |
      * +---+----+----+----------+----------+
      * |1  |null|null|[,]       |[]        |
      * |2  |123 |null|[123,]    |[123]     |
      * |3  |124 |127 |[124, 127]|[124, 127]|
      * +---+----+----+----------+----------+
      */