在 Pyspark 中将任意数量的列组合成一个 Array 类型的新列答案

【问题标题】：Combine arbitrary number of columns into a new column of Array type in Pyspark在 Pyspark 中将任意数量的列组合成一个 Array 类型的新列
【发布时间】：2020-06-18 14:47:07
【问题描述】：

我有一个 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'))

它返回 null （我假设由于初始数据帧中的空值）。从我搜索的内容来看，我想使用.coalesce() 或.fillna() 来处理/删除空值，但我没有设法让它工作。

我的主要要求是我希望新创建的列是 Array 类型，并且我不想枚举我需要连接的所有列名。

【问题讨论】：

标签： python pyspark

【解决方案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)'''))

【讨论】：

【解决方案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]|
      * +---+----+----+----------+----------+
      */

【讨论】：