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?

score 8 · Accepted Answer · answered Apr 14 '18 at 15:14

8

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)

answered Apr 14 '18 at 15:14

Alper t. Turker

34,230
9
83
115

Thanks! I want to call `AFTSurvivalRegression`, and it only works if the "features" columns are of the type `VectorUDT`. – Bruno Apr 14 '18 at 18:11
Then use `pyspark.ml.features` utilities. – Alper t. Turker Apr 14 '18 at 18:22

Adding a Vectors Column to a pyspark DataFrame

1 Answers1

Linked