Repository: sdv-dev/SDV

Issue:
Add helper method for transforming conditions
Problem:
There are several locations in `BaseSingleTableSynthesizer` where the line 

```python
self._data_processor.transform(condition_df, is_condition=True)
```
is used to transform the conditions dataframe. Adding a private helper method `_transform_conditions` that encapsulates this call would make it easier to modify transform behavior for new synthesizers.

Proposal:
Add a `_transform_conditions` method to `BaseSingleTableSynthesizer` and replace existing `self._data_processor.transform(condition_df, is_condition=True)` calls with `_transform_conditions` calls.
Update test cases accordingly and add a new unit test for `_transform_conditions`.


Discussion:

Patch:
diff --git a/sdv/single_table/base.py b/sdv/single_table/base.py
index 8dbd4b1a6..bf19f6ceb 100644
--- a/sdv/single_table/base.py
+++ b/sdv/single_table/base.py
@@ -1208,18 +1208,19 @@ def sample(self, num_rows, max_tries_per_batch=100, batch_size=None, output_file
 
         return sampled_data
 
+    def _transform_conditions(self, condition_df):
+        return self._data_processor.transform(condition_df, is_condition=True)
+
     def _transform_conditions_chained_constraints(self, condition_df):
         try:
             transformed_condition = self._validate_transform_constraints(condition_df)
-            transformed_condition = self._data_processor.transform(
-                transformed_condition, is_condition=True
-            )
+            transformed_condition = self._transform_conditions(transformed_condition)
         except ConstraintNotMetError:
             raise ConstraintNotMetError(
                 'Provided conditions are not valid for the given constraints.'
             )
         except Exception:
-            transformed_condition = self._data_processor.transform(condition_df, is_condition=True)
+            transformed_condition = self._transform_conditions(condition_df)
 
         return transformed_condition
 
@@ -1274,13 +1275,13 @@ def _sample_with_conditions(
 
             condition = dict(zip(condition_columns, group))
             condition_df = dataframe.iloc[0].to_frame().T
+            dtypes = conditions.dtypes.to_dict()
+            condition_df = condition_df.astype(dtypes)
             if hasattr(self, '_chained_constraints'):
                 transformed_condition = self._transform_conditions_chained_constraints(condition_df)
             else:
                 try:
-                    transformed_condition = self._data_processor.transform(
-                        condition_df, is_condition=True
-                    )
+                    transformed_condition = self._transform_conditions(condition_df)
                 except ConstraintsNotMetError as error:
                     raise ConstraintsNotMetError(
                         'Provided conditions are not valid for the given constraints.'


Test patch:
diff --git a/tests/unit/single_table/test_base.py b/tests/unit/single_table/test_base.py
index 1924cf7af..1f597cadb 100644
--- a/tests/unit/single_table/test_base.py
+++ b/tests/unit/single_table/test_base.py
@@ -2341,7 +2341,7 @@ def test__sample_with_conditions_constraints_not_met_backwards_compatability(sel
         conditions = pd.DataFrame({'name': ['Johanna', 'Doe'], 'salary': [100.0, 90.0]})
         instance = Mock()
         delattr(instance, '_chained_constraints')
-        instance._data_processor.transform.side_effect = [ConstraintsNotMetError]
+        instance._transform_conditions.side_effect = [ConstraintsNotMetError]
 
         # Run and Assert
         error_msg = 'Provided conditions are not valid for the given constraints.'
@@ -2407,6 +2407,17 @@ def test__sample_with_conditions_transformed_with_transformed_data(self):
             'keep_extra_columns': False,
         }
 
+    def test__transform_conditions(self):
+        # Setup
+        conditions = pd.DataFrame({'name': ['Johanna']})
+        instance = Mock()
+
+        # Run
+        BaseSingleTableSynthesizer._transform_conditions(instance, conditions)
+
+        # Assert
+        instance._data_processor.transform.assert_called_once_with(conditions, is_condition=True)
+
     def test__transform_conditions_chained_constraints_no_transformed_conditions(self):
         """Test when the conditions are not being transformed.
 
@@ -2417,7 +2428,7 @@ def test__transform_conditions_chained_constraints_no_transformed_conditions(sel
         conditions = pd.DataFrame({'name': ['Johanna']})
         instance = Mock()
         instance._validate_transform_constraints.side_effect = [KeyError]
-        instance._data_processor.transform.side_effect = [
+        instance._transform_conditions.side_effect = [
             pd.DataFrame(),
             pd.DataFrame(),
         ]