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)
-def min_max(arr):
+@numba.njit(error_model="numpy", nogil=True)
+def min_max(arr: np.ndarray):
     """Return (min, max) values for the ndarray."""
-    n = arr.size
-    odd = n % 2
-    if not odd:
-        n -= 1
-    max_val = min_val = arr[0]
-    i = 1
-    while i < n:
+
+    # initialize to first finite value in array.  Normally,
+    # this will exit on the first value.
+    for i in range(arr.size):
+        min_val = max_val = arr[i]
+        if np.isfinite(min_val):
+            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: