BUG: MaskedArray._quantile match non-nullable behavior (#46282)

Author: jbrockmendel

pandas/core/array_algos/quantile.py

@@ -74,6 +74,14 @@ def quantile_with_mask(
 
  Quantile is computed along axis=1.
  """
+ assert values.shape == mask.shape
+ if values.ndim == 1:
+ # unsqueeze, operate, re-squeeze
+ values = np.atleast_2d(values)
+ mask = np.atleast_2d(mask)
+ res_values = quantile_with_mask(values, mask, fill_value, qs, interpolation)
+ return res_values[0]
+
  assert values.ndim == 2
 
  is_empty = values.shape[1] == 0
@@ -189,7 +197,18 @@ def _nanpercentile(
  _nanpercentile_1d(val, m, qs, na_value, interpolation=interpolation)
  for (val, m) in zip(list(values), list(mask))
  ]
- result = np.array(result, dtype=values.dtype, copy=False).T
+ if values.dtype.kind == "f":
+ # preserve itemsize
+ result = np.array(result, dtype=values.dtype, copy=False).T
+ else:
+ result = np.array(result, copy=False).T
+ if (
+ result.dtype != values.dtype
+ and (result == result.astype(values.dtype, copy=False)).all()
+ ):
+ # e.g. values id integer dtype and result is floating dtype,
+ # only cast back to integer dtype if result values are all-integer.
+ result = result.astype(values.dtype, copy=False)
  return result
  else:
  return np.percentile(

pandas/core/arrays/_mixins.py

@@ -473,21 +473,14 @@ def _quantile(
  ) -> NDArrayBackedExtensionArrayT:
  # TODO: disable for Categorical if not ordered?
 
- # asarray needed for Sparse, see GH#24600
  mask = np.asarray(self.isna())
- mask = np.atleast_2d(mask)
-
- arr = np.atleast_2d(self._ndarray)
+ arr = self._ndarray
  fill_value = self._internal_fill_value
 
  res_values = quantile_with_mask(arr, mask, fill_value, qs, interpolation)
- res_values = self._cast_quantile_result(res_values)
- result = self._from_backing_data(res_values)
- if self.ndim == 1:
- assert result.shape == (1, len(qs)), result.shape
- result = result[0]
 
- return result
+ res_values = self._cast_quantile_result(res_values)
+ return self._from_backing_data(res_values)
 
  # TODO: see if we can share this with other dispatch-wrapping methods
  def _cast_quantile_result(self, res_values: np.ndarray) -> np.ndarray:

pandas/core/arrays/base.py

@@ -1586,25 +1586,12 @@ def _quantile(
  -------
  same type as self
  """
- # asarray needed for Sparse, see GH#24600
  mask = np.asarray(self.isna())
- mask = np.atleast_2d(mask)
-
- arr = np.atleast_2d(np.asarray(self))
+ arr = np.asarray(self)
  fill_value = np.nan
 
  res_values = quantile_with_mask(arr, mask, fill_value, qs, interpolation)
-
- if self.ndim == 2:
- # i.e. DatetimeArray
- result = type(self)._from_sequence(res_values)
-
- else:
- # shape[0] should be 1 as long as EAs are 1D
- assert res_values.shape == (1, len(qs)), res_values.shape
- result = type(self)._from_sequence(res_values[0])
-
- return result
+ return type(self)._from_sequence(res_values)
 
  def _mode(self: ExtensionArrayT, dropna: bool = True) -> ExtensionArrayT:
  """

pandas/core/arrays/categorical.py

@@ -351,7 +351,6 @@ class Categorical(NDArrayBackedExtensionArray, PandasObject, ObjectStringArrayMi
  # For comparisons, so that numpy uses our implementation if the compare
  # ops, which raise
  __array_priority__ = 1000
- _internal_fill_value = -1
  # tolist is not actually deprecated, just suppressed in the __dir__
  _hidden_attrs = PandasObject._hidden_attrs | frozenset(["tolist"])
  _typ = "categorical"
@@ -476,6 +475,13 @@ def dtype(self) -> CategoricalDtype:
  """
  return self._dtype
 
+ @property
+ def _internal_fill_value(self) -> int:
+ # using the specific numpy integer instead of python int to get
+ # the correct dtype back from _quantile in the all-NA case
+ dtype = self._ndarray.dtype
+ return dtype.type(-1)
+
  @property
  def _constructor(self) -> type[Categorical]:
  return Categorical
@@ -2300,7 +2306,7 @@ def unique(self):
 
  def _cast_quantile_result(self, res_values: np.ndarray) -> np.ndarray:
  # make sure we have correct itemsize for resulting codes
- res_values = coerce_indexer_dtype(res_values, self.dtype.categories)
+ assert res_values.dtype == self._ndarray.dtype
  return res_values
 
  def equals(self, other: object) -> bool:

pandas/core/arrays/masked.py

@@ -958,7 +958,7 @@ def equals(self, other) -> bool:
 
  def _quantile(
  self: BaseMaskedArrayT, qs: npt.NDArray[np.float64], interpolation: str
- ) -> BaseMaskedArrayT:
+ ) -> BaseMaskedArray:
  """
  Dispatch to quantile_with_mask, needed because we do not have
  _from_factorized.
@@ -967,29 +967,30 @@ def _quantile(
  -----
  We assume that all impacted cases are 1D-only.
  """
- mask = np.atleast_2d(np.asarray(self.isna()))
- npvalues: np.ndarray = np.atleast_2d(np.asarray(self))
-
  res = quantile_with_mask(
- npvalues,
- mask=mask,
- fill_value=self.dtype.na_value,
+ self._data,
+ mask=self._mask,
+ # TODO(GH#40932): na_value_for_dtype(self.dtype.numpy_dtype)
+ # instead of np.nan
+ fill_value=np.nan,
  qs=qs,
  interpolation=interpolation,
  )
- assert res.ndim == 2
- assert res.shape[0] == 1
- res = res[0]
- try:
- out = type(self)._from_sequence(res, dtype=self.dtype)
- except TypeError:
- # GH#42626: not able to safely cast Int64
- # for floating point output
- # error: Incompatible types in assignment (expression has type
- # "ndarray[Any, dtype[floating[_64Bit]]]", variable has type
- # "BaseMaskedArrayT")
- out = np.asarray(res, dtype=np.float64) # type: ignore[assignment]
- return out
+
+ if self._hasna:
+ # Our result mask is all-False unless we are all-NA, in which
+ # case it is all-True.
+ if self.ndim == 2:
+ # I think this should be out_mask=self.isna().all(axis=1)
+ # but am holding off until we have tests
+ raise NotImplementedError
+ elif self.isna().all():
+ out_mask = np.ones(res.shape, dtype=bool)
+ else:
+ out_mask = np.zeros(res.shape, dtype=bool)
+ else:
+ out_mask = np.zeros(res.shape, dtype=bool)
+ return self._maybe_mask_result(res, mask=out_mask)
 
  # ------------------------------------------------------------------
  # Reductions

pandas/core/arrays/period.py

@@ -698,11 +698,16 @@ def fillna(self, value=None, method=None, limit=None) -> PeriodArray:
  return result.view(self.dtype) # type: ignore[return-value]
  return super().fillna(value=value, method=method, limit=limit)
 
- # TODO: alternately could override _quantile like searchsorted
- def _cast_quantile_result(self, res_values: np.ndarray) -> np.ndarray:
- # quantile_with_mask may return float64 instead of int64, in which
- # case we need to cast back
- return res_values.astype(np.int64, copy=False)
+ def _quantile(
+ self: PeriodArray,
+ qs: npt.NDArray[np.float64],
+ interpolation: str,
+ ) -> PeriodArray:
+ # dispatch to DatetimeArray implementation
+ dtres = self.view("M8[ns]")._quantile(qs, interpolation)
+ # error: Incompatible return value type (got "Union[ExtensionArray,
+ # ndarray[Any, Any]]", expected "PeriodArray")
+ return dtres.view(self.dtype) # type: ignore[return-value]
 
  # ------------------------------------------------------------------
  # Arithmetic Methods

pandas/core/arrays/sparse/array.py

@@ -75,6 +75,7 @@
 
 from pandas.core import arraylike
 import pandas.core.algorithms as algos
+from pandas.core.array_algos.quantile import quantile_with_mask
 from pandas.core.arraylike import OpsMixin
 from pandas.core.arrays import ExtensionArray
 from pandas.core.arrays.sparse.dtype import SparseDtype
@@ -890,10 +891,27 @@ def value_counts(self, dropna: bool = True) -> Series:
  return Series(counts, index=keys)
 
  def _quantile(self, qs: npt.NDArray[np.float64], interpolation: str):
+
+ if self._null_fill_value or self.sp_index.ngaps == 0:
+ # We can avoid densifying
+ npvalues = self.sp_values
+ mask = np.zeros(npvalues.shape, dtype=bool)
+ else:
+ npvalues = self.to_numpy()
+ mask = self.isna()
+
+ fill_value = na_value_for_dtype(npvalues.dtype, compat=False)
+ res_values = quantile_with_mask(
+ npvalues,
+ mask,
+ fill_value,
+ qs,
+ interpolation,
+ )
+
  # Special case: the returned array isn't _really_ sparse, so we don't
  # wrap it in a SparseArray
- result = super()._quantile(qs, interpolation)
- return np.asarray(result)
+ return res_values
 
  # --------
  # Indexing

pandas/tests/frame/methods/test_quantile.py

@@ -665,9 +665,14 @@ def test_quantile_ea(self, obj, index):
  qs = [0.5, 0, 1]
  result = self.compute_quantile(obj, qs)
 
+ exp_dtype = index.dtype
+ if index.dtype == "Int64":
+ # match non-nullable casting behavior
+ exp_dtype = "Float64"
+
  # expected here assumes len(index) == 9
  expected = Series(
- [index[4], index[0], index[-1]], dtype=index.dtype, index=qs, name="A"
+ [index[4], index[0], index[-1]], dtype=exp_dtype, index=qs, name="A"
  )
  expected = type(obj)(expected)
 
@@ -712,6 +717,8 @@ def test_quantile_ea_all_na(self, obj, index):
 
  expected = index.take([-1, -1, -1], allow_fill=True, fill_value=index._na_value)
  expected = Series(expected, index=qs, name="A")
+ if expected.dtype == "Int64":
+ expected = expected.astype("Float64")
  expected = type(obj)(expected)
  tm.assert_equal(result, expected)
 
@@ -726,7 +733,11 @@ def test_quantile_ea_scalar(self, obj, index):
  qs = 0.5
  result = self.compute_quantile(obj, qs)
 
- expected = Series({"A": index[4]}, dtype=index.dtype, name=0.5)
+ exp_dtype = index.dtype
+ if index.dtype == "Int64":
+ exp_dtype = "Float64"
+
+ expected = Series({"A": index[4]}, dtype=exp_dtype, name=0.5)
  if isinstance(obj, Series):
  expected = expected["A"]
  assert result == expected

pandas/tests/series/methods/test_quantile.py

@@ -222,4 +222,6 @@ def test_quantile_empty(self):
  def test_quantile_dtypes(self, dtype):
  result = Series([1, 2, 3], dtype=dtype).quantile(np.arange(0, 1, 0.25))
  expected = Series(np.arange(1, 3, 0.5), index=np.arange(0, 1, 0.25))
+ if dtype == "Int64":
+ expected = expected.astype("Float64")
  tm.assert_series_equal(result, expected)