Pyspark：用零填充 Array[Int] 列答案

【问题标题】：Pyspark: Pad Array[Int] column with zerosPyspark：用零填充 Array[Int] 列
【发布时间】：2018-10-27 00:27:47
【问题描述】：

我在 pyspark 数据框中有以下列，类型为 Array[Int]。

+--------------------+
|     feature_indices|
+--------------------+
|                 [0]|
|[0, 1, 4, 10, 11,...|
|           [0, 1, 2]|
|                 [1]|
|                 [0]|
+--------------------+

我试图用零填充数组，然后限制列表长度，以便每一行数组的长度相同。例如，对于 n = 5，我期望：

+--------------------+
|     feature_indices|
+--------------------+
|     [0, 0, 0, 0, 0]|
|   [0, 1, 4, 10, 11]|
|     [0, 1, 2, 0, 0]|
|     [1, 0, 0, 0, 0]|
|     [0, 0, 0, 0, 0]|
+--------------------+

有什么建议吗？我查看了 pyspark rpad 函数，但它只对字符串类型的列起作用。

【问题讨论】：

标签： python dataframe pyspark

【解决方案1】：

你可以写一个udf来做到这一点：

from pyspark.sql.types import ArrayType, IntegerType
import pyspark.sql.functions as F

pad_fix_length = F.udf(
    lambda arr: arr[:5] + [0] * (5 - len(arr[:5])), 
    ArrayType(IntegerType())
)

df.withColumn('feature_indices', pad_fix_length(df.feature_indices)).show()
+-----------------+
|  feature_indices|
+-----------------+
|  [0, 0, 0, 0, 0]|
|[0, 1, 4, 10, 11]|
|  [0, 1, 2, 0, 0]|
|  [1, 0, 0, 0, 0]|
|  [0, 0, 0, 0, 0]|
+-----------------+

【讨论】：

太好了，谢谢！我正在努力正确地编写 udf。
如果我们不在udf中给出ArrayType(IntegerType())怎么办？

【解决方案2】：

我最近在 Keras 中使用了 pad_sequences 函数来做类似的事情。我不确定你的用例，所以这可能是一个不必要的大依赖。

无论如何，这里是函数文档的链接：https://keras.io/preprocessing/sequence/#pad_sequences

from keras.preprocessing.sequence import pad_sequences    

input_sequence =[[1,2,3], [1,2], [1,4]]

padded_sequence = pad_sequences(input_sequence, maxlen=3, padding='post', truncating='post', value=0.0)

print padded_sequence

输出：

[[1 2 3]
 [1 2 0]
 [1 4 0]]

【讨论】：