Type more and fix some bugs

Author: lithomas1

pandas/_libs/algos.pyx

@@ -833,9 +833,20 @@ def backfill_2d_inplace(algos_t[:, :] values,
 
 
 # Fillna logic
+# We have our own fused type since we don't need to support
+# types that can't hold NAs
+ctypedef fused fillna_t:
+    float64_t
+    float32_t
+    object
+    # TODO: Add datetime64 support through viewing data as uint64
+    uint64_t  # Datetime64
+    # TODO: Complex support?
+
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def fillna1d(ndarray arr,
+def fillna1d(fillna_t[:] arr,
              object value,
              limit=None,
              bint inf_as_na=False
@@ -856,15 +867,15 @@ def fillna1d(ndarray arr,
     arr : ndarray
     value : object
         The value to use to replace nans
-    limit : int, default -1
-        The number of elements to fill. If -1, fills all NaN values
+    limit : int, default None
+        The number of elements to fill. If None, fills all NaN values
     inf_as_na:
         Whether to consider INF and NEGINF as NA
     """
     cdef:
         Py_ssize_t i, N, lim
         Py_ssize_t count=0
-        object val
+        fillna_t val
 
     assert arr.ndim == 1, "'arr' must be 1-D."
 
@@ -883,7 +894,7 @@ def fillna1d(ndarray arr,
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def fillna2d(ndarray arr,
+def fillna2d(fillna_t[:, :] arr,
              object value,
              limit=None,
              bint inf_as_na=False
@@ -904,25 +915,26 @@ def fillna2d(ndarray arr,
     arr : ndarray
     value : object
         The value to use to replace nans
-    limit : int, default -1
-        The number of elements to fill. If -1, fills all NaN values
+    limit : int, default None
+        The number of elements to fill. If None, fills all NaN values
     inf_as_na:
         Whether to consider INF and NEGINF as NA
     """
     cdef:
         Py_ssize_t i, j, n, m, lim
         Py_ssize_t count=0
-        object val
+        fillna_t val
 
     assert arr.ndim == 2, "'arr' must be 2-D."
 
     n, m = (<object>arr).shape
-    lim = validate_limit(n, limit)
+    lim = validate_limit(m, limit)
     if inf_as_na:
         check_func = missing.checknull_old
     else:
         check_func = missing.checknull
     for i in range(n):
+        count = 0  # Limit is per axis
         for j in range(m):
             val = arr[i, j]
             if check_func(val) and count < lim:

pandas/core/internals/blocks.py

@@ -436,6 +436,7 @@ def fillna(
         """
         inplace = validate_bool_kwarg(inplace, "inplace")
         arr = self if inplace else self.copy()
+        limit = libalgos.validate_limit(None, limit=limit)
 
         if not self._can_hold_na:
             return [arr]
@@ -461,37 +462,12 @@ def fillna(
                     downcast=None,
                 )
         else:
-            # print("hello")
+            # TODO: Verify that this works for EAs
             return [
                 self.make_block_same_class(
                     values=self.values.fillna(value=value, limit=limit)
                 )
             ]
-        # mask = isna(self.values)
-        # mask, noop = validate_putmask(self.values, mask)
-        #
-        # if limit is not None:
-        #     #limit = libalgos.validate_limit(None, limit=limit)
-        #     #mask[mask.cumsum(self.ndim - 1) > limit] = False
-        #
-        # if self._can_hold_element(value):
-        #     putmask_inplace(arr.values, mask, value)
-        #     return arr._maybe_downcast([arr], downcast)
-        #
-        # if noop:
-        #    # we can't process the value, but nothing to do
-        #    return arr
-        #
-        # elif self.ndim == 1 or self.shape[0] == 1:
-        #     blk = self.coerce_to_target_dtype(value)
-        #     # bc we have already cast, inplace=True may avoid an extra copy
-        #     return blk.fillna(value, limit=limit, inplace=True, downcast=None)
-        #
-        # else:
-        #     # operate column-by-column
-        #     return self.split_and_operate(
-        #         type(self).fillna, value, limit=limit, inplace=inplace, downcast=None
-        #     )
 
     @final
     def _split(self) -> list[Block]: