Fix: Make PolyFit stat robust to missing data

Fixes #2992

Author: mwaskom

seaborn/_stats/regression.py

@@ -38,7 +38,10 @@ def _fit_predict(self, data):
 
     def __call__(self, data, groupby, orient, scales):
 
-        return groupby.apply(data, self._fit_predict)
+        return (
+            groupby
+            .apply(data.dropna(subset=["x", "y"]), self._fit_predict)
+        )
 
 
 @dataclass

tests/_stats/test_regression.py

@@ -4,6 +4,7 @@
 
 import pytest
 from numpy.testing import assert_array_equal, assert_array_almost_equal
+from pandas.testing import assert_frame_equal
 
 from seaborn._core.groupby import GroupBy
 from seaborn._stats.regression import PolyFit
@@ -50,3 +51,11 @@ def test_one_grouper(self, df):
             grid = np.linspace(part["x"].min(), part["x"].max(), gridsize)
             assert_array_equal(part["x"], grid)
             assert part["y"].diff().diff().dropna().abs().gt(0).all()
+
+    def test_missing_data(self, df):
+
+        groupby = GroupBy(["group"])
+        df.iloc[5:10] = np.nan
+        res1 = PolyFit()(df[["x", "y"]], groupby, "x", {})
+        res2 = PolyFit()(df[["x", "y"]].dropna(), groupby, "x", {})
+        assert_frame_equal(res1, res2)