add None test

jitingxu1 · Jun 24, 2024 · 3936fa2 · 3936fa2
1 parent 1e311d6
commit 3936fa2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/ibis_ml/steps/_select_features.py b/ibis_ml/steps/_select_features.py
@@ -67,9 +67,9 @@ def fit_table(self, table: ir.Table, metadata: Metadata) -> None:
                 col_var = results[f"{col_name}_var"]
                 if isinstance(table[col_name], ir.NumericColumn):
                     # Check variance for numeric columns
-                    if not col_var or math.isnan(col_var) or col_var < self.tolerance:
+                    if col_var is None or math.isnan(col_var) or col_var < self.tolerance:
                         cols.append(col_name)
-                elif not col_var or col_var < 2:
+                elif col_var is None or col_var < 2:
                     # Check unique count for non-numeric columns
                     cols.append(col_name)
                     metadata.drop_categories(col_name)

diff --git a/tests/test_select_features.py b/tests/test_select_features.py
@@ -7,7 +7,8 @@
 
 def test_drop_zero_variance():
     zv_numeric_col = [1.0] * 10
-    all_null_numeric_col = [np.nan] * 10
+    all_null_numeric_col = [None] * 10
+    all_nan_numeric_col = [np.nan] * 10
     non_zv_numeric_col = list(range(10))
     zv_string_col = ["String"] * 10
     non_zv_string_col = [f"String_{i}" for i in range(10)]
@@ -22,12 +23,14 @@ def test_drop_zero_variance():
         "zero_variance_string_col",
         "zero_variance_timestamp_col",
         "all_null_numeric_col",
+        "all_nan_numeric_col",
     }
 
     t_train = ibis.memtable(
         {
             "zero_variance_numeric_col": zv_numeric_col,
             "all_null_numeric_col": all_null_numeric_col,
+            "all_nan_numeric_col": all_nan_numeric_col,
             "non_zero_variance_numeric_col": non_zv_numeric_col,
             "zero_variance_string_col": zv_string_col,
             "non_zero_variance_string_col": non_zv_string_col,
@@ -39,6 +42,7 @@ def test_drop_zero_variance():
         {
             "zero_variance_numeric_col": [],
             "all_null_numeric_col": [],
+            "all_nan_numeric_col": [],
             "non_zero_variance_numeric_col": [],
             "zero_variance_string_col": [],
             "non_zero_variance_string_col": [],