train_test_split incorrect column names fix (#998)

* train_test_split incorrect column names fix * fix pep8 issue
uxlfoundation · Apr 13, 2022 · 9764b19 · 9764b19
1 parent 7a39c0a
commit 9764b19
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/daal4py/sklearn/model_selection/_split.py b/daal4py/sklearn/model_selection/_split.py
@@ -237,6 +237,14 @@ def _daal_train_test_split(*arrays, **options):
                 train_arr.index = train
                 test_arr.index = test
 
+            if hasattr(arr, 'columns'):
+                train_arr.columns = arr.columns
+                test_arr.columns = arr.columns
+
+            if hasattr(arr, 'name'):
+                train_arr.name = arr.name
+                test_arr.name = arr.name
+
             res.append(train_arr)
             res.append(test_arr)