add integration_test + _get_output_to_property

rdt/transformers/base.py

@@ -362,7 +362,7 @@ def _fit(self, columns_data):
         raise NotImplementedError()
 
     def _set_seed(self, data):
-        hash_value = self.get_input_column()
+        hash_value = self.columns[0]
         for value in data.head(5):
             hash_value += str(value)
 
@@ -508,6 +508,14 @@ def get_input_columns(self):
         """
         return self.columns
 
+    def _get_output_to_property(self, property_):
+        result = {
+            output_column: properties[property_]
+            for output_column, properties in self.output_properties.items()
+        }
+
+        return result
+
     def _fit(self, columns_data, columns_to_sdtypes):
         """Fit the transformer to the data.
 
@@ -519,6 +527,7 @@ def _fit(self, columns_data, columns_to_sdtypes):
         """
         raise NotImplementedError()
 
+    @random_state
     def fit(self, data, columns_to_sdtypes):
         """Fit the transformer to a ``column`` of the ``data``.
 
@@ -528,9 +537,25 @@ def fit(self, data, columns_to_sdtypes):
             columns_to_sdtypes (dict):
                 Dictionary mapping column names to their sdtypes.
         """
-        column_names = list(columns_to_sdtypes.keys())
+        column_names = tuple(columns_to_sdtypes.keys())
         self._store_columns(column_names, data)
         self._set_seed(data)
         columns_data = self._get_columns_data(data, self.columns)
         self._fit(columns_data, columns_to_sdtypes)
         self._build_output_columns(data)
+
+    def fit_transform(self, data, columns_to_sdtypes):
+        """Fit the transformer to a `column` of the `data` and then transform it.
+
+        Args:
+            data (pandas.DataFrame):
+                The entire table.
+            columns_to_sdtypes (dict):
+                Dictionary mapping column names to their sdtypes.
+
+        Returns:
+            pd.DataFrame:
+                The entire table, containing the transformed data.
+        """
+        self.fit(data, columns_to_sdtypes)
+        return self.transform(data)

tests/integration/transformers/test_base.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pandas as pd
 
-from rdt.transformers.base import BaseTransformer
+from rdt.transformers.base import BaseMultiColumnTransformer, BaseTransformer
 
 
 def test_dummy_transformer_series_output():
@@ -129,3 +129,172 @@ def _reverse_transform(self, data):
     })
     pd.testing.assert_frame_equal(expected_transform, transformed)
     pd.testing.assert_frame_equal(reverse, data)
+
+
+def test_multi_column_transformer_same_number_of_columns_input_output():
+    """Test a multi-column transformer when the same of input and output columns."""
+    # Setup
+    class AdditionTransformer(BaseMultiColumnTransformer):
+        """This transformer takes 3 columns and return the cumulative sum of each row."""
+        def _fit(self, columns_data, columns_to_sdtypes):
+            self.output_properties = {
+                column: {'sdtype': 'numerical'} for column in self.columns
+            }
+            self.dtypes = columns_data.dtypes
+
+        def _transform(self, data):
+            return data.cumsum(axis=1)
+
+        def _reverse_transform(self, data):
+            result = data.diff(axis=1)
+            result.iloc[:, 0] = data.iloc[:, 0]
+
+            return result.astype(self.dtypes)
+
+    data_test = pd.DataFrame({
+        'col_1': [1, 2, 3],
+        'col_2': [10, 20, 30],
+        'col_3': [100, 200, 300]
+    })
+
+    column_to_sdtype = {
+        'col_1': 'numerical',
+        'col_2': 'numerical',
+        'col_3': 'numerical'
+    }
+    transformer = AdditionTransformer()
+
+    # Run
+    transformed = transformer.fit_transform(data_test, column_to_sdtype)
+    reverse = transformer.reverse_transform(transformed)
+
+    # Assert
+    expected_transform = pd.DataFrame({
+        'col_1': [1, 2, 3],
+        'col_2': [11, 22, 33],
+        'col_3': [111, 222, 333]
+    })
+    pd.testing.assert_frame_equal(expected_transform, transformed)
+    pd.testing.assert_frame_equal(reverse, data_test)
+
+
+def test_multi_column_transformer_less_output_than_input_columns():
+    """Test a multi-column transformer when the output has less columns than the input."""
+    class ConcatenateTransformer(BaseMultiColumnTransformer):
+        """This transformer takes 4 columns and concatenate them into 2 columns.
+        The two first and last columns are concatenated together.
+        """
+        def _fit(self, columns_data, columns_to_sdtypes):
+            self.name_1 = self.columns[0] + '#' + self.columns[1]
+            self.name_2 = self.columns[2] + '#' + self.columns[3]
+            self.output_properties = {
+                self.name_1: {'sdtype': 'categorical'},
+                self.name_2: {'sdtype': 'categorical'}
+            }
+            self.dtypes = columns_data.dtypes
+
+        def _transform(self, data):
+            data[self.name_1] = data.iloc[:, 0] + '#' + data.iloc[:, 1]
+            data[self.name_2] = data.iloc[:, 2] + '#' + data.iloc[:, 3]
+
+            return data.drop(columns=self.columns)
+
+        def _reverse_transform(self, data):
+            result = data.copy()
+            column_names = list(data.columns)
+
+            col1, col2 = column_names[0].split('#')
+            result[[col1, col2]] = result[column_names[0]].str.split('#', expand=True)
+
+            col3, col4 = column_names[1].split('#')
+            result[[col3, col4]] = result[column_names[1]].str.split('#', expand=True)
+
+            return result.astype(self.dtypes).drop(columns=column_names)
+
+    data_test = pd.DataFrame({
+        'col_1': ['A', 'B', 'C'],
+        'col_2': ['D', 'E', 'F'],
+        'col_3': ['G', 'H', 'I'],
+        'col_4': ['J', 'K', 'L']
+    })
+
+    column_to_sdtype = {
+        'col_1': 'categorical',
+        'col_2': 'categorical',
+        'col_3': 'categorical',
+        'col_4': 'categorical'
+    }
+    transformer = ConcatenateTransformer()
+
+    # Run
+    transformer.fit(data_test, column_to_sdtype)
+    transformed = transformer.transform(data_test)
+    reverse = transformer.reverse_transform(transformed)
+
+    # Assert
+    expected_transform = pd.DataFrame({
+        'col_1#col_2': ['A#D', 'B#E', 'C#F'],
+        'col_3#col_4': ['G#J', 'H#K', 'I#L']
+    })
+    pd.testing.assert_frame_equal(expected_transform, transformed)
+    pd.testing.assert_frame_equal(reverse, data_test)
+
+
+def test_multi_column_transformer_more_output_than_input_columns():
+    """Test a multi-column transformer when the output has more columns than the input."""
+    class ExpandTransformer(BaseMultiColumnTransformer):
+
+        def _fit(self, columns_data, columns_to_sdtypes):
+            name_1 = self.columns[0] + '.first_part'
+            name_2 = self.columns[0] + '.second_part'
+            name_3 = self.columns[1] + '.first_part'
+            name_4 = self.columns[1] + '.second_part'
+            self.output_properties = {
+                name_1: {'sdtype': 'categorical'},
+                name_2: {'sdtype': 'categorical'},
+                name_3: {'sdtype': 'categorical'},
+                name_4: {'sdtype': 'categorical'}
+            }
+            self.names = [name_1, name_2, name_3, name_4]
+            self.dtypes = columns_data.dtypes
+
+        def _transform(self, data):
+            data[self.names[0]] = data[self.columns[0]].str[0]
+            data[self.names[1]] = data[self.columns[0]].str[1]
+            data[self.names[2]] = data[self.columns[1]].str[0]
+            data[self.names[3]] = data[self.columns[1]].str[1]
+
+            return data.drop(columns=self.columns)
+
+        def _reverse_transform(self, data):
+            result = data.copy()
+            result[self.columns[0]] = result[self.names[0]] + result[self.names[1]]
+            result[self.columns[1]] = result[self.names[2]] + result[self.names[3]]
+
+            return result.astype(self.dtypes).drop(columns=self.names)
+
+    data_test = pd.DataFrame({
+        'col_1': ['AB', 'CD', 'EF'],
+        'col_2': ['GH', 'IJ', 'KL'],
+    })
+
+    column_to_sdtype = {
+        'col_1': 'categorical',
+        'col_2': 'categorical',
+    }
+    transformer = ExpandTransformer()
+
+    # Run
+    transformer.fit(data_test, column_to_sdtype)
+    transformed = transformer.transform(data_test)
+    reverse = transformer.reverse_transform(transformed)
+
+    # Assert
+    expected_transform = pd.DataFrame({
+        'col_1.first_part': ['A', 'C', 'E'],
+        'col_1.second_part': ['B', 'D', 'F'],
+        'col_2.first_part': ['G', 'I', 'K'],
+        'col_2.second_part': ['H', 'J', 'L']
+    })
+    pd.testing.assert_frame_equal(expected_transform, transformed)
+    pd.testing.assert_frame_equal(reverse, data_test)

tests/unit/transformers/test_base.py

@@ -1303,6 +1303,27 @@ def test_get_input_columns(self):
         # Assert
         assert output == ['a', 'b', 'c']
 
+    def test__get_output_to_property(self):
+        """Test the ``_get_output_to_property`` method."""
+        # Setup
+        transformer = BaseMultiColumnTransformer()
+        transformer.output_properties = {
+            'col_1': {'sdtype': 'numerical'},
+            'col_2': {'sdtype': 'categorical'},
+            'col_3': {'sdtype': 'boolean'},
+        }
+
+        # Run
+        output = transformer._get_output_to_property('sdtype')
+
+        # Assert
+        expected_output = {
+            'col_1': 'numerical',
+            'col_2': 'categorical',
+            'col_3': 'boolean',
+        }
+        assert output == expected_output
+
     def test__fit(self):
         """Test the ``_fit`` method.
 
@@ -1341,9 +1362,35 @@ def test_fit(self):
 
         # Assert
         transformer._store_columns.assert_called_once_with(
-            list(columns_to_sdtypes.keys()), data
+            tuple(columns_to_sdtypes.keys()), data
         )
         transformer._set_seed.assert_called_once_with(data)
         transformer._get_columns_data.assert_called_once_with(data, ['a', 'b'])
         transformer._fit.assert_called_once_with(data_transformer, columns_to_sdtypes)
         transformer._build_output_columns.assert_called_once_with(data)
+
+    def test_fit_transform(self):
+        """Test the ``fit_transform`` method."""
+        # Setup
+        transformer = BaseMultiColumnTransformer()
+        columns_to_sdtypes = {
+            'a': 'numerical',
+            'b': 'categorical',
+            'c': 'boolean'
+        }
+        data = pd.DataFrame({
+            'a': [1, 2, 3],
+            'b': ['a', 'b', 'c'],
+        })
+        transformer.columns = ['a', 'b']
+        mock_fit = Mock()
+        mock_transform = Mock(return_value=data)
+        transformer.fit = mock_fit
+        transformer.transform = mock_transform
+
+        # Run
+        transformer.fit_transform(data, columns_to_sdtypes)
+
+        # Assert
+        mock_fit.assert_called_once_with(data, columns_to_sdtypes)
+        mock_transform.assert_called_once_with(data)