simply the logic of ensureIdCol

python/src/spark_rapids_ml/knn.py

@@ -142,19 +142,19 @@ def _ensureIdCol(self, df: DataFrame) -> DataFrame:
         Ensure an id column exists in the input dataframe. Add the column if not exists.
         Overwritten for knn assumption on error for not setting idCol and duplicate exists.
         """
-        if not self.isSet("idCol") and self.getIdCol() in df.columns:
-            raise ValueError(
-                f"Cannot create a default id column since a column with the default name '{self.getIdCol()}' already exists."
-                + "Please specify an id column"
-            )
 
         id_col_name = self.getIdCol()
-        df_withid = (
-            df
-            if id_col_name in df.columns
-            else df.select(monotonically_increasing_id().alias(id_col_name), "*")
-        )
-        return df_withid
+        if id_col_name in df.columns:
+            if self.isSet("idCol"):
+                return df
+            else:
+                raise ValueError(
+                    f"Column '{id_col_name}' already exists."
+                    + f"Please use a different column name or set the column as the idCol."
+                )
+        else:
+            df_withid = df.select(monotonically_increasing_id().alias(id_col_name), "*")
+            return df_withid
 
 
 class NearestNeighbors(