feat: make plotly-express dataframe agnostic via narwhals

packages/python/plotly/_plotly_utils/basevalidators.py

@@ -8,6 +8,7 @@
 import re
 import sys
 import warnings
+import narwhals.stable.v1 as nw
 
 from _plotly_utils.optional_imports import get_module
 
@@ -93,8 +94,19 @@ def copy_to_readonly_numpy_array(v, kind=None, force_numeric=False):
         "O": "object",
     }
 
-    # Handle pandas Series and Index objects
+    if isinstance(v, nw.Series):
+        if nw.dependencies.is_pandas_like_series(v_native := v.to_native()):
+            v = v_native
+        else:
+            v = v.to_numpy()
+    elif isinstance(v, nw.DataFrame):
+        if nw.dependencies.is_pandas_like_dataframe(v_native := v.to_native()):
+            v = v_native
+        else:
+            v = v.to_numpy()
+
     if pd and isinstance(v, (pd.Series, pd.Index)):
+        # Handle pandas Series and Index objects
         if v.dtype.kind in numeric_kinds:
             # Get the numeric numpy array so we use fast path below
             v = v.values
@@ -193,6 +205,7 @@ def is_homogeneous_array(v):
         np
         and isinstance(v, np.ndarray)
         or (pd and isinstance(v, (pd.Series, pd.Index)))
+        or (isinstance(v, nw.Series))
     ):
         return True
     if is_numpy_convertable(v):

packages/python/plotly/optional-requirements.txt

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

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

@@ -4,11 +4,11 @@
 """
 from plotly import optional_imports
 
-pd = optional_imports.get_module("pandas")
-if pd is None:
+np = optional_imports.get_module("numpy")
+if np is None:
     raise ImportError(
         """\
-Plotly express requires pandas to be installed."""
+Plotly express requires numpy to be installed."""
     )
 
 from ._imshow import imshow

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

@@ -503,7 +503,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,7 +321,8 @@ 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

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

@@ -8,9 +8,6 @@
 exposed as part of the public API for documentation purposes.
 """
 
-import pandas as pd
-import numpy as np
-
 __all__ = ["ols", "lowess", "rolling", "ewm", "expanding"]
 
 
@@ -32,6 +29,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 +109,25 @@ 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: Narwhals Series/DataFrame do not support rolling, ewm nor expanding, therefore
+    # it fallbacks to pandas Series independently of the original type.
+    # Plotly issue: https://github.com/plotly/plotly.py/issues/4834
+    # Narwhals issue: https://github.com/narwhals-dev/narwhals/issues/1254
     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,47 @@ 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"], drop_null_keys=True)
+            .__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,10 +472,12 @@ 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"]
-            ),
+            ).to_native(),
             lat=args["lat"],
             lon=args["lon"],
             animation_frame=args["animation_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,43 @@
+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.extend(
+        [
+            pandas_nullable_constructor,
+            pandas_pyarrow_constructor,
+        ]
+    )
+
+
+@pytest.fixture(params=constructors)
+def constructor(request: pytest.FixtureRequest):
+    return request.param  # type: ignore[no-any-return]