Improve handling of edge cases with nullable dtypes (#3394)

* Improve handling of numeric type edge cases

* Ensure numeric data after scaling in Plot

* Remove some back-compat flexibility around pd.NA

Author: mwaskom

seaborn/_core/plot.py

@@ -1392,11 +1392,11 @@ def _setup_scales(
                         spec_error = PlotSpecError._during("Scaling operation", var)
                         raise spec_error from err
 
-            # Now the transformed data series are complete, set update the layer data
+            # Now the transformed data series are complete, update the layer data
             for layer, new_series in zip(layers, transformed_data):
                 layer_df = layer["data"].frame
                 if var in layer_df:
-                    layer_df[var] = new_series
+                    layer_df[var] = pd.to_numeric(new_series)
 
     def _plot_layer(self, p: Plot, layer: Layer) -> None:
 

seaborn/_core/rules.py

@@ -74,6 +74,9 @@ def variable_type(
     if pd.isna(vector).all():
         return VarType("numeric")
 
+    # Now drop nulls to simplify further type inference
+    vector = vector.dropna()
+
     # Special-case binary/boolean data, allow caller to determine
     # This triggers a numpy warning when vector has strings/objects
     # https://github.com/numpy/numpy/issues/6784
@@ -94,7 +97,7 @@ def variable_type(
                 boolean_dtypes = ["bool"]
             boolean_vector = vector.dtype in boolean_dtypes
         else:
-            boolean_vector = bool(np.isin(vector.dropna(), [0, 1]).all())
+            boolean_vector = bool(np.isin(vector, [0, 1]).all())
         if boolean_vector:
             return VarType(boolean_type)
 

seaborn/_oldcore.py

@@ -1128,7 +1128,7 @@ def comp_data(self):
                             # it is similar to GH2419, but more complicated because
                             # supporting `order` in categorical plots is tricky
                             orig = orig[orig.isin(self.var_levels[var])]
-                    comp = pd.to_numeric(converter.convert_units(orig))
+                    comp = pd.to_numeric(converter.convert_units(orig)).astype(float)
                     if converter.get_scale() == "log":
                         comp = np.log10(comp)
                     parts.append(pd.Series(comp, orig.index, name=orig.name))
@@ -1505,6 +1505,9 @@ def variable_type(vector, boolean_type="numeric"):
     if pd.isna(vector).all():
         return VariableType("numeric")
 
+    # At this point, drop nans to simplify further type inference
+    vector = vector.dropna()
+
     # Special-case binary/boolean data, allow caller to determine
     # This triggers a numpy warning when vector has strings/objects
     # https://github.com/numpy/numpy/issues/6784
@@ -1517,7 +1520,7 @@ def variable_type(vector, boolean_type="numeric"):
         warnings.simplefilter(
             action='ignore', category=(FutureWarning, DeprecationWarning)
         )
-        if np.isin(vector.dropna(), [0, 1]).all():
+        if np.isin(vector, [0, 1]).all():
             return VariableType(boolean_type)
 
     # Defer to positive pandas tests

tests/_core/test_rules.py

@@ -38,6 +38,12 @@ def test_variable_type():
     s = pd.Series([pd.NA, pd.NA])
     assert variable_type(s) == "numeric"
 
+    s = pd.Series([1, 2, pd.NA], dtype="Int64")
+    assert variable_type(s) == "numeric"
+
+    s = pd.Series([1, 2, pd.NA], dtype=object)
+    assert variable_type(s) == "numeric"
+
     s = pd.Series(["1", "2", "3"])
     assert variable_type(s) == "categorical"
 

tests/test_core.py

@@ -23,16 +23,9 @@
     categorical_order,
 )
 from seaborn.utils import desaturate
-
 from seaborn.palettes import color_palette
 
 
-try:
-    from pandas import NA as PD_NA
-except ImportError:
-    PD_NA = None
-
-
 @pytest.fixture(params=[
     dict(x="x", y="y"),
     dict(x="t", y="y"),
@@ -1302,28 +1295,23 @@ def test_comp_data_category_order(self):
 
     @pytest.fixture(
         params=itertools.product(
-            [None, np.nan, PD_NA],
-            ["numeric", "category", "datetime"]
+            [None, np.nan, pd.NA],
+            ["numeric", "category", "datetime"],
         )
     )
-    @pytest.mark.parametrize(
-        "NA,var_type",
-    )
+    @pytest.mark.parametrize("NA,var_type")
     def comp_data_missing_fixture(self, request):
 
         # This fixture holds the logic for parameterizing
         # the following test (test_comp_data_missing)
 
         NA, var_type = request.param
 
-        if NA is None:
-            pytest.skip("No pandas.NA available")
-
         comp_data = [0, 1, np.nan, 2, np.nan, 1]
         if var_type == "numeric":
             orig_data = [0, 1, NA, 2, np.inf, 1]
         elif var_type == "category":
-            orig_data = ["a", "b", NA, "c", NA, "b"]
+            orig_data = ["a", "b", NA, "c", pd.NA, "b"]
         elif var_type == "datetime":
             # Use 1-based numbers to avoid issue on matplotlib<3.2
             # Could simplify the test a bit when we roll off that version
@@ -1343,6 +1331,7 @@ def test_comp_data_missing(self, comp_data_missing_fixture):
         ax = plt.figure().subplots()
         p._attach(ax)
         assert_array_equal(p.comp_data["x"], comp_data)
+        assert p.comp_data["x"].dtype == "float"
 
     def test_comp_data_duplicate_index(self):
 
@@ -1352,6 +1341,15 @@ def test_comp_data_duplicate_index(self):
         p._attach(ax)
         assert_array_equal(p.comp_data["x"], x)
 
+    def test_comp_data_nullable_dtype(self):
+
+        x = pd.Series([1, 2, 3, 4], dtype="Int64")
+        p = VectorPlotter(variables={"x": x})
+        ax = plt.figure().subplots()
+        p._attach(ax)
+        assert_array_equal(p.comp_data["x"], x)
+        assert p.comp_data["x"].dtype == "float"
+
     def test_var_order(self, long_df):
 
         order = ["c", "b", "a"]
@@ -1456,7 +1454,12 @@ def test_variable_type(self):
         assert variable_type(s) == "numeric"
 
         s = pd.Series([np.nan, np.nan])
-        # s = pd.Series([pd.NA, pd.NA])
+        assert variable_type(s) == "numeric"
+
+        s = pd.Series([pd.NA, pd.NA])
+        assert variable_type(s) == "numeric"
+
+        s = pd.Series([1, 2, pd.NA], dtype="Int64")
         assert variable_type(s) == "numeric"
 
         s = pd.Series(["1", "2", "3"])