Python Spark DataFrame: заменить ноль на SparseVector

Question

В спарк у меня есть следующий фрейм данных под названием "df" с некоторыми пустыми записями:

+-------+--------------------+--------------------+                     
|     id|           features1|           features2|
+-------+--------------------+--------------------+
|    185|(5,[0,1,4],[0.1,0...|                null|
|    220|(5,[0,2,3],[0.1,0...|(10,[1,2,6],[0.1,...|
|    225|                null|(10,[1,3,5],[0.1,...|
+-------+--------------------+--------------------+

df.features1 и df.features2 являются типом vector (обнуляемый). Затем я попытался использовать следующий код для заполнения пустых записей SparseVectors:

df1 = df.na.fill({"features1":SparseVector(5,{}), "features2":SparseVector(10, {})})

Этот код привел к следующей ошибке:

AttributeError: 'SparseVector' object has no attribute '_get_object_id'

Тогда я нашел следующий абзац в документации искры:

fillna(value, subset=None)
Replace null values, alias for na.fill(). DataFrame.fillna() and DataFrameNaFunctions.fill() are aliases of each other.

Parameters: 
value – int, long, float, string, or dict. Value to replace null values with. If the value is a dict, then subset is ignored and value must be a mapping from column name (string) to replacement value. The replacement value must be an int, long, float, or string.

Объясняет ли это мою неспособность заменить пустые записи на SparseVectors в DataFrame? Или это означает, что в DataFrame нет способа сделать это?

Я могу достичь своей цели, преобразовав DataFrame в RDD и заменив значения None на SparseVectors, но мне будет гораздо удобнее сделать это непосредственно в DataFrame.

Есть ли способ сделать это непосредственно в DataFrame? Спасибо!

Answer 1

Ты можешь использовать udf:

from pyspark.sql.functions import udf, lit
from pyspark.ml.linalg import *

fill_with_vector = udf(
    lambda x, i: x if x is not None else SparseVector(i, {}),
    VectorUDT()
)

df = sc.parallelize([
    (SparseVector(5, {1: 1.0}), SparseVector(10, {1: -1.0})), (None, None)
]).toDF(["features1", "features2"])

(df
    .withColumn("features1", fill_with_vector("features1", lit(5)))
    .withColumn("features2", fill_with_vector("features2", lit(10)))
    .show())

# +-------------+---------------+
# |    features1|      features2|
# +-------------+---------------+
# |(5,[1],[1.0])|(10,[1],[-1.0])|
# |    (5,[],[])|     (10,[],[])|
# +-------------+---------------+