Added categorical encoding function

src/koheesio/categorical_encoding.py

@@ -0,0 +1,30 @@
+# Import necessary libraries
+from sklearn.preprocessing import OneHotEncoder
+import pandas as pd
+from pydantic import BaseModel
+
+# Define the EncodingConfig class
+class EncodingConfig(BaseModel):
+    drop_first: bool = True  # Whether to drop the first dummy column
+
+# Define the categorical_encoding function
+def categorical_encoding(data, columns, config: EncodingConfig):
+    """
+    Encodes categorical columns using one-hot encoding.
+
+    Args:
+        data (pd.DataFrame): Input dataset.
+        columns (list): List of categorical columns to encode.
+        config (EncodingConfig): Configuration for encoding.
+
+    Returns:
+        pd.DataFrame: Dataset with one-hot encoded columns.
+    """
+    encoder = OneHotEncoder(sparse_output=False, drop='first' if config.drop_first else None)
+    encoded_array = encoder.fit_transform(data[columns])
+    encoded_df = pd.DataFrame(
+        encoded_array,
+        columns=encoder.get_feature_names_out(columns),
+        index=data.index,
+    )
+    return pd.concat([data.drop(columns, axis=1), encoded_df], axis=1)

tests/test_categorical_encoding.py

@@ -0,0 +1,42 @@
+import unittest
+import pandas as pd
+from src.koheesio.categorical_encoding import categorical_encoding, EncodingConfig
+
+class TestCategoricalEncoding(unittest.TestCase):
+    def setUp(self):
+        self.data = pd.DataFrame({
+            'color': ['red', 'blue', 'green', 'red'],
+            'size': ['S', 'M', 'L', 'XL'],
+            'price': [10.0, 20.0, 30.0, 40.0]
+        })
+
+    def test_one_hot_encoding(self):
+        # Configure encoding to drop the first column (avoid dummy variable trap)
+        config = EncodingConfig(drop_first=True)
+        result = categorical_encoding(self.data, ['color', 'size'], config)
+        # Check if the encoded columns are present
+        self.assertIn('color_green', result.columns)
+        self.assertIn('color_red', result.columns)
+        self.assertIn('size_M', result.columns)
+        self.assertIn('size_XL', result.columns)
+        # Check if the original categorical columns were removed
+        self.assertNotIn('color', result.columns)
+        self.assertNotIn('size', result.columns)
+
+    def test_no_drop_first_encoding(self):
+        # Configure encoding to retain all dummy columns
+        config = EncodingConfig(drop_first=False)
+        result = categorical_encoding(self.data, ['color', 'size'], config)
+        # Check for all encoded columns
+        self.assertIn('color_blue', result.columns)
+        self.assertIn('color_green', result.columns)
+        self.assertIn('color_red', result.columns)
+        self.assertIn('size_S', result.columns)
+        self.assertIn('size_M', result.columns)
+        self.assertIn('size_XL', result.columns)
+        # Check if the original categorical columns were removed
+        self.assertNotIn('color', result.columns)
+        self.assertNotIn('size', result.columns)
+
+if __name__ == '__main__':
+    unittest.main()