fix: Reuse pd.DataFrame index in _transform

scma2z · Nov 24, 2023 · 9bccf9c · 9bccf9c
1 parent 3754085
commit 9bccf9c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/optbinning/binning/binning_process.py b/optbinning/binning/binning_process.py
@@ -1423,7 +1423,7 @@ def _transform(self, X, metric, metric_special, metric_missing,
             X_transform[:, i] = optb.transform(**tparams)
 
         if isinstance(X, pd.DataFrame):
-            return pd.DataFrame(X_transform, columns=selected_variables)
+            return pd.DataFrame(X_transform, columns=selected_variables, index=X.index)
 
         return X_transform
 

diff --git a/tests/test_binning_process.py b/tests/test_binning_process.py
@@ -580,3 +580,10 @@ def test_verbose():
     with open("tests/results/test_binning_process_verbose.txt", "w") as f:
         with redirect_stdout(f):
             process.fit(X, y, check_input=True)
+
+
+def test_dataframe_index():
+    process = BinningProcess(variable_names)
+    X_train = pd.DataFrame(X, columns=variable_names, index=[2 * i for i in range(len(X))])
+    X_transform = process.fit_transform(X_train, y, metric="indices")
+    pd.testing.assert_index_equal(X_train.index, X_transform.index)