Dot product в фреймах данных pyspark с помощью MLLIB

Question

У меня очень простой фрейм данных в pyspark, что-то вроде этого:

from pyspark.sql import Row
from pyspark.mllib.linalg import DenseVector

row = Row("a", "b")
df = spark.sparkContext.parallelize([
    offer_row(DenseVector([1, 1, 1]), DenseVector([1, 0, 0])),
]).toDF()

, и я хотел бы вычислить скалярное произведение этих векторов, не прибегая к вызову UDF.

Документация spark MLLIB ссылается на метод dot для DenseVectors, но если я попытаюсь применить его следующим образом:

df_offers = df_offers.withColumn("c", col("a").dot(col("b")))

Я получаю ошибки, подобные:

TypeError: 'Column' object is not callable

Кто-нибудь знает, можно ли вызывать эти методы mllib для объектов DataFrame?

Answer 1

Здесь вы применяете метод dot к столбцу, а не к DenseVector, который действительно не работает:

df_offers = df_offers.withColumn("c", col("a").dot(col("b")))

Вам придется использовать udf:

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

def dot_fun(array):
    return array[0].dot(array[1])

dot_udf = udf(dot_fun, DoubleType())

df_offers = df_offers.withColumn("c", dot_udf(array('a', 'b')))

Answer 2

Нет. Вам придется использовать udf:

from pyspark.sql.functions import udf

@udf("double")
def dot(x, y):
    if x is not None and y is not None:
        return float(x.dot(y))