`sns.heatmap` crashes with nullable data types (`Int64` and `Float64`)

[mojones]

Minimal example:

df = pd.DataFrame({
    'foo' : [1,1,1,2,2,2],
    'bar' : [3,3,3,4,4,4],
    'baz' : [1,2,3,4,5,6]
    })

# works as expected
sns.heatmap(df)

# crashes
sns.heatmap(df.astype('Int64'))

Not sure if this is a bug or expected behaviour, I originally encountered this when using a dataset containing missing data that I had manually set the dtype to Int64 on import.

[mwaskom]

If you look at the traceback, this is coming directly out of matplotlib:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In [25], line 1
----> 1 plt.pcolormesh(df.astype('Int64'))

File ~/miniconda/envs/py310/lib/python3.10/site-packages/matplotlib/pyplot.py:2689, in pcolormesh(alpha, norm, cmap, vmin, vmax, shading, antialiased, data, *args, **kwargs)
   2684 @_copy_docstring_and_deprecators(Axes.pcolormesh)
   2685 def pcolormesh(
   2686         *args, alpha=None, norm=None, cmap=None, vmin=None,
   2687         vmax=None, shading=None, antialiased=False, data=None,
   2688         **kwargs):
-> 2689     __ret = gca().pcolormesh(
   2690         *args, alpha=alpha, norm=norm, cmap=cmap, vmin=vmin,
   2691         vmax=vmax, shading=shading, antialiased=antialiased,
   2692         **({"data": data} if data is not None else {}), **kwargs)
   2693     sci(__ret)
   2694     return __ret

File ~/miniconda/envs/py310/lib/python3.10/site-packages/matplotlib/__init__.py:1423, in _preprocess_data.<locals>.inner(ax, data, *args, **kwargs)
   1420 @functools.wraps(func)
   1421 def inner(ax, *args, data=None, **kwargs):
   1422     if data is None:
-> 1423         return func(ax, *map(sanitize_sequence, args), **kwargs)
   1425     bound = new_sig.bind(ax, *args, **kwargs)
   1426     auto_label = (bound.arguments.get(label_namer)
   1427                   or bound.kwargs.get(label_namer))

File ~/miniconda/envs/py310/lib/python3.10/site-packages/matplotlib/axes/_axes.py:6139, in Axes.pcolormesh(self, alpha, norm, cmap, vmin, vmax, shading, antialiased, *args, **kwargs)
   6135 C = C.ravel()
   6137 kwargs.setdefault('snap', mpl.rcParams['pcolormesh.snap'])
-> 6139 collection = mcoll.QuadMesh(
   6140     coords, antialiased=antialiased, shading=shading,
   6141     array=C, cmap=cmap, norm=norm, alpha=alpha, **kwargs)
   6142 collection._scale_norm(norm, vmin, vmax)
   6143 self._pcolor_grid_deprecation_helper()

File ~/miniconda/envs/py310/lib/python3.10/site-packages/matplotlib/collections.py:1988, in QuadMesh.__init__(self, *args, **kwargs)
   1985 self._bbox.update_from_data_xy(self._coordinates.reshape(-1, 2))
   1986 # super init delayed after own init because array kwarg requires
   1987 # self._coordinates and self._shading
-> 1988 super().__init__(**kwargs)
   1989 self.set_mouseover(False)

File ~/miniconda/envs/py310/lib/python3.10/site-packages/matplotlib/_api/deprecation.py:454, in make_keyword_only.<locals>.wrapper(*args, **kwargs)
    448 if len(args) > name_idx:
    449     warn_deprecated(
    450         since, message="Passing the %(name)s %(obj_type)s "
    451         "positionally is deprecated since Matplotlib %(since)s; the "
    452         "parameter will become keyword-only %(removal)s.",
    453         name=name, obj_type=f"parameter of {func.__name__}()")
--> 454 return func(*args, **kwargs)

File ~/miniconda/envs/py310/lib/python3.10/site-packages/matplotlib/collections.py:202, in Collection.__init__(self, edgecolors, facecolors, linewidths, linestyles, capstyle, joinstyle, antialiaseds, offsets, offset_transform, norm, cmap, pickradius, hatch, urls, zorder, **kwargs)
    199 self._offset_transform = offset_transform
    201 self._path_effects = None
--> 202 self._internal_update(kwargs)
    203 self._paths = None

File ~/miniconda/envs/py310/lib/python3.10/site-packages/matplotlib/artist.py:1186, in Artist._internal_update(self, kwargs)
   1179 def _internal_update(self, kwargs):
   1180     """
   1181     Update artist properties without prenormalizing them, but generating
   1182     errors as if calling `set`.
   1183 
   1184     The lack of prenormalization is to maintain backcompatibility.
   1185     """
-> 1186     return self._update_props(
   1187         kwargs, "{cls.__name__}.set() got an unexpected keyword argument "
   1188         "{prop_name!r}")

File ~/miniconda/envs/py310/lib/python3.10/site-packages/matplotlib/artist.py:1162, in Artist._update_props(self, props, errfmt)
   1159             if not callable(func):
   1160                 raise AttributeError(
   1161                     errfmt.format(cls=type(self), prop_name=k))
-> 1162             ret.append(func(v))
   1163 if ret:
   1164     self.pchanged()

File ~/miniconda/envs/py310/lib/python3.10/site-packages/matplotlib/collections.py:2050, in QuadMesh.set_array(self, A)
   2045     if faulty_data:
   2046         raise TypeError(
   2047             f"Dimensions of A {A.shape} are incompatible with "
   2048             f"X ({width}) and/or Y ({height})")
-> 2050 return super().set_array(A)

File ~/miniconda/envs/py310/lib/python3.10/site-packages/matplotlib/cm.py:502, in ScalarMappable.set_array(self, A)
    500 A = cbook.safe_masked_invalid(A, copy=True)
    501 if not np.can_cast(A.dtype, float, "same_kind"):
--> 502     raise TypeError(f"Image data of dtype {A.dtype} cannot be "
    503                     "converted to float")
    505 self._A = A

TypeError: Image data of dtype object cannot be converted to float

I'm sort of inclined to say that this is an issue that needs to get sorted out at the pandas <> matplotlib interface and that seaborn should not do anything specific to work around it.

[mwaskom]

(oops sorry the traceback I copied is what you get with plt.pcolormesh, but it's the same thing you get with heatmap just missing a frame or two within seaborn.)

[mojones]

Thanks, submitted as https://github.com/matplotlib/matplotlib/issues/23991