fix(data): pass non-string column labels to PreTab by position

deeptab/data/datamodule.py

@@ -2,6 +2,7 @@
 
 import lightning as pl
 import numpy as np
+import pandas as pd
 import torch
 from sklearn.model_selection import train_test_split
 from torch.utils.data import DataLoader, WeightedRandomSampler
@@ -43,6 +44,27 @@ def _prepare_regression_labels(y, *, vector_targets: bool = False) -> torch.Tens
     return torch.as_tensor(y_arr, dtype=torch.float32)
 
 
+def _positional_columns(X):
+    """Relabel non-string DataFrame columns by position before they reach PreTab.
+
+    PreTab selects integer-labelled columns by *position* but detects their type
+    by *label*. Integer labels that are not ``0..n-1`` in order (e.g. after
+    ``df.drop(columns=0)`` on a headerless CSV) therefore crash or encode the
+    wrong column. Following scikit-learn, non-string column labels carry no
+    feature names, so such frames are relabelled ``0..n-1`` in column order.
+    String labels and positional integer labels (as produced for NumPy input)
+    are returned unchanged.
+    """
+    if not isinstance(X, pd.DataFrame):
+        return X
+    labels = list(X.columns)
+    if all(isinstance(label, str) for label in labels):
+        return X
+    if all(isinstance(label, int | np.integer) for label in labels) and labels == list(range(len(labels))):
+        return X
+    return X.set_axis(range(len(labels)), axis=1)
+
+
 class TabularDataModule(pl.LightningDataModule):
     """Manage tabular preprocessing, datasets, and Lightning data loaders.
 
@@ -256,6 +278,9 @@ def preprocess_data(
                 self.embeddings_train = None
                 self.embeddings_val = None
 
+        self.X_train = _positional_columns(self.X_train)
+        self.X_val = _positional_columns(self.X_val)
+
         if fit_preprocessing:
             self.preprocessor = fit_preprocessor(self.preprocessor, self.X_train, self.y_train, self.embeddings_train)
 
@@ -445,7 +470,7 @@ def preprocess_new_data(self, X, embeddings=None):
         cat_tensors = []
         num_tensors = []
         emb_tensors = []
-        preprocessed_data = self.preprocessor.transform(X, embeddings)
+        preprocessed_data = self.preprocessor.transform(_positional_columns(X), embeddings)
 
         # Populate tensors for categorical features, if present in processed data
         for key in self.cat_feature_info:  # type: ignore

docs/core_concepts/sklearn_api.md

@@ -100,6 +100,10 @@ X = np.random.randn(1000, 10)
 
 For mixed numerical/categorical data, DataFrames are strongly preferred.
 
+A DataFrame whose column labels are not all strings, for example the integer
+labels of `pd.read_csv(path, header=None)`, is treated like an array: its
+columns are used by position, and `feature_names_in_` is not set.
+
 Targets are required and must be finite. Classifiers require discrete labels;
 point regressors accept one numeric target column. Distributional models retain
 their task-specific target dimensions, including Dirichlet component columns.

tests/data/test_datamodule.py

@@ -803,3 +803,22 @@ def test_weighted_sampler_generator_is_none_when_no_random_state(self, classific
         sampler = dm._build_train_sampler()
         assert isinstance(sampler, WeightedRandomSampler)
         assert sampler.generator is None
+
+
+@pytest.mark.parametrize(
+    "columns, expected",
+    [
+        (["a", "b", "c"], ["a", "b", "c"]),
+        ([0, 1, 2], [0, 1, 2]),
+        ([1, 2, 3], [0, 1, 2]),
+        ([2, 0, 1], [0, 1, 2]),
+        (["a", 0, "c"], [0, 1, 2]),
+    ],
+)
+def test_positional_columns_relabels_only_non_positional_integer_labels(columns, expected):
+    from deeptab.data.datamodule import _positional_columns
+
+    frame = pd.DataFrame(np.arange(6).reshape(2, 3), columns=columns)
+    result = _positional_columns(frame)
+    assert list(result.columns) == expected
+    np.testing.assert_array_equal(result.to_numpy(), frame.to_numpy())

tests/integration/test_estimator_input_validation.py

@@ -119,3 +119,44 @@ def test_all_missing_column_fits_and_predicts(self):
             model.fit(X, y, max_epochs=1, batch_size=16)
         preds = model.predict(X)
         assert preds.shape == (40,)
+
+
+class TestIntegerColumnLabels:
+    """Integer column labels that are not 0..n-1 used to be read by position for
+    selection but by label for type detection, crashing or encoding the wrong column."""
+
+    @staticmethod
+    def _fit_predict(features, targets):
+        from deeptab.configs import TrainerConfig
+        from deeptab.models import MLPRegressor
+
+        model = MLPRegressor(trainer_config=TrainerConfig(max_epochs=1, batch_size=32), random_state=0)
+        model.fit(features, targets, accelerator="cpu", logger=False, enable_progress_bar=False)
+        return model, model.predict(features)
+
+    def test_frame_after_dropping_a_positional_column(self):
+        rng = np.random.default_rng(0)
+        data = pd.DataFrame(rng.normal(size=(96, 4)))
+        features, targets = data.drop(columns=0), data[0].to_numpy()
+        _, predictions = self._fit_predict(features, targets)
+        _, expected = self._fit_predict(features.rename(columns=str), targets)
+        np.testing.assert_allclose(predictions, expected, rtol=1e-5)
+
+    def test_reordered_integer_labels_keep_column_types(self):
+        rng = np.random.default_rng(0)
+        colours = np.array(["red", "green", "blue"])
+        codes = rng.integers(0, 3, 96)
+        features = pd.DataFrame({1: rng.normal(size=96), 0: colours[codes]})
+        targets = features[1].to_numpy() * 3 + codes
+        model, predictions = self._fit_predict(features, targets)
+        string_model, expected = self._fit_predict(features.rename(columns=str), targets)
+        assert model._data_module is not None
+        assert string_model._data_module is not None
+        categorical = model._data_module.cat_feature_info
+        string_categorical = string_model._data_module.cat_feature_info
+        assert categorical
+        assert string_categorical
+        assert [info["categories"] for info in categorical.values()] == [
+            info["categories"] for info in string_categorical.values()
+        ]
+        np.testing.assert_allclose(predictions, expected, rtol=1e-5)