move _table_as_type_by_col from enterprise to sdv public

Author: R-Palazzo

sdv/cag/base.py

@@ -1,9 +1,14 @@
 """Base Multi-Table Pattern."""
 
+import logging
+
+import numpy as np
 import pandas as pd
 
 from sdv.errors import NotFittedError
 
+LOGGER = logging.getLogger(__name__)
+
 
 class BasePattern:
     """Base CAG Pattern Class."""
@@ -121,6 +126,31 @@ def transform(self, data):
     def _reverse_transform(self, data):
         raise NotImplementedError
 
+    def _table_as_type_by_col(self, reverse_transformed, table, table_name):
+        """Cast table to given types on a column by column basis.
+
+        Args:
+            reverse_transformed (dict[str, pd.DataFrame])
+                The reverse transformed data dictionary
+            table (pd.DataFrame)
+                The reverse transformed table
+            table_name (str)
+                The name of the table
+        """
+        for col in table:
+            try:
+                reverse_transformed[table_name][col] = table[col].astype(
+                    self._dtypes[table_name][col]
+                )
+            except pd.errors.IntCastingNaNError:
+                LOGGER.info(
+                    "Column '%s' is being converted to float because it contains NaNs.", col
+                )
+                self._dtypes[table_name][col] = np.dtype('float64')
+                reverse_transformed[table_name][col] = table[col].astype(
+                    self._dtypes[table_name][col]
+                )
+
     def reverse_transform(self, data):
         """Reverse transform the data back into the original space.
 
@@ -135,7 +165,11 @@ def reverse_transform(self, data):
         reverse_transformed = self._reverse_transform(data)
         for table_name, table in reverse_transformed.items():
             table = table[self._original_data_columns[table_name]]
-            reverse_transformed[table_name] = table.astype(self._dtypes[table_name])
+            try:
+                reverse_transformed[table_name] = table.astype(self._dtypes[table_name])
+            except pd.errors.IntCastingNaNError:
+                # iterate over the columns and cast individually
+                self._table_as_type_by_col(reverse_transformed, table, table_name)
 
         if self._single_table:
             return reverse_transformed[self._table_name]