PySpark - 将数组列拆分为更小的块

【问题标题】：PySpark - Split Array Column into smaller chunksPySpark - 将数组列拆分为更小的块
【发布时间】：2021-09-25 08:04:42
【问题描述】：

假设我有一个带有列的数据集

 # Output
 #+-----------------+
 #|         arrayCol|
 #+-----------------+
 #| [1, 2, 3, 4, 5] |
 #+-----------------+

我想知道是否可以在不使用 UDF 的情况下将此列拆分为较小的 max_size 块。

所需的结果如下，max_size = 2：

 # Output
 #+-----------------------+
 #|               arrayCol|
 #+-----------------------+
 #| [[1, 2], [3, 4], [5]] |
 #+-----------------------+

【问题讨论】：

标签： pyspark

【解决方案1】：

使用transform 和filter 的另一种方法是使用if 和使用mod 来决定拆分并使用slice（对数组进行切片）

from pyspark.sql import functions as F
n = 2
df.withColumn("NewCol",F.expr(f""" 
               filter(
         transform(arrayCol,(x,i)-> if (i%{n}=0 ,slice(arrayCol,i+1,{n}), null)),x->
               x is not null)                               
""")).show(truncate=False)


+---------------+---------------------+
|arrayCol       |NewCol               |
+---------------+---------------------+
|[1, 2, 3, 4, 5]|[[1, 2], [3, 4], [5]]|
+---------------+---------------------+

【讨论】：