Pyspark: Pad Array [Int] столбец с нулями

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

Недавно я использовал функцию pad_sequences в Keras, чтобы сделать что-то подобное.Я не уверен в вашем сценарии использования, так что это может быть излишне большая зависимость для добавления.

В любом случае, вот ссылка на документацию для функции: 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]]

Answer 2

Вы можете написать 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]|
+-----------------+