Adding a Vectors Column to a pyspark DataFrame

Question

How do I add a Vectors.dense column to a pyspark dataframe?

import pandas as pd
from pyspark import SparkContext
from pyspark.sql import SQLContext
from pyspark.ml.linalg import DenseVector

py_df = pd.DataFrame.from_dict({"time": [59., 115., 156., 421.], "event": [1, 1, 1, 0]})

sc = SparkContext(master="local")
sqlCtx = SQLContext(sc)
sdf = sqlCtx.createDataFrame(py_df)
sdf.withColumn("features", DenseVector(1))

Gives an error in file anaconda3/lib/python3.6/site-packages/pyspark/sql/dataframe.py, line 1848:

AssertionError: col should be Column

It doesn't like the DenseVector type as a column. Essentially, I have a pandas dataframe that I'd like to transform to a pyspark dataframe and add a column of the type Vectors.dense. Is there another way of doing this?

Answer 1

Constant Vectors cannot be added as literal. You have to use udf:

from pyspark.sql.functions import udf
from pyspark.ml.linalg import VectorUDT

one = udf(lambda: DenseVector([1]), VectorUDT())
sdf.withColumn("features", one()).show()

But I am not sure why you need that at all. If you want to transform existing columns into Vectors use appropriate pyspark.ml tools, like VectorAssembler - Encode and assemble multiple features in PySpark

from pyspark.ml.feature import VectorAssembler

VectorAssembler(inputCols=["time"], outputCol="features").transform(sdf)