PySpark K-means with categorical variables

Question

I started playing with kmeans clustering in pyspark (v 1.6.2) with the following example which includes mixed variable types:

# Import libraries
from pyspark.ml.feature import OneHotEncoder, StringIndexer, VectorAssembler
from pyspark.ml.clustering import KMeans
from pyspark.ml import Pipeline
from pyspark.mllib.clustering import KMeansModel

# Create sample DF
sample = sqlContext.createDataFrame([["a@email.com", 12000,"M"],
                                ["b@email.com", 43000,"M"],
                                ["c@email.com", 5000,"F"],
                                ["d@email.com", 60000,"M"]]
                                , ["email", "income","gender"])

I used StringIndexer, OneHotEncoder, and VectorAssembler to process categorical attributes like this:

# Indexers encode strings
string_indexers=[
StringIndexer(inputCol=x,outputCol="idx_{0}".format(x))
for x in sample.columns if x not in 'income']

encoders=[
OneHotEncoder(inputCol="idx_{0}".format(x),outputCol="enc_{0}".format(x))
for x in sample.columns if x not in 'income']

# Assemble multiple columns into a single vector
assembler=VectorAssembler(
inputCols=["enc_{0}".format(x) for x in sample.columns if x not in 'income'] + ['income'],
outputCol="features")

And this piece makes sure that the conversion works fine:

pipeline= Pipeline(stages=string_indexers+encoders+[assembler])
model=pipeline.fit(sample)
indexed=model.transform(sample)

indexed.show()

I know I can do a k-means on this transformed DF by doing something like this:

kmeans = KMeans()
          .setK(2)
          .setFeaturesCol("features")
          .setPredictionCol("prediction")

kmeans_transformer = kmeans.fit(indexed)
oo = kmeans_transformer.transform(indexed)

oo.select('email', 'income', 'gender', 'features', 
'prediction').show(truncate = False)

+-----------+------+------+-------------------------+----------+
|email      |income|gender|features                 |prediction|
+-----------+------+------+-------------------------+----------+
|a@email.com|12000 |M     |[0.0,1.0,0.0,1.0,12000.0]|1         |
|b@email.com|43000 |M     |(5,[3,4],[1.0,43000.0])  |0         |
|c@email.com|5000  |F     |(5,[0,4],[1.0,5000.0])   |1         |
|d@email.com|60000 |M     |[0.0,0.0,1.0,1.0,60000.0]|0         |
+-----------+------+------+-------------------------+----------+

But I wanted to see:

1) How can I do same thing with pyspark.mllib.clustering.KMeansModel to identify best (least cost) value of K ( aligned with KMeans.train and computeCost functions in pyspark generic example)?

2) How can I get cluster centers in the original scale (meaning "Male" or "Female" labels NOT in encoded scale)?

PySpark version 1.6.2

PySpark K-means with categorical variables

0 Answers