Chuyển đổi cột Spark DataFrame thành danh sách python

Tôi làm việc trên khung dữ liệu có hai cột, mvv và số.

+---+-----+
|mvv|count|
+---+-----+
| 1 |  5  |
| 2 |  9  |
| 3 |  3  |
| 4 |  1  |

tôi muốn lấy hai danh sách chứa các giá trị mvv và giá trị đếm. Cái gì đó như

mvv = [1,2,3,4]
count = [5,9,3,1]

Vì vậy, tôi đã thử đoạn mã sau: Dòng đầu tiên sẽ trả về một danh sách hàng python. Tôi muốn xem giá trị đầu tiên:

mvv_list = mvv_count_df.select('mvv').collect()
firstvalue = mvv_list[0].getInt(0)

Nhưng tôi nhận được thông báo lỗi với dòng thứ hai:

AttributeError: getInt

Hãy xem, tại sao cách bạn đang làm này không hoạt động. Trước tiên, bạn đang cố gắng lấy số nguyên từ Loại hàng , đầu ra của bộ sưu tập của bạn như sau:

>>> mvv_list = mvv_count_df.select('mvv').collect()
>>> mvv_list[0]
Out: Row(mvv=1)

Nếu bạn lấy một cái gì đó như thế này:

>>> firstvalue = mvv_list[0].mvv
Out: 1

Bạn sẽ nhận được mvvgiá trị. Nếu bạn muốn tất cả thông tin của mảng, bạn có thể lấy một cái gì đó như sau:

>>> mvv_array = [int(row.mvv) for row in mvv_list.collect()]
>>> mvv_array
Out: [1,2,3,4]

Nhưng nếu bạn thử tương tự cho cột kia, bạn sẽ nhận được:

>>> mvv_count = [int(row.count) for row in mvv_list.collect()]
Out: TypeError: int() argument must be a string or a number, not 'builtin_function_or_method'

Điều này xảy ra bởi vì countlà một phương thức được tích hợp sẵn. Và cột có tên giống như count. Một giải pháp để thực hiện việc này là thay đổi tên cột countthành _count:

>>> mvv_list = mvv_list.selectExpr("mvv as mvv", "count as _count")
>>> mvv_count = [int(row._count) for row in mvv_list.collect()]

Nhưng cách giải quyết này không cần thiết, vì bạn có thể truy cập cột bằng cú pháp từ điển:

>>> mvv_array = [int(row['mvv']) for row in mvv_list.collect()]
>>> mvv_count = [int(row['count']) for row in mvv_list.collect()]

Và cuối cùng nó sẽ hoạt động!

Sau một lớp lót cung cấp danh sách bạn muốn.

mvv = mvv_count_df.select("mvv").rdd.flatMap(lambda x: x).collect()

Điều này sẽ cung cấp cho bạn tất cả các yếu tố dưới dạng danh sách.

mvv_list = list(
    mvv_count_df.select('mvv').toPandas()['mvv']
)

Đoạn mã sau sẽ giúp bạn

mvv_count_df.select('mvv').rdd.map(lambda row : row[0]).collect()

Trên dữ liệu của mình, tôi nhận được các điểm chuẩn này:

>>> data.select(col).rdd.flatMap(lambda x: x).collect()

0,52 giây

>>> [row[col] for row in data.collect()]

0,271 giây

>>> list(data.select(col).toPandas()[col])

0,427 giây

Kết quả là như nhau

Nếu bạn gặp lỗi bên dưới:

AttributeError: Đối tượng 'list' không có thuộc tính 'collect'

Mã này sẽ giải quyết các vấn đề của bạn:

mvv_list = mvv_count_df.select('mvv').collect()

mvv_array = [int(i.mvv) for i in mvv_list]

Tôi đã chạy một phân tích điểm chuẩn và list(mvv_count_df.select('mvv').toPandas()['mvv'])là phương pháp nhanh nhất. Tôi rất ngạc nhiên.

Tôi đã chạy các phương pháp tiếp cận khác nhau trên tập dữ liệu 100 nghìn / 100 triệu hàng bằng cách sử dụng cụm i3.xlarge 5 nút (mỗi nút có 30,5 GB RAM và 4 lõi) với Spark 2.4.5. Dữ liệu được phân phối đồng đều trên 20 tệp Parquet nén nhanh với một cột duy nhất.

Đây là kết quả đo điểm chuẩn (thời gian chạy tính bằng giây):

+-------------------------------------------------------------+---------+-------------+
|                          Code                               | 100,000 | 100,000,000 |
+-------------------------------------------------------------+---------+-------------+
| df.select("col_name").rdd.flatMap(lambda x: x).collect()    |     0.4 | 55.3        |
| list(df.select('col_name').toPandas()['col_name'])          |     0.4 | 17.5        |
| df.select('col_name').rdd.map(lambda row : row[0]).collect()|     0.9 | 69          |
| [row[0] for row in df.select('col_name').collect()]         |     1.0 | OOM         |
| [r[0] for r in mid_df.select('col_name').toLocalIterator()] |     1.2 | *           |
+-------------------------------------------------------------+---------+-------------+

* cancelled after 800 seconds

Các quy tắc vàng cần tuân theo khi thu thập dữ liệu trên nút trình điều khiển:

• Cố gắng giải quyết vấn đề bằng các cách tiếp cận khác. Việc thu thập dữ liệu vào nút trình điều khiển rất tốn kém, không khai thác được sức mạnh của cụm Spark và nên tránh bất cứ khi nào có thể.
• Thu thập càng ít hàng càng tốt. Tổng hợp, loại bỏ trùng lặp, lọc và cắt bớt các cột trước khi thu thập dữ liệu. Gửi càng ít dữ liệu đến nút trình điều khiển càng tốt.

toPandas đã được cải thiện đáng kể trong Spark 2.3 . Đó có lẽ không phải là cách tốt nhất nếu bạn đang sử dụng phiên bản Spark cũ hơn 2.3.

Xem tại đây để biết thêm chi tiết / kết quả điểm chuẩn.

Một giải pháp khả thi là sử dụng collect_list()hàm from pyspark.sql.functions. Điều này sẽ tổng hợp tất cả các giá trị cột thành một mảng pyspark được chuyển đổi thành danh sách python khi được thu thập:

mvv_list   = df.select(collect_list("mvv")).collect()[0][0]
count_list = df.select(collect_list("count")).collect()[0][0]