BUG: Fixes plotting with nullable integers (#32073)

pandas/core/dtypes/cast.py

@@ -3,7 +3,7 @@
 """
 
 from datetime import date, datetime, timedelta
-from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type, Union
 
 import numpy as np
 
@@ -76,6 +76,7 @@
 if TYPE_CHECKING:
     from pandas import Series
     from pandas.core.arrays import ExtensionArray  # noqa: F401
+    from pandas.core.indexes.base import Index  # noqa: F401
 
 _int8_max = np.iinfo(np.int8).max
 _int16_max = np.iinfo(np.int16).max
@@ -1747,3 +1748,35 @@ def validate_numeric_casting(dtype: np.dtype, value):
     ):
         if is_bool(value):
             raise ValueError("Cannot assign bool to float/integer series")
+
+
+def safe_convert_to_ndarray(values: Union[ArrayLike, "Index"]) -> np.ndarray:
+    """
+    Converts values to ndarray with special handling for extension arrays.
+
+    Cast to ndarray but tries to avoid returning an array of `object` dtype.
+    Nullable integer and boolean arrays will be cast to float, and datetime
+    arrays with timezone information will lose their timezone information.
+
+    Parameters
+    ----------
+    values : Union[ArrayLike, Index]
+        Values to be converted to ndarray.
+
+    Returns
+    -------
+    converted_values : np.ndarray
+        Values cast to np.ndarray.
+    """
+    if hasattr(values, "dtype") and is_extension_array_dtype(values.dtype):
+        if is_integer_dtype(values.dtype):
+            converted_values = values.to_numpy(dtype=float, na_value=np.nan)
+        elif is_bool_dtype(values.dtype):
+            converted_values = values.to_numpy(dtype=float, na_value=np.nan)
+        elif is_datetime64tz_dtype(values.dtype):
+            converted_values = np.asarray(values.tz_localize(tz=None))
+        else:
+            converted_values = np.asarray(values)
+    else:
+        converted_values = np.asarray(values)
+    return converted_values

pandas/core/groupby/groupby.py

@@ -41,12 +41,10 @@ class providing the base-class of operations.
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import Appender, Substitution, cache_readonly, doc
 
-from pandas.core.dtypes.cast import maybe_cast_result
+from pandas.core.dtypes.cast import maybe_cast_result, safe_convert_to_ndarray
 from pandas.core.dtypes.common import (
     ensure_float,
-    is_bool_dtype,
     is_datetime64_dtype,
-    is_extension_array_dtype,
     is_integer_dtype,
     is_numeric_dtype,
     is_object_dtype,
@@ -2052,14 +2050,11 @@ def pre_processor(vals: np.ndarray) -> Tuple[np.ndarray, Optional[Type]]:
 
             inference = None
             if is_integer_dtype(vals.dtype):
-                if is_extension_array_dtype(vals.dtype):
-                    vals = vals.to_numpy(dtype=float, na_value=np.nan)
                 inference = np.int64
-            elif is_bool_dtype(vals.dtype) and is_extension_array_dtype(vals.dtype):
-                vals = vals.to_numpy(dtype=float, na_value=np.nan)
             elif is_datetime64_dtype(vals.dtype):
                 inference = "datetime64[ns]"
-                vals = np.asarray(vals).astype(float)
+
+            vals = safe_convert_to_ndarray(vals)
 
             return vals, inference
 

pandas/plotting/_matplotlib/core.py

@@ -9,6 +9,7 @@
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import cache_readonly
 
+from pandas.core.dtypes.cast import safe_convert_to_ndarray
 from pandas.core.dtypes.common import (
     is_float,
     is_hashable,
@@ -421,7 +422,7 @@ def _compute_plot_data(self):
         # np.ndarray before plot.
         numeric_data = numeric_data.copy()
         for col in numeric_data:
-            numeric_data[col] = np.asarray(numeric_data[col])
+            numeric_data[col] = safe_convert_to_ndarray(numeric_data[col].values)
 
         self.data = numeric_data
 

pandas/tests/dtypes/cast/test_safe_convert_to_ndarray.py

@@ -0,0 +1,51 @@
+import numpy as np
+import pytest
+
+from pandas.core.dtypes.cast import safe_convert_to_ndarray
+
+import pandas as pd
+import pandas._testing as tm
+
+
+@pytest.mark.parametrize(
+    "values, expected",
+    [
+        (pd.Series([1, 2, 3], dtype=int), np.array([1, 2, 3], dtype=int)),
+        (
+            # Nullable integer type cast to float to handle missing values
+            pd.Series([1, np.NaN, 3], dtype="Int64"),
+            np.array([1, np.NaN, 3], dtype=float),
+        ),
+        (
+            # Nullable boolean type cast to float to handle missing values
+            pd.Series([True, np.NaN, False], dtype="boolean"),
+            np.array([1.0, np.NaN, 0.0], dtype=float),
+        ),
+        (
+            # Normal datetime cast not changed
+            pd.to_datetime([2001, None, 2003], format="%Y"),
+            np.array(["2001", "NaT", "2003"], dtype="datetime64").astype(
+                "datetime64[ns]"
+            ),
+        ),
+        (
+            # Extended datetime should be downcast to normal datetime
+            pd.to_datetime([2001, None, 2003], format="%Y", utc=True),
+            np.array(["2001", "NaT", "2003"], dtype="datetime64").astype(
+                "datetime64[ns]"
+            ),
+        ),
+        (
+            # Downcast to naive datetime should result in local dates, not UTC
+            pd.to_datetime([2001, None, 2003], format="%Y").tz_localize(
+                tz="US/Eastern"
+            ),
+            np.array(["2001", "NaT", "2003"], dtype="datetime64").astype(
+                "datetime64[ns]"
+            ),
+        ),
+    ],
+)
+def test_safe_convert_to_ndarray(values, expected):
+    result = safe_convert_to_ndarray(values)
+    tm.assert_numpy_array_equal(result, expected)

pandas/tests/plotting/test_frame.py

@@ -3419,6 +3419,23 @@ def test_xlabel_ylabel_dataframe_subplots(
         assert all(ax.get_ylabel() == str(new_label) for ax in axes)
         assert all(ax.get_xlabel() == str(new_label) for ax in axes)
 
+    def test_nullable_int_plot(self):
+        # GH 32073
+        dates = ["2008", "2009", None, "2011", "2012"]
+        df = pd.DataFrame(
+            {
+                "A": [1, 2, 3, 4, 5],
+                "B": [7, 5, np.nan, 3, 2],
+                "C": pd.to_datetime(dates, format="%Y"),
+                "D": pd.to_datetime(dates, format="%Y", utc=True),
+            }
+        )
+
+        _check_plot_works(df.plot, x="A", y="B")
+        _check_plot_works(df[["A", "B"]].astype("Int64").plot, x="A", y="B")
+        _check_plot_works(df[["A", "C"]].plot, x="A", y="C")
+        _check_plot_works(df[["A", "D"]].plot, x="A", y="D")
+
 
 def _generate_4_axes_via_gridspec():
     import matplotlib.pyplot as plt