[jvm-packages] Don't cast to float if it's already float (#10386)

dmlc · Jun 4, 2024 · bc7643d · bc7643d
1 parent 9b7633c
commit bc7643d
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/...es/xgboost4j-spark-gpu/src/main/scala/ml/dmlc/xgboost4j/scala/rapids/spark/GpuUtils.scala b/...es/xgboost4j-spark-gpu/src/main/scala/ml/dmlc/xgboost4j/scala/rapids/spark/GpuUtils.scala
@@ -89,9 +89,13 @@ private[spark] object GpuUtils {
     val featureNameSet = featureNames.distinct
     validateSchema(dataset.schema, featureNameSet, labelName, weightName, marginName, fitting)
 
-    val castToFloat = (ds: Dataset[_], colName: String) => {
-      val colMeta = ds.schema(colName).metadata
-      ds.withColumn(colName, col(colName).as(colName, colMeta).cast(FloatType))
+    val castToFloat = (df: DataFrame, colName: String) => {
+      if (df.schema(colName).dataType.isInstanceOf[FloatType]) {
+        df
+      } else {
+        val colMeta = df.schema(colName).metadata
+        df.withColumn(colName, col(colName).as(colName, colMeta).cast(FloatType))
+      }
     }
     val colNames = if (fitting) {
       var names = featureNameSet :+ labelName