Cannot create custom constraint with primary key

[amontanez24]

Environment Details

Please indicate the following details about the environment in which you found the bug:

• SDV version: 1.2.1
• Python version: Any
• Operating System: Any

Error Description

If I create a custom constraint that involves a primary key, I get a KeyError during sampling

Steps to reproduce

from sdv.single_table import GaussianCopulaSynthesizer
from sdv.metadata import SingleTableMetadata
from sdv.constraints import create_custom_constraint_class
import pandas as pd

# Add constraints
def is_valid(column_names, data):
    return data['key'] == data['letter'] + '_' + data['number']

def transform(column_names, data):
    new_data = data.drop(['letter', 'number'], axis=1)
    return new_data

def reverse_transform(column_names, data):
    columns = data['key'].str.split('_', expand=True)
    data['letter'] = columns[0]
    data['number'] = columns[1]
    return data

custom_constraint = create_custom_constraint_class(
    is_valid_fn=is_valid,
    transform_fn=transform,
    reverse_transform_fn=reverse_transform
)

data = pd.DataFrame({
    'key': ['a_1', 'b_2', 'c_3'],
    'letter': ['a', 'b', 'c'],
    'number': ['1', '2', '3'],
    'other': [7, 8, 9]
})
metadata = SingleTableMetadata()
metadata.detect_from_dataframe(data)
metadata.update_column('key', sdtype='id', regex_format='\w_\d')
metadata.set_primary_key('key')
synth = GaussianCopulaSynthesizer(metadata)
synth.add_custom_constraint_class(custom_constraint, 'custom')


id_must_match = {
    'constraint_class': 'custom',
    'constraint_parameters': {
        'column_names': ['letter', 'number'],
    }
}

synth.add_constraints([id_must_match])
synth.fit(data)
sampled = synth.sample(100)

Additional context

• This is because we are not adding the key columns to the data that gets reverse transformed by constraints in the data processor. In the following lines, the keys need to be added to reversed_data
  

  SDV/sdv/data_processing/data_processor.py

  Lines 716 to 721 in c4a5230

  	if self._keys and num_rows:
  	generated_keys = self.generate_keys(num_rows, reset_keys)
  	sampled_columns.extend(self._keys)
  	for constraint in reversed(self._constraints_to_reverse):
  	reversed_data = constraint.reverse_transform(reversed_data)