[pyspark] Don't stack for non feature columns (#9088)

2023-04-25 23:09:12 +08:00 · 2023-04-25 23:09:12 +08:00 · 17add4776f
commit 17add4776f
parent a2cc78c1fb
1 changed files with 3 additions and 1 deletions
--- a/python-package/xgboost/spark/data.py
+++ b/python-package/xgboost/spark/data.py
@ -219,6 +219,8 @@ def create_dmatrix_from_partitions(  # pylint: disable=too-many-arguments
                array: Optional[np.ndarray] = part[feature_cols]
            elif part[name].shape[0] > 0:
                array = part[name]
+                if name == alias.data:
+                    # For the array/vector typed case.
                    array = stack_series(array)
            else:
                array = None