有没有更快的方法将 pyspark dataframe 列转换为 python 列表？ （Collect() 非常慢）

Question

I am trying to store a column of pyspark dataframe into python list using collect function. eg

list_a = [row[column_name] for row in dataset_name.collect()]

but this is very slow process and takes more than 10 seconds for a dataframe of 3 columns and 27 rows.

is there a faster way to do so?

I tried caching the data before this step . With this step, the above query is being executed in 2 seconds but cache step itself is taking around 7-8 seconds so my purpose of reducing time is not full filled.

And my code is such that i need to rebuild the dataframe everytime before this step so need to do cache again so this step(caching the dataframe) is not helping a lot in time reduction.

Thanks in advance!

Answer 1

Your code can be slightly optimized by only collecting one column of data:

list_a = [row[column_name] for row in dataset_name.select(column_name).collect()]

This code is cleaner if you use quinn :

import quinn

list_a = quinn.column_to_list(df, col_name)

collect() transfers all the data to the driver node and is expensive. You can only make it faster by collecting less data (eg dataset_name.select(column_name).distinct().collect() would typically be faster).

Spark is optimized for distributing datasets across a cluster and running computations in parallel. The distributed nature of Spark makes computations that collect results on a single node comparatively slow.