Added categorical encoding function

pyproject.toml

@@ -42,6 +42,8 @@ dependencies = [
   "pytz>=2023.3",
   "pyyaml>=6.0",
   "tomli>=2.0.1",
+  "pandas>=1.5.0",
+  "scikit-learn>=1.2.0"
 ]
 
 [project.urls]

src/koheesio/pandas/categorical_encoding.py

@@ -0,0 +1,57 @@
+from sklearn.preprocessing import OneHotEncoder
+import pandas as pd
+from pydantic import BaseModel
+from typing import List, Optional
+from koheesio.steps import Step
+
+from typing import List, Dict
+import pandas as pd
+from pydantic import BaseModel
+
+class PandasCategoricalEncoding(BaseModel):
+    """
+    Encodes categorical columns using one-hot encoding or ordinal encoding.
+
+    Attributes
+    ----------
+    columns : List[str]
+        List of categorical columns to encode.
+    encoding_type : str, optional
+        Type of encoding, either "one-hot" or "ordinal" (default is "one-hot").
+    drop_first : bool, optional
+        Whether to drop the first dummy column for one-hot encoding (default is True).
+    ordinal_mapping : dict, optional
+        A dictionary mapping categorical values to integers for ordinal encoding.
+    """
+
+    columns: List[str]
+    encoding_type: str = "one-hot"
+    drop_first: bool = True
+    ordinal_mapping: Dict[str, Dict] = None
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        if self.encoding_type not in ['one-hot', 'ordinal']:
+            raise ValueError(f"Invalid encoding type: {self.encoding_type}. Expected 'one-hot' or 'ordinal'.")
+
+    def execute(self, data: pd.DataFrame) -> pd.DataFrame:
+        """
+        Executes the categorical encoding transformation on the provided dataset.
+
+        Parameters
+        ----------
+        data : pd.DataFrame
+            The input dataset to encode.
+
+        Returns
+        -------
+        pd.DataFrame
+            The dataset with the specified categorical columns encoded.
+        """
+        if self.encoding_type == 'one-hot':
+            data = pd.get_dummies(data, columns=self.columns, drop_first=self.drop_first)
+        elif self.encoding_type == 'ordinal':
+            for column in self.columns:
+                if column in data.columns and self.ordinal_mapping and column in self.ordinal_mapping:
+                    data[column] = data[column].map(self.ordinal_mapping[column]).fillna(-1).astype(int)
+        return data

tests/test_categorical_encoding.py

@@ -0,0 +1,81 @@
+import unittest
+import pandas as pd
+from src.koheesio.pandas.categorical_encoding  import PandasCategoricalEncoding
+import unittest
+import pandas as pd
+
+class TestPandasCategoricalEncoding(unittest.TestCase):
+
+    """
+    Unit tests for the PandasCategoricalEncoding class.
+
+    Methods
+    -------
+    setUp():
+        Sets up the test environment with sample data.
+    test_one_hot_encoding():
+        Tests the one-hot encoding functionality.
+    test_ordinal_encoding():
+        Tests the ordinal encoding functionality.
+    """
+
+    def setUp(self):
+        """
+        Sets up the test environment with sample data.
+
+        Creates a pandas DataFrame with categorical columns for testing.
+        """
+        self.data = pd.DataFrame({
+            'color': ['red', 'blue', 'green', 'blue', 'red'],
+            'size': ['S', 'M', 'L', 'S', 'XL']
+        })
+
+    def test_one_hot_encoding(self):
+        """
+        Tests one-hot encoding functionality.
+
+        Ensures that the specified categorical columns are correctly 
+        encoded into dummy variables, with an option to drop the first 
+        dummy column.
+        """
+        encoding_step = PandasCategoricalEncoding(
+            columns=["color"],
+            encoding_type="one-hot",
+            drop_first=False  # Adjusted to match expected columns
+        )
+
+        # Apply encoding
+        encoded_data = encoding_step.execute(self.data)
+
+        # Expected columns after one-hot encoding
+        expected_columns = ['color_blue', 'color_green', 'color_red']
+
+        # Check if the encoded data contains expected columns
+        self.assertTrue(all(col in encoded_data.columns for col in expected_columns))
+        print("One-hot encoding passed.")
+
+    def test_ordinal_encoding(self):
+        """
+        Tests ordinal encoding functionality.
+
+        Ensures that the specified categorical columns are correctly 
+        encoded into ordinal values based on a provided mapping.
+        """
+
+        ordinal_mapping = {
+            'size': {'S': 1, 'M': 2, 'L': 3, 'XL': 4}
+        }
+
+        encoding_step = PandasCategoricalEncoding(
+            columns=["size"], 
+            encoding_type="ordinal", 
+            ordinal_mapping=ordinal_mapping
+        )
+
+        # Apply encoding
+        encoded_data = encoding_step.execute(self.data)
+
+        # Check if the "size" column is correctly ordinal encoded
+        expected_values = [1, 2, 3, 1, 4]
+        self.assertTrue(all(encoded_data['size'] == expected_values))
+        print("Ordinal encoding passed.")