Check if Column is Vector Type

Question

I am trying to determine if a column is a vector type, but am running into issues.

After I run a model and create a dataframe called predictions, there is a field called probability.

When I run this code to see the datatype if shows a vector.

predictions.schema['probability'].dataType
Out[128]: VectorUDT

Then when I run this I get a false returned

predictions.schema["probability"].dataType == 'VectorUDT'
Out[129]: False

So I tried this

dict(predictions.dtypes)['probability'] == 'vector'
Out[130]: True

However, when I try to use that in my dataframe I get an error stating TypeError: unhashable type: 'Column'

.withColumn('test',   when(dict(predictions.dtypes)['probability'] == 'vector',1)
                           .otherwise(0)) \

user_2129801 · Answer

predictions.schema['probability'] returns a StructField and StructField.dataType returns a DataType. So

predictions.schema["probability"].dataType == 'VectorUDT'

compares a DataType with a string and this always false.

One way to check for a vector column is to check the type of the dataType using isinstance:

from pyspark.ml.linalg import VectorUDT
isinstance(predictions.schema["probability"].dataType, VectorUDT) #True

Another way to check the column type is to use DataType.simpleString or DataType.typeName. Both return the same string, vector for VectorUDT fields:

predictions.schema['probability'].dataType.simpleString() == 'vector' #True