Fix: Column mismatch during inference with column drift

energy_fault_detector/data_preprocessing/categorical_encoder.py

@@ -105,7 +105,7 @@ def transform(self, x: pd.DataFrame) -> pd.DataFrame:
         numerical_data = x[self.numerical_columns]
         categorical_data = x[self.categorical_columns]
 
-        if not categorical_data.empty:
+        if self.categorical_columns and len(categorical_data) > 0:
             for i, col in enumerate(self.categorical_columns):
                 fitted_cats = set(self.one_hot_encoder.categories_[i])
                 current_cats = set(categorical_data[col].dropna().unique())
@@ -121,6 +121,12 @@ def transform(self, x: pd.DataFrame) -> pd.DataFrame:
             self.feature_names_out_ = self.get_feature_names_out(self.feature_names_in_)
             transformed_data = pd.DataFrame(x_, columns=self.feature_names_out_, index=x.index)
             return transformed_data
+        elif self.categorical_columns:
+            # 0 rows but categorical columns exist — produce empty output with the correct
+            # one-hot encoded columns. OneHotEncoder.transform() raises on 0-row input,
+            # so the result is built manually to keep the column schema stable.
+            self.feature_names_out_ = self.get_feature_names_out(self.feature_names_in_)
+            return pd.DataFrame(columns=self.feature_names_out_, index=x.index)
         else:
             return numerical_data  # Returns input df if no categorical features are specified in config file
 

energy_fault_detector/data_preprocessing/ffill_imputer.py

@@ -6,6 +6,7 @@
 from sklearn.utils.validation import check_is_fitted
 
 from energy_fault_detector.core.data_transformer import DataTransformer
+from energy_fault_detector.data_preprocessing.imputer import Imputer
 
 logger = logging.getLogger('energy_fault_detector')
 
@@ -67,6 +68,7 @@ def __init__(self, ffill_limit: Union[str, pd.Timedelta] = "15min",
         self.numerical_columns: List[str] = []
         self.categorical_columns: List[str] = []
         self.non_declared_categorical_features: List[str] = []
+        self._fallback_values: dict = {}
 
     def fit(self, x: pd.DataFrame, y: Optional[pd.Series] = None) -> "ForwardFillImputer":
         """Fits the imputer to the input data by identifying feature types and storing metadata.
@@ -125,6 +127,21 @@ def fit(self, x: pd.DataFrame, y: Optional[pd.Series] = None) -> "ForwardFillImp
             self.numerical_columns = [col for col in self.numerical_columns if col not in all_nan_cols]
             self.categorical_columns = [col for col in self.categorical_columns if col not in all_nan_cols]
 
+        # Fit a fallback Imputer on the surviving columns so that, during transform, any
+        # column that becomes entirely NaN (e.g. a column missing from inference data) can be
+        # filled with its training mean (numerical) / most-frequent value (categorical) instead
+        # of causing dropna(how="any") to drop every row.
+        if self.numerical_columns or self.categorical_columns:
+            fallback = Imputer(categorical_features=self.categorical_columns)
+            cols = self.numerical_columns + self.categorical_columns
+            fallback.fit(x[cols])
+            for col, stat in zip(fallback.numerical_columns,
+                                 getattr(fallback.numerical_imputer, "statistics_", [])):
+                self._fallback_values[col] = stat
+            for col, stat in zip(fallback.categorical_columns,
+                                 getattr(fallback.categorical_imputer, "statistics_", [])):
+                self._fallback_values[col] = stat
+
         return self
 
     def transform(self, x: pd.DataFrame) -> pd.DataFrame:
@@ -198,6 +215,16 @@ def transform(self, x: pd.DataFrame) -> pd.DataFrame:
             filled[invalid] = np.nan
             df_filled[col] = filled
 
+        all_nan_after_ffill = [col for col in df_filled.columns if df_filled[col].isna().all()]
+        if all_nan_after_ffill:
+            logger.warning(
+                f"Columns entirely NaN after forward-fill during transform: {all_nan_after_ffill}. "
+                f"Filling with training mean/mode to prevent dropping all rows. "
+                f"These columns had valid data during fit."
+            )
+            for col in all_nan_after_ffill:
+                df_filled[col] = self._fallback_values.get(col, 0.0)
+
         df_cleaned = df_filled.drop_duplicates(keep="first")
         df_final = df_cleaned.dropna(how="any")
 

energy_fault_detector/data_preprocessing/imputer.py

@@ -127,9 +127,9 @@ def fit(self, x: pd.DataFrame, y=None) -> "Imputer":
         logger.debug(f"Categorical columns: {self.categorical_columns}")
 
         # Fit the imputers
-        if not numerical_data.empty:
+        if self.numerical_columns:
             self.numerical_imputer.fit(numerical_data)
-        if not categorical_data.empty:
+        if self.categorical_columns:
             self.categorical_imputer.fit(categorical_data)
 
         return self
@@ -173,20 +173,30 @@ def transform(self, x: pd.DataFrame) -> pd.DataFrame:
         categorical_data = x.loc[:, self.categorical_columns]
 
         # Transform the data
-        if not numerical_data.empty:
+        if self.numerical_columns and len(numerical_data) > 0:
             numerical_transformed = pd.DataFrame(
                 self.numerical_imputer.transform(numerical_data),
                 columns=self.numerical_columns,
                 index=x.index
             )
+        elif self.numerical_columns:
+            # 0 rows — SimpleImputer.transform() raises on 0-sample input, so build
+            # the output manually to keep the column schema stable.
+            numerical_transformed = pd.DataFrame(
+                columns=self.numerical_columns, index=x.index
+            )
         else:
             numerical_transformed = pd.DataFrame(index=x.index)
-        if not categorical_data.empty:
+        if self.categorical_columns and len(categorical_data) > 0:
             categorical_transformed = pd.DataFrame(
                 self.categorical_imputer.transform(categorical_data),
                 columns=self.categorical_columns,
                 index=x.index
             )
+        elif self.categorical_columns:
+            categorical_transformed = pd.DataFrame(
+                columns=self.categorical_columns, index=x.index
+            )
         else:
             categorical_transformed = pd.DataFrame(index=x.index)  # Empty DataFrame for consistency
 

tests/data_preprocessing/test_categorical_encoder.py

@@ -194,6 +194,29 @@ def test_empty_dataframe_handling(self):
         result = encoder.transform(empty_df)
 
         self.assertEqual(result.shape[0], 0)
+        # Regression for EFD 0.7.0 Bug 1: a 0-row DataFrame must still produce the
+        # full one-hot encoded column schema, not just the numerical columns.
+        self.assertListEqual(list(result.columns), encoder.get_feature_names_out())
+        for col in self.data_clean_encoded_features:
+            self.assertIn(col, result.columns)
+
+    def test_transform_zero_rows_preserves_one_hot_columns(self):
+        """Regression for EFD 0.7.0 Bug 1: a 0-row DataFrame must still yield the
+        full one-hot encoded column schema so downstream pipeline steps don't hit
+        a column-length mismatch (ValueError from sklearn _set_output wrapper)."""
+        encoder = CategoricalEncoder(categorical_features=['category', 'region'])
+        encoder.fit(self.data_clean)
+
+        empty_df = pd.DataFrame(columns=self.data_clean.columns, index=pd.DatetimeIndex([]))
+        result = encoder.transform(empty_df)
+
+        # Must contain the one-hot encoded columns, not just the 2 numerical ones.
+        self.assertEqual(result.shape[0], 0)
+        self.assertEqual(result.shape[1], len(encoder.get_feature_names_out()))
+        self.assertListEqual(list(result.columns), encoder.get_feature_names_out())
+        self.assertIn('category_A', result.columns)
+        self.assertIn('region_North', result.columns)
+        self.assertNotEqual(result.shape[1], len(encoder.numerical_columns))
 
     def test_transform_preserves_index_type(self):
         """Test that index type is preserved after transformation."""

tests/data_preprocessing/test_ffill_imputer.py

@@ -276,6 +276,59 @@ def test_all_nan_column_handling(self):
         self.assertListEqual(list(result.columns), ['value1'])
         self.assertEqual(result.shape[0], 3)
 
+    def test_transform_column_becomes_all_nan_keeps_rows(self):
+        """Regression for EFD 0.7.0 Bug 2: a tracked column that becomes entirely NaN
+        during transform (e.g. a column missing from inference data and filled with NaN)
+        must be filled with its training mean rather than causing dropna(how='any') to
+        drop every row. Categorical columns are filled with their training mode."""
+        df_train = pd.DataFrame(
+            {"temp": [20.0, 21.0, 22.0], "extra": [1.0, 2.0, 3.0]},
+            index=pd.date_range("2024-01-01", periods=3, freq="1h"),
+        )
+        imputer = ForwardFillImputer(ffill_limit="1h")
+        imputer.fit(df_train)
+
+        expected_extra_mean = df_train["extra"].mean()  # 2.0
+
+        # 'extra' is entirely NaN during inference (no valid value to forward-fill).
+        df_infer = pd.DataFrame(
+            {"temp": [20.0, 21.0, 22.0], "extra": [np.nan, np.nan, np.nan]},
+            index=pd.date_range("2024-02-01", periods=3, freq="1h"),
+        )
+        result = imputer.transform(df_infer)
+
+        # All rows must be retained; none should be dropped because 'extra' was all-NaN.
+        self.assertEqual(result.shape[0], 3)
+        self.assertListEqual(list(result.columns), ['temp', 'extra'])
+        # 'extra' should be filled with the training mean and hold no NaNs.
+        self.assertFalse(result['extra'].isna().any())
+        self.assertTrue((result['extra'] == expected_extra_mean).all())
+
+    def test_transform_categorical_column_becomes_all_nan_filled_with_mode(self):
+        """Regression for EFD 0.7.0 Bug 2 (categorical): a categorical column that becomes
+        entirely NaN during transform is filled with its training most-frequent value
+        via the package Imputer, keeping all rows."""
+        df_train = pd.DataFrame(
+            {"temp": [20.0, 21.0, 22.0],
+             "status": ["A", "A", "B"]},
+            index=pd.date_range("2024-01-01", periods=3, freq="1h"),
+        )
+        imputer = ForwardFillImputer(ffill_limit="1h", categorical_features=["status"])
+        imputer.fit(df_train)
+
+        # 'status' is entirely NaN during inference.
+        df_infer = pd.DataFrame(
+            {"temp": [20.0, 21.0, 22.0],
+             "status": [np.nan, np.nan, np.nan]},
+            index=pd.date_range("2024-02-01", periods=3, freq="1h"),
+        )
+        result = imputer.transform(df_infer)
+
+        self.assertEqual(result.shape[0], 3)
+        self.assertFalse(result['status'].isna().any())
+        # Most-frequent training value for 'status' is 'A'.
+        self.assertTrue((result['status'] == "A").all())
+
     def test_numerical_conversion_error(self):
         """Test error handling for non-convertible numerical columns."""
 

tests/data_preprocessing/test_imputer.py

@@ -325,5 +325,31 @@ def test_transform_with_only_categorical(self):
         self.assertEqual(result.shape, self.data_only_categorical.shape)
         self.assertFalse(result.isna().any().any())
 
+    def test_transform_zero_rows_preserves_columns(self):
+        """Regression for EFD 0.7.0 Bug 1 pattern: a 0-row DataFrame must still
+        produce the correct output column schema. The old 'not numerical_data.empty'
+        check skipped SimpleImputer.transform() on 0 rows, producing a 0-column
+        DataFrame that relied on a downstream reindex to recover the columns."""
+        imputer = Imputer(categorical_features=['status', 'region'])
+        imputer.fit(self.data_mixed)
+
+        empty_df = pd.DataFrame(columns=self.data_mixed.columns, index=pd.DatetimeIndex([]))
+        result = imputer.transform(empty_df)
+
+        self.assertEqual(result.shape[0], 0)
+        self.assertListEqual(list(result.columns), imputer.get_feature_names_out())
+
+    def test_transform_zero_rows_numerical_only_preserves_columns(self):
+        """Regression for EFD 0.7.0 Bug 1 pattern (numerical-only): 0-row input
+        must preserve the numerical column schema."""
+        imputer = Imputer()
+        imputer.fit(self.data_numerical)
+
+        empty_df = pd.DataFrame(columns=self.data_numerical.columns, index=pd.DatetimeIndex([]))
+        result = imputer.transform(empty_df)
+
+        self.assertEqual(result.shape[0], 0)
+        self.assertListEqual(list(result.columns), ['temperature', 'humidity', 'pressure'])
+
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()