spark - Converting dataframe to list提高性能

Question

我需要将Spark dataframe的一列转换为列表，以便稍后用于matplotlib

df.toPandas()[col_name].values.tolist()

看起来有很高的性能开销，这个操作大约需要18秒，有没有其他方法可以做到这一点或提高性能？

Answer 1

如果您确实需要一个本地列表，那么在这里您可以做的并不多，但是一个改进是只收集单个列，而不是整个DataFrame

df.select(col_name).flatMap(lambda x: x).collect()

Answer 2

你可以这样做：

>>> [list(row) for row in df.collect()]

示例：

>>> d = [['Alice', 1], ['Bob', 2]]

>>> df = spark.createDataFrame(d, ['name', 'age'])

>>> df.show()

+-----+---+

| name|age|

+-----+---+

|Alice| 1|

| Bob| 2|

+-----+---+

>>> to_list = [list(row) for row in df.collect()]

print list

结果：[[u'Alice', 1], [u'Bob', 2]]

Answer 3

您可以使用迭代器来节省内存toLocalIterator。迭代器将消耗与此分区中最大分区一样多的内存。如果你只需要使用一次结果，那么迭代器就是完美的例子。

d = [['Bender', 12], ['Flex', 123],['Fry', 1234]]
df = spark.createDataFrame(d, ['name', 'value'])
df.show()
+------+-----+
|  name|value|
+------+-----+
|Bender|   12|
|  Flex|  123|
|   Fry| 1234|
+------+-----+`
values = [row.value for row in df.toLocalIterator()]

print(values)
>>> [12, 123, 1234]

此外，只有当生成的Pandas的DataFrame很小时，才应该使用toPandas()方法，因为所有数据都加载到驱动程序的内存中。