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

There is a bug determining the column index.