[pyspark] rework transform to reuse same code #9292
Conversation
|
@WeichenXu123 @trivialfis could you please help to review this PR? Previously, SparkXGBClassifier totally overrides the whole _transform function which causes duplicated common code, so this PR tries to unify them. |
| result = data[pred.prediction] | ||
| if pred_contrib_col_name: | ||
| contribs = pred_contribs(model, X, base_margin) | ||
| data[pred.pred_contrib] = pd.Series(list(contribs)) |
There was a problem hiding this comment.
Not sure how it works. Wouldn't this be super slow?
There was a problem hiding this comment.
any thoughts to rework this? using contribs.tolist() ?
python-package/xgboost/spark/core.py
Outdated
| def _post_transform(self, dataset: DataFrame, pred_col: Column) -> DataFrame: | ||
| """Post process of transform""" | ||
| prediction_col_name = self.getOrDefault(self.predictionCol) | ||
| single_pred = "," not in self._out_schema() |
There was a problem hiding this comment.
That's a bit, hmm, unconventional. Is there a way to refactor this to make it more "standard"?
There was a problem hiding this comment.
The _out_schema has typing hint with str returned which is a DDL formatted string, so if there're many columns, it must have at least a ",". So it's conventional when we check if it is a single column according to if "," is in the schema.
There was a problem hiding this comment.
We are meta-programming by manipulating strings here.
There was a problem hiding this comment.
Maybe just use the pred_contrib_col_name as a predicate?
python-package/xgboost/spark/core.py
Outdated
| pred_contrib_col_name = self._get_pred_contrib_col_name() | ||
|
|
||
| def _predict( | ||
| model: XGBModel, X: ArrayLike, base_margin: Optional[np.ndarray] |
There was a problem hiding this comment.
base_margin is not necessarily np.ndarray. Let's stick with ArrayLike.
python-package/xgboost/spark/core.py
Outdated
| def _post_transform(self, dataset: DataFrame, pred_col: Column) -> DataFrame: | ||
| """Post process of transform""" | ||
| prediction_col_name = self.getOrDefault(self.predictionCol) | ||
| single_pred = "," not in self._out_schema() |
There was a problem hiding this comment.
Maybe just use the pred_contrib_col_name as a predicate?
wbo4958
force-pushed
the
rework-transform
branch
from
62aebbd
to
c8d353b
Compare
Co-authored-by: Bobby Wang <[email protected]>
To fix #9170