ENH: speed up very-wide DataFrame line plots via single-pass LineCollection

Author: evgenii

pandas/plotting/_matplotlib/core.py

@@ -51,6 +51,7 @@
     ABCMultiIndex,
     ABCPeriodIndex,
     ABCSeries,
+    ABCTimedeltaIndex,
 )
 from pandas.core.dtypes.missing import isna
 
@@ -99,7 +100,7 @@
         Series,
     )
 
-import itertools
+from itertools import islice
 
 from matplotlib.collections import LineCollection
 
@@ -1538,6 +1539,7 @@ def _make_legend(self) -> None:
 
 class LinePlot(MPLPlot):
     _default_rot = 0
+    _wide_line_threshold: int = 200
 
     @property
     def orientation(self) -> PlottingOrientation:
@@ -1547,97 +1549,101 @@ def orientation(self) -> PlottingOrientation:
     def _kind(self) -> Literal["line", "area", "hist", "kde", "box"]:
         return "line"
 
-    def __init__(self, data, **kwargs) -> None:
-        MPLPlot.__init__(self, data, **kwargs)
+    def __init__(self, data, **kwargs):
+        super().__init__(data, **kwargs)
         if self.stacked:
-            self.data = self.data.fillna(value=0)
+            self.data = self.data.fillna(0)
 
-    def _make_plot(self, fig: Figure) -> None:
-        threshold = 200  # switch when DataFrame has more than this many columns
-        can_use_lc = (
-            not self._is_ts_plot()  # not a TS plot
-            and not self.stacked  # stacking not requested
-            and not com.any_not_none(*self.errors.values())  # no error bars
-            and len(self.data.columns) > threshold
-        )
-        if can_use_lc:
-            ax = self._get_ax(0)
-            x = self._get_xticks()
-            segments = [
-                np.column_stack((x, self.data[col].values)) for col in self.data.columns
-            ]
-            base_colors = mpl.rcParams["axes.prop_cycle"].by_key()["color"]
-            colors = list(itertools.islice(itertools.cycle(base_colors), len(segments)))
-            lc = LineCollection(
-                segments,
-                colors=colors,
-                linewidths=self.kwds.get("linewidth", mpl.rcParams["lines.linewidth"]),
-            )
-            ax.add_collection(lc)
-            ax.margins(0.05)
-            return  # skip the per-column Line2D loop
+    def _make_plot(self, fig):
+        is_ts = self._is_ts_plot()
 
-        if self._is_ts_plot():
+        # establish X, iterator, plot function
+        if is_ts:
             data = maybe_convert_index(self._get_ax(0), self.data)
-
-            x = data.index  # dummy, not used
-            plotf = self._ts_plot
-            it = data.items()
+            x_vals, iterator, plotf = data.index, data.items(), self._ts_plot
         else:
-            x = self._get_xticks()
-            # error: Incompatible types in assignment (expression has type
-            # "Callable[[Any, Any, Any, Any, Any, Any, KwArg(Any)], Any]", variable has
-            # type "Callable[[Any, Any, Any, Any, KwArg(Any)], Any]")
-            plotf = self._plot  # type: ignore[assignment]
-            # error: Incompatible types in assignment (expression has type
-            # "Iterator[tuple[Hashable, ndarray[Any, Any]]]", variable has
-            # type "Iterable[tuple[Hashable, Series]]")
-            it = self._iter_data(data=self.data)  # type: ignore[assignment]
+            x_vals = self._get_xticks()
+            iterator, plotf = self._iter_data(self.data), self._plot
 
+        n_cols = len(self.data.columns)
         stacking_id = self._get_stacking_id()
         is_errorbar = com.any_not_none(*self.errors.values())
+        default_lw = self.kwds.get("linewidth", mpl.rcParams["lines.linewidth"])
+        colours_seq = self._get_colors()
+        colours = list(islice(colours_seq, n_cols))
+
+        # fast aggregate-LC path (numeric, wide, simple)
+        idx = self.data.index
+        numeric_idx = is_any_real_numeric_dtype(idx.dtype)
+
+        complex_idx = isinstance(
+            idx, (ABCTimedeltaIndex, ABCPeriodIndex, ABCDatetimeIndex)
+        ) or getattr(idx, "_is_all_dates", False)
+
+        lc_fast_path = (
+            n_cols > self._wide_line_threshold
+            and numeric_idx
+            and not complex_idx
+            and not is_ts
+            and not self.stacked
+            and not is_errorbar
+            and self.style is None
+            and "marker" not in self.kwds
+            and "linestyle" not in self.kwds
+        )
 
-        colors = self._get_colors()
-        for i, (label, y) in enumerate(it):
-            ax = self._get_ax(i)
+        if lc_fast_path:
+            ax = self._get_ax(0)
+
+            segs = np.empty((n_cols, len(x_vals), 2), dtype=float)
+            segs[:, :, 0] = np.asarray(x_vals, dtype=float)
+            segs[:, :, 1] = self.data.values.T
+
+            lc = LineCollection(segs, colors=colours, linewidths=default_lw)
+            ax.add_collection(lc)
+            ax.relim()
+            ax.autoscale()
+
+            if self.legend:
+                for i, col in enumerate(self.data.columns):
+                    h = mpl.lines.Line2D(
+                        [], [], color=colours[i], label=pprint_thing(col)
+                    )
+                    self._append_legend_handles_labels(h, h.get_label())
+
+            self._post_plot_logic(ax, self.data)
+            return
+
+        # unified per-column loop (complex or narrow frames)
+        for i, (label, y) in enumerate(iterator):
+            ax: Axes = self._get_ax(i)
             kwds = self.kwds.copy()
             if self.color is not None:
                 kwds["color"] = self.color
-            style, kwds = self._apply_style_colors(
-                colors,
-                kwds,
-                i,
-                # error: Argument 4 to "_apply_style_colors" of "MPLPlot" has
-                # incompatible type "Hashable"; expected "str"
-                label,  # type: ignore[arg-type]
-            )
 
-            errors = self._get_errorbars(label=label, index=i)
-            kwds = dict(kwds, **errors)
-
-            label = pprint_thing(label)
-            label = self._mark_right_label(label, index=i)
-            kwds["label"] = label
+            style, kwds = self._apply_style_colors(colours_seq, kwds, i, label)
+            kwds.update(self._get_errorbars(label, i))
+            kwds["label"] = self._mark_right_label(pprint_thing(label), index=i)
 
             newlines = plotf(
                 ax,
-                x,
+                x_vals,
                 y,
                 style=style,
                 column_num=i,
                 stacking_id=stacking_id,
                 is_errorbar=is_errorbar,
                 **kwds,
             )
-            self._append_legend_handles_labels(newlines[0], label)
+            self._append_legend_handles_labels(newlines[0], kwds["label"])
 
-            if self._is_ts_plot():
-                # reset of xlim should be used for ts data
-                # TODO: GH28021, should find a way to change view limit on xaxis
-                lines = get_all_lines(ax)
-                left, right = get_xlim(lines)
+            if is_ts:  # keep TS rescale
+                left, right = get_xlim(get_all_lines(ax))
                 ax.set_xlim(left, right)
 
+        # shared clean-up
+        self._post_plot_logic(self._get_ax(0), self.data)
+
     # error: Signature of "_plot" incompatible with supertype "MPLPlot"
     @classmethod
     def _plot(  # type: ignore[override]