Skip to content

Commit ba9b4cb

Browse files
authored
Fix pylint. (dmlc#10296)
1 parent 835e59e commit ba9b4cb

File tree

1 file changed

+5
-3
lines changed
  • python-package/xgboost/spark

1 file changed

+5
-3
lines changed

python-package/xgboost/spark/core.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -527,7 +527,8 @@ def _validate_and_convert_feature_col_as_array_col(
527527
(DoubleType, FloatType, LongType, IntegerType, ShortType),
528528
):
529529
raise ValueError(
530-
"If feature column is array type, its elements must be number type."
530+
"If feature column is array type, its elements must be number type, "
531+
f"got {features_col_datatype.elementType}."
531532
)
532533
features_array_col = features_col.cast(ArrayType(FloatType())).alias(alias.data)
533534
elif isinstance(features_col_datatype, VectorUDT):
@@ -1379,15 +1380,15 @@ def _transform(self, dataset: DataFrame) -> DataFrame:
13791380
# to avoid the `self` object to be pickled to remote.
13801381
xgb_sklearn_model = self._xgb_sklearn_model
13811382

1382-
has_base_margin = False
1383+
base_margin_col = None
13831384
if (
13841385
self.isDefined(self.base_margin_col)
13851386
and self.getOrDefault(self.base_margin_col) != ""
13861387
):
1387-
has_base_margin = True
13881388
base_margin_col = col(self.getOrDefault(self.base_margin_col)).alias(
13891389
alias.margin
13901390
)
1391+
has_base_margin = base_margin_col is not None
13911392

13921393
features_col, feature_col_names = self._get_feature_col(dataset)
13931394
enable_sparse_data_optim = self.getOrDefault(self.enable_sparse_data_optim)
@@ -1472,6 +1473,7 @@ def to_gpu_if_possible(data: ArrayLike) -> ArrayLike:
14721473
yield predict_func(model, X, base_margin)
14731474

14741475
if has_base_margin:
1476+
assert base_margin_col is not None
14751477
pred_col = predict_udf(struct(*features_col, base_margin_col))
14761478
else:
14771479
pred_col = predict_udf(struct(*features_col))

0 commit comments

Comments
 (0)