Fix pyarrow for old pandas

sdv-dev · Oct 9, 2024 · cbe1abf · cbe1abf
1 parent 3e9bd7a
commit cbe1abf
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/rdt/transformers/utils.py b/rdt/transformers/utils.py
@@ -252,8 +252,7 @@ def learn_rounding_digits(data):
         int or None:
             Number of digits to round to.
     """
-    name = data.name
-    if isinstance(data.dtype, pd.ArrowDtype):
+    if str(data.dtype).endswith("[pyarrow]"):
         data = data.to_numpy()
 
     # check if data has any decimals
@@ -276,7 +275,7 @@ def learn_rounding_digits(data):
     # Can't round, not equal after MAX_DECIMALS digits of precision
     LOGGER.info(
         "No rounding scheme detected for column '%s'. Data will not be rounded.",
-        name,
+        data.name,
     )
     return None