Exposing conversion to Float32 in the Cleaner (#1440)

Author: rcap107

CHANGES.rst

@@ -83,6 +83,10 @@ Changes
 
 - The :class:`TableReport` now supports Series in addition to Dataframes. :pr:`1420` by :user:`Vitor Pohlenz<vitorpohlenz>`.
 
+- The :class:`Cleaner` now exposes a parameter to convert numerical values to float32. :pr:`1440` by
+  :user:`Riccardo Cappuzzo<rcap107>`.
+
+
 Bugfixes
 --------
 - Fixed a bug that caused the :class:`StringEncoder` and :class:`TextEncoder` to raise an exception if the

skrub/_table_vectorizer.py

@@ -182,6 +182,10 @@ class Cleaner(TransformerMixin, BaseEstimator):
     datetime_format : str, default=None
         The format to use when parsing dates. If None, the format is inferred.
 
+    numeric_dtype : "float32" or None, default=None
+        If set to ``float32``, convert numeric columns to ``np.float32`` dtype. If
+        ``None``, numerical dtypes are not modified.
+
     n_jobs : int, default=None
         Number of jobs to run in parallel.
         ``None`` means 1 unless in a joblib ``parallel_backend`` context.
@@ -225,11 +229,10 @@ class Cleaner(TransformerMixin, BaseEstimator):
     - ``ToStr()``: convert columns to strings, unless they are numerical,
       categorical, or datetime.
 
-    The ``Cleaner`` object should only be used for preliminary sanitizing of
-    the data because it does not perform any transformations on numeric columns.
-    On the other hand, the ``TableVectorizer`` converts numeric columns to ``float32``
-    and ensures that null values are represented with NaNs, which can be handled
-    correctly by downstream scikit-learn estimators.
+    If ``numeric_dtype`` is set to ``float32``, the ``Cleaner`` will also convert
+    numeric columns to ``np.float32`` dtype, ensuring a consistent representation
+    of numbers and missing values. This can be useful if the ``Cleaner``
+    is used as a preprocessing step in a skrub pipeline.
 
     Examples
     --------
@@ -288,6 +291,11 @@ class Cleaner(TransformerMixin, BaseEstimator):
     TableVectorizer :
         Process columns of a dataframe and convert them to a numeric (vectorized)
         representation.
+
+    ToFloat32 :
+        Convert numeric columns to ``np.float32``, to have consistent numeric
+        types and representation of missing values. More informative columns (e.g.,
+        categorical or datetime) are not converted.
     """
 
     def __init__(
@@ -296,12 +304,14 @@ def __init__(
         drop_if_constant=False,
         drop_if_unique=False,
         datetime_format=None,
+        numeric_dtype=None,
         n_jobs=1,
     ):
         self.drop_null_fraction = drop_null_fraction
         self.drop_if_constant = drop_if_constant
         self.drop_if_unique = drop_if_unique
         self.datetime_format = datetime_format
+        self.numeric_dtype = numeric_dtype
         self.n_jobs = n_jobs
 
     def fit_transform(self, X, y=None):
@@ -322,13 +332,21 @@ def fit_transform(self, X, y=None):
         dataframe
             The transformed input.
         """
+
+        add_tofloat32 = self.numeric_dtype == "float32"
+        if self.numeric_dtype not in (None, "float32"):
+            raise ValueError(
+                "`numeric_dtype` must be one of"
+                f"[`None`, `'float32'`]. Found {self.numeric_dtype}."
+            )
+
         all_steps = _get_preprocessors(
             cols=s.all(),
             drop_null_fraction=self.drop_null_fraction,
             drop_if_constant=self.drop_if_constant,
             drop_if_unique=self.drop_if_unique,
             n_jobs=self.n_jobs,
-            add_tofloat32=False,
+            add_tofloat32=add_tofloat32,
             datetime_format=self.datetime_format,
         )
         self._pipeline = make_pipeline(*all_steps)
@@ -539,7 +557,7 @@ class TableVectorizer(TransformerMixin, BaseEstimator):
 
     Cleaner :
         Preprocesses each column of a dataframe with consistency checks and
-        sanitization, eg of null values or dates.
+        sanitization, e.g., of null values or dates.
 
     Examples
     --------

skrub/tests/test_table_vectorizer.py

@@ -389,21 +389,35 @@ def test_cleaner_dtypes(X, dict_expected_types):
                     assert dict_expected_types[col] == X_trans[col].dtype
 
 
-def test_convert_float32():
+def test_convert_float32(df_module):
     """
     Test that the TableVectorizer converts float64 to float32
     when using the default parameters.
     """
     X = _get_clean_dataframe()
     vectorizer = TableVectorizer()
     out = vectorizer.fit_transform(X)
-    assert out.dtypes["float"] == "float32"
-    assert out.dtypes["int"] == "float32"
+    assert sbd.dtype(out["float"]) == "float32"
+    assert sbd.dtype(out["int"]) == "float32"
 
+    # default behavior: keep numeric type
     vectorizer = Cleaner()
     out = vectorizer.fit_transform(X)
-    assert sbd.is_float(out["float"])
-    assert sbd.is_integer(out["int"])
+    assert sbd.dtype(out["float"]) == sbd.dtype(X["float"])
+    assert sbd.dtype(out["int"]) == sbd.dtype(X["int"])
+    assert sbd.dtype(out["float"]) == sbd.dtype(X["float"])
+    assert sbd.dtype(out["int"]) == sbd.dtype(X["int"])
+
+    vectorizer = Cleaner(numeric_dtype="float32")
+    out = vectorizer.fit_transform(X)
+    assert out.dtypes["float"] == "float32"
+    assert out.dtypes["int"] == "float32"
+
+
+def test_cleaner_invalid_numeric_dtype():
+    X = _get_clean_dataframe()
+    with pytest.raises(ValueError, match="numeric_dtype.*must be one of"):
+        Cleaner(numeric_dtype="wrong").fit_transform(X)
 
 
 def test_auto_cast_missing_categories():