Question

我在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函数，但它只对字符串类型的列进行操作。

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

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

Pyspark：填充数组[Int]列，带零

2 个答案: