Correctly handle non-finite numbers in heuristic determination of X distribution (#2342)

* handle non-finites explicitly

* improve and test edge case handling for distribution estimation

* revert debugging changes

* code readability

backend/common/compute/estimate_distribution.py

@@ -5,27 +5,45 @@ from scipy import sparse
 from backend.common.constants import XApproximateDistribution
 
 
-@numba.njit(fastmath=True, error_model="numpy", nogil=True)
+@numba.njit(error_model="numpy", nogil=True)
-def min_max(arr):
+def min_max(arr: np.ndarray):
     """Return (min, max) values for the ndarray."""
-    n = arr.size
+
-    odd = n % 2
+    # initialize to first finite value in array.  Normally,
-    if not odd:
+    # this will exit on the first value.
-        n -= 1
+    for i in range(arr.size):
-    max_val = min_val = arr[0]
+        min_val = max_val = arr[i]
-    i = 1
+        if np.isfinite(min_val):
-    while i < n:
+            break
+
+    # now find min/max, unrolled by two
+    odd = arr.size % 2
+    unrolled_loop_limit = arr.size - 1 if odd else arr.size
+    i = 0
+    while i < unrolled_loop_limit:
         x = arr[i]
         y = arr[i + 1]
+
+        # ignore non-finites
+        x = x if np.isfinite(x) else min_val
+        y = y if np.isfinite(y) else min_val
+
         if x > y:
             x, y = y, x
         min_val = min(x, min_val)
         max_val = max(y, max_val)
         i += 2
-    if not odd:
+
-        x = arr[n]
+    # handle the tail if any
+    if odd:
+        x = arr[arr.size - 1]
+
+        # ignore non-finites
+        x = x if np.isfinite(x) else min_val
+
         min_val = min(x, min_val)
         max_val = max(x, max_val)
+
     return min_val, max_val
 
 
@@ -38,6 +56,13 @@ def estimate_approximate_distribution(X) -> XApproximateDistribution:
     any (max-min) range in excess of 24 is implies tens of millions of
     observations of a single feature and so is extremely unlikely.
     """
+    if X.dtype.kind not in ["i", "u", "f"]:
+        raise TypeError(f"Unsupported matrix dtype: {X.dtype.name}")
+
+    if X.size == 0:
+        # default for empty array
+        return XApproximateDistribution.NORMAL
+
     if sparse.isspmatrix_csc(X) or sparse.isspmatrix_csr(X):
         Xdata = X.data
     elif type(X) is np.ndarray:
@@ -45,7 +70,7 @@ def estimate_approximate_distribution(X) -> XApproximateDistribution:
             X.size,
         )
     else:
-        raise TypeError(f"Unsupported matrix type: {str(type(X))}")
+        raise TypeError(f"Unsupported matrix format: {str(type(X))}")
 
     CHUNKSIZE = 1 << 24
     if Xdata.size > CHUNKSIZE:

backend/test/test_server/unit/compute/test_est_dist.py

@@ -25,6 +25,9 @@ class EstDistTest(unittest.TestCase):
     def test_estimate_approximate_distribution(self):
         raw = np.random.exponential(scale=1000, size=(100, 40))
 
+        # empty
+        self.assertEqual(estimate_approximate_distribution(np.zeros((0,))), XApproximateDistribution.NORMAL)
+
         # ndarray
         self.assertEqual(estimate_approximate_distribution(raw), XApproximateDistribution.COUNT)
         self.assertEqual(estimate_approximate_distribution(np.log1p(raw)), XApproximateDistribution.NORMAL)
@@ -35,7 +38,83 @@ class EstDistTest(unittest.TestCase):
             estimate_approximate_distribution(sparse.csr_matrix(np.log1p(raw))), XApproximateDistribution.NORMAL
         )
 
+        # csc_matrix
+        self.assertEqual(estimate_approximate_distribution(sparse.csc_matrix(raw)), XApproximateDistribution.COUNT)
+        self.assertEqual(
+            estimate_approximate_distribution(sparse.csc_matrix(np.log1p(raw))), XApproximateDistribution.NORMAL
+        )
+
         # BIG (ie, trigger MT)
         big = np.random.exponential(scale=100, size=(1_000_000, 100))
         self.assertEqual(estimate_approximate_distribution(big), XApproximateDistribution.COUNT)
         self.assertEqual(estimate_approximate_distribution(np.log1p(big)), XApproximateDistribution.NORMAL)
+
+    def test_unsupported_throws(self):
+        # dtypes and matrix formats we do not support
+        with self.assertRaises(TypeError):
+            estimate_approximate_distribution(np.array(["a", "b"]))
+        with self.assertRaises(TypeError):
+            estimate_approximate_distribution(sparse.coo_matrix(np.array([[0, 1, 2], [3, 0, 2]])))
+
+    def test_nonfinites(self):
+        def put(arr, ind, vals):
+            # like np.put, but creates and returns a modified copy of original array
+            a = arr.copy()
+            np.put(a, ind, vals)
+            return a
+
+        # non-finites
+        self.assertEqual(estimate_approximate_distribution(np.array([np.nan])), XApproximateDistribution.NORMAL)
+        self.assertEqual(estimate_approximate_distribution(np.array([np.PINF])), XApproximateDistribution.NORMAL)
+        self.assertEqual(estimate_approximate_distribution(np.array([np.NINF])), XApproximateDistribution.NORMAL)
+        self.assertEqual(
+            estimate_approximate_distribution(np.array([np.PINF, np.NINF, 0])), XApproximateDistribution.NORMAL
+        )
+        self.assertEqual(
+            estimate_approximate_distribution(np.array([np.nan, np.PINF, np.NINF])), XApproximateDistribution.NORMAL
+        )
+
+        raw = np.random.exponential(scale=1000, size=(50, 3))
+        logged = np.log1p(raw)
+
+        self.assertEqual(
+            estimate_approximate_distribution(put(raw, [1], [np.nan])),
+            XApproximateDistribution.COUNT,
+        )
+        self.assertEqual(
+            estimate_approximate_distribution(put(raw, [1], [np.PINF])),
+            XApproximateDistribution.COUNT,
+        )
+        self.assertEqual(
+            estimate_approximate_distribution(put(raw, [1], [np.NINF])),
+            XApproximateDistribution.COUNT,
+        )
+        self.assertEqual(
+            estimate_approximate_distribution(put(raw, [1, 3, 88], [np.nan, np.PINF, np.NINF])),
+            XApproximateDistribution.COUNT,
+        )
+        self.assertEqual(
+            estimate_approximate_distribution(put(raw, [0, 1], [np.nan, np.nan])),
+            XApproximateDistribution.COUNT,
+        )
+
+        self.assertEqual(
+            estimate_approximate_distribution(put(logged, [1], [np.nan])),
+            XApproximateDistribution.NORMAL,
+        )
+        self.assertEqual(
+            estimate_approximate_distribution(put(logged, [1], [np.PINF])),
+            XApproximateDistribution.NORMAL,
+        )
+        self.assertEqual(
+            estimate_approximate_distribution(put(logged, [1], [np.NINF])),
+            XApproximateDistribution.NORMAL,
+        )
+        self.assertEqual(
+            estimate_approximate_distribution(put(logged, [1, 3, 88], [np.nan, np.PINF, np.NINF])),
+            XApproximateDistribution.NORMAL,
+        )
+        self.assertEqual(
+            estimate_approximate_distribution(put(logged, [0, 1], [np.nan, np.nan])),
+            XApproximateDistribution.NORMAL,
+        )