Handle a None encoding in load_data when assertions are disabled

causalml/features.py

@@ -263,13 +263,19 @@ def load_data(data, features, transformations={}):
         logger.info("Applying one-hot-encoding to {}".format(cat_cols))
         ohe = OneHotEncoder(min_obs=df.shape[0] * 0.001)
         try:
-            X_cat = ohe.fit_transform(df[cat_cols]).toarray()
+            X_cat = ohe.fit_transform(df[cat_cols])
         except AssertionError:
             # OneHotEncoder asserts when no column produced a non-empty
-            # encoding. That happens when every categorical column holds a
-            # single value, or when all of their levels fall below min_obs,
-            # so there is nothing to append.
+            # encoding: every categorical column holds a single value, or all
+            # of their levels fall below min_obs. Under python -O that
+            # assertion is stripped and transform returns None instead of
+            # raising, so both outcomes end up here.
+            X_cat = None
+
+        if X_cat is None:
             logger.info("No categorical column produced an encoding")
+        else:
+            X_cat = X_cat.toarray()
 
     if X_cat is None:
         X = df[num_cols].values

tests/test_features.py

@@ -55,6 +55,20 @@ def test_load_data_categorical_column_that_encodes_to_nothing():
     np.testing.assert_array_equal(features, df[["num1"]].values)
 
 
+def test_load_data_when_transform_returns_none(monkeypatch):
+    # Under python -O the assertion inside OneHotEncoder.transform is stripped
+    # and it returns None rather than raising, so load_data has to handle the
+    # value as well as the exception.
+    monkeypatch.setattr(OneHotEncoder, "fit_transform", lambda self, X, y=None: None)
+    df = pd.DataFrame({"const_cat": ["x"] * 7, "num1": [1, 2, 1, 2, 1, 1, 1]})
+
+    features = load_data(df, ["const_cat", "num1"])
+
+    assert isinstance(features, np.ndarray) and not isinstance(features, np.matrix)
+    assert features.shape == (7, 1)
+    np.testing.assert_array_equal(features, df[["num1"]].values)
+
+
 def test_load_data_returns_ndarray_with_categorical_features(generate_categorical_data):
     df = generate_categorical_data()