include train and test indices in iter_cv_splits output (#2012)

Author: jeromedockes

CHANGES.rst

@@ -14,6 +14,9 @@ New Features
 
 Changes
 -------
+- The row indices of training and testing samples are now also included in the
+  dictionaries produced by :meth:`DataOp.skb.iter_cv_splits`. :pr:`2012` by
+  :user:`Jérôme Dockès <jeromedockes>`.
 
 Bugfixes
 --------

skrub/_data_ops/_estimator.py

@@ -842,6 +842,8 @@ def iter_cv_splits(data_op, environment, *, keep_subsampling=False, cv=None):
             "test": test_env,
             "X_train": X_train,
             "X_test": X_test,
+            "row_indices_train": train_idx,
+            "row_indices_test": test_idx,
         }
         if y is not None:
             y_train, y_test = (

skrub/_data_ops/_skrub_namespace.py

@@ -1939,6 +1939,8 @@ def iter_cv_splits(self, environment=None, *, keep_subsampling=False, cv=None):
             - y_test: the value of the variable marked with ``skb.mark_as_y()`` in
               the test environment, if there is one (may not be the case for
               unsupervised learning).
+            - row_indices_train: the row indices (in X and y) of the training samples.
+            - row_indices_test: the row indices (in X and y) of the testing samples.
 
         Examples
         --------

skrub/_data_ops/tests/test_estimators.py

@@ -606,9 +606,13 @@ def test_iter_cv_splits():
     s = next(splits)
     assert list(s["X_train"]) == list(s["train"]["_skrub_X"]) == [10, 20, 30, 40]
     assert list(s["X_test"]) == list(s["test"]["_skrub_X"]) == [0]
+    assert list(s["row_indices_train"]) == [1, 2, 3, 4]
+    assert list(s["row_indices_test"]) == [0]
     s = next(splits)
     assert list(s["X_train"]) == list(s["train"]["_skrub_X"]) == [0, 20, 30, 40]
     assert list(s["X_test"]) == list(s["test"]["_skrub_X"]) == [10]
+    assert list(s["row_indices_train"]) == [0, 2, 3, 4]
+    assert list(s["row_indices_test"]) == [1]
 
     X = skrub.X(np.arange(4) * 10)
     y = skrub.y(np.arange(4) * -10)
@@ -618,11 +622,15 @@ def test_iter_cv_splits():
     assert list(s["X_test"]) == list(s["test"]["_skrub_X"]) == [0]
     assert list(s["y_train"]) == list(s["train"]["_skrub_y"]) == [-10, -20, -30]
     assert list(s["y_test"]) == list(s["test"]["_skrub_y"]) == [0]
+    assert list(s["row_indices_train"]) == [1, 2, 3]
+    assert list(s["row_indices_test"]) == [0]
     s = next(splits)
     assert list(s["X_train"]) == list(s["train"]["_skrub_X"]) == [0, 20, 30]
     assert list(s["X_test"]) == list(s["test"]["_skrub_X"]) == [10]
     assert list(s["y_train"]) == list(s["train"]["_skrub_y"]) == [0, -20, -30]
     assert list(s["y_test"]) == list(s["test"]["_skrub_y"]) == [-10]
+    assert list(s["row_indices_train"]) == [0, 2, 3]
+    assert list(s["row_indices_test"]) == [1]
 
 
 def test_train_test_split_splitter_renaming():