feat: make plotly-express dataframe agnostic via narwhals

packages/python/plotly/optional-requirements.txt

@@ -39,6 +39,7 @@ ipython
 
 ## pandas deps for some matplotlib functionality ##
 pandas
+narwhals>=1.9.2
 
 ## scipy deps for some FigureFactory functions ##
 scipy

packages/python/plotly/plotly/express/__init__.py

@@ -2,15 +2,6 @@
 `plotly.express` is a terse, consistent, high-level wrapper around `plotly.graph_objects`
 for rapid data exploration and figure generation. Learn more at https://plotly.com/python/plotly-express/
 """
-from plotly import optional_imports
-
-pd = optional_imports.get_module("pandas")
-if pd is None:
-    raise ImportError(
-        """\
-Plotly express requires pandas to be installed."""
-    )
-
 from ._imshow import imshow
 from ._chart_types import (  # noqa: F401
     scatter,

packages/python/plotly/plotly/express/_doc.py

@@ -502,7 +502,7 @@
     zoom=["int (default `8`)", "Between 0 and 20.", "Sets map zoom level."],
     orientation=[
         "str, one of `'h'` for horizontal or `'v'` for vertical. ",
-        "(default `'v'` if `x` and `y` are provided and both continous or both categorical, ",
+        "(default `'v'` if `x` and `y` are provided and both continuous or both categorical, ",
         "otherwise `'v'`(`'h'`) if `x`(`y`) is categorical and `y`(`x`) is continuous, ",
         "otherwise `'v'`(`'h'`) if only `x`(`y`) is provided) ",
     ],

packages/python/plotly/plotly/express/_imshow.py

@@ -2,7 +2,7 @@
 from _plotly_utils.basevalidators import ColorscaleValidator
 from ._core import apply_default_cascade, init_figure, configure_animation_controls
 from .imshow_utils import rescale_intensity, _integer_ranges, _integer_types
-import pandas as pd
+import narwhals.stable.v1 as nw
 import numpy as np
 import itertools
 from plotly.utils import image_array_to_data_uri
@@ -321,10 +321,12 @@ def imshow(
             aspect = "equal"
 
     # --- Set the value of binary_string (forbidden for pandas)
-    if isinstance(img, pd.DataFrame):
+    img = nw.from_native(img, strict=False)
+    if isinstance(img, nw.DataFrame):
         if binary_string:
             raise ValueError("Binary strings cannot be used with pandas arrays")
         is_dataframe = True
+        img = img.to_numpy()
     else:
         is_dataframe = False
 

packages/python/plotly/plotly/express/trendline_functions/__init__.py

@@ -8,8 +8,7 @@
 exposed as part of the public API for documentation purposes.
 """
 
-import pandas as pd
-import numpy as np
+import narwhals.stable.v1 as nw
 
 __all__ = ["ols", "lowess", "rolling", "ewm", "expanding"]
 
@@ -32,6 +31,8 @@ def ols(trendline_options, x_raw, x, y, x_label, y_label, non_missing):
     respect to the base 10 logarithm of the input. Note that this means no zeros can
     be present in the input.
     """
+    import numpy as np
+
     valid_options = ["add_constant", "log_x", "log_y"]
     for k in trendline_options.keys():
         if k not in valid_options:
@@ -110,11 +111,22 @@ def lowess(trendline_options, x_raw, x, y, x_label, y_label, non_missing):
 
 
 def _pandas(mode, trendline_options, x_raw, y, non_missing):
+    import numpy as np
+
+    try:
+        import pandas as pd
+    except ImportError:
+        msg = "Trendline requires pandas to be installed"
+        raise ImportError(msg)
+
     modes = dict(rolling="Rolling", ewm="Exponentially Weighted", expanding="Expanding")
     trendline_options = trendline_options.copy()
     function_name = trendline_options.pop("function", "mean")
     function_args = trendline_options.pop("function_args", dict())
-    series = pd.Series(y, index=x_raw)
+
+    series = pd.Series(np.copy(y), index=x_raw.to_pandas())
+
+    # TODO: If narwhals were to support rolling, ewm and expanding then we could go around these
     agg = getattr(series, mode)  # e.g. series.rolling
     agg_obj = agg(**trendline_options)  # e.g. series.rolling(**opts)
     function = getattr(agg_obj, function_name)  # e.g. series.rolling(**opts).mean

packages/python/plotly/plotly/figure_factory/_hexbin_mapbox.py

@@ -1,8 +1,8 @@
 from plotly.express._core import build_dataframe
 from plotly.express._doc import make_docstring
 from plotly.express._chart_types import choropleth_mapbox, scatter_mapbox
+import narwhals.stable.v1 as nw
 import numpy as np
-import pandas as pd
 
 
 def _project_latlon_to_wgs84(lat, lon):
@@ -231,6 +231,7 @@ def _compute_wgs84_hexbin(
     nx=None,
     agg_func=None,
     min_count=None,
+    native_namespace=None,
 ):
     """
     Computes the lat-lon aggregation at hexagonal bin level.
@@ -263,7 +264,7 @@ def _compute_wgs84_hexbin(
         Lat coordinates of each hexagon (shape M x 6)
     np.ndarray
         Lon coordinates of each hexagon (shape M x 6)
-    pd.Series
+    nw.Series
         Unique id for each hexagon, to be used in the geojson data (shape M)
     np.ndarray
         Aggregated value in each hexagon (shape M)
@@ -288,7 +289,14 @@ def _compute_wgs84_hexbin(
 
     # Create unique feature id based on hexagon center
     centers = centers.astype(str)
-    hexagons_ids = pd.Series(centers[:, 0]) + "," + pd.Series(centers[:, 1])
+    hexagons_ids = (
+        nw.from_dict(
+            {"x1": centers[:, 0], "x2": centers[:, 1]},
+            native_namespace=native_namespace,
+        )
+        .select(hexagons_ids=nw.concat_str([nw.col("x1"), nw.col("x2")], separator=","))
+        .get_column("hexagons_ids")
+    )
 
     return hexagons_lats, hexagons_lons, hexagons_ids, agreggated_value
 
@@ -344,22 +352,40 @@ def create_hexbin_mapbox(
     Returns a figure aggregating scattered points into connected hexagons
     """
     args = build_dataframe(args=locals(), constructor=None)
-
+    native_namespace = nw.get_native_namespace(args["data_frame"])
     if agg_func is None:
         agg_func = np.mean
 
-    lat_range = args["data_frame"][args["lat"]].agg(["min", "max"]).values
-    lon_range = args["data_frame"][args["lon"]].agg(["min", "max"]).values
+    lat_range = (
+        args["data_frame"]
+        .select(
+            nw.min(args["lat"]).name.suffix("_min"),
+            nw.max(args["lat"]).name.suffix("_max"),
+        )
+        .to_numpy()
+        .squeeze()
+    )
+
+    lon_range = (
+        args["data_frame"]
+        .select(
+            nw.min(args["lon"]).name.suffix("_min"),
+            nw.max(args["lon"]).name.suffix("_max"),
+        )
+        .to_numpy()
+        .squeeze()
+    )
 
     hexagons_lats, hexagons_lons, hexagons_ids, count = _compute_wgs84_hexbin(
-        lat=args["data_frame"][args["lat"]].values,
-        lon=args["data_frame"][args["lon"]].values,
+        lat=args["data_frame"].get_column(args["lat"]).to_numpy(),
+        lon=args["data_frame"].get_column(args["lon"]).to_numpy(),
         lat_range=lat_range,
         lon_range=lon_range,
         color=None,
         nx=nx_hexagon,
         agg_func=agg_func,
         min_count=min_count,
+        native_namespace=native_namespace,
     )
 
     geojson = _hexagons_to_geojson(hexagons_lats, hexagons_lons, hexagons_ids)
@@ -381,41 +407,43 @@ def create_hexbin_mapbox(
         center = dict(lat=lat_range.mean(), lon=lon_range.mean())
 
     if args["animation_frame"] is not None:
-        groups = args["data_frame"].groupby(args["animation_frame"]).groups
+        groups = dict(args["data_frame"].group_by(args["animation_frame"]).__iter__())
     else:
-        groups = {0: args["data_frame"].index}
+        groups = {(0,): args["data_frame"]}
 
     agg_data_frame_list = []
-    for frame, index in groups.items():
-        df = args["data_frame"].loc[index]
+    for key, df in groups.items():
         _, _, hexagons_ids, aggregated_value = _compute_wgs84_hexbin(
-            lat=df[args["lat"]].values,
-            lon=df[args["lon"]].values,
+            lat=df.get_column(args["lat"]).to_numpy(),
+            lon=df.get_column(args["lon"]).to_numpy(),
             lat_range=lat_range,
             lon_range=lon_range,
-            color=df[args["color"]].values if args["color"] else None,
+            color=df.get_column(args["color"]).to_numpy() if args["color"] else None,
             nx=nx_hexagon,
             agg_func=agg_func,
             min_count=min_count,
+            native_namespace=native_namespace,
         )
         agg_data_frame_list.append(
-            pd.DataFrame(
-                np.c_[hexagons_ids, aggregated_value], columns=["locations", "color"]
+            nw.from_dict(
+                {
+                    "frame": [key[0]] * len(hexagons_ids),
+                    "locations": hexagons_ids,
+                    "color": aggregated_value,
+                },
+                native_namespace=native_namespace,
             )
         )
-    agg_data_frame = (
-        pd.concat(agg_data_frame_list, axis=0, keys=groups.keys())
-        .rename_axis(index=("frame", "index"))
-        .reset_index("frame")
-    )
 
-    agg_data_frame["color"] = pd.to_numeric(agg_data_frame["color"])
+    agg_data_frame = nw.concat(agg_data_frame_list, how="vertical").with_columns(
+        color=nw.col("color").cast(nw.Int64)
+    )
 
     if range_color is None:
         range_color = [agg_data_frame["color"].min(), agg_data_frame["color"].max()]
 
     fig = choropleth_mapbox(
-        data_frame=agg_data_frame,
+        data_frame=agg_data_frame.to_native(),
         geojson=geojson,
         locations="locations",
         color="color",
@@ -440,7 +468,9 @@ def create_hexbin_mapbox(
     if show_original_data:
         original_fig = scatter_mapbox(
             data_frame=(
-                args["data_frame"].sort_values(by=args["animation_frame"])
+                args["data_frame"].sort(
+                    by=args["animation_frame"], descending=False, nulls_last=True
+                )
                 if args["animation_frame"] is not None
                 else args["data_frame"]
             ),

packages/python/plotly/plotly/tests/test_optional/test_figure_factory/test_figure_factory.py

@@ -4474,7 +4474,7 @@ def test_build_dataframe(self):
         lon = np.random.randn(N)
         color = np.ones(N)
         frame = np.random.randint(0, n_frames, N)
-        df = pd.DataFrame(
+        df = pd.DataFrame(  # TODO: Test other constructors?
             np.c_[lat, lon, color, frame],
             columns=["Latitude", "Longitude", "Metric", "Frame"],
         )

packages/python/plotly/plotly/tests/test_optional/test_px/conftest.py

@@ -0,0 +1,41 @@
+import pandas as pd
+import polars as pl
+import pyarrow as pa
+import pytest
+
+from narwhals.typing import IntoDataFrame
+from narwhals.utils import parse_version
+
+
+def pandas_constructor(obj) -> IntoDataFrame:
+    return pd.DataFrame(obj)  # type: ignore[no-any-return]
+
+
+def pandas_nullable_constructor(obj) -> IntoDataFrame:
+    return pd.DataFrame(obj).convert_dtypes(dtype_backend="numpy_nullable")  # type: ignore[no-any-return]
+
+
+def pandas_pyarrow_constructor(obj) -> IntoDataFrame:
+    return pd.DataFrame(obj).convert_dtypes(dtype_backend="pyarrow")  # type: ignore[no-any-return]
+
+
+def polars_eager_constructor(obj) -> IntoDataFrame:
+    return pl.DataFrame(obj)
+
+
+def pyarrow_table_constructor(obj) -> IntoDataFrame:
+    return pa.table(obj)  # type: ignore[no-any-return]
+
+
+constructors = [polars_eager_constructor, pyarrow_table_constructor, pandas_constructor]
+
+if parse_version(pd.__version__) >= parse_version("2.0.0"):
+    constructors = [
+        pandas_nullable_constructor,
+        pandas_pyarrow_constructor,
+    ]
+
+
+@pytest.fixture(params=constructors)
+def constructor(request: pytest.FixtureRequest):
+    return request.param  # type: ignore[no-any-return]

packages/python/plotly/plotly/tests/test_optional/test_px/test_facets.py

@@ -1,12 +1,13 @@
-import pandas as pd
 import plotly.express as px
 from pytest import approx
 import pytest
 import random
 
 
-def test_facets():
-    df = px.data.tips()
+def test_facets(constructor):
+    data = px.data.tips().to_dict(orient="list")
+    df = constructor(data)
+
     fig = px.scatter(df, x="total_bill", y="tip")
     assert "xaxis2" not in fig.layout
     assert "yaxis2" not in fig.layout
@@ -46,8 +47,9 @@ def test_facets():
     assert fig.layout.yaxis4.domain[0] - fig.layout.yaxis.domain[1] == approx(0.08)
 
 
-def test_facets_with_marginals():
-    df = px.data.tips()
+def test_facets_with_marginals(constructor):
+    data = px.data.tips().to_dict(orient="list")
+    df = constructor(data)
 
     fig = px.histogram(df, x="total_bill", facet_col="sex", marginal="rug")
     assert len(fig.data) == 4
@@ -93,12 +95,11 @@ def test_facets_with_marginals():
     assert len(fig.data) == 2  # ignore all marginals
 
 
-@pytest.fixture
-def bad_facet_spacing_df():
+def bad_facet_spacing_df(constructor_func):
     NROWS = 101
     NDATA = 1000
     categories = [n % NROWS for n in range(NDATA)]
-    df = pd.DataFrame(
+    df = constructor_func(
         {
             "x": [random.random() for _ in range(NDATA)],
             "y": [random.random() for _ in range(NDATA)],
@@ -108,8 +109,8 @@ def bad_facet_spacing_df():
     return df
 
 
-def test_bad_facet_spacing_eror(bad_facet_spacing_df):
-    df = bad_facet_spacing_df
+def test_bad_facet_spacing_error(constructor):
+    df = bad_facet_spacing_df(constructor_func=constructor)
     with pytest.raises(
         ValueError, match="Use the facet_row_spacing argument to adjust this spacing."
     ):

packages/python/plotly/plotly/tests/test_optional/test_px/test_marginals.py

@@ -5,8 +5,9 @@
 @pytest.mark.parametrize("px_fn", [px.scatter, px.density_heatmap, px.density_contour])
 @pytest.mark.parametrize("marginal_x", [None, "histogram", "box", "violin"])
 @pytest.mark.parametrize("marginal_y", [None, "rug"])
-def test_xy_marginals(px_fn, marginal_x, marginal_y):
-    df = px.data.tips()
+def test_xy_marginals(constructor, px_fn, marginal_x, marginal_y):
+    data = px.data.tips().to_dict(orient="list")
+    df = constructor(data)
 
     fig = px_fn(
         df, x="total_bill", y="tip", marginal_x=marginal_x, marginal_y=marginal_y
@@ -17,8 +18,9 @@ def test_xy_marginals(px_fn, marginal_x, marginal_y):
 @pytest.mark.parametrize("px_fn", [px.histogram, px.ecdf])
 @pytest.mark.parametrize("marginal", [None, "rug", "histogram", "box", "violin"])
 @pytest.mark.parametrize("orientation", ["h", "v"])
-def test_single_marginals(px_fn, marginal, orientation):
-    df = px.data.tips()
+def test_single_marginals(constructor, px_fn, marginal, orientation):
+    data = px.data.tips().to_dict(orient="list")
+    df = constructor(data)
 
     fig = px_fn(
         df, x="total_bill", y="total_bill", marginal=marginal, orientation=orientation