如何使用逗号分隔值拆分列并存储在 PySpark Dataframe 中的数组中？如下所示

Posted 2023-04-15

【中文标题】如何使用逗号分隔值拆分列并存储在 PySpark Dataframe 中的数组中？如下所示

【英文标题】：How to split a column with comma separated values and store in array in PySpark's Dataframe? As given below

【发布时间】：2020-03-10 21:43:44

【问题描述】：

我有一个 PySpark 数据框，其中有一列包含逗号分隔值。值的数量 该列包含是固定的（比如 4）。示例：

+------------------------+
|col1                    |
+------------------------+
|1,val1, val4            |
|2,val1                  |
|3,val1, val2, val3      |
|4,val1, val2, val3, val4|
+------------------------+

现在我希望它被分成 2 列，如下所示

+----+------------------------+
|col1|col2                    |
+----+------------------------+
|   1|[val1, val4]            |
|   2|[val1]                  |
|   3|[val1, val2, val3]      |
|   4|[val1, val2, val3, val4]|
+----+------------------------+

如何做到这一点？

【参考方案1】：

您可以使用slice 和split 实现此功能：

from pyspark.sql.functions import col, split, slice

array_len = 4
df.withColumn("ar", split(col("col1"), ",")) \
  .select(
     col("ar")[0].alias("col1"), 
     slice(col("ar"), 2, array_len).alias("col2")
  )

# +----+---------------------------+
# |col1|col2                       |
# +----+---------------------------+
# |1   |[val1,  val4]              |
# |2   |[val1]                     |
# |3   |[val1,  val2,  val3]       |
# |4   |[val1,  val2,  val3,  val4]|
# +----+---------------------------+

首先我们将数组拆分并存储到ar，接下来我们使用select检索数组的第一项col("ar")[0]和数组的其余部分slice(col("ar"), 2, array_len)，这将返回除第一项之外的所有项目一个。

【讨论】：

切片中的2是什么意思？

我们保留2-4项，切片函数从1开始计数

好的明白了，但为什么要手动定义array_len？有没有办法像我们在 pandas iloc[10:] 中那样指定范围，这意味着第十行及以后。

不，切片只接受整数@Sid_K（甚至不接受列，即size(col("ar"))）

酷。谢谢亚历克斯！