adding key columns before reversing constraints

sdv-dev · Jul 28, 2023 · 794c5e5 · 794c5e5
1 parent bfed446
commit 794c5e5
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/sdv/data_processing/data_processor.py b/sdv/data_processing/data_processor.py
@@ -716,6 +716,7 @@ def reverse_transform(self, data, reset_keys=False):
         if self._keys and num_rows:
             generated_keys = self.generate_keys(num_rows, reset_keys)
             sampled_columns.extend(self._keys)
+            reversed_data[generated_keys.columns] = generated_keys[generated_keys.notna()]
 
         for constraint in reversed(self._constraints_to_reverse):
             reversed_data = constraint.reverse_transform(reversed_data)