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

Question

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

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

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

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

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

Answer 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]]|
+---------------+---------------------+