Skip to content

Commit

Permalink
fix: use PyCapsule Interface instead of Dataframe Interchange Protocol
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed Nov 9, 2024
1 parent b4e5f8d commit cb86e7a
Show file tree
Hide file tree
Showing 5 changed files with 36 additions and 39 deletions.
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ dev = [
"mypy",
"pandas-stubs",
"pre-commit",
"pyarrow",
"flit",
]
docs = [
Expand Down
66 changes: 31 additions & 35 deletions seaborn/_core/data.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

from collections.abc import Mapping, Sized
from typing import cast
import warnings

import pandas as pd
from pandas import DataFrame
Expand Down Expand Up @@ -269,9 +268,9 @@ def _assign_variables(

def handle_data_source(data: object) -> pd.DataFrame | Mapping | None:
"""Convert the data source object to a common union representation."""
if isinstance(data, pd.DataFrame) or hasattr(data, "__dataframe__"):
if isinstance(data, pd.DataFrame) or hasattr(data, "__arrow_c_stream__"):
# Check for pd.DataFrame inheritance could be removed once
# minimal pandas version supports dataframe interchange (1.5.0).
# minimal pandas version supports PyCapsule Interface (2.2).
data = convert_dataframe_to_pandas(data)
elif data is not None and not isinstance(data, Mapping):
err = f"Data source must be a DataFrame or Mapping, not {type(data)!r}."
Expand All @@ -285,35 +284,32 @@ def convert_dataframe_to_pandas(data: object) -> pd.DataFrame:
if isinstance(data, pd.DataFrame):
return data

if not hasattr(pd.api, "interchange"):
msg = (
"Support for non-pandas DataFrame objects requires a version of pandas "
"that implements the DataFrame interchange protocol. Please upgrade "
"your pandas version or coerce your data to pandas before passing "
"it to seaborn."
)
raise TypeError(msg)

if _version_predates(pd, "2.0.2"):
msg = (
"DataFrame interchange with pandas<2.0.2 has some known issues. "
f"You are using pandas {pd.__version__}. "
"Continuing, but it is recommended to carefully inspect the results and to "
"consider upgrading."
)
warnings.warn(msg, stacklevel=2)

try:
# This is going to convert all columns in the input dataframe, even though
# we may only need one or two of them. It would be more efficient to select
# the columns that are going to be used in the plot prior to interchange.
# Solving that in general is a hard problem, especially with the objects
# interface where variables passed in Plot() may only be referenced later
# in Plot.add(). But noting here in case this seems to be a bottleneck.
return pd.api.interchange.from_dataframe(data)
except Exception as err:
msg = (
"Encountered an exception when converting data source "
"to a pandas DataFrame. See traceback above for details."
)
raise RuntimeError(msg) from err
if hasattr(data, '__arrow_c_stream__'):
try:
import pyarrow
except ImportError as err:
msg = "PyArrow is required for non-pandas Dataframe support."
raise RuntimeError(msg) from err
if _version_predates(pyarrow, '14.0.0'):
msg = "PyArrow>=14.0.0 is required for non-pandas Dataframe support."
raise RuntimeError(msg)
try:
# This is going to convert all columns in the input dataframe, even though
# we may only need one or two of them. It would be more efficient to select
# the columns that are going to be used in the plot prior to interchange.
# Solving that in general is a hard problem, especially with the objects
# interface where variables passed in Plot() may only be referenced later
# in Plot.add(). But noting here in case this seems to be a bottleneck.
return pyarrow.table(data).to_pandas()
except Exception as err:
msg = (
"Encountered an exception when converting data source "
"to a pandas DataFrame. See traceback above for details."
)
raise RuntimeError(msg) from err

msg = (
"Expected object which implements '__arrow_c_stream__' from the "
f"PyCapsule Interface, got: {type(data)}"
)
raise TypeError(msg)
2 changes: 1 addition & 1 deletion seaborn/_core/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,7 @@ def _resolve_positionals(

if (
isinstance(args[0], (abc.Mapping, pd.DataFrame))
or hasattr(args[0], "__dataframe__")
or hasattr(args[0], "__arrow_c_stream__")
):
if data is not None:
raise TypeError("`data` given by both name and position.")
Expand Down
2 changes: 1 addition & 1 deletion tests/_core/test_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -425,7 +425,7 @@ def test_data_interchange(self, mock_long_df, long_df):
)
def test_data_interchange_failure(self, mock_long_df):

mock_long_df._data = None # Break __dataframe__()
mock_long_df.__arrow_c_stream__ = lambda _x: 1 / 0 # Break __arrow_c_stream__()
with pytest.raises(RuntimeError, match="Encountered an exception"):
PlotData(mock_long_df, {"x": "x"})

Expand Down
4 changes: 2 additions & 2 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,8 +188,8 @@ class MockInterchangeableDataFrame:
def __init__(self, data):
self._data = data

def __dataframe__(self, *args, **kwargs):
return self._data.__dataframe__(*args, **kwargs)
def __arrow_c_stream__(self, *args, **kwargs):
return self._data.__arrow_c_stream__()


@pytest.fixture
Expand Down

0 comments on commit cb86e7a

Please sign in to comment.