Repository: sdv-dev/SDV

Issue:
[OneHotEncoding constraint] Allow me to specify whether to keep the one-hot columns or collapse them into one categorical column
### Problem Description
The [OneHotEncoding constraint](https://docs.sdv.dev/sdv/concepts/constraint-augmented-generation-cag/predefined-constraints/onehotencoding) is meant to be used when you already have several one-hot encoded columns in your data. The constraint guarantees that the one-hot scheme will also be carried over into the synthetic data -- i.e. within each row, exactly 1 of the columns should contain a 1 while the rest should be 0.

Currently, the constraint works by keeping all the 0/1 columns for the ML model to learn. When creating synthetic data, it modifies the data so that only one column contains a `1` (and others contain a `0`). This strategy may not be optimal for a few reasons:

- **Performance**: If there are a lot of one-hot columns, the ML model's fitting and sampling time will be slow
- **Quality**: Some synthesizers are able to learn this type of data much better when it's presented as a single, discrete variable (categorical) rather than a set of one-hot columns

### Expected behavior
Add a parameter to this synthesizer called `learning_strategy`. This parameter can take on one of two values:
- (default) `'one_hot'`: The status quo. The AI will learn each of the one-hot columns, and the constraint will make sure that only one of the synthesized values is a `1`
- `'categorical'`: The constraint will collapse the one-hot columns into a single, categorical column for the AI model to learn. This may improve the performance (if you have a lot of one-hot columns), and quality (for certain synthesizers like GaussianCopula)
    - _Note that the final, outputted synthetic data should have the same 0/1 columns as the original data. (The promise is always that the synthetic data matches the format of the original.) The categorical column will just be created internally for the purposes of modeling; it should be reversed back into 0/1 columns when sampling._

```python
from sdv.cag import OneHotEncoding

constraint = OneHotEncoding(
  column_names=['status_new', 'status_on_hold', 'status_active', 'status_lapsed', 'status_churned'],
  learning_strategy='categorical'
)
```


Discussion:

Patch:
diff --git a/sdv/cag/one_hot_encoding.py b/sdv/cag/one_hot_encoding.py
index bd86cf09b..a38979bef 100644
--- a/sdv/cag/one_hot_encoding.py
+++ b/sdv/cag/one_hot_encoding.py
@@ -4,10 +4,12 @@
 
 import numpy as np
 
+from sdv._utils import _create_unique_name
 from sdv.cag._errors import ConstraintNotMetError
 from sdv.cag._utils import (
     _get_is_valid_dict,
     _is_list_of_type,
+    _remove_columns_from_metadata,
     _validate_table_and_column_names,
     _validate_table_name_if_defined,
 )
@@ -30,20 +32,30 @@ class OneHotEncoding(BaseConstraint):
         table_name (str, optional):
             The name of the table that contains the columns. Optional if the
             data is only a single table. Defaults to None.
+        learning_strategy (str, optional):
+            Strategy for how the model should learn the one-hot fields. Supported values:
+            - 'one_hot' (default): Learn each one-hot column separately.
+            - 'categorical': Internally collapse the one-hot columns into a single categorical
+              column for the model to learn, then expand back to one-hot at sampling time.
     """
 
     @staticmethod
-    def _validate_init_inputs(column_names, table_name):
+    def _validate_init_inputs(column_names, table_name, learning_strategy):
         if not _is_list_of_type(column_names):
             raise ValueError('`column_names` must be a list of strings.')
 
         _validate_table_name_if_defined(table_name)
 
-    def __init__(self, column_names, table_name=None):
+        if learning_strategy not in ['one_hot', 'categorical']:
+            raise ValueError("`learning_strategy` must be either 'one_hot' or 'categorical'.")
+
+    def __init__(self, column_names, table_name=None, learning_strategy='one_hot'):
         super().__init__()
-        self._validate_init_inputs(column_names, table_name)
+        self._validate_init_inputs(column_names, table_name, learning_strategy)
         self._column_names = column_names
         self.table_name = table_name
+        self.learning_strategy = learning_strategy
+        self._categorical_column = '#'.join(self._column_names)
 
     def _validate_constraint_with_metadata(self, metadata):
         """Validate the constraint is compatible with the provided metadata.
@@ -94,12 +106,25 @@ def _fit(self, data, metadata):
 
     def _get_updated_metadata(self, metadata):
         table_name = self._get_single_table_name(metadata)
-        metadata = deepcopy(metadata)
-        for column in self._column_names:
-            if metadata.tables[table_name].columns[column]['sdtype'] in ['categorical', 'boolean']:
-                metadata.tables[table_name].columns[column]['sdtype'] = 'numerical'
-
-        return metadata
+        if self.learning_strategy == 'categorical':
+            self._categorical_column = _create_unique_name(
+                self._categorical_column, metadata.tables[table_name].columns
+            )
+            md = metadata.to_dict()
+            md['tables'][table_name]['columns'][self._categorical_column] = {
+                'sdtype': 'categorical'
+            }
+            return _remove_columns_from_metadata(md, table_name, columns_to_drop=self._column_names)
+
+        else:
+            metadata = deepcopy(metadata)
+            for column in self._column_names:
+                if metadata.tables[table_name].columns[column]['sdtype'] in [
+                    'categorical',
+                    'boolean',
+                ]:
+                    metadata.tables[table_name].columns[column]['sdtype'] = 'numerical'
+            return metadata
 
     def _transform(self, data):
         """Transform the data.
@@ -113,9 +138,15 @@ def _transform(self, data):
                 Transformed data.
         """
         table_name = self._get_single_table_name(self.metadata)
-        one_hot_data = data[table_name][self._column_names]
-        one_hot_data = np.where(one_hot_data == 0, EPSILON, 1 - EPSILON)
-        data[table_name][self._column_names] = one_hot_data
+        if self.learning_strategy == 'categorical':
+            table_data = data[table_name]
+            categories = table_data[self._column_names].idxmax(axis=1)
+            table_data[self._categorical_column] = categories
+            data[table_name] = table_data.drop(self._column_names, axis=1)
+        else:
+            one_hot_data = data[table_name][self._column_names]
+            one_hot_data = np.where(one_hot_data == 0, EPSILON, 1 - EPSILON)
+            data[table_name][self._column_names] = one_hot_data
 
         return data
 
@@ -134,13 +165,28 @@ def _reverse_transform(self, data):
         """
         table_name = self._get_single_table_name(self.metadata)
         table_data = data[table_name]
-        one_hot_data = table_data[self._column_names]
-        transformed_data = np.zeros_like(one_hot_data.to_numpy())
-        max_category_indices = np.argmax(one_hot_data.to_numpy(), axis=1)
-        transformed_data[np.arange(len(one_hot_data)), max_category_indices] = 1
-        table_data[self._column_names] = transformed_data
-        data[table_name] = table_data
 
+        if self.learning_strategy == 'categorical':
+            categories = table_data.pop(self._categorical_column)
+            num_rows = len(table_data)
+            num_cols = len(self._column_names)
+            transformed = np.zeros((num_rows, num_cols), dtype=float)
+
+            column_to_index = {name: idx for idx, name in enumerate(self._column_names)}
+            indices = categories.map(lambda x: column_to_index[x]).to_numpy()
+            transformed[np.arange(num_rows), indices] = 1
+
+            for idx, col in enumerate(self._column_names):
+                table_data[col] = transformed[:, idx]
+
+        else:
+            one_hot_data = table_data[self._column_names]
+            transformed_data = np.zeros_like(one_hot_data.to_numpy())
+            max_category_indices = np.argmax(one_hot_data.to_numpy(), axis=1)
+            transformed_data[np.arange(len(one_hot_data)), max_category_indices] = 1
+            table_data[self._column_names] = transformed_data
+
+        data[table_name] = table_data
         return data
 
     def _is_valid(self, data, metadata):


Test patch:
diff --git a/tests/integration/cag/test_one_hot_encoding.py b/tests/integration/cag/test_one_hot_encoding.py
index 5570d4499..5f100ea31 100644
--- a/tests/integration/cag/test_one_hot_encoding.py
+++ b/tests/integration/cag/test_one_hot_encoding.py
@@ -8,7 +8,7 @@
 from sdv.cag._errors import ConstraintNotMetError
 from sdv.metadata import Metadata
 from sdv.single_table import GaussianCopulaSynthesizer
-from tests.utils import run_copula, run_hma
+from tests.utils import run_constraint, run_copula, run_hma
 
 
 @pytest.fixture()
@@ -201,3 +201,144 @@ def test_end_to_end_boolean():
     assert (samples.sum(axis=1) == 1).all()
     for col in columns:
         assert sorted(samples[col].unique().tolist()) == [0, 1]
+
+
+def test_end_to_end_categorical_single(data, metadata):
+    """End-to-end with learning_strategy='categorical' for single-table data."""
+    # Setup
+    constraint = OneHotEncoding(column_names=['a', 'b', 'c'], learning_strategy='categorical')
+
+    # Run
+    synthesizer = run_copula(data, metadata, [constraint])
+    synthetic_data = synthesizer.sample(200)
+    synthesizer.validate_constraints(synthetic_data=synthetic_data)
+
+    # Assert
+    assert set(synthetic_data.columns) == {'a', 'b', 'c'}
+    for col in ['a', 'b', 'c']:
+        assert set(synthetic_data[col]) == {0, 1}
+    assert (synthetic_data[['a', 'b', 'c']].sum(axis=1) == 1).all()
+
+
+def test_end_to_end_categorical_single_raises(data, metadata):
+    """Invalid synthetic data should raise with learning_strategy='categorical'."""
+    # Setup
+    invalid_data = pd.DataFrame({
+        'a': [1, 2, 0],
+        'b': [0, 1, np.nan],
+        'c': [0, 0, 3],
+    })
+    constraint = OneHotEncoding(column_names=['a', 'b', 'c'], learning_strategy='categorical')
+
+    # Run and Assert
+    msg = re.escape(
+        "Data is not valid for the 'OneHotEncoding' constraint in table 'table':\n"
+        '   a    b  c\n'
+        '1  2  1.0  0\n'
+        '2  0  NaN  3'
+    )
+    with pytest.raises(ConstraintNotMetError, match=msg):
+        run_copula(invalid_data, metadata, [constraint])
+
+    # Run and Assert
+    msg = re.escape('The one hot encoding requirement is not met for row indices: 1, 2')
+    with pytest.raises(ConstraintNotMetError, match=msg):
+        synthesizer = run_copula(data, metadata, [constraint])
+        synthesizer.validate_constraints(synthetic_data=invalid_data)
+
+
+def test_end_to_end_categorical_multi(data_multi, metadata_multi):
+    """End-to-end with learning_strategy='categorical' for multi-table data."""
+    # Setup
+    constraint = OneHotEncoding(
+        column_names=['a', 'b', 'c'], table_name='table1', learning_strategy='categorical'
+    )
+
+    # Run
+    synthesizer = run_hma(data_multi, metadata_multi, [constraint])
+    synthetic = synthesizer.sample(200)
+    synthesizer.validate_constraints(synthetic_data=synthetic)
+
+    # Assert
+    assert set(synthetic['table1'].columns) == {'a', 'b', 'c'}
+    for col in ['a', 'b', 'c']:
+        assert set(synthetic['table1'][col]) == {0, 1}
+    assert (synthetic['table1'][['a', 'b', 'c']].sum(axis=1) == 1).all()
+
+
+def test_end_to_end_categorical_multi_raises(data_multi, metadata_multi):
+    """Invalid multi-table synthetic data should raise with learning_strategy='categorical'."""
+    # Setup
+    constraint = OneHotEncoding(
+        column_names=['a', 'b', 'c'], table_name='table1', learning_strategy='categorical'
+    )
+    invalid = {
+        'table1': pd.DataFrame({
+            'a': [1, 2, 0],
+            'b': [0, 1, np.nan],
+            'c': [0, 0, 3],
+        }),
+        'table2': pd.DataFrame({'id': range(5)}),
+    }
+
+    # Run and Assert
+    msg = re.escape(
+        "Data is not valid for the 'OneHotEncoding' constraint in table 'table1':\n   "
+        'a    b  c\n1  2  1.0  0\n2  0  NaN  3'
+    )
+    with pytest.raises(ConstraintNotMetError, match=msg):
+        run_hma(invalid, metadata_multi, [constraint])
+
+    # Run and Assert
+    msg = "Table 'table1': The one hot encoding requirement is not met for row indices: 1, 2."
+    with pytest.raises(ConstraintNotMetError, match=msg):
+        synthesizer = run_hma(data_multi, metadata_multi, [constraint])
+        synthesizer.validate_constraints(synthetic_data=invalid)
+
+
+def test_constraint_pipeline_categorical_single(data, metadata):
+    """Constraint pipeline behavior for categorical strategy (single table)."""
+    # Setup
+    constraint = OneHotEncoding(column_names=['a', 'b', 'c'], learning_strategy='categorical')
+
+    # Run
+    updated_metadata, transformed, reverse_transformed = run_constraint(constraint, data, metadata)
+
+    # Assert metadata
+    assert updated_metadata.get_column_names() == ['a#b#c']
+
+    # Assert transform
+    assert transformed.shape[1] == 1
+    assert not any(col in transformed.columns for col in ['a', 'b', 'c'])
+    assert set(transformed.columns) == {'a#b#c'}
+
+    # Assert reverse_transform
+    assert set(reverse_transformed.columns) == {'a', 'b', 'c'}
+    assert (reverse_transformed[['a', 'b', 'c']].sum(axis=1) == 1).all()
+    assert set(reverse_transformed.columns) == {'a', 'b', 'c'}
+
+
+def test_constraint_pipeline_categorical_multi(data_multi, metadata_multi):
+    """Constraint pipeline behavior for categorical strategy (multi table)."""
+    # Setup
+    orig_cols = ['a', 'b', 'c']
+    constraint = OneHotEncoding(
+        column_names=orig_cols, table_name='table1', learning_strategy='categorical'
+    )
+
+    # Run
+    updated_metadata, transformed, reverse_transformed = run_constraint(
+        constraint, data_multi, metadata_multi
+    )
+
+    # Assert metadata
+    assert updated_metadata.tables['table1'].get_column_names() == ['a#b#c']
+
+    # Assert transform
+    assert list(transformed['table1'].columns) != orig_cols
+    assert transformed['table1'].shape[1] == 1
+    assert list(transformed['table2'].columns) == list(data_multi['table2'].columns)
+
+    # Assert reverse_transform
+    assert set(reverse_transformed['table1'].columns) == set(orig_cols)
+    assert (reverse_transformed['table1'][orig_cols].sum(axis=1) == 1).all()
diff --git a/tests/unit/cag/test_one_hot_encoding.py b/tests/unit/cag/test_one_hot_encoding.py
index 3cdef6b08..17650aa2c 100644
--- a/tests/unit/cag/test_one_hot_encoding.py
+++ b/tests/unit/cag/test_one_hot_encoding.py
@@ -287,3 +287,136 @@ def test_is_valid(self):
         }
         for table, series in out.items():
             pd.testing.assert_series_equal(expected_out[table], series)
+
+    def test___init___invalid_learning_strategy(self):
+        """Test it raises if learning_strategy is invalid."""
+        # Run and Assert
+        err_msg = "`learning_strategy` must be either 'one_hot' or 'categorical'."
+        with pytest.raises(ValueError, match=err_msg):
+            OneHotEncoding(column_names=['a', 'b', 'c'], learning_strategy='bad')
+
+    def test_get_updated_metadata_categorical_single(self):
+        """Test updated metadata drops one-hot columns and adds a categorical column."""
+        # Setup
+        metadata = Metadata.load_from_dict({
+            'columns': {
+                'a': {'sdtype': 'numerical'},
+                'b': {'sdtype': 'numerical'},
+                'c': {'sdtype': 'numerical'},
+            }
+        })
+        instance = OneHotEncoding(column_names=['a', 'b', 'c'], learning_strategy='categorical')
+
+        # Run
+        updated = instance.get_updated_metadata(metadata)
+
+        # Assert
+        cols = updated.get_column_names()
+        assert cols == ['a#b#c']
+
+        table_name = updated._get_single_table_name()
+        assert updated.tables[table_name].columns[cols[0]]['sdtype'] == 'categorical'
+
+        assert 'a' not in cols
+        assert 'b' not in cols
+        assert 'c' not in cols
+
+    def test_get_updated_metadata_categorical_repeated_column_names(self):
+        """Test updated metadata drops one-hot columns and adds a categorical column."""
+        # Setup
+        metadata = Metadata.load_from_dict({
+            'columns': {
+                'a': {'sdtype': 'numerical'},
+                'b': {'sdtype': 'numerical'},
+                'c': {'sdtype': 'numerical'},
+                'a#b#c': {'sdtype': 'numerical'},
+            }
+        })
+        instance = OneHotEncoding(column_names=['a', 'b', 'c'], learning_strategy='categorical')
+
+        # Run
+        updated = instance.get_updated_metadata(metadata)
+
+        # Assert
+        cols = updated.get_column_names()
+        assert cols == ['a#b#c', 'a#b#c_']
+
+        table_name = updated._get_single_table_name()
+        assert updated.tables[table_name].columns[cols[0]]['sdtype'] == 'numerical'
+        assert updated.tables[table_name].columns[cols[1]]['sdtype'] == 'categorical'
+
+        assert 'a' not in cols
+        assert 'b' not in cols
+        assert 'c' not in cols
+
+    def test_transform_and_reverse_categorical_single(self):
+        """Test transform collapses to categorical and reverse restores one-hot (single table)."""
+        # Setup
+        data = pd.DataFrame({
+            'a': [1, 0, 0, 1],
+            'b': [0, 1, 0, 0],
+            'c': [0, 0, 1, 0],
+        })
+        metadata = Metadata.load_from_dict({
+            'columns': {
+                'a': {'sdtype': 'numerical'},
+                'b': {'sdtype': 'numerical'},
+                'c': {'sdtype': 'numerical'},
+            }
+        })
+        instance = OneHotEncoding(column_names=['a', 'b', 'c'], learning_strategy='categorical')
+
+        # Run
+        instance.fit(data, metadata)
+        transformed = instance.transform(data)
+        reversed_df = instance.reverse_transform(transformed)
+
+        # Assert
+        pd.testing.assert_frame_equal(transformed, pd.DataFrame({'a#b#c': ['a', 'b', 'c', 'a']}))
+        pd.testing.assert_frame_equal(data, reversed_df)
+
+    def test_transform_and_reverse_categorical_multi(self):
+        """Test transform/reverse with categorical strategy (multi table)."""
+        # Setup
+        data = {
+            'table1': pd.DataFrame({
+                'a': [1, 0, 0, 1],
+                'b': [0, 1, 0, 0],
+                'c': [0, 0, 1, 0],
+            }),
+            'table2': pd.DataFrame({'id': range(4)}),
+        }
+        metadata = Metadata.load_from_dict({
+            'tables': {
+                'table1': {
+                    'columns': {
+                        'a': {'sdtype': 'numerical'},
+                        'b': {'sdtype': 'numerical'},
+                        'c': {'sdtype': 'numerical'},
+                    }
+                },
+                'table2': {
+                    'columns': {
+                        'id': {'sdtype': 'id'},
+                    }
+                },
+            }
+        })
+        instance = OneHotEncoding(
+            column_names=['a', 'b', 'c'], table_name='table1', learning_strategy='categorical'
+        )
+
+        # Run
+        instance.fit(data, metadata)
+        transformed = instance.transform(data)
+        reversed_data = instance.reverse_transform(transformed)
+
+        # Assert
+        pd.testing.assert_frame_equal(
+            transformed['table1'], pd.DataFrame({'a#b#c': ['a', 'b', 'c', 'a']})
+        )
+        pd.testing.assert_frame_equal(data['table2'], transformed['table2'])
+
+        # Assert
+        pd.testing.assert_frame_equal(data['table1'], reversed_data['table1'])
+        pd.testing.assert_frame_equal(data['table2'], reversed_data['table2'])