如何将Pyspark数据框列转换为Numpy数组

Question

我试图将一个包含约9000万行的pyspark dataframe列转换为numpy数组。

我需要将该数组作为scipy.optimize.minimize函数的输入。

我尝试了将其转换为Pandas并使用collect()，但这些方法非常耗时。

我是PySpark的新手，如果有更快更好的方法，请帮忙指导。

谢谢

这是我的数据框的样子。

+----------+
|Adolescent|
+----------+
|       0.0|
|       0.0|
|       0.0|
|       0.0|
|       0.0|
|       0.0|
|       0.0|
|       0.0|
|       0.0|
|       0.0|
+----------+

- Vaibhav Rathi

---

你尝试过 df['Adolescent'].to_numpy() 或者 df['Adolescent'].array 吗？ - Nils Werner

似乎 to_numpy() 只适用于 pandas 数据框，而不适用于 pyspark。我尝试了 df["Adolescent"].array，它给出的输出是："Column<b'Adolescent[array]'>"。我不知道如何将其用作数组。 - Vaibhav Rathi

Answer 1

#1

你需要以任何方式调用 .collect()。要从Pyspark DataFrame创建Numpy数组，可以使用以下方法：

adoles = np.array(df.select("Adolescent").collect()) #.reshape(-1) for 1-D array

#2

你可以使用toPandas()将其转换为pandas dataframe，然后使用.values将其转换为numpy数组。

pdf = df.toPandas()
adoles = df["Adolescent"].values

或者简单地：

adoles = df.select("Adolescent").toPandas().values #.reshape(-1) for 1-D array

#3

对于分布式数组，您可以尝试 Dask Arrays

我没有测试过这个方法，但是假设它与 numpy 的工作方式相同（可能存在不一致之处）：

import dask.array as da
adoles = da.array(df.select("Adolescent").collect()) #.reshape(-1) for 1-D array

- pissall

---

我尝试使用toPandas()，但是它花费了很多时间。 - Vaibhav Rathi

你应该看一下Dask Arrays。 - pissall

谢谢你的帮助。我现在正在尝试使用Dask数组。 - Vaibhav Rathi

Answer 2

另一种方法是将所选列转换为RDD，然后通过提取每个Row的值（可以滥用.keys()）来展平，然后转换为numpy数组：

x = df.select("colname").rdd.map(lambda r: r[0]).collect()  # python list
np.array(x)  # numpy array

- qwr