[SPARK-2839][MLlib] Addressed comments