From 17add4776fabc93de2c305f917cc530466e19db3 Mon Sep 17 00:00:00 2001
From: Bobby Wang <wbo4958@gmail.com>
Date: Tue, 25 Apr 2023 23:09:12 +0800
Subject: [PATCH] [pyspark] Don't stack for non feature columns (#9088)

---
 python-package/xgboost/spark/data.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/python-package/xgboost/spark/data.py b/python-package/xgboost/spark/data.py
index f2c5e1197..8f84459d7 100644
--- a/python-package/xgboost/spark/data.py
+++ b/python-package/xgboost/spark/data.py
@@ -219,7 +219,9 @@ 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]
-                array = stack_series(array)
+                if name == alias.data:
+                    # For the array/vector typed case.
+                    array = stack_series(array)
             else:
                 array = None