Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Bobby Wang 2024-06-04 18:01:51 +08:00 committed by GitHub
parent 9b7633c01d
commit bc7643d35e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
jvm-packages/xgboost4j-spark-gpu/src/main/scala/ml/dmlc/xgboost4j/scala/rapids/spark/GpuUtils.scala
View File

@ -89,9 +89,13 @@ private[spark] object GpuUtils {
val featureNameSet = featureNames.distinct val featureNameSet = featureNames.distinct
validateSchema(dataset.schema, featureNameSet, labelName, weightName, marginName, fitting) validateSchema(dataset.schema, featureNameSet, labelName, weightName, marginName, fitting)
val castToFloat = (ds: Dataset[_], colName: String) => { val castToFloat = (df: DataFrame, colName: String) => {
val colMeta = ds.schema(colName).metadata if (df.schema(colName).dataType.isInstanceOf[FloatType]) {
ds.withColumn(colName, col(colName).as(colName, colMeta).cast(FloatType)) df
} else {
val colMeta = df.schema(colName).metadata
df.withColumn(colName, col(colName).as(colName, colMeta).cast(FloatType))
}
} }
val colNames = if (fitting) { val colNames = if (fitting) {
var names = featureNameSet :+ labelName var names = featureNameSet :+ labelName