Added DropNullColumn transformer to remove columns that contain only nulls

CHANGES.rst

@@ -60,6 +60,10 @@ Minor changes
   contained datetimes with time zones and missing values; this has been fixed in
   :pr:`1114` by :user:`Jérôme Dockès <jeromedockes>`.
 
+* Added a `DropNullColumn` transformer that drops columns that contain only null
+  values. :pr:`1115` by :user: `Riccardo Cappuzzo <riccardocappuzzo>`
+
+
 Release 0.3.1
 =============
 

skrub/_dataframe/_common.py

@@ -73,6 +73,7 @@
     "to_datetime",
     "is_categorical",
     "to_categorical",
+    "is_all_null",
     #
     # Inspecting, selecting and modifying values
     #
@@ -825,6 +826,27 @@
     return _cast_polars(col, pl.Categorical())
 
 
+@dispatch
+def is_all_null(col):
+    raise NotImplementedError()
+
+
+@is_all_null.specialize("pandas", argument_type="Column")
+def _is_all_null_pandas(col):
+    return bool(col.isna().all())
+
+
+@is_all_null.specialize("polars", argument_type="Column")
+def _is_all_null_polars(col):
+    if col.dtype == pl.Null:
+        return True
+    elif col.dtype.is_numeric() and col.is_nan().all():
+        return True
+    # col is non numeric
+    elif col.null_count() == col.len():
+        return True
+
+
 #
 # Inspecting, selecting and modifying values
 # ==========================================

skrub/_drop_null.py

@@ -0,0 +1,54 @@
+# drop columns that contain all null values
+from sklearn.utils.validation import check_is_fitted
+
+from . import _dataframe as sbd
+from ._on_each_column import SingleColumnTransformer
+
+__all__ = ["DropNullColumn"]
+
+
+class DropNullColumn(SingleColumnTransformer):
+    """Drop a single column if it contains only null values."""
+
+    def __init__(self):
+        super().__init__()
+        self._is_fitted = False
+
+    def fit_transform(self, column, y=None):
+        """Fit the encoder and transform a column.
+
+        Args:
+            column : Pandas or Polars series. The input column to check.
+            y : None. Ignored.
+
+        Returns:
+            The input column, or an empty list if the column contains only null values.
+        """
+        del y
+
+        self._is_fitted = True
+        return self.transform(column)
+
+    def transform(self, column):
+        """Transform a column.
+
+        Args:
+            column : Pandas or Polars series. The input column to check.
+
+        Returns:
+            The input column, or an empty list if the column contains only null values.
+        """
+        check_is_fitted(
+            self,
+        )
+
+        if sbd.is_all_null(column):
+            return []
+        else:
+            return column
+
+    def __sklearn_is_fitted__(self):
+        """
+        Check fitted status and return a Boolean value.
+        """
+        return hasattr(self, "_is_fitted") and self._is_fitted

skrub/_table_vectorizer.py

@@ -16,6 +16,7 @@
 from ._clean_categories import CleanCategories
 from ._clean_null_strings import CleanNullStrings
 from ._datetime_encoder import DatetimeEncoder
+from ._drop_null import DropNullColumn
 from ._gap_encoder import GapEncoder
 from ._on_each_column import SingleColumnTransformer
 from ._select_cols import Drop
@@ -191,6 +192,9 @@ class TableVectorizer(TransformerMixin, BaseEstimator):
         similar functionality to what is offered by scikit-learn's
         :class:`~sklearn.compose.ColumnTransformer`.
 
+    drop_null_columns : bool, default=False
+        If set to `True`, columns that contain only null values are dropped.
+
     n_jobs : int, default=None
         Number of jobs to run in parallel.
         ``None`` means 1 unless in a joblib ``parallel_backend`` context.
@@ -412,6 +416,7 @@ def __init__(
         numeric=NUMERIC_TRANSFORMER,
         datetime=DATETIME_TRANSFORMER,
         specific_transformers=(),
+        drop_null_columns=False,
         n_jobs=None,
     ):
         self.cardinality_threshold = cardinality_threshold
@@ -425,6 +430,7 @@ def __init__(
         self.datetime = _utils.clone_if_default(datetime, DATETIME_TRANSFORMER)
         self.specific_transformers = specific_transformers
         self.n_jobs = n_jobs
+        self.drop_null_columns = drop_null_columns
 
     def fit(self, X, y=None):
         """Fit transformer.
@@ -536,6 +542,9 @@ def add_step(steps, transformer, cols, allow_reject=False):
         cols = s.all() - self._specific_columns
 
         self._preprocessors = [CheckInputDataFrame()]
+        if self.drop_null_columns:
+            add_step(self._preprocessors, DropNullColumn(), cols, allow_reject=True)
+
         for transformer in [
             CleanNullStrings(),
             ToDatetime(),

skrub/tests/test_drop_nulls.py

@@ -0,0 +1,86 @@
+import numpy as np
+import pytest
+
+from skrub import TableVectorizer
+from skrub import _dataframe as sbd
+from skrub._drop_null import DropNullColumn
+
+
+@pytest.fixture
+def drop_null_table(df_module):
+    return df_module.make_dataframe(
+        {
+            "idx": [
+                1,
+                2,
+                3,
+            ],
+            "value_nan": [
+                np.nan,
+                np.nan,
+                np.nan,
+            ],
+            "value_null": [
+                None,
+                None,
+                None,
+            ],
+            "value_almost_nan": [
+                2.5,
+                np.nan,
+                np.nan,
+            ],
+            "value_almost_null": [
+                "almost",
+                None,
+                None,
+            ],
+        }
+    )
+
+
+def test_single_column(drop_null_table, df_module):
+    """Check that null columns are dropped and non-null columns are kept."""
+    dn = DropNullColumn()
+    assert dn.fit_transform(drop_null_table["value_nan"]) == []
+    assert dn.fit_transform(drop_null_table["value_null"]) == []
+
+    df_module.assert_column_equal(
+        sbd.col(drop_null_table, "idx"), df_module.make_column("idx", [1, 2, 3])
+    )
+
+    df_module.assert_column_equal(
+        sbd.col(drop_null_table, "value_almost_nan"),
+        df_module.make_column("value_almost_nan", [2.5, np.nan, np.nan]),
+    )
+
+    df_module.assert_column_equal(
+        sbd.col(drop_null_table, "value_almost_null"),
+        df_module.make_column("value_almost_null", ["almost", None, None]),
+    )
+
+
+def test_drop_null_column(drop_null_table):
+    """Check that all null columns are dropped, and no more."""
+    # Don't drop null columns
+    tv = TableVectorizer(drop_null_columns=False)
+    transformed = tv.fit_transform(drop_null_table)
+
+    assert sbd.shape(transformed) == sbd.shape(drop_null_table)
+
+    # Drop null columns
+    tv = TableVectorizer(drop_null_columns=True)
+    transformed = tv.fit_transform(drop_null_table)
+    assert sbd.shape(transformed) == (sbd.shape(drop_null_table)[0], 3)
+
+
+def test_is_all_null(drop_null_table):
+    """Check that is_all_null is evaluating null counts correctly."""
+    # Check that all null columns are marked as "all null"
+    assert sbd.is_all_null(drop_null_table["value_nan"])
+    assert sbd.is_all_null(drop_null_table["value_null"])
+
+    # Check that the other columns are *not* marked as "all null"
+    assert not sbd.is_all_null(drop_null_table["value_almost_null"])
+    assert not sbd.is_all_null(drop_null_table["value_almost_nan"])
+    assert not sbd.is_all_null(drop_null_table["idx"])