Added DropNullColumn transformer to remove columns that contain only nulls

skrub/tests/test_drop_column_if_null.py

@@ -3,6 +3,7 @@
 
 from skrub import _dataframe as sbd
 from skrub._drop_column_if_null import DropColumnIfNull
+from skrub._on_each_column import RejectColumn
 
 
 @pytest.fixture
@@ -39,9 +40,9 @@ def drop_null_table(df_module):
     )
 
 
-def test_single_column(drop_null_table, df_module):
+def test_single_column_drop(drop_null_table, df_module):
     """Check that null columns are dropped and non-null columns are kept."""
-    dn = DropColumnIfNull()
+    dn = DropColumnIfNull(null_column_strategy="drop")
     assert dn.fit_transform(sbd.col(drop_null_table, "value_nan")) == []
     assert dn.fit_transform(sbd.col(drop_null_table, "value_null")) == []
     assert dn.fit_transform(sbd.col(drop_null_table, "mixed_null")) == []
@@ -60,3 +61,68 @@ def test_single_column(drop_null_table, df_module):
         dn.fit_transform(sbd.col(drop_null_table, "value_almost_null")),
         df_module.make_column("value_almost_null", ["almost", None, None]),
     )
+
+
+def test_single_column_keep(drop_null_table, df_module):
+    """Check that all columns are kept."""
+    dn = DropColumnIfNull(null_column_strategy="keep")
+
+    df_module.assert_column_equal(
+        dn.fit_transform(sbd.col(drop_null_table, "idx")),
+        df_module.make_column("idx", [1, 2, 3]),
+    )
+
+    df_module.assert_column_equal(
+        dn.fit_transform(sbd.col(drop_null_table, "value_null")),
+        df_module.make_column(
+            "value_null",
+            [
+                None,
+                None,
+                None,
+            ],
+        ),
+    )
+
+    df_module.assert_column_equal(
+        dn.fit_transform(sbd.col(drop_null_table, "value_nan")),
+        df_module.make_column(
+            "value_nan",
+            [
+                np.nan,
+                np.nan,
+                np.nan,
+            ],
+        ),
+    )
+
+    df_module.assert_column_equal(
+        dn.fit_transform(sbd.col(drop_null_table, "mixed_null")),
+        df_module.make_column("mixed_null", [None, np.nan, None]),
+    )
+
+    df_module.assert_column_equal(
+        dn.fit_transform(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(
+        dn.fit_transform(sbd.col(drop_null_table, "value_almost_null")),
+        df_module.make_column("value_almost_null", ["almost", None, None]),
+    )
+
+
+def test_single_column_raise(drop_null_table, df_module):
+    """Check that an exception is raised if a null column is detected."""
+    dn = DropColumnIfNull(null_column_strategy="raise")
+    with pytest.raises(RejectColumn):
+        dn.fit_transform(sbd.col(drop_null_table, "value_nan"))
+    with pytest.raises(RejectColumn):
+        dn.fit_transform(sbd.col(drop_null_table, "value_null"))
+    with pytest.raises(RejectColumn):
+        dn.fit_transform(sbd.col(drop_null_table, "mixed_null"))
+
+
+def test_incorrect_argument():
+    with pytest.raises(ValueError):
+        DropColumnIfNull(null_column_strategy="wrong value")

skrub/tests/test_table_vectorizer.py

@@ -19,7 +19,6 @@
 from skrub._datetime_encoder import DatetimeEncoder
 from skrub._gap_encoder import GapEncoder
 from skrub._minhash_encoder import MinHashEncoder
-from skrub._on_each_column import RejectColumn
 from skrub._table_vectorizer import TableVectorizer
 
 MSG_PANDAS_DEPRECATED_WARNING = "Skip deprecation warning"
@@ -531,7 +530,7 @@ def test_changing_types(X_train, X_test, expected_X_out):
         # only extract the total seconds
         datetime=DatetimeEncoder(resolution=None),
         # True by default
-        null_column_strategy=False,
+        null_column_strategy="keep",
     )
 
     table_vec.fit(X_train)
@@ -766,7 +765,7 @@ def test_drop_null_column():
     """Check that all null columns are dropped, and no more."""
     # Don't drop null columns
     X = _get_missing_values_dataframe()
-    tv = TableVectorizer(null_column_strategy="ignore")
+    tv = TableVectorizer(null_column_strategy="keep")
     transformed = tv.fit_transform(X)
 
     assert sbd.shape(transformed) == sbd.shape(X)
@@ -778,11 +777,7 @@ def test_drop_null_column():
 
     # Raise exception if a null column is found
     with pytest.raises(
-        RejectColumn, match="Column all_null contains only null values."
+        ValueError,
     ):
         tv = TableVectorizer(null_column_strategy="raise")
         transformed = tv.fit_transform(X)
-
-    # # Raise an exception if an unknown parameter is found
-    # tv = TableVectorizer(null_column_strategy="wrong_parameter")
-    # transformed = tv.fit_transform(X)