python - 分解多行中具有密集向量的列

Question

我有一个包含两列的数据框:BrandWatchErwaehnungID 和 word_counts。word_counts 列是“CountVectorizer”(稀疏向量)的输出。删除空行后，我创建了两列新列，一列包含稀疏向量的索引，另一列包含它们的值。

help0 = countedwords_text['BrandWatchErwaehnungID','word_counts'].rdd\
    .filter(lambda x : x[1].indices.size!=0)\
    .map(lambda x : (x[0],x[1],DenseVector(x[1].indices) , DenseVector(x[1].values))).toDF()\
    .withColumnRenamed("_1", "BrandWatchErwaenungID").withColumnRenamed("_2", "word_counts")\
    .withColumnRenamed("_3", "word_indices").withColumnRenamed("_4", "single_word_counts")

由于 Spark 不接受 numpy.ndarray，我需要在添加到我的 Dataframe 之前将它们转换为密集向量。我的问题是，我现在想要在 word_indices 列上分解该 Dataframe，但 pyspark.sql.functions 中的 explode 方法仅支持数组或 map 作为输入。

我已经尝试过:

help1 = help0.withColumn('b' , explode(help0.word_indices))

并出现以下错误:

cannot resolve 'explode(`word_indices')' due to data type mismatch: input to function explode should be array or map type

后来我尝试了:

help1 = help0.withColumn('b' , explode(help0.word_indices.toArray()))

这也不起作用......有什么建议吗？

Answer 1

您必须使用udf:

from pyspark.sql.functions import udf, explode
from pyspark.sql.types import *
from pyspark.ml.linalg import *

@udf("array<integer>")
def indices(v):
   if isinstance(v, DenseVector):
      return list(range(len(v)))
   if isinstance(v, SparseVector):
      return v.indices.tolist()

df = spark.createDataFrame([
   (1, DenseVector([1, 2, 3])), (2, SparseVector(5, {4: 42}))], 
   ("id", "v"))

df.select("id", explode(indices("v"))).show()

# +---+---+
# | id|col|
# +---+---+
# |  1|  0|
# |  1|  1|
# |  1|  2|
# |  2|  4|
# +---+---+