如何从 pyspark SparseVector 获取密钥答案

【问题标题】：How to get keys from pyspark SparseVector如何从 pyspark SparseVector 获取密钥
【发布时间】：2019-05-28 13:21:50
【问题描述】：

我进行了 tf-idf 转换，现在我想从结果中获取键和值。

我正在使用以下 udf 代码来获取值：

def extract_values_from_vector(vector):
    return vector.values.tolist()

extract_values_from_vector_udf = udf(lambda vector:extract_values_from_vector(vector), ArrayType(DoubleType()))

extract = rescaledData.withColumn("extracted_keys", extract_keys_from_vector_udf("features"))

所以如果稀疏向量看起来像： features=SparseVector(123241, {20672: 4.4233, 37393: 0.0, 109847: 3.7096, 118474: 5.4042}))

extracted_keys 在我的提取中将如下所示： [4.4233, 0.0, 3.7096, 5.4042]

我的问题是，如何获取 SparseVector 字典中的键？比如keys = [20672, 37393, 109847, 118474]？

我正在尝试以下代码，但它不起作用

def extract_keys_from_vector(vector):
    return vector.indices.tolist()
extract_keys_from_vector_udf = spf.udf(lambda vector:extract_keys_from_vector(vector), ArrayType(DoubleType()))

它给我的结果是：[null,null,null,null]

有人可以帮忙吗？非常感谢！

【问题讨论】：

我不这么认为，那是 RDD
@Astory-teller 索引是整数值，但您的 UDF 返回一个双精度数组。我猜你只是希望它是IntegerType()
@Sergey Khudyakov 我认为你是对的！你想回答这个问题，我可以标记接受吗？

标签： pyspark tf-idf

【解决方案1】：

既然答案就在上面的cmets中，我想我会花时间（当然是在等待写拼花的时候）写下答案。

from pyspark.sql.types import *
from pyspark.sql import functions as F

def extract_keys_from_vector(vector):
    return vector.indices.tolist()

feature_extract = F.UserDefinedFunction(lambda vector: extract_keys_from_vector(vector), ArrayType(IntegerType()))

df = df.withColumn("features", feature_extract(F.col("features")))

【讨论】：