[SPARK-3807][MLLIB] fix col indexing bug and add a check for number of distinct values

There is a bug determining the column index. @dorx
Should use `zipWithIndex` before slicing.