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

@@ -254,14 +254,16 @@ def learn_rounding_digits(data):
     """
     # check if data has any decimals
     name = data.name
+    if str(data.dtype).endswith('[pyarrow]'):
+        data = data.to_numpy()
     roundable_data = data[~(np.isinf(data.astype(float)) | pd.isna(data))]
 
     # Doesn't contain numbers
     if len(roundable_data) == 0:
         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

@@ -44,6 +44,19 @@ 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
+        try:
+            data = pd.Series(range(10), dtype='int64[pyarrow]')
+        except TypeError:
+            pytest.skip("Skipping as old numpy/pandas versions don't support arrow")
+        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

@@ -1,7 +1,7 @@
 import sre_parse
 import warnings
 from sre_constants import MAXREPEAT
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import numpy as np
 import pandas as pd
@@ -225,6 +225,36 @@ def test_learn_rounding_digits_less_than_15_decimals():
     assert output == 3
 
 
+def test_learn_rounding_digits_pyarrow():
+    """Test it works with pyarrow."""
+    # Setup
+    try:
+        data = pd.Series(range(10), dtype='int64[pyarrow]')
+    except TypeError:
+        pytest.skip("Skipping as old numpy/pandas versions don't support arrow")
+
+    # Run
+    output = learn_rounding_digits(data)
+
+    # Assert
+    assert output == 0
+
+
+def test_learn_rounding_digits_pyarrow_float():
+    """Test it learns the proper amount of digits with pyarrow."""
+    # Setup
+    try:
+        data = pd.Series([0.5, 0.19, 3], dtype='float64[pyarrow]')
+    except TypeError:
+        pytest.skip("Skipping as old numpy/pandas versions don't support arrow")
+
+    # Run
+    output = learn_rounding_digits(data)
+
+    # Assert
+    assert output == 2
+
+
 def test_learn_rounding_digits_negative_decimals_float():
     """Test the learn_rounding_digits method with floats multiples of powers of 10.
 
@@ -299,6 +329,20 @@ def test_learn_rounding_digits_nullable_numerical_pandas_dtypes():
         assert output == expected_output[column]
 
 
+def test_learn_rounding_digits_pyarrow_to_numpy():
+    """Test that ``learn_rounding_digits`` works with pyarrow to numpy conversion."""
+    # Setup
+    data = Mock()
+    data.dtype = 'int64[pyarrow]'
+    data.to_numpy.return_value = np.array([1, 2, 3])
+
+    # Run
+    learn_rounding_digits(data)
+
+    # Assert
+    assert data.to_numpy.called
+
+
 def test_warn_dict():
     """Test that ``WarnDict`` will raise a warning when called with `text`."""
     # Setup