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

dmlc · Apr 25, 2023 · 17add47 · 17add47
1 parent a2cc78c
commit 17add47
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/python-package/xgboost/spark/data.py b/python-package/xgboost/spark/data.py
@@ -219,7 +219,9 @@ def append_m(part: pd.DataFrame, name: str, is_valid: bool) -> None:
                 array: Optional[np.ndarray] = part[feature_cols]
             elif part[name].shape[0] > 0:
                 array = part[name]
-                array = stack_series(array)
+                if name == alias.data:
+                    # For the array/vector typed case.
+                    array = stack_series(array)
             else:
                 array = None