PySpark：如何将带有稀疏向量列的Spark数据框写入CSV文件？

Question

我有一个Spark DataFrame，其中有一列的类型为spark.mllib.linalg.SparseVector：

1) 如何将它写入CSV文件？

2) 如何打印出所有向量？

- Alex

Answer 1

要将数据框写入csv文件，您可以使用标准的df.write.csv(output_path)。但是，如果您仅使用以上内容，则可能会在具有SparseVector类型的列上收到“java.lang.UnsupportedOperationException：CSV数据源不支持struct<type:tinyint,size:int,indices:array<int>,values:array<double>>”错误。

有两种方法可以打印SparseVector并避免该错误：稀疏格式或密集格式。

如果您想以密集格式打印，则可以像这样定义udf：

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType
from pyspark.sql.functions import col

dense_format_udf = udf(lambda x: ','.join([str(elem) for elem in x], StringType())

df = df.withColumn('column_name', dense_format_udf(col('column_name')))

df.write.option("delimiter", "\t").csv(output_path)

在紧凑格式下，该列的输出如下：1.0,0.0,5.0,0.0

如果您想以稀疏格式打印，可以利用SparseVector类的OOB __str__函数，或者创造性地定义自己的输出格式。这里我将使用OOB函数。

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType
from pyspark.sql.functions import col

sparse_format_udf = udf(lambda x: str(x), StringType())

df = df.withColumn('column_name', sparse_format_udf(col('column_name')))

df.write.option("delimiter", "\t").csv(output_path)

这一列以稀疏格式打印，类似于(4,[0,2],[1.0,5.0])

注意我之前尝试过这个方法：df = df.withColumn("column_name", col("column_name").cast("string")) 但是该列只会打印出类似于 [0,5,org.apache.spark.sql.catalyst.expressions.UnsafeArrayData@6988050,org.apache.spark.sql.catalyst.expressions.UnsafeArrayData@ec4ae6ab] 这样并不理想。

- cindyxiaoxiaoli

---

谢谢，这很有用，但是您在UDF名称中是否混淆了稀疏和密集？即1.0,0.0,5.0,0.0应该使用dense_format_udf进行密集格式，而(4,[0,2],[1.0,5.0])应该使用sparse_format_udf进行稀疏格式。 - Artem Trunov

Answer 2

1. https://github.com/databricks/spark-csv

2. df2 = df1.map(lambda row: row.yourVectorCol)

   或者 df1.map(lambda row: row[1])

   如果你有一个命名的列，可以直接引用它，否则可以通过位置引用。

   然后，要打印它，可以使用df2.collect()

没有更多信息的话，这可能对你有帮助，也可能不够详细。请稍微详细说明一下。

- Kristian