BUG: Fix sparse where

Author: Licht-T

pandas/core/internals.py

@@ -27,6 +27,7 @@
     is_bool_dtype,
     is_object_dtype,
     is_datetimelike_v_numeric,
+    is_complex_dtype,
     is_float_dtype, is_numeric_dtype,
     is_numeric_v_string_like, is_extension_type,
     is_list_like,
@@ -46,7 +47,6 @@
     find_common_type)
 from pandas.core.dtypes.missing import (
     isna, notna, array_equivalent,
-    na_value_for_dtype,
     _isna_compat,
     is_null_datelike_scalar)
 import pandas.core.dtypes.concat as _concat
@@ -444,18 +444,14 @@ def make_a_block(nv, ref_loc):
                 except (AttributeError, NotImplementedError):
                     pass
 
-                if isinstance(self, SparseBlock):
-                    block = self.make_block_same_class(values=nv,
-                                                       placement=ref_loc,
-                                                       fastpath=True)
-                else:
-                    block = self.make_block(values=nv,
-                                            placement=ref_loc,
-                                            fastpath=True)
+                block = self.make_block(values=nv,
+                                        placement=ref_loc,
+                                        fastpath=True)
+
             return block
 
         # ndim == 1
-        if self.ndim == 1 or isinstance(self, SparseBlock):
+        if self.ndim == 1:
             if mask.any():
                 nv = f(mask, new_values, None)
             else:
@@ -1392,6 +1388,11 @@ def where(self, other, cond, align=True, raise_on_error=True,
         if hasattr(other, 'reindex_axis'):
             other = other.values
 
+        if is_scalar(other) or is_list_like(other):
+            fill_value = other
+        else:
+            fill_value = None
+
         if hasattr(cond, 'reindex_axis'):
             cond = cond.values
 
@@ -1405,7 +1406,8 @@ def where(self, other, cond, align=True, raise_on_error=True,
             raise ValueError("where must have a condition that is ndarray "
                              "like")
         else:
-            cond = cond.reshape(values.shape)
+            if values.ndim == 1:
+                cond = cond.flatten()
 
         # our where function
         def func(cond, values, other):
@@ -1443,7 +1445,8 @@ def func(cond, values, other):
                                  transpose=transpose)
             return self._maybe_downcast(blocks, 'infer')
 
-        if self._can_hold_na or self.ndim == 1:
+        if self._can_hold_element(fill_value) or \
+           self.is_sparse or self.ndim == 1:
 
             if transpose:
                 result = result.T
@@ -1452,13 +1455,12 @@ def func(cond, values, other):
             if try_cast:
                 result = self._try_cast_result(result)
 
-            if isinstance(self, SparseBlock):
-                fill_value = na_value_for_dtype(result.dtype)
-                return self.make_block_same_class(result,
-                                                  self.mgr_locs,
-                                                  fill_value=fill_value)
+            if isinstance(result, np.ndarray):
+                ndim = result.ndim
             else:
-                return self.make_block(result)
+                ndim = None
+
+            return self.make_block(result, ndim=ndim, fill_value=fill_value)
 
         # might need to separate out blocks
         axis = cond.ndim - 1
@@ -1582,9 +1584,6 @@ def _nanpercentile(values, q, axis, **kw):
         if is_scalar(result):
             return ax, self.make_block_scalar(result)
 
-        if isinstance(self, SparseBlock):
-            result = SparseArray(result.flatten())
-
         return ax, make_block(result,
                               placement=np.arange(len(result)),
                               ndim=ndim)
@@ -1736,6 +1735,7 @@ class FloatBlock(FloatOrComplexBlock):
     is_float = True
     _downcast_dtype = 'int64'
 
+    @classmethod
     def _can_hold_element(self, element):
         if is_list_like(element):
             element = np.asarray(element)
@@ -1786,6 +1786,7 @@ class ComplexBlock(FloatOrComplexBlock):
     __slots__ = ()
     is_complex = True
 
+    @classmethod
     def _can_hold_element(self, element):
         if is_list_like(element):
             element = np.array(element)
@@ -1948,6 +1949,7 @@ class BoolBlock(NumericBlock):
     is_bool = True
     _can_hold_na = False
 
+    @classmethod
     def _can_hold_element(self, element):
         if is_list_like(element):
             element = np.asarray(element)
@@ -2632,10 +2634,62 @@ class SparseBlock(NonConsolidatableMixIn, Block):
     is_sparse = True
     is_numeric = True
     _box_to_block_values = False
-    _can_hold_na = True
     _ftype = 'sparse'
     _holder = SparseArray
 
+    def __init__(self, values, placement, ndim=None, fastpath=False, **kwargs):
+        super(SparseBlock, self).__init__(values, placement,
+                                          ndim, fastpath,
+                                          **kwargs)
+
+        dtype = self.values.sp_values.dtype
+
+        if is_float_dtype(dtype):
+            self.is_float = True
+            self._can_hold_na = True
+        elif is_complex_dtype(dtype):
+            self.is_complex = True
+            self._can_hold_na = True
+        elif is_integer_dtype(dtype):
+            self.is_integer = True
+            self._can_hold_na = False
+        elif is_bool_dtype(dtype):
+            self.is_bool = True
+            self._can_hold_na = False
+        elif is_object_dtype(dtype):
+            self.is_object = True
+            self._can_hold_na = True
+        else:
+            self._can_hold_na = False
+
+    def _can_hold_element(self, element):
+        """ require the same dtype as ourselves """
+        dtype = self.values.sp_values.dtype
+
+        if is_bool_dtype(dtype):
+            return BoolBlock._can_hold_element(element)
+        elif is_integer_dtype(dtype):
+            if is_list_like(element):
+                element = np.array(element)
+                tipo = element.dtype.type
+                return (issubclass(tipo, np.integer) and
+                        not issubclass(tipo,
+                                       (np.datetime64,
+                                        np.timedelta64)) and
+                        dtype.itemsize >= element.dtype.itemsize)
+            return is_integer(element)
+        elif is_float_dtype(dtype):
+            return FloatBlock._can_hold_element(element)
+        elif is_complex_dtype(dtype):
+            return ComplexBlock._can_hold_element(element)
+        elif is_object_dtype(dtype):
+            return True
+        else:
+            return False
+
+    def coerce_to_target_dtype(self, other, copy=True):
+        return super(SparseBlock, self).coerce_to_target_dtype(other, copy)
+
     @property
     def shape(self):
         return (len(self.mgr_locs), self.sp_index.length)
@@ -2696,6 +2750,20 @@ def copy(self, deep=True, mgr=None):
                                           kind=self.kind, copy=deep,
                                           placement=self.mgr_locs)
 
+    def make_block(self, values, placement=None,
+                   ndim=None, fill_value=None, **kwargs):
+        """
+        Create a new block, with type inference propagate any values that are
+        not specified
+        """
+        if fill_value is not None and isinstance(values, SparseArray):
+            values = SparseArray(values.to_dense(), fill_value=fill_value,
+                                 kind=values.kind, dtype=values.dtype)
+
+        return super(SparseBlock, self).make_block(values, placement=placement,
+                                                   ndim=ndim, fill_value=None,
+                                                   **kwargs)
+
     def make_block_same_class(self, values, placement, sparse_index=None,
                               kind=None, dtype=None, fill_value=None,
                               copy=False, fastpath=True, **kwargs):
@@ -2726,7 +2794,7 @@ def make_block_same_class(self, values, placement, sparse_index=None,
 
         new_values = SparseArray(values, sparse_index=sparse_index,
                                  kind=kind or self.kind, dtype=dtype,
-                                 fill_value=fill_value, copy=copy)
+                                 fill_value=fill_value, copy=copy).flatten()
         return self.make_block(new_values, fastpath=fastpath,
                                placement=placement)
 
@@ -2792,9 +2860,15 @@ def sparse_reindex(self, new_index):
         return self.make_block_same_class(values, sparse_index=new_index,
                                           placement=self.mgr_locs)
 
+    def _try_coerce_result(self, result):
+        """ reverse of try_coerce_args """
+        if isinstance(result, np.ndarray):
+            result = SparseArray(result.flatten(), kind=self.kind)
+        return result
+
 
 def make_block(values, placement, klass=None, ndim=None, dtype=None,
-               fastpath=False):
+               fastpath=False, **kwargs):
     if klass is None:
         dtype = dtype or values.dtype
         vtype = dtype.type

pandas/core/sparse/frame.py

@@ -334,6 +334,15 @@ def copy(self, deep=True):
         result._default_kind = self._default_kind
         return result
 
+    def where(self, cond, other=np.nan, inplace=False, axis=None, level=None,
+              try_cast=False, raise_on_error=True):
+        result = super(SparseDataFrame, self).where(cond, other,
+                                                    inplace, axis,
+                                                    level, try_cast,
+                                                    raise_on_error)
+        result._default_fill_value = other
+        return result
+
     @property
     def default_fill_value(self):
         return self._default_fill_value