将int列转换为列表类型pyspark

时间:2019-01-07 22:14:52

标签: pyspark

我的DataFrame有一列num_of_items。这是一个计数字段。现在,我想将其从int类型转换为列表类型。

我尝试使用array(col)甚至通过将int值作为输入来创建函数以返回列表。没用

from pyspark.sql.types import ArrayType
from array import array

def to_array(x):
    return [x]

df=df.withColumn("num_of_items", monotonically_increasing_id())

df

col_1    | num_of_items
A        |  1
B        |  2

预期产量

col_1    | num_of_items
A        | [23]
B        | [43]

1 个答案:

答案 0 :(得分:2)

  

我尝试使用array(col)

使用pyspark.sql.functions.array似乎对我有用。

from pyspark.sql.functions import array
df.withColumn("num_of_items", array("num_of_items")).show()
#+-----+------------+
#|col_1|num_of_items|
#+-----+------------+
#|    A|         [1]|
#|    B|         [2]|
#+-----+------------+
  

,甚至创建一个函数,以int值作为输入返回列表。

如果要使用创建的函数,则必须将其设为udf并指定返回类型:

from pyspark.sql.types import ArrayType, IntegerType
from pyspark.sql.functions import udf, col

to_array_udf = udf(to_array, ArrayType(IntegerType()))
df.withColumn("num_of_items", to_array_udf(col("num_of_items"))).show()
#+-----+------------+
#|col_1|num_of_items|
#+-----+------------+
#|    A|         [1]|
#|    B|         [2]|
#+-----+------------+

但是最好尽可能避免使用udf:请参见Spark functions vs UDF performance?