Spark MLlib: Including categorical features

Question

What is the correct or best method for including categorical variables (both string and int) into a feature for an MLlib algorithm?

Is it correct to use OneHotEncoders on the categorical variables and then include the output columns with other columns in a VectorAssembler like in the code below?

The reason is that I end up with a data frame with rows like this where it looks like feature3 and feature4 combined look like they are on the same 'level' of importance as the two categorical features singly.

+------------------+-----------------------+---------------------------+
|prediction        |actualVal |features                                |
+------------------+-----------------------+---------------------------+
|355416.44924898935|990000.0  |(17,[0,1,2,3,4,5,10,15],[1.0,206.0])    |
|358917.32988024893|210000.0  |(17,[0,1,2,3,4,5,10,15,16],[1.0,172.0]) |
|291313.84175674635|4600000.0 |(17,[0,1,2,3,4,5,12,15,16],[1.0,239.0]) |

Here is my code:

val indexer = new StringIndexer()
  .setInputCol("stringFeatureCode")
  .setOutputCol("stringFeatureCodeIndex")
  .fit(data)
val indexed = indexer.transform(data)

val encoder = new OneHotEncoder()
  .setInputCol("stringFeatureCodeIndex")
  .setOutputCol("stringFeatureCodeVec")

var encoded = encoder.transform(indexed)

encoded = encoded.withColumn("intFeatureCodeTmp", encoded.col("intFeatureCode")
  .cast(DoubleType))
  .drop("intFeatureCode")
  .withColumnRenamed("intFeatureCodeTmp", "intFeatureCode")

val intFeatureCodeEncoder = new OneHotEncoder()
  .setInputCol("intFeatureCode")
  .setOutputCol("intFeatureCodeVec")

encoded = intFeatureCodeEncoder.transform(encoded)

val assemblerDeparture =
  new VectorAssembler()
    .setInputCols(
      Array("stringFeatureCodeVec", "intFeatureCodeVec", "feature3", "feature4"))
    .setOutputCol("features")
var data2 = assemblerDeparture.transform(encoded)

val Array(trainingData, testData) = data2.randomSplit(Array(0.7, 0.3))

val rf = new RandomForestRegressor()
  .setLabelCol("actualVal")
  .setFeaturesCol("features")
  .setNumTrees(100)

Unknown Unknown · Accepted Answer · 2016-11-10T22:01:28

In general this is a recommended method.
When working tree models it unnecessary and should be avoided. You can use StringIndexer only.

Spark MLlib: Including categorical features

1 Answers