如何访问Spark DataFrame中VectorUDT列的元素？

Question

我有一个数据帧df有VectorUDT指定的列features.如何获取列的元素,比如第一个元素？

我尝试过以下操作

from pyspark.sql.functions import udf
first_elem_udf = udf(lambda row: row.values[0])
df.select(first_elem_udf(df.features)).show()

但是我收到了一个net.razorvine.pickle.PickleException: expected zero arguments for construction of ClassDict(for numpy.dtype)错误.如果我first_elem_udf = first_elem_udf(lambda row: row.toArray()[0])改为相同的错误.

我也试过,explode()但我得到一个错误,因为它需要一个数组或地图类型.

我认为这应该是一种常见的操作.

Answer 1

将输出转换为float:

from pyspark.sql.types import DoubleType
from pyspark.sql.functions import lit, udf

def ith_(v, i):
    try:
        return float(v[i])
    except ValueError:
        return None

ith = udf(ith_, DoubleType())

用法示例:

from pyspark.ml.linalg import Vectors

df = sc.parallelize([
    (1, Vectors.dense([1, 2, 3])),
    (2, Vectors.sparse(3, [1], [9]))
]).toDF(["id", "features"])

df.select(ith("features", lit(1))).show()

## +-----------------+
## |ith_(features, 1)|
## +-----------------+
## |              2.0|
## |              9.0|
## +-----------------+

说明:

输出值必须重新序列化为等效的Java对象.如果你想访问values(提防SparseVectors)你应该使用item方法:

v.values.item(0)

返回标准的Python标量.同样,如果要将所有值作为密集结构访问:

v.toArray().tolist()