Cast finite quantile boundaries to float64 to avoid precision errors.

PiperOrigin-RevId: 296235385

Author: paulgc17

tensorflow_data_validation/utils/quantiles_util.py

@@ -238,6 +238,13 @@ def generate_equi_width_buckets(quantiles: List[float],
     # the (+inf, +inf) buckets.
     last_bucket_count -= sample_count
 
+  # Cast finite boundaries from float32 to float64 to avoid precision errors.
+  # This error can happen when both min and max in the boundaries are valid
+  # float32 values, but when we compute (max-min) to compute the width it can
+  # result in an overflow.
+  # Example: min=-3.4e+38, max=3.4e+38
+  finite_values = np.array(finite_values, dtype=np.float64)
+
   # Check if the finite quantile boundaries are sorted.
   assert np.all(np.diff(finite_values) >= 0), (
       'Quantiles output not sorted %r'  % ','.join(map(str, finite_values)))
@@ -262,7 +269,7 @@ def generate_equi_width_buckets(quantiles: List[float],
 
 
 def _generate_equi_width_buckets_from_finite_boundaries(
-    quantiles: List[float], sample_count: float, num_buckets: int):
+    quantiles: np.ndarray, sample_count: float, num_buckets: int):
   """Generates equi width buckets from finite quantiles boundaries."""
 
   def _compute_count(start_index, end_index, start_pos):

tensorflow_data_validation/utils/quantiles_util_test.py

@@ -192,6 +192,22 @@ def test_find_median(self):
             quantiles_util.Bucket(2.33333333, 3.66666666, 1.33333333),
             quantiles_util.Bucket(3.66666666, 5, 2.33333333)],
     },
+    {
+        'testcase_name': 'finite_values_float_boundaries_float32_overflow',
+        'quantiles': list(np.array(
+            [-3.4e+38, 1.0, 2.0, 3.0, 4.0, 5.0, 3.4e+38], dtype=np.float32)),
+        'finite_min': 1,
+        'finite_max': 5,
+        'total_count': 6,
+        'num_buckets': 3,
+        'expected_buckets': [
+            quantiles_util.Bucket(
+                -3.3999999521443642e+38, -1.1333333173814546e+38, 0.66666666),
+            quantiles_util.Bucket(
+                -1.1333333173814546e+38, 1.133333317381455e+38, 4.66666666),
+            quantiles_util.Bucket(
+                1.133333317381455e+38, 3.3999999521443642e+38, 0.66666666)],
+    },
     {
         'testcase_name': 'same_min_max',
         'quantiles': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],