从带有标签的特征向量创建 PySpark 数据框

Question

我有一个使用 Pipeline 对象创建的数据框，如下所示：

df.show()

+--------------------+-----+
|            features|label|
+--------------------+-----+
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
|[-0.0775219322931...|    0|
+--------------------+-----+

我已经成功提取了这样的特征向量：

df_table = df.rdd.map(lambda x: [float(y) for y in x['features']]).toDF(cols)

上面的问题是它没有保留标签列。作为一种解决方法，我成功地使用了 Join 来恢复该标签列，但我发现它太复杂了。

我将如何使用像上面这样的单线来提取特征向量并从中制作 Spark DF，同时将标签列附加到它上面？

Answer 1

here 有很好的选择，尤其是如果你有一个版本 Spark >= 3.0.0

假设您没有这样的最新版本，您的问题来自于您丢失了地图中的密钥这一事实。你可以这样做：

df_table = df.rdd.map(lambda l: tuple([l['label']] + [float(y) for y in l['features']])).toDF()

您最终会得到一个宽格式的数据框。如果你想要一个长格式的矢量，你有更多的选择。

如果你想要长格式的数据

首先，用rdd：

df.rdd.flatMapValues(lambda l: l).toDF(['label','feature'])

或者，甚至更好，直接使用DataFrame API:（未经测试的解决方案）

import pyspark.sql.functions as psf
df.select('label', psf.explode(psf.col('label')))