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

Question

我有一个 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]|
+----+------------------------+

如何做到这一点？

Answer 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)，这将返回除第一项之外的所有项目一个。