Support transforming and reversing a subset of the training columns

rdt/hyper_transformer.py

@@ -174,18 +174,19 @@ def transform(self, data):
             data = data.copy()
 
         for column_name, transformer in self._transformers.items():
-            column = data.pop(column_name)
-            transformed = transformer.transform(column)
+            if column_name in data:
+                column = data.pop(column_name)
+                transformed = transformer.transform(column)
 
-            shape = transformed.shape
+                shape = transformed.shape
 
-            if len(shape) == 2:
-                for index in range(shape[1]):
-                    new_column = '{}#{}'.format(column_name, index)
-                    data[new_column] = transformed[:, index]
+                if len(shape) == 2:
+                    for index in range(shape[1]):
+                        new_column = '{}#{}'.format(column_name, index)
+                        data[new_column] = transformed[:, index]
 
-            else:
-                data[column_name] = transformed
+                else:
+                    data[column_name] = transformed
 
         return data
 
@@ -224,7 +225,7 @@ def _get_columns(data, column_name):
         regex = r'{}(#[0-9]+)?$'.format(re.escape(column_name))
         columns = data.columns[data.columns.str.match(regex)]
         if columns.empty:
-            raise ValueError('No columns match_ {}'.format(column_name))
+            return None
 
         values = [data.pop(column).values for column in columns]
 
@@ -249,6 +250,7 @@ def reverse_transform(self, data):
 
         for column_name, transformer in self._transformers.items():
             columns = self._get_columns(data, column_name)
-            data[column_name] = transformer.reverse_transform(columns)
+            if columns is not None:
+                data[column_name] = transformer.reverse_transform(columns)
 
         return data

tests/integration/test_hyper_transformer.py

@@ -159,3 +159,20 @@ def test_empty_transformers():
 
     pd.testing.assert_frame_equal(data, transformed)
     pd.testing.assert_frame_equal(data, reverse)
+
+
+def test_subset_of_columns():
+    """HyperTransform should be able to transform a subset of the training columns.
+
+    See https://github.com/sdv-dev/RDT/issues/152
+    """
+    data = get_input_data()
+
+    ht = HyperTransformer()
+    ht.fit(data)
+
+    subset = data[[data.columns[0]]]
+    transformed = ht.transform(subset)
+    reverse = ht.reverse_transform(transformed)
+
+    pd.testing.assert_frame_equal(subset, reverse)

tests/test_hyper_transformer.py

@@ -152,13 +152,14 @@ def test__get_columns_two(self):
         ])
         np.testing.assert_equal(returned, expected)
 
-    def test__get_columns_error(self):
+    def test__get_columns_none(self):
         data = pd.DataFrame({
             'a': [1, 2, 3],
         })
 
-        with pytest.raises(ValueError):
-            HyperTransformer._get_columns(data, 'b')
+        returned = HyperTransformer._get_columns(data, 'b')
+
+        assert returned is None
 
     def test__get_columns_regex(self):
         data = pd.DataFrame({