PySpark - 将数组列分割成更小的块

Question

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

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

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

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

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

Answer 1

transform使用and的另一种方法filter是使用 if 和 using 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]]|
+---------------+---------------------+