Fix NotImplementedError when using FloatFormatter with numerical data types during fit due to PyArrow

pyproject.toml

@@ -49,7 +49,9 @@ rdt = { main = 'rdt.cli.__main__:main' }
 
 [project.optional-dependencies]
 copulas = ['copulas>=0.11.0',]
+pyarrow = ['pyarrow>=17.0.0']
 test = [
+    'rdt[pyarrow]',
     'rdt[copulas]',
 
     'pytest>=3.4.2',

rdt/transformers/numerical.py

@@ -104,7 +104,7 @@ def _raise_out_of_bounds_error(self, value, name, bound_type, min_bound, max_bou
 
     def _validate_values_within_bounds(self, data):
         if not self.computer_representation.startswith('Float'):
-            fractions = data[~data.isna() & data % 1 != 0]
+            fractions = data[~data.isna() & (data != (data // 1))]
             if not fractions.empty:
                 raise ValueError(
                     f"The column '{data.name}' contains float values {fractions.tolist()}. "

rdt/transformers/utils.py

@@ -261,7 +261,7 @@ def learn_rounding_digits(data):
         return None
 
     # Doesn't contain decimal digits
-    if ((roundable_data % 1) == 0).all():
+    if (roundable_data == roundable_data.astype(int)).all():
         return 0
 
     # Try to round to fewer digits

tests/unit/transformers/test_numerical.py

@@ -5,6 +5,7 @@
 import copulas
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 import pytest
 from copulas import univariate
 from pandas.api.types import is_float_dtype
@@ -44,6 +45,16 @@ def test__validate_values_within_bounds(self):
         # Run
         transformer._validate_values_within_bounds(data)
 
+    def test__validate_values_within_bounds_pyarrow(self):
+        """Test it works with pyarrow."""
+        # Setup
+        data = pd.Series(range(10), dtype='int64[pyarrow]')
+        transformer = FloatFormatter()
+        transformer.computer_representation = 'UInt8'
+
+        # Run
+        transformer._validate_values_within_bounds(data)
+
     def test__validate_values_within_bounds_under_minimum(self):
         """Test the ``_validate_values_within_bounds`` method.
 

tests/unit/transformers/test_utils.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 import pytest
 
 from rdt.transformers.utils import (
@@ -225,6 +226,18 @@ def test_learn_rounding_digits_less_than_15_decimals():
     assert output == 3
 
 
+def test_learn_rounding_digits_pyarrow():
+    """Test it works with pyarrow."""
+    # Setup
+    data = pd.Series(range(10), dtype='int64[pyarrow]')
+
+    # Run
+    output = learn_rounding_digits(data)
+
+    # Assert
+    assert output == 0
+
+
 def test_learn_rounding_digits_negative_decimals_float():
     """Test the learn_rounding_digits method with floats multiples of powers of 10.