ordered_columns + prefix

rdt/transformers/base.py

@@ -486,8 +486,19 @@ class BaseMultiColumnTransformer(BaseTransformer):
 
     The ``BaseMultiColumnTransformer`` class contains methods that must be implemented
     in order to create a new multi column transformer.
+
+    Attributes:
+        ordered_columns (tuple):
+            Order of the columns to be used for the transformer.
+        prefixes (dict):
+            Dictionary mapping each output column to its prefix.
     """
 
+    def __init__(self):
+        super().__init__()
+        self.ordered_columns = None
+        self.prefixes = {}
+
     def get_input_column(self):
         """Override ``get_input_column`` method from ``BaseTransformer``.
 
@@ -510,52 +521,85 @@ def get_input_columns(self):
 
     def _get_output_to_property(self, property_):
         result = {
-            output_column: properties[property_]
+            f'{self.prefixes[output_column]}.{output_column}': properties[property_]
             for output_column, properties in self.output_properties.items()
         }
 
         return result
 
-    def _fit(self, columns_data, columns_to_sdtypes):
+    def _validate_ordered_columns(self, data, ordered_columns):
+        """Check that all the columns in ``ordered_columns`` are present in the data."""
+        missing = set(ordered_columns) - set(data.columns)
+        if missing:
+            missing_to_print = ', '.join(missing)
+            raise KeyError(f'Columns ({missing_to_print}) are not present in the data.')
+
+    def _generate_prefixes(self, ordered_columns):
+        """Generate prefixes for the output columns to precised which column they come from.
+
+        Returns:
+            dict:
+                Dictionary mapping each output column to its prefix.
+                The key is the output column name and the value is the prefix.
+        """
+        raise NotImplementedError()
+
+    def _validate_prefixes(self, ordered_columns):
+        """Check that the prefixes are valid.
+
+        Every prefix must include the name of at least one column in the data.
+        """
+        for prefix in self.prefixes.values():
+            if not any(column in prefix for column in ordered_columns):
+                raise ValueError(
+                    f"The prefix '{prefix}' does not include the name of any column in the data."
+                )
+
+    def _fit(self, data, ordered_columns):
         """Fit the transformer to the data.
 
         Args:
-            columns_data (pandas.DataFrame):
+            data (pandas.DataFrame):
                 Data to transform.
-            columns_to_sdtypes (dict):
-                Dictionary mapping column names to their sdtypes.
+            ordered_columns (tuple):
+                Order of the columns to be used for the transformer.
         """
         raise NotImplementedError()
 
     @random_state
-    def fit(self, data, columns_to_sdtypes):
+    def fit(self, data, ordered_columns):
         """Fit the transformer to a ``column`` of the ``data``.
 
         Args:
             data (pandas.DataFrame):
                 The entire table.
-            columns_to_sdtypes (dict):
-                Dictionary mapping column names to their sdtypes.
+            ordered_columns (tuple):
+                Order of the columns to be used for the transformer.
         """
-        column_names = tuple(columns_to_sdtypes.keys())
-        self._store_columns(column_names, data)
+        self._validate_ordered_columns(data, ordered_columns)
+        self.ordered_columns = ordered_columns
+        self._store_columns(ordered_columns, data)
         self._set_seed(data)
+
         columns_data = self._get_columns_data(data, self.columns)
-        self._fit(columns_data, columns_to_sdtypes)
+        self._fit(columns_data, ordered_columns)
+
+        self.prefixes = self._generate_prefixes(ordered_columns)
+        self._validate_prefixes(ordered_columns)
         self._build_output_columns(data)
 
-    def fit_transform(self, data, columns_to_sdtypes):
+    def fit_transform(self, data, ordered_columns):
         """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.
+            ordered_columns (tuple):
+                Order of the columns to be used for the transformer.
 
         Returns:
             pd.DataFrame:
                 The entire table, containing the transformed data.
         """
-        self.fit(data, columns_to_sdtypes)
+        self.fit(data, ordered_columns)
         return self.transform(data)

tests/integration/transformers/test_base.py

@@ -136,11 +136,17 @@ def test_multi_column_transformer_same_number_of_columns_input_output():
     # 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):
+        def _fit(self, columns_data, ordered_columns):
             self.output_properties = {
                 column: {'sdtype': 'numerical'} for column in self.columns
             }
-            self.dtypes = columns_data.dtypes
+
+        def _generate_prefixes(self, ordered_columns):
+            prefixes = {}
+            for idx, column in enumerate(self.output_properties):
+                prefixes[column] = '#'.join(ordered_columns[:idx + 1])
+
+            return prefixes
 
         def _transform(self, data):
             return data.cumsum(axis=1)
@@ -149,30 +155,26 @@ def _reverse_transform(self, data):
             result = data.diff(axis=1)
             result.iloc[:, 0] = data.iloc[:, 0]
 
-            return result.astype(self.dtypes)
+            return result.astype(int)
 
     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'
-    }
+    order_columns = ('col_1', 'col_2', 'col_3')
     transformer = AdditionTransformer()
 
     # Run
-    transformed = transformer.fit_transform(data_test, column_to_sdtype)
+    transformed = transformer.fit_transform(data_test, order_columns)
     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]
+        'col_1.col_1': [1, 2, 3],
+        'col_1#col_2.col_2': [11, 22, 33],
+        'col_1#col_2#col_3.col_3': [111, 222, 333]
     })
     pd.testing.assert_frame_equal(expected_transform, transformed)
     pd.testing.assert_frame_equal(reverse, data_test)
@@ -184,14 +186,20 @@ 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):
+        def _fit(self, columns_data, ordered_columns):
             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'}
+                'concatenate_1': {'sdtype': 'categorical'},
+                'concatenate_2': {'sdtype': 'categorical'}
             }
-            self.dtypes = columns_data.dtypes
+
+        def _generate_prefixes(self, ordered_columns):
+            prefixes = {}
+            for idx, column in enumerate(self.output_properties):
+                prefixes[column] = self.name_1 if idx == 0 else self.name_2
+
+            return prefixes
 
         def _transform(self, data):
             data[self.name_1] = data.iloc[:, 0] + '#' + data.iloc[:, 1]
@@ -209,7 +217,7 @@ def _reverse_transform(self, data):
             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)
+            return result.drop(columns=column_names)
 
     data_test = pd.DataFrame({
         'col_1': ['A', 'B', 'C'],
@@ -218,23 +226,18 @@ def _reverse_transform(self, data):
         'col_4': ['J', 'K', 'L']
     })
 
-    column_to_sdtype = {
-        'col_1': 'categorical',
-        'col_2': 'categorical',
-        'col_3': 'categorical',
-        'col_4': 'categorical'
-    }
+    ordered_columns = ('col_1', 'col_2', 'col_3', 'col_4')
     transformer = ConcatenateTransformer()
 
     # Run
-    transformer.fit(data_test, column_to_sdtype)
+    transformer.fit(data_test, ordered_columns)
     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']
+        'col_1#col_2.concatenate_1': ['A#D', 'B#E', 'C#F'],
+        'col_3#col_4.concatenate_2': ['G#J', 'H#K', 'I#L']
     })
     pd.testing.assert_frame_equal(expected_transform, transformed)
     pd.testing.assert_frame_equal(reverse, data_test)
@@ -244,57 +247,61 @@ 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'
+        def _fit(self, columns_data, ordered_columns):
             self.output_properties = {
-                name_1: {'sdtype': 'categorical'},
-                name_2: {'sdtype': 'categorical'},
-                name_3: {'sdtype': 'categorical'},
-                name_4: {'sdtype': 'categorical'}
+                'first_part_1': {'sdtype': 'categorical'},
+                'second_part_1': {'sdtype': 'categorical'},
+                'first_part_2': {'sdtype': 'categorical'},
+                'second_part_2': {'sdtype': 'categorical'}
             }
-            self.names = [name_1, name_2, name_3, name_4]
-            self.dtypes = columns_data.dtypes
+
+        def _generate_prefixes(self, ordered_columns):
+            list_prefixes = [
+                self.columns[0], self.columns[0],
+                self.columns[1], self.columns[1]
+            ]
+            prefixes = {}
+            for idx, column in enumerate(self.output_properties):
+                prefixes[column] = list_prefixes[idx]
+
+            return prefixes
 
         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]
+            data[self.output_columns[0]] = data[self.columns[0]].str[0]
+            data[self.output_columns[1]] = data[self.columns[0]].str[1]
+            data[self.output_columns[2]] = data[self.columns[1]].str[0]
+            data[self.output_columns[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]]
+            reverse_1 = result[self.output_columns[0]] + result[self.output_columns[1]]
+            reverse_2 = result[self.output_columns[2]] + result[self.output_columns[3]]
+            result[self.columns[0]] = reverse_1
+            result[self.columns[1]] = reverse_2
 
-            return result.astype(self.dtypes).drop(columns=self.names)
+            return result.drop(columns=self.output_columns)
 
     data_test = pd.DataFrame({
         'col_1': ['AB', 'CD', 'EF'],
         'col_2': ['GH', 'IJ', 'KL'],
     })
 
-    column_to_sdtype = {
-        'col_1': 'categorical',
-        'col_2': 'categorical',
-    }
+    ordered_columns = ('col_1', 'col_2')
     transformer = ExpandTransformer()
 
     # Run
-    transformer.fit(data_test, column_to_sdtype)
+    transformer.fit(data_test, ordered_columns)
     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']
+        'col_1.first_part_1': ['A', 'C', 'E'],
+        'col_1.second_part_1': ['B', 'D', 'F'],
+        'col_2.first_part_2': ['G', 'I', 'K'],
+        'col_2.second_part_2': ['H', 'J', 'L']
     })
     pd.testing.assert_frame_equal(expected_transform, transformed)
     pd.testing.assert_frame_equal(reverse, data_test)