Skip NaN in quantile() and median()

changelog.d/quantile-skipna.yaml

@@ -0,0 +1,7 @@
+- bump: minor
+  changes:
+    fixed:
+    - quantile() and median() now skip NaN values by default, so NaN weight no
+      longer inflates the cumulative distribution and pushes the cutoff up.
+    added:
+    - skipna argument on MicroSeries.quantile and MicroSeries.median.

microdf/microseries.py

@@ -306,7 +306,7 @@ def corr(self, other, *args, **kwargs):
         )
         return super().corr(other, *args, **kwargs)
 
-    def quantile(self, q: np.array) -> pd.Series:
+    def quantile(self, q: np.array, skipna: bool = True) -> pd.Series:
         """Calculates weighted quantiles of the MicroSeries.
 
         Uses the inverse CDF method: the q-th quantile is the smallest
@@ -315,13 +315,25 @@ def quantile(self, q: np.array) -> pd.Series:
 
         :param q: Quantile(s) to calculate, must be in [0, 1].
         :type q: float or np.array
+        :param skipna: Exclude NaN values (default True). NaN sorts to the
+            end of the array, so leaving NaN rows in would let their weight
+            inflate the cumulative distribution and push the cutoff upward.
+            If False, NaN is returned whenever any value is NaN.
+        :type skipna: bool
 
         :return: Weighted quantile value(s).
         :rtype: float or pd.Series
         """
         values = np.array(self._values)
         quantiles = np.atleast_1d(q)
         sample_weight = np.array(self.weights)
+        na_mask = pd.isna(values)
+        if not skipna and na_mask.any():
+            return (
+                np.nan
+                if np.array(q).shape == ()
+                else pd.Series(np.full(len(quantiles), np.nan), index=quantiles)
+            )
         assert np.all(quantiles >= 0) and np.all(quantiles <= 1), (
             "quantiles should be in [0, 1]"
         )
@@ -330,7 +342,10 @@ def quantile(self, q: np.array) -> pd.Series:
         # that should have been skipped by the inverse CDF. E.g.
         # MicroSeries([10, 20, 30], weights=[0, 1, 1]).quantile(0)
         # returned 10 instead of 20.
-        nonzero = sample_weight > 0
+        # Drop NaN rows for the same reason: NaN sorts last, so its weight
+        # would inflate the cumulative distribution and push the cutoff up
+        # (median of [1, nan, 3] returned 3.0 instead of 1.0).
+        nonzero = (sample_weight > 0) & ~na_mask
         if not nonzero.any():
             return (
                 np.nan
@@ -355,13 +370,15 @@ def quantile(self, q: np.array) -> pd.Series:
         return pd.Series(result, index=quantiles)
 
     @scalar_function
-    def median(self) -> float:
+    def median(self, skipna: bool = True) -> float:
         """Calculates the weighted median of the MicroSeries.
 
+        :param skipna: Exclude NaN values (default True).
+        :type skipna: bool
         :returns: The weighted median of a DataFrame's column.
         :rtype: float
         """
-        return self.quantile(0.5)
+        return self.quantile(0.5, skipna=skipna)
 
     @scalar_function
     def gini(self, negatives: Optional[str] = None) -> float:

microdf/tests/test_microseries_dataframe.py

@@ -815,3 +815,34 @@ def test_rank_ties_share_bucket() -> None:
     # existing ``test_rank`` expectations hold.
     s = mdf.MicroSeries([1, 2, 3], weights=[4, 5, 6])
     np.testing.assert_array_equal(s.rank().values, [4, 9, 15])
+
+
+def test_quantile_skips_nan():
+    """NaN weight must not inflate the cumulative distribution.
+
+    Dropping a NaN row should give the same answer as never having had
+    it: the inverse-CDF quantile of [1, nan, 3] equals that of [1, 3].
+    """
+    with_nan = mdf.MicroSeries([1.0, np.nan, 3.0], weights=[1, 1, 1])
+    without_nan = mdf.MicroSeries([1.0, 3.0], weights=[1, 1])
+    assert with_nan.median() == without_nan.median()
+    assert with_nan.quantile(0.5) == without_nan.quantile(0.5)
+
+    q = [0.25, 0.5, 0.75]
+    np.testing.assert_array_equal(
+        mdf.MicroSeries([1.0, np.nan, 3.0, 5.0], weights=[1, 1, 1, 1]).quantile(q),
+        mdf.MicroSeries([1.0, 3.0, 5.0], weights=[1, 1, 1]).quantile(q),
+    )
+
+
+def test_quantile_skipna_false_propagates_nan():
+    """Skipna=False returns NaN when any value is NaN, like mean/var."""
+    s = mdf.MicroSeries([1.0, np.nan, 3.0], weights=[1, 1, 1])
+    assert np.isnan(s.quantile(0.5, skipna=False))
+    assert np.isnan(s.median(skipna=False))
+    assert s.quantile([0.25, 0.75], skipna=False).isna().all()
+
+
+def test_quantile_all_nan_returns_nan():
+    s = mdf.MicroSeries([np.nan, np.nan], weights=[1, 1])
+    assert np.isnan(s.median())