fix(transformers): handle DataFrame input in ContinuousOrdinalTransformer

pretab/transformers/encoders/continuous_ordinal.py

@@ -2,6 +2,8 @@
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.utils.validation import check_is_fitted
 
+from ...core.exceptions import PretabDataError
+
 
 class ContinuousOrdinalTransformer(TransformerMixin, BaseEstimator):
     """Encode categorical features as continuous integer values.
@@ -45,16 +47,28 @@ def fit(self, X, y=None):
         self : object
             Fitted transformer.
         """
+        # Coerce to a 2D object array first. Iterating ``X.T`` directly yields
+        # the *column labels* of the transpose for a DataFrame -- i.e. the
+        # original row index -- which silently produced one mapping per row.
+        X = self._as_2d(X)
         # Fit should determine the mapping from original categories to sequential integers starting from 0
         self.mapping_ = [
-            {category: i + 1 for i, category in enumerate(np.unique(col))}
-            for col in X.T
+            {category: i + 1 for i, category in enumerate(np.unique(X[:, j]))}
+            for j in range(X.shape[1])
         ]
         for mapping in self.mapping_:
             mapping[None] = 0  # Assign 0 to unknown values
         self.n_features_in_ = len(self.mapping_)
         return self
 
+    @staticmethod
+    def _as_2d(X) -> np.ndarray:
+        """Return ``X`` as a 2D object ndarray, accepting frames, arrays and lists."""
+        array = np.asarray(X, dtype=object)
+        if array.ndim == 1:
+            array = array.reshape(-1, 1)
+        return array
+
     def transform(self, X):
         """Apply the learned category-to-integer mapping.
 
@@ -69,13 +83,19 @@ def transform(self, X):
             The transformed data with integer values.
         """
         check_is_fitted(self, "mapping_")
-        # Transform the categories to their mapped integer values
-        X_transformed = np.array(
-            [
-                [self.mapping_[col].get(value, 0) for col, value in enumerate(row)]
-                for row in X
-            ]
-        )
+        # As in ``fit``: iterating a DataFrame directly yields column *names*,
+        # not rows. Normalize first, then index by position.
+        X = self._as_2d(X)
+        if X.shape[1] != len(self.mapping_):
+            raise PretabDataError(
+                f"X has {X.shape[1]} features, but {type(self).__name__} "
+                f"is expecting {len(self.mapping_)} features as input."
+            )
+        # Allocating up front keeps the output 2D even for zero rows, where the
+        # comprehension used to collapse to shape ``(0,)``.
+        X_transformed = np.zeros(X.shape, dtype=int)
+        for col, mapping in enumerate(self.mapping_):
+            X_transformed[:, col] = [mapping.get(value, 0) for value in X[:, col]]
         return X_transformed
 
     def get_feature_names_out(self, input_features=None):

tests/test_continuous_ordinal_transformer.py

@@ -0,0 +1,76 @@
+"""``ContinuousOrdinalTransformer`` must treat a DataFrame like the equivalent array.
+
+``fit`` iterated ``X.T``, which for a DataFrame yields the *column labels of the
+transpose* -- i.e. the original row index -- and ``transform`` iterated ``X``,
+which yields column names. Frames therefore produced one mapping per row and an
+all-zero result of the wrong shape, with no error raised.
+"""
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from pretab.core.exceptions import PretabDataError
+from pretab.transformers import ContinuousOrdinalTransformer
+
+
+@pytest.fixture
+def frame():
+    return pd.DataFrame({"c1": ["a", "b", "a", "c"], "c2": ["x", "y", "x", "y"]})
+
+
+def test_fit_maps_one_dict_per_column(frame):
+    transformer = ContinuousOrdinalTransformer().fit(frame)
+
+    assert len(transformer.mapping_) == 2
+    assert transformer.n_features_in_ == 2
+    assert set(transformer.mapping_[0]) == {"a", "b", "c", None}
+    assert set(transformer.mapping_[1]) == {"x", "y", None}
+
+
+def test_transform_of_dataframe_keeps_shape_and_codes(frame):
+    out = ContinuousOrdinalTransformer().fit(frame).transform(frame)
+
+    assert out.shape == (4, 2)
+    np.testing.assert_array_equal(out[:, 0], [1, 2, 1, 3])
+    np.testing.assert_array_equal(out[:, 1], [1, 2, 1, 2])
+
+
+def test_dataframe_and_ndarray_agree(frame):
+    array = frame.to_numpy(dtype=object)
+
+    from_frame = ContinuousOrdinalTransformer().fit(frame).transform(frame)
+    from_array = ContinuousOrdinalTransformer().fit(array).transform(array)
+
+    np.testing.assert_array_equal(from_frame, from_array)
+
+
+def test_list_input_is_accepted():
+    rows = [["a", "x"], ["b", "y"], ["a", "x"]]
+    out = ContinuousOrdinalTransformer().fit(rows).transform(rows)
+
+    assert out.shape == (3, 2)
+
+
+def test_empty_transform_keeps_two_dimensions():
+    X = np.array([["a", "x"], ["b", "y"]], dtype=object)
+    transformer = ContinuousOrdinalTransformer().fit(X)
+
+    assert transformer.transform(np.empty((0, 2), dtype=object)).shape == (0, 2)
+
+
+def test_transform_rejects_wrong_feature_count():
+    X = np.array([["a", "x"], ["b", "y"]], dtype=object)
+    transformer = ContinuousOrdinalTransformer().fit(X)
+
+    with pytest.raises(PretabDataError, match="expecting 2 features"):
+        transformer.transform(np.array([["a"], ["b"]], dtype=object))
+
+
+def test_unknown_categories_still_map_to_zero():
+    X = np.array([["a"], ["b"]], dtype=object)
+    transformer = ContinuousOrdinalTransformer().fit(X)
+
+    np.testing.assert_array_equal(
+        transformer.transform(np.array([["a"], ["zzz"]], dtype=object)).ravel(), [1, 0]
+    )